Update main.py

main.py

@@ -1,107 +1,122 @@
-"""
-Agent with RAG memory using Qdrant and Ollama.
-
-This is a minimal example that demonstrates:
-1. Connecting to a local Qdrant instance.
-2. Creating two tools: `search_knowledge_base` and `add_to_knowledge_base`.
-3. Using LangChain's RecursiveCharacterTextSplitter to chunk documents.
-4. Building an agent with the tools via `create_agent`.
-5. A simple REPL that accepts `/add`, `/search` and `/quit` commands.
-
-To run:
-    pip install -r requirements.txt
-    python main.py
-
-Make sure a Qdrant instance is running locally (default port 6333) and Ollama is available at http://localhost:11434.
-"""
+"""Interactive LangChain agent with local RAG memory on Qdrant and Ollama."""
 
 import os
-from pathlib import Path
-from typing import List, Dict
+from functools import lru_cache
+from typing import Any
+from uuid import uuid4
 
-from langchain_community.document_loaders import DirectoryLoader
+from langchain.agents import create_agent
+from langchain_core.documents import Document
+from langchain_core.tools import tool
+from langchain_ollama import ChatOllama, OllamaEmbeddings
+from langchain_qdrant import QdrantVectorStore
 from langchain_text_splitters import RecursiveCharacterTextSplitter
-from langchain.embeddings.ollama import OllamaEmbeddings
-from langchain.vectorstores.qdrant import QdrantVectorStore
-from langchain.agents import create_agent, AgentExecutor, Tool
-from langchain.schema import Document
+from qdrant_client import QdrantClient
+from qdrant_client.models import Distance, VectorParams
 
-# Configuration
 QDRANT_URL = os.getenv("QDRANT_URL", "http://localhost:6333")
-OLLAMA_MODEL = os.getenv("OLLAMA_MODEL", "llama2")
-COLLECTION_NAME = "rag_collection"
+QDRANT_COLLECTION = os.getenv("QDRANT_COLLECTION", "rag_memory")
+OLLAMA_CHAT_MODEL = os.getenv("OLLAMA_CHAT_MODEL", "llama3")
+OLLAMA_EMBED_MODEL = os.getenv("OLLAMA_EMBED_MODEL", "nomic-embed-text")
+EMBEDDING_SIZE = int(os.getenv("OLLAMA_EMBEDDING_SIZE", "768"))
 
-# Initialize embeddings and vector store
-embeddings = OllamaEmbeddings(model=OLLAMA_MODEL)
-vector_store = QdrantVectorStore(
-    url=QDRANT_URL,
-    collection_name=COLLECTION_NAME,
-    embedding_function=embeddings,
-)
 
-# Tool: search knowledge base
+@lru_cache(maxsize=1)
+def get_embeddings() -> OllamaEmbeddings:
+    return OllamaEmbeddings(model=OLLAMA_EMBED_MODEL)
 
-def search_knowledge_base(query: str) -> str:
-    """Return the top 3 relevant snippets for a query."""
-    docs = vector_store.similarity_search_with_score(query, k=3)
-    if not docs:
+
+@lru_cache(maxsize=1)
+def get_vector_store() -> QdrantVectorStore:
+    client = QdrantClient(url=QDRANT_URL)
+    collections = {item.name for item in client.get_collections().collections}
+    if QDRANT_COLLECTION not in collections:
+        client.create_collection(
+            collection_name=QDRANT_COLLECTION,
+            vectors_config=VectorParams(size=EMBEDDING_SIZE, distance=Distance.COSINE),
+        )
+    return QdrantVectorStore(
+        client=client,
+        collection_name=QDRANT_COLLECTION,
+        embedding=get_embeddings(),
+    )
+
+
+def chunk_document(content: str, title: str) -> list[Document]:
+    splitter = RecursiveCharacterTextSplitter(chunk_size=700, chunk_overlap=100)
+    return splitter.create_documents([content], metadatas=[{"title": title}])
+
+
+@tool
+def search_knowledge_base(query: str, max_results: int = 3) -> str:
+    """Search relevant chunks in the local Qdrant knowledge base."""
+    results = get_vector_store().similarity_search_with_score(query, k=max_results)
+    if not results:
         return "No relevant documents found."
