Spaces:
Sleeping
Sleeping
| """Web search agent for online information.""" | |
| import logging | |
| from typing import Optional | |
| from src.agents.base_agent import BaseAgent | |
| from src.tools.web_search import get_web_search | |
| logger = logging.getLogger(__name__) | |
| class SearchAgent(BaseAgent): | |
| """Agent specialized in web search and online information.""" | |
| def __init__(self, use_planning: bool = True): | |
| """Initialize search agent.""" | |
| web_search = get_web_search() | |
| tools = [web_search.get_tool_schema()] | |
| super().__init__( | |
| name="search_agent", | |
| description=( | |
| "You are a specialized agent for searching the web and finding " | |
| "online information. You can search the internet to answer questions " | |
| "that require current or external information." | |
| ), | |
| tools=tools, | |
| use_memory=True, | |
| use_planning=use_planning, | |
| planning_type="react", | |
| ) | |
| # Register tool function (async wrapper) | |
| async def web_search_tool(query: str, max_results: int = 5): | |
| return await web_search.search(query, max_results) | |
| self.add_tool( | |
| tool=web_search.get_tool_schema(), | |
| tool_function=web_search_tool, | |
| ) | |
| self.web_search = web_search | |
| async def retrieve_context(self, query: str) -> str: | |
| """ | |
| Retrieve relevant context from web search. | |
| Args: | |
| query: User query | |
| Returns: | |
| Context string from web search results | |
| """ | |
| try: | |
| # Perform web search | |
| search_results = await self.web_search.search(query, max_results=5) | |
| if not search_results.get("success") or not search_results.get("results"): | |
| return "No relevant information found from web search." | |
| # Format results | |
| context_parts = ["Web search results:"] | |
| for i, result in enumerate(search_results["results"], 1): | |
| title = result.get("title", "No title") | |
| url = result.get("url", "") | |
| content = result.get("content", "")[:300] # Truncate | |
| context_parts.append(f"\n[{i}] {title}") | |
| context_parts.append(f"URL: {url}") | |
| context_parts.append(f"Content: {content}...") | |
| return "\n".join(context_parts) | |
| except Exception as e: | |
| logger.error(f"Error retrieving web context: {e}") | |
| return f"Error performing web search: {str(e)}" | |
| async def process( | |
| self, | |
| query: str, | |
| session_id: Optional[str] = None, | |
| context: Optional[str] = None, | |
| ) -> dict: | |
| """ | |
| Process query with web search. | |
| Args: | |
| query: User query | |
| session_id: Optional session ID | |
| context: Optional additional context | |
| Returns: | |
| Response dictionary | |
| """ | |
| # Retrieve web context | |
| web_context = await self.retrieve_context(query) | |
| # Combine with provided context | |
| full_context = web_context | |
| if context: | |
| full_context = f"{context}\n\n{web_context}" | |
| # Process using base agent (which will use planning if enabled) | |
| return await super().process(query, session_id, full_context) | |