main.py written

main.py

@@ -1,27 +1,11 @@
-# main.py
+"""Main module implementing a RAG-enabled agent with Qdrant and Ollama.
-# Полностью рабочий пример агента с RAG‑памятью на базе Qdrant и Ollama.
+
-# Весь код написан в одном файле для простоты демонстрации.
+The code follows the assignment specification and uses only the required libraries.
-#
+"""
-# Требования:
-#   - Python 3.10+
-#   - Qdrant (работает по умолчанию на localhost:6333)
-#   - Ollama (llama3 + nomic-embed-text)
-#   - LangChain 1.2.10+ и связанные пакеты
-#
-# Запуск:
-#   1. Убедитесь, что Qdrant и Ollama запущены.
-#   2. pip install -r requirements.txt
-#   3. python main.py
-#
-# После запуска появится интерактивный клиент с командами:
-#   /add <path_to_file>   – добавить документ в базу знаний.
-#   /search <query>       – выполнить поиск.
-#   /quit                 – выйти.
-#   Любой другой ввод – агент обработает как обычный запрос.
 
 import os
 import sys
-import textwrap
+import json
 from pathlib import Path
 from typing import List, Dict, Any
 
@@ -29,170 +13,161 @@ from langchain_ollama import ChatOllama, OllamaEmbeddings
 from langchain_qdrant import QdrantVectorStore
 from langchain_text_splitters import RecursiveCharacterTextSplitter
 from langchain.tools import tool
-from langchain.agents import create_agent, AgentExecutor, AgentType
+from langchain.agents import create_agent, AgentExecutor
-from langchain_core.messages import HumanMessage
+from langchain.schema import AgentAction, AgentFinish
+from langchain_core.prompts import ChatPromptTemplate
 
 # ---------------------------------------------------------------------------
-# Конфигурация
+# Configuration
 # ---------------------------------------------------------------------------
-QDRANT_URL = "http://localhost:6333"
+QDRANT_URL = os.getenv("QDRANT_URL", "http://localhost:6333")
-COLLECTION_NAME = "knowledge"
+QDRANT_COLLECTION = "knowledge"
 EMBEDDING_MODEL = "nomic-embed-text"
 LLM_MODEL = "llama3"
+CHUNK_SIZE = 1000  # characters
+CHUNK_OVERLAP = 200
 
 # ---------------------------------------------------------------------------
-# Векторное хранилище
+# Vector store helper
 # ---------------------------------------------------------------------------
-# Инициализируем эмбеддер и клиент Qdrant
+class KnowledgeBase:
-embeddings = OllamaEmbeddings(model=EMBEDDING_MODEL)
+    """Wrapper around QdrantVectorStore providing add/search helpers."""
-vector_store = QdrantVectorStore(
+
-    url=QDRANT_URL,
+    def __init__(self, url: str = QDRANT_URL, collection: str = QDRANT_COLLECTION):
-    collection_name=COLLECTION_NAME,
+        self.embeddings = OllamaEmbeddings(model=EMBEDDING_MODEL)
-    embedding=embeddings,
+        self.store = QdrantVectorStore.from_existing_collection(
+            collection_name=collection,
+            url=url,
+            embedding=self.embeddings,
         )
 
-# ---------------------------------------------------------------------------
+    def add_documents(self, documents: List[str], titles: List[str]):
-# Чанкинг
+        """Adds a list of documents with corresponding titles to the store.
-# ---------------------------------------------------------------------------
+
-text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
+        Each document is split into chunks before being stored.
+        """
+        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,
+                "title": meta.get("title", "unknown"),
+                "distance": score,
+            }
+            for text, score, meta in results
+        ]
 
 # ---------------------------------------------------------------------------
-# Инструменты
+# Global knowledge base instance
 # ---------------------------------------------------------------------------
