diff --git a/tools.py b/tools.py new file mode 100644 index 0000000..e6bdad4 --- /dev/null +++ b/tools.py @@ -0,0 +1,82 @@ +"""Tools for 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. +""" + +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 + + +@tool("search_knowledge_base") +def search_knowledge_base(query: str, max_results: int = 5) -> List[dict]: + """Semantic search in the knowledge base. + + Parameters + ---------- + query: str + The search query. + max_results: int + Maximum number of results to return. + + Returns + ------- + List[dict] + List of matching documents with ``content`` and ``metadata``. + """ + 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 + ] + + +@tool("add_to_knowledge_base") +def add_to_knowledge_base(content: str, title: str) -> str: + """Add a document to the knowledge base. + + The content is split into chunks before being added. + """ + 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."