module_agent_web_search.py

"""
File: web_app/module_agent_web_search.py
Description: Gradio module for the Agent Web Search functionality.
Author: Didier Guillevic
Date: 2025-10-20
"""

import gradio as gr

from google.adk.agents import Agent
from google.adk.runners import Runner
from google.adk.sessions import InMemorySessionService
from google.adk.tools import google_search
from google.genai import types

import asyncio
import uuid

APP_NAME="google_search_agent"
SESSION_ID="1234"

model = "gemini-2.5-flash"

#
# ===== agent =====
#
root_agent = Agent(
    name="basic_search_agent",
    model=model,
    description=(
        "Agent to answer questions with the option to call Google Search "
        "if needed for up-to-date information."
    ),
    instruction=(
        "I can answer your questions from my own knowledge or by searching the "
        "web using Google Search. Just ask me anything!"
    ),
    # google_search: pre-built tool allows agent to perform Google searches.
    tools=[google_search]
)

#
# ===== Session and Runner =====
#
async def setup_session_and_runner(user_id: str):
    session_service = InMemorySessionService()
    session = await session_service.create_session(
        app_name=APP_NAME,
        user_id=user_id,
        session_id=SESSION_ID
    )
    runner = Runner(
        agent=root_agent,
        app_name=APP_NAME,
        session_service=session_service
    )
    return session, runner


#
# ===== Call Agent Asynchronously =====
#
async def call_agent_async(query: str, user_id: str):
    content = types.Content(role='user', parts=[types.Part(text=query)])
    session, runner = await setup_session_and_runner(user_id=user_id)
    events = runner.run_async(
        user_id=user_id,
        session_id=SESSION_ID,
        new_message=content
    )

    final_response = ""
    rendered_content = ""

    async for event in events:
        if event.is_final_response():
            final_response = event.content.parts[0].text

            # Check if the event has grounding metadata and rendered content
            if (
                event.grounding_metadata and
                event.grounding_metadata.search_entry_point and
                event.grounding_metadata.search_entry_point.rendered_content
            ):
                rendered_content = event.grounding_metadata.search_entry_point.rendered_content
            else:
                rendered_content = None

    return final_response, rendered_content


#
# ===== Call Agent Asynchronously with Streaming =====
#
async def call_agent_streaming(query: str, user_id: str):
    content = types.Content(role='user', parts=[types.Part(text=query)])
    session, runner = await setup_session_and_runner(user_id=user_id)
    events = runner.run_async(
        user_id=user_id,
        session_id=SESSION_ID,
        new_message=content
    )

    accumulated_response = ""
    rendered_content = None # Initialize to None

    async for event in events:
        # Check for intermediate text parts to stream
        if event.content and event.content.parts and event.content.parts[0].text:
            # Accumulate and yield the new text
            new_text = event.content.parts[0].text
            accumulated_response += new_text
            yield accumulated_response, None, user_id # Yield the current text and empty grounding

        # When the final response event is received, capture the grounding content
        if event.is_final_response():
            # The final response text should already be in accumulated_response from earlier yields,
            # but we can ensure it's fully captured here.
            # accumulated_response = event.content.parts[0].text # The final text

            # Capture the rendered_content from grounding_metadata
            if (
                event.grounding_metadata and
                event.grounding_metadata.search_entry_point and
                event.grounding_metadata.search_entry_point.rendered_content
            ):
                rendered_content = event.grounding_metadata.search_entry_point.rendered_content
            
            # After the final response, yield one last time with the accumulated text AND the grounding content
            # This final yield updates the grounding block.
            yield accumulated_response, rendered_content, user_id 

    # If the grounding content wasn't in the final event (e.g., if no search was performed), 
    # make sure to yield the final accumulated text.
    if rendered_content is None:
        yield accumulated_response, None, user_id


#
# ===== User interface Block =====
#
def agent_web_search(query: str, user_id=None):
    """Calls a language model agent with Google Search tool to answer the query.

    Args:
        query (str): The user query.
        user_id (str, optional): The user ID for session management. If None, a new ID is generated. Defaults to None.

    Returns:
        tuple: A tuple containing the agent's response (str), rendered grounding content (str or None), and user_id (str).
    """
    if user_id is None:
        user_id = str(uuid.uuid4())  # Generate a unique user ID

    response, rendered_content = asyncio.run(call_agent_async(query, user_id))
    return response, rendered_content, user_id


async def agent_web_search_streaming(query: str, current_user_id: str | None):
    # If the user ID state is None (first run), generate a new one
    if current_user_id is None:
        user_id = str(uuid.uuid4())
    else:
        user_id = current_user_id

    # The user_id is passed as part of the yield from the generator
    # but we need to ensure the Gradio state is updated initially for the generator to use the correct ID.
    
    # Gradio handles the asynchronous generator return and streams the output to the UI.
    return call_agent_streaming(query, user_id)


with gr.Blocks() as demo:
    gr.Markdown(
        """
        **Agent with Google Search tool**: be patient :-) Currently looking into (async) streaming support...
        """
    )

    with gr.Row():
        input_text = gr.Textbox(
            lines=2,
            placeholder="Enter your query here...",
            label="Query",
            render=True
        )

    user_id = gr.State(None)

    with gr.Row():
        submit_button = gr.Button("Submit", variant="primary")
        clear_button = gr.Button("Clear", variant="secondary")
    
    with gr.Row():
        output_text = gr.Markdown(
            label="Agent Response",
            render=True
        )
    
    with gr.Row():
        grounding = gr.HTML(
            label="Grounding Content",
            render=True
        )
    
    with gr.Accordion("Examples", open=False):
        examples = gr.Examples(
            examples=[
                ["What is the prime number factorization of 15?",], # no need got Google Search
                ["Who won the Nobel Peace Prize in 2025?",],
                ["What is the weather like tomorrow in Montreal, Canada?",],
                ["What are the latest news about Graph Neural Networks?",],
            ],
            inputs=[input_text,],
            cache_examples=False,
            label="Click to use an example"
        )

    # ===== Button Actions =====
    submit_button.click(
        fn=agent_web_search,
        inputs=[input_text, user_id],
        outputs=[output_text, grounding, user_id]
    )
    clear_button.click(
        fn=lambda : ('', '', None),
        inputs=None,
        outputs=[input_text, output_text, grounding]
    )


if __name__ == "__main__":
    demo.launch(mcp_server=True)