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:

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:

@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:

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:

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())