From ad4fe7c6dd6b24fb4de16dcce56b74f10beff0d8 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=D0=A0=D0=B8=D0=BD=D0=B0=D1=80=20=D0=9C=D0=B8=D1=80=D0=B7?= =?UTF-8?q?=D0=B0=D0=B3=D0=B8=D1=82=D0=BE=D0=B2?= Date: Wed, 27 May 2026 08:02:42 +0000 Subject: [PATCH] Update main.py --- main.py | 108 +++++++++++++++++++++++++++++++++++++++++++++++++++++++- 1 file changed, 107 insertions(+), 1 deletion(-) diff --git a/main.py b/main.py index c353358..a435270 100644 --- a/main.py +++ b/main.py @@ -1 +1,107 @@ -# main.py content \ No newline at end of file +""" +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 , /search , /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", ""))