-    results = [f"{idx+1}. {doc[0].page_content[:200]}… (score: {doc[1]:.4f})" for idx, doc in enumerate(docs)]
-    return "\n".join(results)
 
-# Tool: add to knowledge base
+    lines: list[str] = []
+    for index, (document, score) in enumerate(results, start=1):
+        title = document.metadata.get("title", "untitled")
+        snippet = document.page_content.replace("\n", " ")[:300]
+        lines.append(f"{index}. {title} (score={score:.4f}): {snippet}")
+    return "\n".join(lines)
 
-def add_to_knowledge_base(file_path: str) -> str:
-    """Load a text file, chunk it and add to Qdrant."""
-    loader = DirectoryLoader(Path(file_path).parent.as_posix(), glob=Path(file_path).name)
-    docs = loader.load()
-    if not docs:
-        return f"No documents found in {file_path}."
-    splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=50)
-    chunks: List[Document] = []
-    for doc in docs:
-        chunks.extend(splitter.split_documents([doc]))
-    vector_store.add_documents(chunks)
-    return f"Added {len(chunks)} chunks from {file_path} to the knowledge base."
 
-# Define tools list
-TOOLS: List[Tool] = [
-    Tool(
-        name="search_knowledge_base",
-        func=search_knowledge_base,
-        description="Search the local Qdrant knowledge base for relevant information.",
-    ),
-    Tool(
-        name="add_to_knowledge_base",
-        func=add_to_knowledge_base,
-        description="Add a text file to the knowledge base. Provide full path.",
-    ),
-]
+@tool
+def add_to_knowledge_base(content: str, title: str) -> str:
+    """Split content into chunks and store it in the local Qdrant knowledge base."""
+    documents = chunk_document(content, title)
+    ids = [str(uuid4()) for _ in documents]
+    get_vector_store().add_documents(documents, ids=ids)
+    return f"Added {len(documents)} chunk(s) from '{title}' to the knowledge base."
 
-# Build agent
-agent = create_agent(TOOLS, llm=embeddings)  # embeddings can act as LLM via Ollama
-executor = AgentExecutor(agent=agent, tools=TOOLS, verbose=True)
 
-# Simple REPL
-if __name__ == "__main__":
-    print("Welcome to the RAG agent. Commands: /add <file>, /search <query>, /quit")
+def build_agent() -> Any:
+    llm = ChatOllama(model=OLLAMA_CHAT_MODEL)
+    return create_agent(
+        model=llm,
+        tools=[search_knowledge_base, add_to_knowledge_base],
+        system_prompt=(
+            "You are a RAG assistant. Use search_knowledge_base before answering "
+            "questions that may depend on stored knowledge, and use "
+            "add_to_knowledge_base when the user asks to remember information."
+        ),
+    )
+
+
+def run_cli() -> None:
+    print("RAG agent. Commands: /add <title> | <content>, /search <query>, /quit")
+    agent = None
+
     while True:
         try:
-            inp = input(">>> ")
+            user_input = input(">>> ").strip()
         except EOFError:
             break
-        if not inp:
+
+        if not user_input:
             continue
-        if inp.startswith("/quit"):
-            print("Goodbye!")
+        if user_input == "/quit":
             break
-        elif inp.startswith("/add "):
-            path = inp.split(maxsplit=1)[1]
-            print(add_to_knowledge_base(path))
-        elif inp.startswith("/search "):
-            query = inp.split(maxsplit=1)[1]
-            print(search_knowledge_base(query))
-        else:
-            # Treat as normal agent prompt
-            result = executor.invoke({"input": inp})
-            print(result.get("output", ""))
+        if user_input.startswith("/add "):
+            payload = user_input[5:]
+            if "|" in payload:
+                title, content = [part.strip() for part in payload.split("|", 1)]
+            else:
+                title, content = "note", payload.strip()
+            print(add_to_knowledge_base.invoke({"content": content, "title": title}))
+            continue
+        if user_input.startswith("/search "):
+            query = user_input[8:].strip()
+            print(search_knowledge_base.invoke({"query": query, "max_results": 3}))
+            continue
+
+        if agent is None:
+            agent = build_agent()
+        response = agent.invoke({"messages": [{"role": "user", "content": user_input}]})
+        messages = response.get("messages", [])
+        print(messages[-1].content if messages else response)
+
+
+if __name__ == "__main__":
+    run_cli()