74 lines
2.4 KiB
Python
Executable File
74 lines
2.4 KiB
Python
Executable File
"""
|
|
RAG Tool
|
|
Search the knowledge base for relevant documents.
|
|
"""
|
|
from typing import Dict, Any, Optional
|
|
from loguru import logger
|
|
|
|
from tools.base import BaseTool, ToolResult
|
|
|
|
|
|
class RAGTool(BaseTool):
|
|
"""Search the RAG knowledge base."""
|
|
|
|
def __init__(self, rag_store, config: Optional[Dict] = None):
|
|
self.rag_store = rag_store
|
|
super().__init__(config)
|
|
|
|
@property
|
|
def name(self) -> str:
|
|
return "search_knowledge_base"
|
|
|
|
@property
|
|
def description(self) -> str:
|
|
return "Search uploaded documents for relevant information. Use this for information from uploaded files, documents, or custom knowledge."
|
|
|
|
@property
|
|
def parameters(self) -> Dict[str, Any]:
|
|
return {
|
|
"type": "object",
|
|
"properties": {
|
|
"query": {
|
|
"type": "string",
|
|
"description": "The search query"
|
|
},
|
|
"top_k": {
|
|
"type": "integer",
|
|
"description": "Number of results to return (default: 5)",
|
|
"default": 5
|
|
}
|
|
},
|
|
"required": ["query"]
|
|
}
|
|
|
|
async def execute(self, query: str, top_k: int = 5, **kwargs) -> ToolResult:
|
|
"""Execute RAG search."""
|
|
self._log_execution({"query": query, "top_k": top_k})
|
|
|
|
try:
|
|
results = await self.rag_store.search(query, top_k=top_k)
|
|
|
|
if not results:
|
|
return ToolResult(
|
|
success=True,
|
|
data="No relevant documents found in the knowledge base."
|
|
)
|
|
|
|
# Format results
|
|
formatted_results = []
|
|
for i, result in enumerate(results, 1):
|
|
formatted_results.append(
|
|
f"{i}. From '{result.get('document_name', 'Unknown')}':\n"
|
|
f" {result.get('content', '')}\n"
|
|
f" Relevance: {result.get('score', 0):.2f}"
|
|
)
|
|
|
|
output = f"Knowledge base results for '{query}':\n\n" + "\n\n".join(formatted_results)
|
|
|
|
self._log_success(output[:100])
|
|
return ToolResult(success=True, data=output)
|
|
|
|
except Exception as e:
|
|
self._log_error(str(e))
|
|
return ToolResult(success=False, error=str(e))
|