Updated main.py with Qdrant-based RAG implementation.

main.py

@@ -1,16 +1,23 @@
 # main.py
 # Полностью рабочий пример агента с RAG‑памятью на базе Qdrant и Ollama.
-# Используется LangChain 1.x, create_agent, инструменты @tool, и Ollama‑LLM/embeddings.
+# Весь код написан в одном файле для простоты демонстрации.
+#
+# Требования:
+#   - Python 3.10+
+#   - Qdrant (работает по умолчанию на localhost:6333)
+#   - Ollama (llama3 + nomic-embed-text)
+#   - LangChain 1.2.10+ и связанные пакеты
 #
 # Запуск:
-#   python main.py
+#   1. Убедитесь, что Qdrant и Ollama запущены.
-# После запуска можно использовать команды:
+#   2. pip install -r requirements.txt
-#   /add <title> <content>   – добавить документ
+#   3. python main.py
-#   /search <query> <max>    – семантический поиск
-#   /quit                    – выйти
 #
-# Для загрузки документов из директории используйте функцию load_documents_from_dir.
+# После запуска появится интерактивный клиент с командами:
-#"""
+#   /add <path_to_file>   – добавить документ в базу знаний.
+#   /search <query>       – выполнить поиск.
+#   /quit                 – выйти.
+#   Любой другой ввод – агент обработает как обычный запрос.
 
 import os
 import sys
@@ -22,14 +29,14 @@ 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, AgentToolkit, Tool
+from langchain.agents import create_agent, AgentExecutor, AgentType
 from langchain_core.messages import HumanMessage
 
 # ---------------------------------------------------------------------------
 # Конфигурация
 # ---------------------------------------------------------------------------
-QDRANT_URL = os.getenv("QDRANT_URL", "http://localhost:6333")
+QDRANT_URL = "http://localhost:6333"
-QDRANT_COLLECTION = "knowledge_base"
+COLLECTION_NAME = "knowledge"
 EMBEDDING_MODEL = "nomic-embed-text"
 LLM_MODEL = "llama3"
 
@@ -40,141 +47,152 @@ LLM_MODEL = "llama3"
 embeddings = OllamaEmbeddings(model=EMBEDDING_MODEL)
 vector_store = QdrantVectorStore(
     url=QDRANT_URL,
-    collection_name=QDRANT_COLLECTION,
+    collection_name=COLLECTION_NAME,
     embedding=embeddings,
 )
 
 # ---------------------------------------------------------------------------
 # Чанкинг
 # ---------------------------------------------------------------------------
-text_splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=50)
+text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
 
 # ---------------------------------------------------------------------------
 # Инструменты
 # ---------------------------------------------------------------------------
-@tool("search_knowledge_base", "Semantic search in the knowledge base.")
+@tool("search_knowledge_base", "Semantic search in the knowledge base")
 def search_knowledge_base(query: str, max_results: int = 5) -> str:
-    """Return top‑k relevant documents for a query.
+    """Возвращает топ‑N релевантных фрагментов из Qdrant."""
-    The function returns a formatted string with titles and snippets.
-    """
     results = vector_store.similarity_search_with_score(query, k=max_results)
     if not results:
         return "No relevant documents found."
+    # Формируем читаемый ответ
     formatted = []
-    for doc, score in results:
+    for i, (doc, score) in enumerate(results, 1):
-        title = doc.metadata.get("title", "Untitled")
+        formatted.append(f"{i}. (score={score:.3f})\n{doc.page_content[:500]}\n---")
-        snippet = doc.page_content[:200].replace("\n", " ")
-        formatted.append(f"{title} (score: {score:.3f}): {snippet}...")
     return "\n".join(formatted)
 
