diff --git a/main.py b/main.py index a435270..f2542dd 100644 --- a/main.py +++ b/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 , /search , /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 | <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()