Update src/tools.py

src/tools.py

@@ -1,22 +1,24 @@
-"""Agent tools for interacting with the knowledge base.
+"""Tools for the RAG agent.
 
-This module defines two tools that can be used by the LangChain agent:
+This module defines two tools that the agent can call:
 
 * ``search_knowledge_base`` – performs a semantic search in the Qdrant vector store.
-* ``add_to_knowledge_base`` – adds a new document (title + content) to the store.
+* ``add_to_knowledge_base`` – adds a document to the store.
 
-Both tools are decorated with ``@tool`` from ``langchain.tools`` so that they can be
-exposed to the agent.
+Both tools are decorated with ``@tool`` from ``langchain.tools`` so that the LLM can invoke them.
 """
 
+from __future__ import annotations
+
 from typing import List, Dict, Any
 
 from langchain.tools import tool
 
+# Import the knowledge base implementation.
 from .vector_store import KnowledgeBase
 
-# Create a single global knowledge base instance that all tools will use.
-# In a real deployment you might want to inject this via dependency injection.
+# Create a global knowledge base instance. In a real application you might
+# want to inject this via a dependency injection container.
 kb = KnowledgeBase()
 
 @tool("search_knowledge_base")
@@ -26,32 +28,32 @@ def search_knowledge_base(query: str, max_results: int = 5) -> List[Dict[str, An
     Parameters
     ----------
     query: str
-        The search query.
-    max_results: int, optional
-        Number of top results to return. Defaults to 5.
+        Search query.
+    max_results: int
+        Number of results to return.
 
     Returns
     -------
-    List[Dict[str, Any]]
-        A list of dictionaries containing ``content``, ``title``, ``chunk_index`` and ``score``.
+    list[dict]
+        List of dictionaries with ``content``, ``title`` and ``chunk_index``.
     """
-    return kb.search(query, max_results)
+    return kb.search(query, limit=max_results)
 
 @tool("add_to_knowledge_base")
 def add_to_knowledge_base(content: str, title: str) -> str:
-    """Add a new document to the knowledge base.
+    """Add a document to the knowledge base.
 
     Parameters
     ----------
     content: str
         Full text of the document.
     title: str
-        Title or name of the document.
+        Title or identifier for the document.
 
     Returns
     -------
     str
         Confirmation message.
     """
-    kb.add_document(content, title)
+    kb.add_document(title=title, content=content)
     return f"Document '{title}' added to the knowledge base."