Coverage for src / dataknobs_bots / tools / knowledge_search.py: 54%

1"""Knowledge search tool for RAG integration.""" 

2 

3from typing import Any 

4 

5from dataknobs_llm.tools import Tool 

6 

7 

8class KnowledgeSearchTool(Tool): 

9 """Tool for searching the knowledge base. 

10 

11 This tool allows LLMs to search the bot's knowledge base 

12 for relevant information during conversations. 

13 

14 Example: 

15 ```python 

16 # Create tool with knowledge base 

17 tool = KnowledgeSearchTool(knowledge_base=kb) 

18 

19 # Register with bot 

20 bot.tool_registry.register_tool(tool) 

21 

22 # LLM can now call the tool 

23 results = await tool.execute( 

24 query="How do I configure the database?", 

25 max_results=3 

26 ) 

27 ``` 

28 """ 

29 

30 def __init__(self, knowledge_base: Any, name: str = "knowledge_search"): 

31 """Initialize knowledge search tool. 

32 

33 Args: 

34 knowledge_base: RAGKnowledgeBase instance to search 

35 name: Tool name (default: knowledge_search) 

36 """ 

37 super().__init__( 

38 name=name, 

39 description="Search the knowledge base for relevant information. " 

40 "Use this when you need to find documentation, examples, or " 

41 "specific information to answer user questions.", 

42 ) 

43 self.knowledge_base = knowledge_base 

44 

45 @property 

46 def schema(self) -> dict[str, Any]: 

47 """Get JSON schema for tool parameters. 

48 

49 Returns: 

50 JSON Schema for the tool parameters 

51 """ 

52 return { 

53 "type": "object", 

54 "properties": { 

55 "query": { 

56 "type": "string", 

57 "description": "The search query or question to find information about", 

58 }, 

59 "max_results": { 

60 "type": "integer", 

61 "description": "Maximum number of results to return", 

62 "default": 3, 

63 "minimum": 1, 

64 "maximum": 10, 

65 }, 

66 }, 

67 "required": ["query"], 

68 } 

69 

70 async def execute(self, query: str, max_results: int = 3, **kwargs: Any) -> dict[str, Any]: 

71 """Execute knowledge base search. 

72 

73 Args: 

74 query: Search query text 

75 max_results: Maximum number of results (default: 3) 

76 **kwargs: Additional arguments (ignored) 

77 

78 Returns: 

79 Dictionary with search results: 

80 - query: Original query 

81 - results: List of relevant chunks 

82 - num_results: Number of results found 

83 

84 Example: 

85 ```python 

86 result = await tool.execute( 

87 query="How do I configure the database?", 

88 max_results=3 

89 ) 

90 for chunk in result['results']: 

91 print(f"{chunk['heading_path']}: {chunk['text']}") 

92 ``` 

93 """ 

94 # Clamp max_results to valid range 

95 max_results = max(1, min(10, max_results)) 

96 

97 # Search knowledge base 

98 results = await self.knowledge_base.query(query, k=max_results) 

99 

100 # Format response 

101 return { 

102 "query": query, 

103 "results": [ 

104 { 

105 "text": r["text"], 

106 "source": r["source"], 

107 "heading": r["heading_path"], 

108 "similarity": round(r["similarity"], 3), 

109 } 

110 for r in results 

111 ], 

112 "num_results": len(results), 

113 }