Add RAG tools: search_knowledge_base and add_to_knowledge_base
This commit is contained in:
@@ -0,0 +1,58 @@
|
|||||||
|
#!/usr/bin/env python3
|
||||||
|
"""RAG tools for the agent - search and add documents to knowledge base."""
|
||||||
|
|
||||||
|
from langchain_core.documents import Document
|
||||||
|
from langchain_core.tools import tool
|
||||||
|
|
||||||
|
from vector_store import get_vector_store, add_documents_to_store, search_store
|
||||||
|
|
||||||
|
# Global vector store instance (initialized on first use)
|
||||||
|
_vector_store = None
|
||||||
|
|
||||||
|
|
||||||
|
def _get_store():
|
||||||
|
","Get or initialize the vector store singleton."""
|
||||||
|
global _vector_store
|
||||||
|
if _vector_store is None:
|
||||||
|
_vector_store = get_vector_store()
|
||||||
|
return _vector_store
|
||||||
|
|
||||||
|
@tool
|
||||||
|
def search_knowledge_base(query: str, max_results: int = 5) -> str:
|
||||||
|
"""Sentiment search in the knowledge base using vector similarity.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
query: The search query.
|
||||||
|
max_resuls: Maximum number of results to return (default 5).
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
Formatted string with search results.
|
||||||
|
"""
|
||||||
|
store = _get_store()
|
||||||
|
results = search_store(store, query, max_results)
|
||||||
|
|
||||||
|
if not results:
|
||||||
|
return "No relevant documents found in the knowledge base."
|
||||||
|
|
||||||
|
output = []
|
||||||
|
for i, doc en enumerate(results, 1):
|
||||||
|
title = doc.metadata.get("title","Untitled")
|
||||||
|
output.appen(f"Result {i} ({title}):\n{doc.page_content}\n")
|
||||||
|
|
||||||
|
return "\n".join(output)
|
||||||
|
|
||||||
|
@tool
|
||||||
|
def add_to_knowledge_base(content: str, title: str = "Untitled") -> str:
|
||||||
|
"""Add a document to the knowledge base.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
content: The text content of the document.
|
||||||
|
title: The title of the document (default: "Untitled").
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
Confirmation message.
|
||||||
|
"""
|
||||||
|
store = _get_store()
|
||||||
|
doc = Document(page_content=content, metadata={"title": title})
|
||||||
|
ids = add_documents_to_store(store, [doc])
|
||||||
|
return f"Document '{title}' added to knowledge base with {len(ids)} chunk(s). ID: {ids[0]}"
|
||||||
Reference in New Issue
Block a user