Update tools.py

This commit is contained in:
2026-06-03 09:04:35 +00:00
parent 03e2d76001
commit 63c222876b
+26 -57
View File
@@ -1,82 +1,51 @@
"""Tools for RAG agent.
"""Tools for the RAG agent.
This module defines two LangChain tools:
* ``search_knowledge_base`` semantic search in Qdrant.
* ``add_to_knowledge_base`` add a document to Qdrant.
The tools use the global ``vector_store`` instance defined in this module.
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_ollama import OllamaEmbeddings
from langchain_qdrant import QdrantVectorStore
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain.tools import tool
from langchain.schema import Document
# Global vector store instance
# We initialise it lazily the first call to the tools will create the store.
_vector_store = None
# Embedding model
_embedding = OllamaEmbeddings(model="nomic-embed-text")
# Text splitter
_splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=50)
def _get_vector_store() -> QdrantVectorStore:
global _vector_store
if _vector_store is None:
# Create a new collection named "knowledge_base"
_vector_store = QdrantVectorStore.from_texts(
texts=[],
embedding=_embedding,
url="http://localhost:6333",
collection_name="knowledge_base",
)
return _vector_store
from vector_store import store
@tool("search_knowledge_base")
def search_knowledge_base(query: str, max_results: int = 5) -> List[dict]:
"""Semantic search in the 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
Maximum number of results to return.
Number of top results to return.
Returns
-------
List[dict]
List of matching documents with ``content`` and ``metadata``.
List[Document]
List of documents returned by Qdrant similarity search.
"""
store = _get_vector_store()
results = store.similarity_search_with_score(query, k=max_results)
# Return only content and metadata for simplicity
return [
{"content": doc.page_content, "metadata": doc.metadata}
for doc, _ in results
]
return store.search(query, max_results)
@tool("add_to_knowledge_base")
def add_to_knowledge_base(content: str, title: str) -> str:
"""Add a document to the knowledge base.
async def add_to_knowledge_base(content: str, title: str) -> str:
"""Add a new document to the knowledge base.
The content is split into chunks before being added.
Parameters
----------
content: str
Full text of the document.
title: str
Title or identifier for the document.
Returns
-------
str
Confirmation message.
"""
store = _get_vector_store()
chunks = _splitter.split_text(content)
docs: List[Document] = []
for i, chunk in enumerate(chunks):
docs.append(
Document(page_content=chunk, metadata={"title": title, "chunk_index": i})
)
store.add_documents(docs)
return f"Document '{title}' added with {len(chunks)} chunks."
store.add_document(content, title)
return f"Document '{title}' added to the knowledge base."