Update src/tools.py
This commit is contained in:
+23
-27
@@ -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
|
||||
-------
|
||||
|
||||
Reference in New Issue
Block a user