async def overwrite_refactored_search_results(callback_context: CallbackContext) -> Optional[types.Content]:

    """
    Owerwrite the response from the refactored_search_agent.
    :param callback_context: the context
    :return: a confirmation message that the search was successful or not.
    """

    # Get the name of the agent calling this callback
    agent_name = callback_context.agent_name
    print(f"[Callback] After model call for agent: {agent_name}")

    # List the existing state variables for debugging
    print(f"[Callback] Existing state variables: {callback_context.state.to_dict().keys()}")

    # check if the 'refactored_search_results' variable can be found in the state
    if 'refactored_search_results' in callback_context.state and 'search_results' in callback_context.state:
        return types.Content(
            parts=[types.Part(text="The original search results and the refactored search results have been preserved")],
            role="model"
        )
    else:
        return types.Content(
            parts=[types.Part(text="Sorry, I was not able to retrieve anything intersting from the Web.")],
            role="model"
        )

class informationChunk(BaseModel):
    companies: List[str] = Field(description="The companies mentioned in the news") 
    news: str = Field(description="The news content")
    date: str = Field(description="The date of the news")
    origin: str = Field(description="The origin of the news")

class informationChunks(BaseModel):
    informations: List[informationChunk] = Field(description="List of information chunks")
    date: str = Field(description="Today's date")

refactor_search_agent = LlmAgent(
    model='gemini-2.5-flash',
    name='refactor_search_agent',
    description='A helpful assistant for refactoring the information found on the web following instructions received from the user',
    instruction=REFACTOR_SEARCH_PROMPT,
    output_schema=informationChunks,
    output_key='refactored_search_results',
    after_agent_callback=overwrite_refactored_search_results # type: ignore
)