diff --git a/src/tools.py b/src/tools.py index 94bd0bf..27bfa4e 100644 --- a/src/tools.py +++ b/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." \ No newline at end of file