Update src/tools.py

This commit is contained in:
2026-06-05 11:29:27 +00:00
parent 9c65a8ca4e
commit 686525b3de
+23 -27
View File
@@ -1,54 +1,50 @@
"""Tools for the RAG agent. """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. The tools are decorated with ``@tool`` from ``langchain.tools`` so that
* ``add_to_knowledge_base`` adds a document to the store. the agent can invoke them automatically.
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 from langchain.tools import tool
from .vector_store import kb
# 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()
@tool("search_knowledge_base") @tool("search_knowledge_base")
def search_knowledge_base(query: str, max_results: int = 5) -> List[Dict[str, Any]]: def search_knowledge_base(query: str, max_results: int = 5) -> str:
"""Search the knowledge base for relevant chunks. """Search the local knowledge base.
Parameters Parameters
---------- ----------
query: str query: str
Search query. The search query.
max_results: int max_results: int, optional
Number of results to return. Limit of results to return.
Returns Returns
------- -------
list[dict] str
List of dictionaries with ``content``, ``title`` and ``chunk_index``. 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") @tool("add_to_knowledge_base")
def add_to_knowledge_base(content: str, title: str) -> str: 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 Parameters
---------- ----------
content: str content: str
Full text of the document. The full text of the document.
title: str title: str
Title or identifier for the document. A short title for the document.
Returns Returns
------- -------