Обновить solutions/6a02e23da6fe2e4ac16acf65_Агент_с_RAG_памятью/solution.py

solutions/6a02e23da6fe2e4ac16acf65_Агент_с_RAG_памятью/solution.py

@@ -2,43 +2,46 @@
 # -*- coding: utf-8 -*-
 
 """
-solution.py – основной скрипт с реализацией RAG‑агента и инструментов.
+solution.py – Агент с RAG‑памятью на базе Qdrant и Ollama.
 """
 
-import os
 import sys
 from pathlib import Path
 from typing import List, Dict
 
-from langchain_ollama import OllamaEmbeddings, OllamaLLM
+from langchain_ollama import OllamaEmbeddings
 from langchain_qdrant import QdrantVectorStore
 from langchain.schema import Document
 from langchain.text_splitter import RecursiveCharacterTextSplitter
 from langchain.tools import tool
-from langchain.agents import create_agent, AgentExecutor, Tool
+from langchain.agents import create_agent
 
 # --------------------------------------------------------------------------- #
 # Конфигурация
 # --------------------------------------------------------------------------- #
 
-QDRANT_HOST = "localhost"
+QDRANT_HOST      = "localhost"
-QDRANT_PORT = 6333
+QDRANT_PORT      = 6333
-COLLECTION_NAME = "knowledge_base"
+COLLECTION_NAME  = "knowledge_base"
 
-EMBEDDING_MODEL = "nomic-embed-text"
+EMBEDDING_MODEL  = "nomic-embed-text"
-LLM_MODEL = "llama3"
+LLM_MODEL        = "ollama:llama3"   # формат model-string для create_agent
 
-MAX_CHUNK_SIZE = 1000   # символов
+MAX_CHUNK_SIZE   = 1000   # символов
-CHUNK_OVERLAP = 200     # символов
+CHUNK_OVERLAP    = 200    # символов
 
 
 # --------------------------------------------------------------------------- #
-# Векторное хранилище (обёртка над QdrantVectorStore)
+# Векторное хранилище
 # --------------------------------------------------------------------------- #
 
 class KnowledgeBase:
     """
-    Класс, инкапсулирующий работу с Qdrant и Ollama.
+    Инкапсулирует работу с Qdrant и эмбеддингами Ollama.
+
+    Инициализация через QdrantVectorStore.from_existing_collection —
+    подключается к уже существующей коллекции (или создаёт при первом
+    add_documents, когда Qdrant сам управляет коллекцией через URL).
     """
 
     def __init__(self,
@@ -46,35 +49,31 @@ class KnowledgeBase:
                  port: int = QDRANT_PORT,
                  collection_name: str = COLLECTION_NAME):
         self.embeddings = OllamaEmbeddings(model=EMBEDDING_MODEL)
