Web Search Agent
==============

This example demonstrates how to create an agent that can search the web for information. We'll use a simple web search API to enable the agent to find up-to-date information.

Implementing a Web Search Tool
---------------------------

First, let's create a web search tool using a hypothetical search API:

.. code-block:: python

    from brain.agents.tool import tool
    from pydantic import BaseModel, Field
    import aiohttp
    import os
    from typing import List

    class SearchResult(BaseModel):
        title: str
        snippet: str
        url: str

    class SearchResponse(BaseModel):
        results: List[SearchResult]
        total_results: int

    @tool()
    async def search_web(query: str, num_results: int = 5) -> SearchResponse:
        """
        Search the web for information
        
        Args:
            query: The search query
            num_results: Number of results to return (max 10)
            
        Returns:
            Search results including titles, snippets, and URLs
        """
        # In a real implementation, you would use a search API
        # For this example, we'll use a mock implementation
        
        # Limit number of results to a reasonable range
        num_results = min(max(1, num_results), 10)
        
        # Example using a hypothetical search API
        async with aiohttp.ClientSession() as session:
            params = {
                "q": query,
                "num": num_results,
                "api_key": os.environ.get("SEARCH_API_KEY")
            }
            async with session.get("https://api.example.com/search", params=params) as response:
                if response.status != 200:
                    # For the example, we'll return mock data instead of raising an exception
                    return SearchResponse(
                        results=[
                            SearchResult(
                                title="Mock search result",
                                snippet="This is a mock search result since the actual API call failed.",
                                url="https://example.com"
                            )
                        ],
                        total_results=1
                    )
                
                data = await response.json()
                
                # Transform API response to our model
                results = [
                    SearchResult(
                        title=item["title"],
                        snippet=item["snippet"],
                        url=item["url"]
                    )
                    for item in data["items"][:num_results]
                ]
                
                return SearchResponse(
                    results=results,
                    total_results=data["totalResults"]
                )

Creating a Mock Implementation
----------------------------

For testing without a real API, here's a mock implementation:

.. code-block:: python

    @tool()
    async def search_web(query: str, num_results: int = 5) -> SearchResponse:
        """
        Search the web for information
        
        Args:
            query: The search query
            num_results: Number of results to return (max 10)
            
        Returns:
            Search results including titles, snippets, and URLs
        """
        # Limit number of results
        num_results = min(max(1, num_results), 10)
        
        # Mock data based on the query
        mock_results = []
        for i in range(num_results):
            mock_results.append(
                SearchResult(
                    title=f"Result {i+1} for {query}",
                    snippet=f"This is a snippet of information about {query}. This is just mock data for demonstration purposes.",
                    url=f"https://example.com/result/{i+1}"
                )
            )
        
        return SearchResponse(
            results=mock_results,
            total_results=num_results
        )

Building the Web Search Agent
---------------------------

Now, let's create an agent that uses the web search tool:

.. code-block:: python

    import asyncio
    import os
    
    from brain.agents.agent import Agent
    from brain.agents.llm.openai import OpenAIBaseLLM
    from brain.agents.callback import callback
    
    @callback("message_stream.assistant")
    async def stream_to_console(agent, event, stream):
        print("\nAssistant: ", end="", flush=True)
        async for chunk in stream:
            if hasattr(chunk, "chunk"):
                print(chunk.chunk, end="", flush=True)
    
    async def main():
        # Initialize the LLM
        llm = OpenAIBaseLLM(
            api_key=os.environ.get("OPENAI_API_KEY"),
            default_model="gpt-4o"  # Using a more capable model for better search queries
        )
        
        # Create the agent with the web search tool
        agent = Agent(
            llm=llm,
            tools=[search_web],
            instructions="""
            You are a helpful web search assistant. You can search the web for information 
            to answer user questions. Use the search_web tool to find relevant information.
            
            Follow these guidelines:
            1. For factual questions, always search the web to provide up-to-date information
            2. When searching, use precise and specific search queries
            3. Include URLs in your responses so users can verify the information
            4. If the search doesn't return useful results, try reformulating the search query
            """,
            callbacks=[stream_to_console]
        )
        
        # Run a conversation loop
        print("Web Search Agent (type 'exit' to quit)")
        while True:
            # Get user input
            user_input = input("\nYou: ")
            if user_input.lower() == "exit":
                break
                
            # Process with the agent
            await agent.run(user_input)
            print()  # Add a newline after the response
    
    if __name__ == "__main__":
        asyncio.run(main())

Complete Example
--------------

Here's the complete example with the mock web search implementation:

.. code-block:: python

    import asyncio
    import os
    from typing import List
    
    from brain.agents.agent import Agent
    from brain.agents.callback import callback
    from brain.agents.llm.openai import OpenAIBaseLLM
    from brain.agents.tool import tool
    from pydantic import BaseModel
    
    class SearchResult(BaseModel):
        title: str
        snippet: str
        url: str
    
    class SearchResponse(BaseModel):
        results: List[SearchResult]
        total_results: int
    
    @tool()
    async def search_web(query: str, num_results: int = 5) -> SearchResponse:
        """
        Search the web for information
        
        Args:
            query: The search query
            num_results: Number of results to return (max 10)
            
        Returns:
            Search results including titles, snippets, and URLs
        """
        # Limit number of results
        num_results = min(max(1, num_results), 10)
        
        # Mock data based on the query
        mock_results = []
        for i in range(num_results):
            mock_results.append(
                SearchResult(
                    title=f"Result {i+1} for {query}",
                    snippet=f"This is a snippet of information about {query}. This is just mock data for demonstration purposes.",
                    url=f"https://example.com/result/{i+1}"
                )
            )
        
        return SearchResponse(
            results=mock_results,
            total_results=num_results
        )
    
    @callback("message_stream.assistant")
    async def stream_to_console(agent, event, stream):
        print("\nAssistant: ", end="", flush=True)
        async for chunk in stream:
            if hasattr(chunk, "chunk"):
                print(chunk.chunk, end="", flush=True)
    
    async def main():
        # Initialize the LLM
        llm = OpenAIBaseLLM(
            api_key=os.environ.get("OPENAI_API_KEY"),
            default_model="gpt-4o"  # Using a more capable model for better search queries
        )
        
        # Create the agent with the web search tool
        agent = Agent(
            llm=llm,
            tools=[search_web],
            instructions="""
            You are a helpful web search assistant. You can search the web for information 
            to answer user questions. Use the search_web tool to find relevant information.
            
            Follow these guidelines:
            1. For factual questions, always search the web to provide up-to-date information
            2. When searching, use precise and specific search queries
            3. Include URLs in your responses so users can verify the information
            4. If the search doesn't return useful results, try reformulating the search query
            """,
            callbacks=[stream_to_console]
        )
        
        # Run a conversation loop
        print("Web Search Agent (type 'exit' to quit)")
        while True:
            # Get user input
            user_input = input("\nYou: ")
            if user_input.lower() == "exit":
                break
                
            # Process with the agent
            await agent.run(user_input)
            print()  # Add a newline after the response
    
    if __name__ == "__main__":
        asyncio.run(main())