add main.py

main.py

@@ -0,0 +1,112 @@
+"""
+Agent with RAG memory using Qdrant and Ollama.
+"""
+
+import os
+import sys
+from pathlib import Path
+from typing import List, Dict, Any
+
+from langchain_ollama import OllamaEmbeddings, Ollama
+from langchain_qdrant import QdrantVectorStore
+from langchain.text_splitter import RecursiveCharacterTextSplitter
+from langchain.tools import tool
+from langchain.agents import create_agent, AgentExecutor, AgentType
+from langchain.schema import AgentAction, AgentFinish
+from langchain.callbacks import get_openai_callback
+
+# Configuration
+QDRANT_URL = os.getenv("QDRANT_URL", "http://localhost:6333")
+QDRANT_COLLECTION = os.getenv("QDRANT_COLLECTION", "rag_agent")
+Ollama_EMBED_MODEL = os.getenv("Ollama_EMBED_MODEL", "nomic-embed-text")
+Ollama_MODEL = os.getenv("Ollama_MODEL", "llama3")
+
+# Initialize embeddings and vector store
+embeddings = OllamaEmbeddings(model=Ollama_EMBED_MODEL)
+vector_store = QdrantVectorStore(
+    url=QDRANT_URL,
+    collection_name=QDRANT_COLLECTION,
+    embeddings=embeddings,
+)
+
+# Text splitter
+splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=50)
+
+# Tools
+@tool
+def search_knowledge_base(query: str, max_results: int = 5) -> str:
+    """Semantic search in the knowledge base."""
+    docs = vector_store.similarity_search(query, k=max_results)
+    if not docs:
+        return "No relevant documents found."
+    return "\n\n".join([f"{i+1}. {doc.page_content[:200]}..." for i, doc in enumerate(docs)])
+
+@tool
+def add_to_knowledge_base(content: str, title: str) -> str:
+    """Add a document to the knowledge base."""
+    # Split content into chunks
+    chunks = splitter.split_text(content)
+    # Create documents with metadata
+    docs = []
+    for i, chunk in enumerate(chunks):
+        docs.append(
+            {
+                "page_content": chunk,
+                "metadata": {"title": title, "chunk_index": i},
+            }
+        )
+    # Add to vector store
+    vector_store.add_documents(docs)
+    return f"Document '{title}' added with {len(chunks)} chunks."
+
+# Agent setup
+SYSTEM_PROMPT = (
+    "You are an AI assistant with access to a local knowledge base. "
+    "Use the provided tools to search and add information. "
+    "When answering, rely on the knowledge base and the LLM."
+)
+
+# Create agent with tools
+agent = create_agent(
+    llm=Ollama(model=Ollama_MODEL),
+    tools=[search_knowledge_base, add_to_knowledge_base],
+    system_prompt=SYSTEM_PROMPT,
+    agent_type=AgentType.ZERO_SHOT_REACT_DESCRIPTION,
+)
+
+executor = AgentExecutor(agent=agent, tools=[search_knowledge_base, add_to_knowledge_base], verbose=True)
+
+# CLI
+def main():
+    print("RAG Agent CLI. Commands: /add <title> <content>, /search <query>, /quit")
+    while True:
+        try:
+            inp = input("> ")
+        except EOFError:
+            break
+        if not inp:
+            continue
+        if inp.startswith("/quit"):
+            print("Bye!")
+            break
+        if inp.startswith("/add"):
+            parts = inp.split(maxsplit=2)
+            if len(parts) < 3:
+                print("Usage: /add <title> <content>")
+                continue
+            title, content = parts[1], parts[2]
+            print(add_to_knowledge_base(content, title))
+            continue
+        if inp.startswith("/search"):
+            query = inp[len("/search"):].strip()
+            if not query:
+                print("Usage: /search <query>")
+                continue
+            print(search_knowledge_base(query))
+            continue
+        # Treat as normal user query
+        response = executor.invoke({"input": inp})
+        print(response.get("output", ""))
+
+if __name__ == "__main__":
+    main()