base.py
 1  """Base web search tool interface."""
 2  
 3  from abc import ABC, abstractmethod
 4  from typing import Literal, Optional
 5  from .models import SearchResult
 6  
 7  
 8  class WebSearchTool(ABC):
 9      """Base class for web search tools."""
10      
11      def __init__(self, api_key: Optional[str] = None, **kwargs):
12          """Initialize the search tool.
13          
14          Args:
15              api_key: API key for the search provider (if required)
16              **kwargs: Additional provider-specific configuration
17          """
18          self.api_key = api_key
19          self.config = kwargs
20      
21      @abstractmethod
22      async def search(
23          self,
24          query: str,
25          max_results: int = 5,
26          search_depth: Literal["basic", "advanced"] = "basic",
27          **kwargs
28      ) -> SearchResult:
29          """Execute a search query.
30          
31          Args:
32              query: The search query string
33              max_results: Maximum number of results to return (1-10)
34              search_depth: Search depth - 'basic' for quick results, 'advanced' for comprehensive
35              **kwargs: Additional provider-specific parameters
36              
37          Returns:
38              SearchResult object containing the search results
39              
40          Raises:
41              Exception: If the search fails
42          """
43          raise NotImplementedError("Subclasses must implement search()")
44      
45      def get_tool_definition(self) -> dict:
46          """Get the tool definition for LLM function calling.
47          
48          Returns:
49              Dictionary containing the tool definition
50          """
51          return {
52              "type": "function",
53              "function": {
54                  "name": "web_search",
55                  "description": (
56                      "Search the web for current information. Use this when you need "
57                      "up-to-date information, facts, news, or data that may not be in "
58                      "your training data. Always cite sources using the citation format."
59                  ),
60                  "parameters": {
61                      "type": "object",
62                      "properties": {
63                          "query": {
64                              "type": "string",
65                              "description": "The search query"
66                          },
67                          "max_results": {
68                              "type": "integer",
69                              "description": "Maximum number of results (1-10)",
70                              "minimum": 1,
71                              "maximum": 10,
72                              "default": 5
73                          },
74                          "search_depth": {
75                              "type": "string",
76                              "enum": ["basic", "advanced"],
77                              "description": "Search depth: 'basic' for quick results, 'advanced' for comprehensive",
78                              "default": "basic"
79                          }
80                      },
81                      "required": ["query"]
82                  }
83              }
84          }