From 63c222876b1523aa805d1d3437ce27bf378098dc Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=D0=9A=D0=B8=D1=80=D0=B8=D0=BB=D0=BB=20=D0=A0=D0=BE=D0=BC?= =?UTF-8?q?=D0=B0=D0=BD=D0=BE=D0=B2?= Date: Wed, 3 Jun 2026 09:04:35 +0000 Subject: [PATCH] Update tools.py --- tools.py | 83 ++++++++++++++++++-------------------------------------- 1 file changed, 26 insertions(+), 57 deletions(-) diff --git a/tools.py b/tools.py index e6bdad4..7342633 100644 --- a/tools.py +++ b/tools.py @@ -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."