Spaces:
Sleeping
Sleeping
| """Tavily search integration for web-based research.""" | |
| from typing import List, Dict, Any, Optional | |
| from tavily import TavilyClient | |
| class TavilySearcher: | |
| """Wrapper for Tavily search API.""" | |
| def __init__(self, api_key: Optional[str] = None): | |
| """Initialize Tavily client. | |
| Args: | |
| api_key: Tavily API key | |
| """ | |
| self.api_key = api_key | |
| self.client = TavilyClient(api_key=api_key) if api_key else None | |
| def search( | |
| self, | |
| query: str, | |
| max_results: int = 5, | |
| search_depth: str = "advanced" | |
| ) -> Dict[str, Any]: | |
| """Perform a search using Tavily. | |
| Args: | |
| query: Search query | |
| max_results: Maximum number of results to return | |
| search_depth: Search depth ('basic' or 'advanced') | |
| Returns: | |
| Search results dict with 'results' and 'answer' keys | |
| """ | |
| if not self.client: | |
| raise ValueError("Tavily API key not configured") | |
| try: | |
| response = self.client.search( | |
| query=query, | |
| max_results=max_results, | |
| search_depth=search_depth, | |
| include_answer=True, | |
| include_raw_content=False | |
| ) | |
| return response | |
| except Exception as e: | |
| raise Exception(f"Tavily search error: {str(e)}") | |
| def format_search_context(self, search_results: Dict[str, Any]) -> str: | |
| """Format Tavily search results into readable context. | |
| Args: | |
| search_results: Raw Tavily search results | |
| Returns: | |
| Formatted context string | |
| """ | |
| if not search_results: | |
| return "" | |
| context_parts = [] | |
| # Add Tavily's AI-generated answer if available | |
| if "answer" in search_results and search_results["answer"]: | |
| context_parts.append(f"**Tavily AI Summary:**\n{search_results['answer']}\n") | |
| # Add individual search results | |
| if "results" in search_results and search_results["results"]: | |
| context_parts.append("**Web Search Results:**\n") | |
| for i, result in enumerate(search_results["results"], 1): | |
| title = result.get("title", "No title") | |
| url = result.get("url", "") | |
| content = result.get("content", "") | |
| context_parts.append( | |
| f"{i}. **{title}**\n" | |
| f" URL: {url}\n" | |
| f" {content}\n" | |
| ) | |
| return "\n".join(context_parts) | |
| async def async_search( | |
| self, | |
| query: str, | |
| max_results: int = 5, | |
| search_depth: str = "advanced" | |
| ) -> Dict[str, Any]: | |
| """Async wrapper for search. | |
| Args: | |
| query: Search query | |
| max_results: Maximum number of results | |
| search_depth: Search depth | |
| Returns: | |
| Search results dict | |
| """ | |
| import asyncio | |
| loop = asyncio.get_event_loop() | |
| return await loop.run_in_executor( | |
| None, | |
| lambda: self.search(query, max_results, search_depth) | |
| ) | |