Add tools.py

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."