Add rag_tools.py
This commit is contained in:
@@ -0,0 +1,48 @@
|
|||||||
|
"""
|
||||||
|
Tools for the RAG agent.
|
||||||
|
"""
|
||||||
|
|
||||||
|
from typing import List
|
||||||
|
|
||||||
|
from langchain.tools import tool
|
||||||
|
from qdrant_store import search, add_document
|
||||||
|
|
||||||
|
@tool("search_knowledge_base")
|
||||||
|
async def search_knowledge_base(query: str, max_results: int = 5) -> List[str]:
|
||||||
|
"""Perform semantic search in the knowledge base.
|
||||||
|
|
||||||
|
Parameters
|
||||||
|
----------
|
||||||
|
query: str
|
||||||
|
Search query.
|
||||||
|
max_results: int
|
||||||
|
Number of results to return.
|
||||||
|
|
||||||
|
Returns
|
||||||
|
-------
|
||||||
|
List[str]
|
||||||
|
List of retrieved document snippets.
|
||||||
|
"""
|
||||||
|
docs = search(query, max_results)
|
||||||
|
return [doc.page_content for doc in docs]
|
||||||
|
|
||||||
|
@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.
|
||||||
|
"""
|
||||||
|
add_document(content, title)
|
||||||
|
return f"Document '{title}' added successfully."
|
||||||
|
|
||||||
|
__all__ = ["search_knowledge_base", "add_to_knowledge_base"]
|
||||||
Reference in New Issue
Block a user