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: Two tools are defined:
1. search_knowledge_base performs semantic search in the vector store.
* ``search_knowledge_base`` semantic search in Qdrant. 2. add_to_knowledge_base adds a new document to the vector store.
* ``add_to_knowledge_base`` add a document to Qdrant.
The tools use the global ``vector_store`` instance defined in this module.
""" """
from typing import List 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.tools import tool
from langchain.schema import Document from langchain.schema import Document
# Global vector store instance from vector_store import store
# 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
@tool("search_knowledge_base") @tool("search_knowledge_base")
def search_knowledge_base(query: str, max_results: int = 5) -> List[dict]: async def search_knowledge_base(query: str, max_results: int = 5) -> List[Document]:
"""Semantic search in the knowledge base. """Search the knowledge base for relevant chunks.
Parameters Parameters
---------- ----------
query: str query: str
The search query. The search query.
max_results: int max_results: int
Maximum number of results to return. Number of top results to return.
Returns Returns
------- -------
List[dict] List[Document]
List of matching documents with ``content`` and ``metadata``. List of documents returned by Qdrant similarity search.
""" """
store = _get_vector_store() return store.search(query, max_results)
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
]
@tool("add_to_knowledge_base") @tool("add_to_knowledge_base")
def add_to_knowledge_base(content: str, title: str) -> str: async def add_to_knowledge_base(content: str, title: str) -> str:
"""Add a document to the knowledge base. """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() store.add_document(content, title)
chunks = _splitter.split_text(content) return f"Document '{title}' added to the knowledge base."
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."