diff --git a/tools.py b/tools.py index 64e4c16..1967392 100644 --- a/tools.py +++ b/tools.py @@ -1,33 +1,31 @@ """ -RAG tools for the agent. +Tools for the RAG agent. -search_knowledge_base and add_to_knowledge_base are decorated with @tool. +search_knowledge_base and add_to_knowledge_base are implemented using VectorStore. """ -import os from typing import List, Dict - from langchain.tools import tool -from vector_store import vector_store -from chunker import split_text +from vector_store import VectorStore -@tool("Search knowledge base") +# Instantiate a global store +store = VectorStore() + +@tool def search_knowledge_base(query: str, max_results: int = 5) -> str: - """Semantic search in the vector store.""" - results = vector_store.similarity_search(query, k=max_results) + """Semantic search in the knowledge base.""" + results = store.similarity_search(query, k=max_results) if not results: return "No relevant documents found." out_lines = [] - for i, res in enumerate(results, 1): - out_lines.append(f"{i}. {res['content'][:200]}... (distance: {res['distance']:.3f})") + for i, r in enumerate(results, 1): + out_lines.append(f"{i}. {r['content'][:200]}... (source: {r['metadata'].get('title', 'unknown')})") return "\n".join(out_lines) -@tool("Add document to knowledge base") +@tool def add_to_knowledge_base(content: str, title: str = "document") -> str: - """Adds a text chunk to the vector store.""" - # Split content into chunks - chunks = split_text(content) - docs = [] - for idx, chunk in enumerate(chunks): - docs.append({"content": chunk, "metadata": {"title": title, "chunk_index": idx}}) - vector_store.add_documents(docs) - return f"Added {len(chunks)} chunks to the knowledge base." + """Add a document to the knowledge base. + + The content is split into chunks and stored with metadata. + """ + store.add_documents([content], [{"title": title}]) + return f"Document '{title}' added to knowledge base."