diff --git a/src/tools.py b/src/tools.py index 27bfa4e..d15e3f9 100644 --- a/src/tools.py +++ b/src/tools.py @@ -1,54 +1,50 @@ """Tools for the RAG agent. -This module defines two tools that the agent can call: +This module defines two LangChain tools that interact with the +``KnowledgeBase`` defined in :mod:`src.vector_store`. -* ``search_knowledge_base`` – performs a semantic search in the Qdrant vector store. -* ``add_to_knowledge_base`` – adds a document to the store. - -Both tools are decorated with ``@tool`` from ``langchain.tools`` so that the LLM can invoke them. +The tools are decorated with ``@tool`` from ``langchain.tools`` so that +the agent can invoke them automatically. """ -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 global knowledge base instance. In a real application you might -# want to inject this via a dependency injection container. -kb = KnowledgeBase() +from .vector_store import kb @tool("search_knowledge_base") -def search_knowledge_base(query: str, max_results: int = 5) -> List[Dict[str, Any]]: - """Search the knowledge base for relevant chunks. +def search_knowledge_base(query: str, max_results: int = 5) -> str: + """Search the local knowledge base. Parameters ---------- query: str - Search query. - max_results: int - Number of results to return. + The search query. + max_results: int, optional + Limit of results to return. Returns ------- - list[dict] - List of dictionaries with ``content``, ``title`` and ``chunk_index``. + str + A formatted string with the search results. """ - return kb.search(query, limit=max_results) + results = kb.search(query, limit=max_results) + if not results: + return "No relevant documents found." + lines = [] + for i, res in enumerate(results, 1): + title = res["metadata"].get("title", "Untitled") + lines.append(f"{i}. Title: {title}\nContent: {res['page_content']}\n") + return "\n".join(lines) @tool("add_to_knowledge_base") def add_to_knowledge_base(content: str, title: str) -> str: - """Add a document to the knowledge base. + """Add a new document to the knowledge base. Parameters ---------- content: str - Full text of the document. + The full text of the document. title: str - Title or identifier for the document. + A short title for the document. Returns -------