52 lines
1.2 KiB
Python
52 lines
1.2 KiB
Python
"""Tools for the RAG agent.
|
||
|
||
Two tools are defined:
|
||
1. search_knowledge_base – performs semantic search in the vector store.
|
||
2. add_to_knowledge_base – adds a new document to the vector store.
|
||
"""
|
||
|
||
from typing import List
|
||
from langchain.tools import tool
|
||
from langchain.schema import Document
|
||
|
||
from vector_store import store
|
||
|
||
|
||
@tool("search_knowledge_base")
|
||
async def search_knowledge_base(query: str, max_results: int = 5) -> List[Document]:
|
||
"""Search the knowledge base for relevant chunks.
|
||
|
||
Parameters
|
||
----------
|
||
query: str
|
||
The search query.
|
||
max_results: int
|
||
Number of top results to return.
|
||
|
||
Returns
|
||
-------
|
||
List[Document]
|
||
List of documents returned by Qdrant similarity search.
|
||
"""
|
||
return store.search(query, max_results)
|
||
|
||
|
||
@tool("add_to_knowledge_base")
|
||
async 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 identifier for the document.
|
||
|
||
Returns
|
||
-------
|
||
str
|
||
Confirmation message.
|
||
"""
|
||
store.add_document(content, title)
|
||
return f"Document '{title}' added to the knowledge base."
|