add tools.py

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 vector_store import VectorStore
-from chunker import split_text
 
-@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."""
+    """Semantic search in the knowledge base."""
-    results = vector_store.similarity_search(query, k=max_results)
+    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):
+    for i, r in enumerate(results, 1):
-        out_lines.append(f"{i}. {res['content'][:200]}... (distance: {res['distance']:.3f})")
+        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."""
+    """Add a document to the knowledge base.
-    # Split content into chunks
+
-    chunks = split_text(content)
+    The content is split into chunks and stored with metadata.
-    docs = []
+    """
-    for idx, chunk in enumerate(chunks):
+    store.add_documents([content], [{"title": title}])
-        docs.append({"content": chunk, "metadata": {"title": title, "chunk_index": idx}})
+    return f"Document '{title}' added to knowledge base."
-    vector_store.add_documents(docs)
-    return f"Added {len(chunks)} chunks to the knowledge base."