Update main.py

main.py

@@ -1 +1,107 @@
-# main.py content
+"""
+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.
+"""
+
+import os
+from pathlib import Path
+from typing import List, Dict
+
+from langchain_community.document_loaders import DirectoryLoader
+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
+
+# Configuration
+QDRANT_URL = os.getenv("QDRANT_URL", "http://localhost:6333")
+OLLAMA_MODEL = os.getenv("OLLAMA_MODEL", "llama2")
+COLLECTION_NAME = "rag_collection"
+
+# 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
+
+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:
+        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
+
+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.",
+    ),
+]
+
+# 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")
+    while True:
+        try:
+            inp = input(">>> ")
+        except EOFError:
+            break
+        if not inp:
+            continue
+        if inp.startswith("/quit"):
+            print("Goodbye!")
+            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", ""))