-@tool("search_knowledge_base", "Semantic search in the knowledge base")
+kb = KnowledgeBase()
+
+# ---------------------------------------------------------------------------
+# Tools
+# ---------------------------------------------------------------------------
+@tool("search_knowledge_base")
 def search_knowledge_base(query: str, max_results: int = 5) -> str:
-    """Возвращает топ‑N релевантных фрагментов из Qdrant."""
+    """Search the knowledge base for a query and return formatted results."""
-    results = vector_store.similarity_search_with_score(query, k=max_results)
+    results = kb.search(query, max_results)
     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]
-    formatted = []
+    return "\n\n".join(formatted)
-    for i, (doc, score) in enumerate(results, 1):
-        formatted.append(f"{i}. (score={score:.3f})\n{doc.page_content[:500]}\n---")
-    return "\n".join(formatted)
 
-@tool("add_to_knowledge_base", "Add a document to the knowledge base")
+@tool("add_to_knowledge_base")
-def add_to_knowledge_base(content: str, title: str = "") -> str:
+def add_to_knowledge_base(content: str, title: str) -> str:
-    """Разбивает документ на чанки, эмбеддит и сохраняет в Qdrant."""
+    """Add a new document to the knowledge base."""
-    # Разбиваем на чанки
+    kb.add_documents([content], [title])
-    chunks = text_splitter.split_text(content)
+    return f"Document '{title}' added successfully."
-    # Создаём Document объекты с метаданными
-    from langchain.schema import Document
-    docs = [Document(page_content=chunk, metadata={"title": title, "source": title}) for chunk in chunks]
-    # Добавляем в хранилище
-    vector_store.add_documents(docs)
-    return f"Added {len(docs)} chunks from '{title}'."
 
 # ---------------------------------------------------------------------------
-# Агент
+# Agent definition
 # ---------------------------------------------------------------------------