-@tool("add_to_knowledge_base", "Add a document to the 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 = "") -> str:
-    """Chunk the content, embed, and store in Qdrant.
+    """Разбивает документ на чанки, эмбеддит и сохраняет в Qdrant."""
-    Returns a confirmation message.
+    # Разбиваем на чанки
-    """
     chunks = text_splitter.split_text(content)
-    docs = []
+    # Создаём Document объекты с метаданными
-    for i, chunk in enumerate(chunks):
+    from langchain.schema import Document
-        docs.append(
+    docs = [Document(page_content=chunk, metadata={"title": title, "source": title}) for chunk in chunks]
-            {
+    # Добавляем в хранилище
-                "page_content": chunk,
-                "metadata": {"title": title, "chunk_index": i},
-            }
-        )
     vector_store.add_documents(docs)
-    return f"Added {len(chunks)} chunks of '{title}' to the knowledge base."
+    return f"Added {len(docs)} chunks from '{title}'."
 
 # ---------------------------------------------------------------------------
 # Агент
 # ---------------------------------------------------------------------------
+# Системный промпт, который подсказывает агенту использовать инструменты
+SYSTEM_PROMPT = textwrap.dedent("""
+You are an intelligent assistant with access to a knowledge base.
+Use the following tools when you need to retrieve or store information:
+- search_knowledge_base: Perform a semantic search in the knowledge base.
+- 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)
 
-# Список инструментов
-tools = [search_knowledge_base, add_to_knowledge_base]
-
 # Создаём агент
 agent = create_agent(
     llm=llm,
-    tools=tools,
+    tools=[search_knowledge_base, add_to_knowledge_base],
-    system_message="You are an assistant that can search and add documents to a local knowledge base. Use the provided tools.",
+    system_message=SYSTEM_PROMPT,
-    verbose=True,
+    agent_type=AgentType.ZERO_SHOT_REACT_DESCRIPTION,
 )
 
-# Обёртка для выполнения
+# Обёртка для удобного вызова
-agent_executor = AgentExecutor(agent=agent, tools=tools, verbose=True)
+agent_executor = AgentExecutor(agent=agent, tools=[search_knowledge_base, add_to_knowledge_base], verbose=True)
 
 # ---------------------------------------------------------------------------
-# Загрузка документов из директории
+# Загрузка документов из директории (инициализация)
 # ---------------------------------------------------------------------------
 
 def load_documents_from_dir(directory: str) -> None:
-    """Load all .txt files from a directory into the knowledge base.
+    """Загружает все .txt и .md файлы из указанной папки в базу."""
-    Each file becomes a separate document with its filename as title.
+    p = Path(directory)
-    """
+    if not p.is_dir():
-    path = Path(directory)
-    if not path.is_dir():
         print(f"Directory {directory} does not exist.")
         return
-    for file in path.glob("*.txt"):
+    for file_path in p.rglob("*.txt") | p.rglob("*.md"):
-        title = file.stem
+        try:
-        content = file.read_text(encoding="utf-8")
+            content = file_path.read_text(encoding="utf-8")
-        print(f"Adding {title}...", end=" ")
+            title = file_path.stem
-        result = add_to_knowledge_base(content, title)
+            add_to_knowledge_base(content, title)
-        print(result)
+            print(f"Loaded {file_path}")
+        except Exception as e:
+            print(f"Failed to load {file_path}: {e}")
 
 # ---------------------------------------------------------------------------
-# CLI
+# Интерактивный клиент
 # ---------------------------------------------------------------------------
 
-def main():
+def interactive_client():
-    # Если пользователь передал путь к директории, загрузим документы
+    print("Welcome to the RAG agent. Type /help for commands.")
-    if len(sys.argv) > 1:
-        load_documents_from_dir(sys.argv[1])
-
-    print("\n--- RAG Agent CLI ---")
-    print("Commands:")
-    print("  /add <title> <content>   – add a document")
-    print("  /search <query> <max>    – search knowledge base")
-    print("  /quit                    – exit")
-
     while True:
         try:
             user_input = input("\n> ")
         except (EOFError, KeyboardInterrupt):
             print("\nExiting.")
             break
-
+        if not user_input:
-        if not user_input.strip():
             continue
+        if user_input.startswith("/"):
+            # Команды
+            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":
+                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})
+                print(response["output"])
+            except Exception as e:
+                print(f"Agent error: {e}")
 
-        if user_input.startswith("/quit"):
+# ---------------------------------------------------------------------------
-            print("Goodbye!")
+# Точка входа
-            break
+# ---------------------------------------------------------------------------
-
-        if user_input.startswith("/add"):
-            parts = user_input.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 user_input.startswith("/search"):
-            parts = user_input.split(maxsplit=2)
-            if len(parts) < 2:
-                print("Usage: /search <query> [max_results]")
-                continue
-            query = parts[1]
-            max_results = int(parts[2]) if len(parts) > 2 else 5
-            print(search_knowledge_base(query, max_results))
-            continue
-
-        # Любой другой ввод – передаём агенту
-        response = agent_executor.invoke({"input": user_input})
-        print(response.get("output", ""))
-
-
 if __name__ == "__main__":
-    main()
+    # Если передан аргумент – загрузить документы из указанной папки
+    if len(sys.argv) > 1:
+        load_documents_from_dir(sys.argv[1])
+    interactive_client()