Spaces:
Running
Running
File size: 4,140 Bytes
afefaea | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 | """Search tool wrapper for search engine providers."""
from typing import Any, Optional
from dataclasses import dataclass
from app.utils.logging import get_logger
logger = get_logger(__name__)
@dataclass
class SearchResult:
"""Individual search result."""
title: str
url: str
snippet: str
position: int
source: str # Provider name
metadata: dict[str, Any] | None = None
@dataclass
class SearchResponse:
"""Response from a search query."""
query: str
results: list[SearchResult]
total_results: int
provider: str
success: bool
error: Optional[str] = None
class SearchTool:
"""
Search tool that wraps search engine providers.
Provides a unified interface for searching across different
search engine providers.
"""
def __init__(self, default_provider: str = "duckduckgo") -> None:
self.default_provider = default_provider
self._engine: Any = None
self._initialized: bool = False
async def initialize(self, engine: Any = None) -> None:
"""
Initialize the search tool with a search engine.
Args:
engine: SearchEngineRouter instance to use
"""
logger.info("Initializing SearchTool")
self._engine = engine
self._initialized = True
logger.info("SearchTool initialized")
async def shutdown(self) -> None:
"""Shutdown the search tool."""
logger.info("Shutting down SearchTool")
self._engine = None
self._initialized = False
async def search(
self,
query: str,
max_results: int = 10,
provider: Optional[str] = None,
) -> SearchResponse:
"""
Perform a search query.
Args:
query: Search query string
max_results: Maximum number of results to return
provider: Specific provider to use (optional)
Returns:
SearchResponse with results
"""
logger.info(f"Searching for: {query}")
provider_name = provider or self.default_provider
if not self._initialized or self._engine is None:
logger.warning("SearchTool not properly initialized, using stub response")
return SearchResponse(
query=query,
results=[],
total_results=0,
provider=provider_name,
success=False,
error="Search engine not initialized",
)
try:
# Delegate to search engine router
results = await self._engine.search(
query=query,
max_results=max_results,
provider=provider_name,
)
return SearchResponse(
query=query,
results=results,
total_results=len(results),
provider=provider_name,
success=True,
)
except Exception as e:
logger.error(f"Search failed: {e}")
return SearchResponse(
query=query,
results=[],
total_results=0,
provider=provider_name,
success=False,
error=str(e),
)
async def get_results(
self,
query: str,
max_results: int = 10,
provider: Optional[str] = None,
) -> list[SearchResult]:
"""
Get search results as a list.
Args:
query: Search query string
max_results: Maximum number of results to return
provider: Specific provider to use (optional)
Returns:
List of SearchResult objects
"""
response = await self.search(query, max_results, provider)
return response.results
def health_check(self) -> bool:
"""Check if the search tool is healthy."""
return self._initialized and self._engine is not None
@property
def is_initialized(self) -> bool:
"""Check if the search tool has been initialized."""
return self._initialized
|