-# Системный промпт, который подсказывает агенту использовать инструменты
+SYSTEM_PROMPT = (
-SYSTEM_PROMPT = textwrap.dedent("""
+    "You are an assistant that can search and add documents to a local knowledge base. "
-You are an intelligent assistant with access to a knowledge base.
+    "Use the provided tools to perform semantic search and store new information. "
-Use the following tools when you need to retrieve or store information:
+    "When answering user queries, first determine if the user needs a search or an addition. "
-- search_knowledge_base: Perform a semantic search in the knowledge base.
+    "If no relevant information is found, suggest adding new content."
-- add_to_knowledge_base: Add new content to the knowledge base.
-When you answer a user query, first decide if you need to search the base.
-If you do, call search_knowledge_base with a relevant query.
-If you need to add new information, call add_to_knowledge_base.
-Otherwise, answer directly using your internal knowledge.
-""")
-
-# Создаём LLM
-llm = ChatOllama(model=LLM_MODEL, temperature=0.2)
-
-# Создаём агент
-agent = create_agent(
-    llm=llm,
-    tools=[search_knowledge_base, add_to_knowledge_base],
-    system_message=SYSTEM_PROMPT,
-    agent_type=AgentType.ZERO_SHOT_REACT_DESCRIPTION,
 )
 
-# Обёртка для удобного вызова
+prompt = ChatPromptTemplate.from_messages([
-agent_executor = AgentExecutor(agent=agent, tools=[search_knowledge_base, add_to_knowledge_base], verbose=True)
+    ("system", SYSTEM_PROMPT),
+    ("user", "{input}"),
+])
+
+agent = create_agent(
+    llm=ChatOllama(model=LLM_MODEL),
+    tools=[search_knowledge_base, add_to_knowledge_base],
+    prompt=prompt,
+    verbose=True,
+)
+
+agent_executor = AgentExecutor(agent=agent, tools=[search_knowledge_base, add_to_knowledge_base])
 
 # ---------------------------------------------------------------------------
-# Загрузка документов из директории (инициализация)
+# CLI client
 # ---------------------------------------------------------------------------
+def load_documents_from_dir(directory: str):
+    """Load all .txt files from a directory and add them to the knowledge base."""
+    docs = []
+    titles = []
+    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 load_documents_from_dir(directory: str) -> None:
-    """Загружает все .txt и .md файлы из указанной папки в базу."""
-    p = Path(directory)
-    if not p.is_dir():
-        print(f"Directory {directory} does not exist.")
-        return
-    for file_path in p.rglob("*.txt") | p.rglob("*.md"):
-        try:
-            content = file_path.read_text(encoding="utf-8")
-            title = file_path.stem
-            add_to_knowledge_base(content, title)
-            print(f"Loaded {file_path}")
-        except Exception as e:
-            print(f"Failed to load {file_path}: {e}")
 
-# ---------------------------------------------------------------------------
+def interactive_loop():
-# Интерактивный клиент
+    print("RAG Agent CLI. Commands: /add <title> <content>, /search <query>, /load <dir>, /quit")
-# ---------------------------------------------------------------------------
-
-def interactive_client():
-    print("Welcome to the RAG agent. Type /help for commands.")
     while True:
         try:
-            user_input = input("\n> ")
+            user_input = input("> ")
-        except (EOFError, KeyboardInterrupt):
+        except (KeyboardInterrupt, EOFError):
             print("\nExiting.")
             break
         if not user_input:
             continue
-        if user_input.startswith("/"):
+        if user_input.startswith("/quit"):
-            # Команды
-            if user_input.startswith("/add "):
-                path = user_input[5:].strip()
-                if not path:
-                    print("Usage: /add <path_to_file>")
-                    continue
-                try:
-                    content = Path(path).read_text(encoding="utf-8")
-                    title = Path(path).stem
-                    result = add_to_knowledge_base(content, title)
-                    print(result)
-                except Exception as e:
-                    print(f"Error reading file: {e}")
-            elif user_input.startswith("/search "):
-                query = user_input[8:].strip()
-                if not query:
-                    print("Usage: /search <query>")
-                    continue
-                result = search_knowledge_base(query)
-                print(result)
-            elif user_input in {"/quit", "/exit"}:
             print("Goodbye!")
             break
-            elif user_input == "/help":
+        elif user_input.startswith("/add "):
-                print(textwrap.dedent("""
-Commands:
-  /add <path>    – add a document to the knowledge base
-  /search <q>    – search the knowledge base
-  /quit /exit    – exit the program
-  /help          – show this help
-Any other input is treated as a normal user query for the agent.
-"""))
-            else:
-                print("Unknown command. Type /help for list of commands.")
-        else:
-            # Передаём запрос агенту
             try:
-                response = agent_executor.invoke({"input": user_input})
+                _, rest = user_input.split("/add ", 1)
-                print(response["output"])
+                title, content = rest.split(" ", 1)
-            except Exception as e:
+                response = add_to_knowledge_base(content, title)
-                print(f"Agent error: {e}")
+                print(response)
+            except ValueError:
+                print("Usage: /add <title> <content>")
+        elif user_input.startswith("/search "):
+            query = user_input.split("/search ", 1)[1]
+            results = search_knowledge_base(query)
+            print(results)
+        elif user_input.startswith("/load "):
+            dir_path = user_input.split("/load ", 1)[1]
+            load_documents_from_dir(dir_path)
+        else:
+            # Treat as normal user query
+            result = agent_executor.invoke({"input": user_input})
+            print(result.get("output", ""))
 
-# ---------------------------------------------------------------------------
-# Точка входа
-# ---------------------------------------------------------------------------
 if __name__ == "__main__":
-    # Если передан аргумент – загрузить документы из указанной папки
+    if len(sys.argv) > 1 and sys.argv[1] == "cli":
-    if len(sys.argv) > 1:
+        interactive_loop()
-        load_documents_from_dir(sys.argv[1])
+    else:
-    interactive_client()
+        print("Usage: python main.py cli")
+        print("Run the interactive CLI with: python main.py cli")