add tools.py

tools.py

@@ -1,38 +1,33 @@
 """
-Tools for the RAG agent.
+RAG tools for the agent.
 
-Two tools: search_knowledge_base and add_to_knowledge_base.
+search_knowledge_base and add_to_knowledge_base are decorated with @tool.
 """
-
+import os
 from typing import List, Dict
+
 from langchain.tools import tool
 from vector_store import vector_store
 from chunker import split_text
 
-@tool("search_knowledge_base")
+@tool("Search knowledge base")
 def search_knowledge_base(query: str, max_results: int = 5) -> str:
-    """Semantic search in the knowledge base.
-
-    Returns a formatted string of results.
-    """
-    hits = vector_store.search(query, k=max_results)
-    if not hits:
+    """Semantic search in the vector store."""
+    results = vector_store.similarity_search(query, k=max_results)
+    if not results:
         return "No relevant documents found."
-    lines: List[str] = []
-    for i, hit in enumerate(hits, 1):
-        title = hit["metadata"].get("title", f"doc_{hit['id']}")
-        snippet = hit["document"][:200]
-        lines.append(f"{i}. {title}: {snippet}...")
-    return "\n".join(lines)
+    out_lines = []
+    for i, res in enumerate(results, 1):
+        out_lines.append(f"{i}. {res['content'][:200]}... (distance: {res['distance']:.3f})")
+    return "\n".join(out_lines)
 
-@tool("add_to_knowledge_base")
+@tool("Add document to knowledge base")
 def add_to_knowledge_base(content: str, title: str = "document") -> str:
-    """Add a document to the knowledge base.
-
-    Splits content into chunks and stores each with metadata.
-    """
+    """Adds a text chunk to the vector store."""
+    # Split content into chunks
     chunks = split_text(content)
+    docs = []
     for idx, chunk in enumerate(chunks):
-        doc_id = f"{title}_{idx}"
-        vector_store.add_document(doc_id=doc_id, text=chunk, metadata={"title": title})
-    return f"Added {len(chunks)} chunks from '{title}'."
+        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."