Updated main.py with Qdrant-based implementation

main.py

@@ -1,8 +1,4 @@
-"""Main module implementing a RAG-enabled agent with Qdrant and Ollama.
+# Main script implementing Qdrant-based RAG agent
-
-The code follows the assignment specification and uses only the required libraries.
-"""
-
 import os
 import sys
 import json
@@ -12,162 +8,120 @@ from typing import List, Dict, Any
 from langchain_ollama import ChatOllama, OllamaEmbeddings
 from langchain_qdrant import QdrantVectorStore
 from langchain_text_splitters import RecursiveCharacterTextSplitter
-from langchain.tools import tool
+from langchain.tools import tool, BaseTool
-from langchain.agents import create_agent, AgentExecutor
+from langchain.agents import create_agent, AgentExecutor, AgentType
-from langchain.schema import AgentAction, AgentFinish
-from langchain_core.prompts import ChatPromptTemplate
 
-# ---------------------------------------------------------------------------
 # Configuration
-# ---------------------------------------------------------------------------
+QDRANT_HOST = os.getenv("QDRANT_HOST", "localhost")
-QDRANT_URL = os.getenv("QDRANT_URL", "http://localhost:6333")
+QDRANT_PORT = int(os.getenv("QDRANT_PORT", "6333"))
-QDRANT_COLLECTION = "knowledge"
+COLLECTION_NAME = "knowledge"
 EMBEDDING_MODEL = "nomic-embed-text"
 LLM_MODEL = "llama3"
-CHUNK_SIZE = 1000  # characters
-CHUNK_OVERLAP = 200
 
