Add rag_tools.py
This commit is contained in:
@@ -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
|
||||
Reference in New Issue
Block a user