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

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

2

3 from typing import Any

4

5 from dataknobs_llm.tools import Tool

6

7

8 class 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 }