-# ---------------------------------------------------------------------------
+# Vector store wrapper
-# Vector store helper
+class QdrantStore:
-# ---------------------------------------------------------------------------
+    def __init__(self, host: str, port: int, collection: str):
-class KnowledgeBase:
+        self.store = QdrantVectorStore(
-    """Wrapper around QdrantVectorStore providing add/search helpers."""
+            url=f"http://{host}:{port}",
-
-    def __init__(self, url: str = QDRANT_URL, collection: str = QDRANT_COLLECTION):
-        self.embeddings = OllamaEmbeddings(model=EMBEDDING_MODEL)
-        self.store = QdrantVectorStore.from_existing_collection(
             collection_name=collection,
-            url=url,
+            embedding=OllamaEmbeddings(model=EMBEDDING_MODEL),
-            embedding=self.embeddings,
         )
 
-    def add_documents(self, documents: List[str], titles: List[str]):
+    def add_documents(self, documents: List[str], metadatas: List[Dict[str, Any]]):
-        """Adds a list of documents with corresponding titles to the store.
+        self.store.add_texts(documents, metadatas=metadatas)
 
-        Each document is split into chunks before being stored.
+    def similarity_search(self, query: str, k: int = 5) -> List[Dict[str, Any]]:
-        """
+        results = self.store.similarity_search(query, k=k)
-        splitter = RecursiveCharacterTextSplitter(chunk_size=CHUNK_SIZE, chunk_overlap=CHUNK_OVERLAP)
-        for doc, title in zip(documents, titles):
-            chunks = splitter.split_text(doc)
-            metadatas = [{"title": title, "chunk_index": i} for i in range(len(chunks))]
-            self.store.add_texts(chunks, metadatas=metadatas)
-
-    def search(self, query: str, max_results: int = 5) -> List[Dict[str, Any]]:
-        """Semantic search in the knowledge base.
-
-        Returns a list of dicts with keys: text, title, distance.
-        """
-        results = self.store.similarity_search_with_score(query, k=max_results)
         return [
             {
-                "text": text,
+                "content": doc.page_content,
-                "title": meta.get("title", "unknown"),
+                "metadata": doc.metadata,
-                "distance": score,
+                "score": doc.metadata.get("score", 0),
             }
-            for text, score, meta in results
+            for doc in results
         ]
 
-# ---------------------------------------------------------------------------
+# Text splitter
-# Global knowledge base instance
+text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
-# ---------------------------------------------------------------------------
+
-kb = KnowledgeBase()
+# Store instance
+store = QdrantStore(QDRANT_HOST, QDRANT_PORT, COLLECTION_NAME)
 
-# ---------------------------------------------------------------------------
 # Tools
-# ---------------------------------------------------------------------------
+@tool("search_knowledge_base", "Semantic search in the knowledge base")
-@tool("search_knowledge_base")
 def search_knowledge_base(query: str, max_results: int = 5) -> str:
-    """Search the knowledge base for a query and return formatted results."""
+    results = store.similarity_search(query, k=max_results)
-    results = kb.search(query, max_results)
+    return json.dumps(results, ensure_ascii=False)
-    if not results:
-        return "No relevant documents found."
-    formatted = [f"Title: {r['title']}\nSnippet: {r['text'][:200]}...\nDistance: {r['distance']:.4f}" for r in results]
-    return "\n\n".join(formatted)
 
-@tool("add_to_knowledge_base")
+@tool("add_to_knowledge_base", "Add a document to the knowledge base")
-def add_to_knowledge_base(content: str, title: str) -> str:
+def add_to_knowledge_base(content: str, title: str = "Untitled") -> str:
-    """Add a new document to the knowledge base."""
+    chunks = text_splitter.split_text(content)
-    kb.add_documents([content], [title])
+    metadatas = [{"title": title, "chunk_idx": i} for i in range(len(chunks))]
-    return f"Document '{title}' added successfully."
+    store.add_documents(chunks, metadatas)
-
+    return f"Added {len(chunks)} chunks titled '{title}'."
-# ---------------------------------------------------------------------------
-# Agent definition
-# ---------------------------------------------------------------------------
-SYSTEM_PROMPT = (
-    "You are an assistant that can search and add documents to a local knowledge base. "
-    "Use the provided tools to perform semantic search and store new information. "
-    "When answering user queries, first determine if the user needs a search or an addition. "
-    "If no relevant information is found, suggest adding new content."
-)
-
-prompt = ChatPromptTemplate.from_messages([
-    ("system", SYSTEM_PROMPT),
-    ("user", "{input}"),
-])
 
+# Agent
+llm = ChatOllama(model=LLM_MODEL)
 agent = create_agent(
-    llm=ChatOllama(model=LLM_MODEL),
+    llm=llm,
     tools=[search_knowledge_base, add_to_knowledge_base],
-    prompt=prompt,
+    agent_type=AgentType.ZERO_SHOT_REACT_DESCRIPTION,
-    verbose=True,
+    system_message="You are an assistant that can search and add information to a local knowledge base. Use the tools when appropriate.",
 )
+executor = AgentExecutor(agent=agent, tools=[search_knowledge_base, add_to_knowledge_base], verbose=True)
 
-agent_executor = AgentExecutor(agent=agent, tools=[search_knowledge_base, add_to_knowledge_base])
+# CLI helpers
 
-# ---------------------------------------------------------------------------
+def load_documents_from_dir(dir_path: str):
-# CLI client
+    for path in Path(dir_path).glob("**/*"):
-# ---------------------------------------------------------------------------
+        if path.suffix.lower() in {".txt", ".md"}:
-def load_documents_from_dir(directory: str):
+            content = path.read_text(encoding="utf-8")
-    """Load all .txt files from a directory and add them to the knowledge base."""
+            title = path.stem
-    docs = []
+            add_to_knowledge_base(content, title)
-    titles = []
+    print("Loading complete.")
-    for path in Path(directory).glob("**/*.txt"):
-        text = path.read_text(encoding="utf-8")
-        docs.append(text)
-        titles.append(path.stem)
-    if docs:
-        kb.add_documents(docs, titles)
-        print(f"Loaded {len(docs)} documents from {directory}.")
-    else:
-        print("No .txt files found.")
 
-
+def main():
-def interactive_loop():
+    if len(sys.argv) > 1 and sys.argv[1] == "load":
-    print("RAG Agent CLI. Commands: /add <title> <content>, /search <query>, /load <dir>, /quit")
+        if len(sys.argv) < 3:
+            print("Usage: python main.py load <directory>")
+            sys.exit(1)
+        load_documents_from_dir(sys.argv[2])
+        sys.exit(0)
+    print("Interactive mode. Commands: /add <title> <file>, /search <query>, /quit")
     while True:
         try:
-            user_input = input("> ")
+            user_input = input("\n> ")
-        except (KeyboardInterrupt, EOFError):
+        except (EOFError, KeyboardInterrupt):
-            print("\nExiting.")
             break
         if not user_input:
             continue
         if user_input.startswith("/quit"):
-            print("Goodbye!")
             break
-        elif user_input.startswith("/add "):
+        if user_input.startswith("/add"):
+            parts = user_input.split(maxsplit=2)
+            if len(parts) != 3:
+                print("Usage: /add <title> <file_path>")
+                continue
+            title, file_path = parts[1], parts[2]
             try:
-                _, rest = user_input.split("/add ", 1)
+                content = Path(file_path).read_text(encoding="utf-8")
-                title, content = rest.split(" ", 1)
+            except Exception as e:
-                response = add_to_knowledge_base(content, title)
+                print(f"Error reading file: {e}")
-                print(response)
+                continue
-            except ValueError:
+            print(add_to_knowledge_base(content, title))
-                print("Usage: /add <title> <content>")
+            continue
-        elif user_input.startswith("/search "):
+        if user_input.startswith("/search"):
-            query = user_input.split("/search ", 1)[1]
+            query = user_input[len("/search"):].strip()
+            if not query:
+                print("Provide a query.")
+                continue
             results = search_knowledge_base(query)
-            print(results)
+            print("Search results:")
-        elif user_input.startswith("/load "):
+            for r in json.loads(results):
-            dir_path = user_input.split("/load ", 1)[1]
+                print(f"- {r['metadata'].get('title', 'Untitled')} (chunk {r['metadata'].get('chunk_idx')})\n  {r['content'][:200]}...")
-            load_documents_from_dir(dir_path)
+            continue
-        else:
+        response = executor.invoke({"input": user_input})
-            # Treat as normal user query
+        print(response["output"])
-            result = agent_executor.invoke({"input": user_input})
-            print(result.get("output", ""))
 
 if __name__ == "__main__":
-    if len(sys.argv) > 1 and sys.argv[1] == "cli":
+    main()
-        interactive_loop()
-    else:
-        print("Usage: python main.py cli")
-        print("Run the interactive CLI with: python main.py cli")