diff --git a/tools.py b/tools.py new file mode 100644 index 0000000..44fd120 --- /dev/null +++ b/tools.py @@ -0,0 +1,28 @@ +import os +from langchain.tools import tool +from langchain_text_splitters import RecursiveCharacterTextSplitter +from langchain_core.documents import Document +from rag_store import KnowledgeBase + +kb = KnowledgeBase() + +splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200) + +@tool +def search_knowledge_base(query: str, max_results: int = 3) -> str: + """Search the knowledge base for relevant information.""" + docs = kb.similarity_search(query, k=max_results) + if not docs: + return "No results found." + return "\n\n".join( + f"Title: {doc.metadata.get('title', 'N/A')}\nContent: {doc.page_content}" + for doc in docs + ) + +@tool +def add_to_knowledge_base(content: str, title: str = "Untitled") -> str: + """Add content to the knowledge base.""" + chunks = splitter.split_text(content) + docs = [Document(page_content=chunk, metadata={"title": title}) for chunk in chunks] + kb.add_documents(docs) + return f"Added document '{title}'." \ No newline at end of file