add tools.py

tools.py

@@ -0,0 +1,71 @@
+"""
+LangChain tools for adding and searching a Qdrant knowledge base.
+
+The module exposes two @tool‑decorated functions:
+* add_to_knowledge_base(content: str, title: str = "document")
+* search_knowledge_base(query: str, max_results: int = 5)
+
+Both tools use an in‑memory Qdrant client and Ollama embeddings.
+"""
+import os
+from typing import List
+
+from langchain_core.documents import Document
+from langchain_text_splitters import RecursiveCharacterTextSplitter
+from langchain_ollama import ChatOllama, OllamaEmbeddings
+from langchain_qdrant import QdrantVectorStore
+from langchain.tools import tool
+
+# --- Configuration -------------------------------------------------------
+QDRANT_COLLECTION = "rag_agent"
+EMBEDDING_MODEL = "nomic-embed-text"
+LLM_MODEL = "llama3"
+
+# Initialize embeddings and vector store (in‑memory Qdrant)
+embeddings = OllamaEmbeddings(model=EMBEDDING_MODEL, base_url="http://localhost:11434")
+vector_store = QdrantVectorStore(client=None, collection_name=QDRANT_COLLECTION, embedding=embeddings)
+
+# Ensure the collection exists (creates if not present)
+try:
+    vector_store.client.create_collection(
+        name=QDRANT_COLLECTION,
+        vectors_config=dict(size=embeddings.embed_query("test").shape[0], distance="cosine"),
+    )
+except Exception:
+    # Collection already exists – ignore
+    pass
+
+# Text splitter for chunking documents
+splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=50)
+
+@tool("Add a document to the knowledge base")
+def add_to_knowledge_base(content: str, title: str = "document") -> str:
+    """
+    Splits *content* into chunks, embeds them via Ollama and stores in Qdrant.
+    Returns a confirmation string with number of chunks added.
+    """
+    # Split content
+    docs: List[Document] = splitter.split_text(content)
+    for i, chunk in enumerate(docs):
+        doc = Document(page_content=chunk, metadata={"title": title, "chunk_index": i})
+        vector_store.add_documents([doc])
+    return f"Added {len(docs)} chunks from '{title}'."
+
+@tool("Search the knowledge base")
+def search_knowledge_base(query: str, max_results: int = 5) -> str:
+    """
+    Performs a semantic similarity search in Qdrant.
+    Returns a formatted string with top results and their metadata.
+    """
+    docs = vector_store.similarity_search(query, k=max_results)
+    if not docs:
+        return "No relevant documents found."
+    lines: List[str] = []
+    for i, doc in enumerate(docs, 1):
+        title = doc.metadata.get("title", "unknown")
+        idx = doc.metadata.get("chunk_index", "?")
+        lines.append(f"{i}. [{title} - chunk {idx}]\n{doc.page_content[:200]}...")
+    return "\n\n".join(lines)
+
+# Expose tool names for agent creation
+TOOLS = [add_to_knowledge_base, search_knowledge_base]