-        self.store = QdrantVectorStore(
+        # Подключаемся к существующей коллекции (или используем URL-режим,
-            url=f"http://{host}:{port}",
+        # при котором QdrantVectorStore создаёт коллекцию автоматически).
+        self.store = QdrantVectorStore.from_existing_collection(
+            embedding=self.embeddings,
             collection_name=collection_name,
-            embedding=self.embeddings
+            url=f"http://{host}:{port}",
         )
-        # Создаём коллекцию, если её ещё нет
-        if not self.store.collection_exists():
-            self.store.create_collection()
 
     def add_documents(self, documents: List[Document]) -> None:
-        """
+        """Добавляет список Document в хранилище."""
-        Добавляет список документов в хранилище.
-        """
         self.store.add_documents(documents)
 
-    def similarity_search(
+    def similarity_search(self, query: str, k: int = 5) -> List[Dict]:
-        self,
-        query: str,
-        k: int = 5
-    ) -> List[Dict]:
         """
-        Семантический поиск по запросу. Возвращает списки словарей с полями:
+        Семантический поиск с метрикой релевантности.
-        'content', 'metadata'.
+        Возвращает список словарей с полями 'content', 'metadata', 'score'.
         """
-        results = self.store.similarity_search(query, k=k)
+        results = self.store.similarity_search_with_score(query, k=k)
         return [
-            {"content": doc.page_content,
+            {
-             "metadata": doc.metadata}
+                "content":  doc.page_content,
-            for doc in results
+                "metadata": doc.metadata,
+                "score":    score,
+            }
+            for doc, score in results
         ]
 
 
@@ -84,157 +83,164 @@ class KnowledgeBase:
 
 def split_text_to_documents(text: str, title: str) -> List[Document]:
     """
-    Делит текст на чанки и возвращает список Document.
+    Разбивает текст на чанки через RecursiveCharacterTextSplitter.
+    Каждый чанк получает метаданные: title и порядковый номер чанка.
     """
     splitter = RecursiveCharacterTextSplitter(
         chunk_size=MAX_CHUNK_SIZE,
-        chunk_overlap=CHUNK_OVERLAP
+        chunk_overlap=CHUNK_OVERLAP,
     )
     chunks = splitter.split_text(text)
     return [
-        Document(page_content=chunk, metadata={"title": title})
+        Document(
-        for chunk in chunks
+            page_content=chunk,
+            metadata={"title": title, "chunk_index": i},
+        )
+        for i, chunk in enumerate(chunks)
     ]
 
 
 # --------------------------------------------------------------------------- #
-# Инструменты для агента
+# Глобальный экземпляр базы знаний + инструменты агента (@tool)
 # --------------------------------------------------------------------------- #
 
 kb = KnowledgeBase()
 
 
-@tool("search_knowledge_base")
+@tool
 def search_knowledge_base(query: str, max_results: int = 5) -> List[Dict]:
     """
     Семантический поиск в базе знаний.
-    Возвращает список найденных документов с их содержимым и метаданными.
+    Возвращает список найденных документов с содержимым, метаданными и оценкой релевантности.
     """
     return kb.similarity_search(query, k=max_results)
 
 
-@tool("add_to_knowledge_base")
+@tool
 def add_to_knowledge_base(content: str, title: str) -> str:
     """
     Добавляет новый документ в базу знаний после разбиения на чанки.
-    Возвращает сообщение об успешном добавлении.
+    Возвращает подтверждение успешного добавления.
     """
     docs = split_text_to_documents(content, title)
     kb.add_documents(docs)
-    return f"Документ '{title}' успешно добавлен в базу."
+    return f"Документ '{title}' успешно добавлен в базу ({len(docs)} чанков)."
 
 
 # --------------------------------------------------------------------------- #
-# Создание агента
+# Агент с RAG‑интеграцией
 # --------------------------------------------------------------------------- #
 
-llm = OllamaLLM(model=LLM_MODEL)
+system_prompt = (
-
+    "Ты – интеллектуальный агент с доступом к локальной базе знаний.\n"
-system_prompt = """
+    "При ответе на запросы сначала ищи релевантную информацию через инструмент "
-Ты – интеллектуальный агент с доступом к локальной базе знаний.
+    "`search_knowledge_base`. Если нужно сохранить новую информацию — используй "
-При ответе на запросы используй только информацию из базы, если она есть.
+    "`add_to_knowledge_base`. Не выдумывай факты, которых нет в базе."
-Если необходима дополнительная информация, сначала выполни поиск в базе,
-затем сформируй ответ. Не выдавай вымышленных фактов.
-"""
-
-tools = [
-    Tool.from_function(
-        func=search_knowledge_base,
-        name="search_knowledge_base",
-        description="Используется для поиска информации в базе знаний."
-    ),
-    Tool.from_function(
-        func=add_to_knowledge_base,
-        name="add_to_knowledge_base",
-        description="Добавляет новый документ в базу знаний."
-    )
-]
-
-agent = create_agent(
-    llm=llm,
-    tools=tools,
-    system_prompt=system_prompt
 )
 
-executor = AgentExecutor(agent=agent, tools=tools, verbose=True)
+# create_agent(model, tools, system_prompt) — актуальный API LangChain.
+# Первый аргумент — строка модели в формате "provider:model".
+# Возвращает граф-агент с .invoke() / .stream() — AgentExecutor не нужен.
+agent = create_agent(
+    LLM_MODEL,
+    tools=[search_knowledge_base, add_to_knowledge_base],
+    system_prompt=system_prompt,
+)
 
 
 # --------------------------------------------------------------------------- #
-# Инициализация: загрузка документов из директории
+# Клиент инициализации: загрузка документов из директории
 # --------------------------------------------------------------------------- #
 
 def load_documents_from_dir(directory: Path) -> None:
     """
-    Загружает все текстовые файлы из указанной папки в базу знаний.
+    Загружает все .txt и .md файлы из указанной папки в базу знаний.
     """
+    loaded = 0
     for file_path in directory.rglob("*"):
         if file_path.is_file() and file_path.suffix.lower() in {".txt", ".md"}:
-            text = file_path.read_text(encoding="utf-8")
+            text  = file_path.read_text(encoding="utf-8")
             title = file_path.stem
-            docs = split_text_to_documents(text, title)
+            docs  = split_text_to_documents(text, title)
             kb.add_documents(docs)
-            print(f"Загружен документ: {title}")
+            print(f"Загружен: {file_path}  ({len(docs)} чанков)")
+            loaded += 1
+    print(f"\nИтого загружено файлов: {loaded}")
 
 
 # --------------------------------------------------------------------------- #
-# Интерактивный клиент
+# Интерактивный тестовый клиент
 # --------------------------------------------------------------------------- #
 
 def interactive_cli() -> None:
     """
-    Простая CLI с командами /add, /search и /quit.
+    REPL с командами /add, /search, /quit.
+    Любой другой ввод передаётся агенту.
     """
     print("=== RAG‑Агент ===")
     print("Команды:")
-    print("/add <title> – добавить новый документ (будет запрошен ввод текста)")
+    print("  /add <title>    – добавить документ (ввод текста, конец — пустая строка)")
-    print("/search <query> – поиск в базе знаний")
+    print("  /search <query> – поиск в базе знаний")
-    print("/quit – выйти")
+    print("  /quit           – выйти")
+    print("  <любой текст>   – запрос к агенту\n")
 
     while True:
         try:
-            user_input = input("\n> ").strip()
+            user_input = input("> ").strip()
         except EOFError:
             break
 
         if not user_input:
             continue
 
+        # /quit
         if user_input.lower() == "/quit":
             print("До свидания!")
             break
 
-        if user_input.startswith("/add"):
+        # /add <title>
+        elif user_input.startswith("/add"):
             parts = user_input.split(maxsplit=1)
-            title = parts[1] if len(parts) > 1 else "Untitled"
+            title = parts[1].strip() if len(parts) > 1 else "Untitled"
-            print(f"Введите текст для документа '{title}'. Завершите ввод пустой строкой.")
+            print(f"Вводите текст для «{title}». Пустая строка — конец ввода.")
-            lines = []
+            lines: List[str] = []
             while True:
                 line = input()
                 if line == "":
                     break
                 lines.append(line)
             content = "\n".join(lines)
-            result = add_to_knowledge_base(content, title)
+            # Вызов @tool через .invoke() — правильный способ вызова инструмента
+            result = add_to_knowledge_base.invoke({"content": content, "title": title})
             print(result)
 
+        # /search <query>
         elif user_input.startswith("/search"):
             query = user_input[len("/search"):].strip()
             if not query:
-                print("Введите запрос после /search")
+                print("Укажите запрос: /search <query>")
                 continue
-            results = search_knowledge_base(query, max_results=3)
+            results = search_knowledge_base.invoke({"query": query, "max_results": 3})
             if not results:
                 print("Ничего не найдено.")
-                continue
+            else:
-            for i, res in enumerate(results, 1):
+                for i, res in enumerate(results, 1):
-                print(f"\nРезультат {i}:")
+                    title = res["metadata"].get("title", "—")
-                print(f"Титул: {res['metadata'].get('title', 'Неизвестно')}")
+                    score = res.get("score", 0)
-                print(res["content"][:500] + ("..." if len(res["content"]) > 500 else ""))
+                    print(f"\nРезультат {i} | {title} | score={score:.4f}")
+                    snippet = res["content"]
+                    print(snippet[:500] + ("..." if len(snippet) > 500 else ""))
 
+        # Запрос к агенту
         else:
-            # Любой другой ввод считается запросом к агенту
+            # Актуальный API: invoke принимает {"messages": [...]},
-            response = executor.invoke({"input": user_input})
+            # ответ — в result["messages"][-1].content
-            print(response.get("output", "Нет ответа."))
+            result   = agent.invoke({"messages": [{"role": "user", "content": user_input}]})
+            messages = result.get("messages", [])
+            if messages:
+                print(messages[-1].content)
+            else:
+                print("Нет ответа.")
 
 
 # --------------------------------------------------------------------------- #
@@ -243,10 +249,10 @@ def interactive_cli() -> None:
 
 if __name__ == "__main__":
     if len(sys.argv) > 1 and sys.argv[1] == "--init":
-        # Инициализация: загрузка документов из папки docs/
+        # python solution.py --init   →  загрузить документы из ./docs/
-        docs_dir = Path("docs")
+        docs_dir = Path(sys.argv[2]) if len(sys.argv) > 2 else Path("docs")
         if not docs_dir.exists():
-            print(f"Папка {docs_dir} не найдена.")
+            print(f"Папка '{docs_dir}' не найдена.")
             sys.exit(1)
         load_documents_from_dir(docs_dir)
     else: