Add src/tools.py
This commit is contained in:
@@ -0,0 +1,57 @@
|
||||
"""Agent tools for interacting with the knowledge base.
|
||||
|
||||
This module defines two tools that can be used by the LangChain agent:
|
||||
|
||||
* ``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.
|
||||
|
||||
Both tools are decorated with ``@tool`` from ``langchain.tools`` so that they can be
|
||||
exposed to the agent.
|
||||
"""
|
||||
|
||||
from typing import List, Dict, Any
|
||||
|
||||
from langchain.tools import tool
|
||||
|
||||
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.
|
||||
kb = KnowledgeBase()
|
||||
|
||||
@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.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
query: str
|
||||
The search query.
|
||||
max_results: int, optional
|
||||
Number of top results to return. Defaults to 5.
|
||||
|
||||
Returns
|
||||
-------
|
||||
List[Dict[str, Any]]
|
||||
A list of dictionaries containing ``content``, ``title``, ``chunk_index`` and ``score``.
|
||||
"""
|
||||
return kb.search(query, 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.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
content: str
|
||||
Full text of the document.
|
||||
title: str
|
||||
Title or name of the document.
|
||||
|
||||
Returns
|
||||
-------
|
||||
str
|
||||
Confirmation message.
|
||||
"""
|
||||
kb.add_document(content, title)
|
||||
return f"Document '{title}' added to the knowledge base."
|
||||
Reference in New Issue
Block a user