Coverage for baibu/chat/tools.py: 100%

1"""Tools the assistant can call during a reply. 

2 

3A tool has a stable name, a description and a JSON schema for its arguments 

4(the format LiteLLM passes to models), and returns text for the model plus 

5optional sources to show under the reply. Tools are offered only when they 

6are configured; with none, replies are plain model calls. 

7""" 

8 

9import json 

10from dataclasses import dataclass 

11from dataclasses import field 

12 

13from django.conf import settings 

14 

15from . import search 

16 

17 

18class ToolError(Exception): 

19 """The tool failed. The model is told, and may answer without it.""" 

20 

21 

22@dataclass 

23class ToolResult: 

24 content: str 

25 sources: list[dict] = field(default_factory=list) 

26 provider: str = "" 

27 

28 

29class Tool: 

30 name = "" 

31 description = "" 

32 parameters: dict = {} 

33 

34 def is_available(self) -> bool: 

35 return True 

36 

37 def schema(self) -> dict: 

38 return { 

39 "type": "function", 

40 "function": {"name": self.name, "description": self.description, "parameters": self.parameters}, 

41 } 

42 

43 def run(self, arguments: dict) -> ToolResult: 

44 raise NotImplementedError 

45 

46 

47class InternetSearchTool(Tool): 

48 name = "internet_search" 

49 description = ( 

50 "Search the web for current or local information. Use it when the answer depends on recent events, " 

51 "places, opening times, prices or facts you are not sure of. Cite the sources you use." 

52 ) 

53 parameters = { 

54 "type": "object", 

55 "properties": {"query": {"type": "string", "description": "What to search for."}}, 

56 "required": ["query"], 

57 } 

58 

59 def is_available(self) -> bool: 

60 return bool(settings.CHAT_SEARCH_PROVIDER) 

61 

62 def run(self, arguments: dict) -> ToolResult: 

63 query = str(arguments.get("query") or "").strip()[:300] 

64 if not query: 

65 msg = "A search query is required." 

66 raise ToolError(msg) 

67 provider = search.get_provider() 

68 try: 

69 results = provider.search(query, max_results=settings.CHAT_SEARCH_MAX_RESULTS) 

70 except search.SearchError as exc: 

71 raise ToolError(str(exc)) from exc 

72 except Exception as exc: 

73 msg = f"Search failed: {type(exc).__name__}" 

74 raise ToolError(msg) from exc 

75 results = results[: settings.CHAT_SEARCH_MAX_RESULTS] 

76 payload = [r.as_dict() for r in results] 

77 content = json.dumps({"query": query, "results": payload}, ensure_ascii=False) 

78 sources = [{"title": r.title, "url": r.url} for r in results if r.url.startswith(("https://", "http://"))] 

79 return ToolResult(content=content, sources=sources, provider=provider.name) 

80 

81 

82TOOLS: dict[str, Tool] = {tool.name: tool for tool in (InternetSearchTool(),)} 

83 

84 

85def available_tools() -> list[Tool]: 

86 return [tool for tool in TOOLS.values() if tool.is_available()]