Files
task-6a02e23da6fe2e4ac16acf65/src/tools.py
T
2026-06-04 22:57:36 +00:00

57 lines
1.6 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
"""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."