diff --git a/rag_tools.py b/rag_tools.py new file mode 100644 index 0000000..045dece --- /dev/null +++ b/rag_tools.py @@ -0,0 +1,41 @@ +from langchain.tools import tool +from .vector_store import VectorStore +from .splitter import chunk_text + +# Shared store instance +_store: VectorStore | None = None + +@tool +def add_to_knowledge_base(content: str, title: str = "Document") -> str: + """Add content to the knowledge base. + + Parameters + ---------- + content: str + Text content to add. + title: str + Optional title for the document. + """ + global _store + if _store is None: + _store = VectorStore() + chunks = chunk_text(content) + _store.add_documents(chunks) + return f"Added {len(chunks)} chunks to the knowledge base under title '{title}'." + +@tool +def search_knowledge_base(query: str, max_results: int = 5) -> list[tuple[str, float]]: + """Search the knowledge base for relevant chunks. + + Parameters + ---------- + query: str + Search query. + max_results: int + Number of top results to return. + """ + global _store + if _store is None: + _store = VectorStore() + results = _store.search(query, max_results=max_results) + return results