Агент с RAG‑памятью: solution.py

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

@@ -2,19 +2,20 @@
 # -*- coding: utf-8 -*-
 
 """
-solution.py – Агент с RAG‑памятью на базе Qdrant и Ollama.
+solution.py – основной скрипт с реализацией RAG‑агента и инструментов.
 """
 
-import argparse
+import os
+import sys
 from pathlib import Path
 from typing import List, Dict
 
-from langchain_ollama import OllamaEmbeddings
+from langchain_ollama import OllamaEmbeddings, OllamaLLM
 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
-from langchain.schema import Document
+from langchain.agents import create_agent, AgentExecutor, Tool
 
 # --------------------------------------------------------------------------- #
 # Конфигурация
@@ -24,16 +25,20 @@ QDRANT_HOST = "localhost"
 QDRANT_PORT = 6333
 COLLECTION_NAME = "knowledge_base"
 
-EMBEDDING_MODEL = "nomic-embed-text"   # Ollama embedding model
-LLM_MODEL = "ollama:llama3"            # Ollama LLM (формат для create_agent)
+EMBEDDING_MODEL = "nomic-embed-text"
+LLM_MODEL = "llama3"
+
+MAX_CHUNK_SIZE = 1000   # символов
+CHUNK_OVERLAP = 200     # символов
+
 
 # --------------------------------------------------------------------------- #
-# Векторное хранилище и вспомогательные функции
+# Векторное хранилище (обёртка над QdrantVectorStore)
 # --------------------------------------------------------------------------- #
 
 class KnowledgeBase:
     """
-    Обёртка над QdrantVectorStore с поддержкой чанкинга и инструментов.
+    Класс, инкапсулирующий работу с Qdrant и Ollama.
     """
 
     def __init__(self,
@@ -41,108 +46,121 @@ class KnowledgeBase:
                  port: int = QDRANT_PORT,
                  collection_name: str = COLLECTION_NAME):
         self.embeddings = OllamaEmbeddings(model=EMBEDDING_MODEL)
-        # Создаём/подключаем коллекцию
         self.store = QdrantVectorStore(
-            embedding=self.embeddings,
             url=f"http://{host}:{port}",
             collection_name=collection_name,
+            embedding=self.embeddings
         )
-        # Чанкер для разбивки документов
-        self.splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=50)
+        # Создаём коллекцию, если её ещё нет
+        if not self.store.collection_exists():
+            self.store.create_collection()
 
-    def add_document(self, content: str, title: str) -> None:
+    def add_documents(self, documents: List[Document]) -> None:
         """
-        Добавляет документ в базу после разбиения на чанки.
+        Добавляет список документов в хранилище.
         """
-        chunks = self.splitter.split_text(content)
-        docs = [
-            Document(
-                page_content=chunk,
-                metadata={"title": title, "source": f"{title} (chunk {i})"}
-            )
-            for i, chunk in enumerate(chunks)
+        self.store.add_documents(documents)
+
+    def similarity_search(
+        self,
+        query: str,
+        k: int = 5
+    ) -> List[Dict]:
+        """
+        Семантический поиск по запросу. Возвращает списки словарей с полями:
+        'content', 'metadata'.
+        """
+        results = self.store.similarity_search(query, k=k)
+        return [
+            {"content": doc.page_content,
+             "metadata": doc.metadata}
+            for doc in results
         ]
-        self.store.add_documents(docs)
-
-    def search(self, query: str, max_results: int = 5) -> List[Dict]:
-        """
-        Семантический поиск по базе знаний.
-        Возвращает список словарей с полями 'title', 'source' и 'content'.
-        """
-        results = self.store.similarity_search_with_score(query, k=max_results)
-        output = []
-        for doc, score in results:
-            output.append({
-                "score": score,
-                "title": doc.metadata.get("title", ""),
-                "source": doc.metadata.get("source", ""),
-                "content": doc.page_content,
-            })
-        return output
 
 
 # --------------------------------------------------------------------------- #
-# Глобальный экземпляр базы знаний
+# Чанкинг документов
 # --------------------------------------------------------------------------- #
 
-kb = KnowledgeBase()
+def split_text_to_documents(text: str, title: str) -> List[Document]:
+    """
+    Делит текст на чанки и возвращает список Document.
+    """
+    splitter = RecursiveCharacterTextSplitter(
+        chunk_size=MAX_CHUNK_SIZE,
+        chunk_overlap=CHUNK_OVERLAP
+    )
+    chunks = splitter.split_text(text)
+    return [
+        Document(page_content=chunk, metadata={"title": title})
+        for chunk in chunks
+    ]
+
 
 # --------------------------------------------------------------------------- #
 # Инструменты для агента
 # --------------------------------------------------------------------------- #
 
-@tool
-def search_knowledge_base(query: str, max_results: int = 5) -> str:
-    """
-    Поиск в базе знаний.
-    Возвращает строку с найденными результатами в формате Markdown.
-    """
-    results = kb.search(query, max_results)
-    if not results:
-        return "Ничего не найдено."
-    lines = []
-    for r in results:
-        lines.append(
-            f"**{r['title']}** ({r['source']}) – {r['score']:.2f}\n> {r['content'][:200]}..."
-        )
-    return "\n\n".join(lines)
+kb = KnowledgeBase()
 
 
-@tool
+@tool("search_knowledge_base")
+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")
 def add_to_knowledge_base(content: str, title: str) -> str:
     """
-    Добавление нового документа в базу знаний.
+    Добавляет новый документ в базу знаний после разбиения на чанки.
+    Возвращает сообщение об успешном добавлении.
     """
-    kb.add_document(content, title)
-    return f"Документ «{title}» успешно добавлен."
+    docs = split_text_to_documents(content, title)
+    kb.add_documents(docs)
+    return f"Документ '{title}' успешно добавлен в базу."
 
 
 # --------------------------------------------------------------------------- #
 # Создание агента
 # --------------------------------------------------------------------------- #
 
-def build_agent():
-    """
-    Возвращает готовый агент с инструментами RAG‑памяти.
-    Используется актуальный API: create_agent(model, tools, system_prompt).
-    """
-    system_prompt = (
-        "Ты – интеллектуальный ассистент, который использует локальную базу знаний. "
-        "При ответе на запросы сначала ищи релевантную информацию в базе через инструмент "
-        "`search_knowledge_base`. Если нужно сохранить новую информацию, добавляй её через "
-        "`add_to_knowledge_base`."
-    )
+llm = OllamaLLM(model=LLM_MODEL)
 
-    agent = create_agent(
-        LLM_MODEL,
-        tools=[search_knowledge_base, add_to_knowledge_base],
-        system_prompt=system_prompt,
+system_prompt = """
+Ты – интеллектуальный агент с доступом к локальной базе знаний.
+При ответе на запросы используй только информацию из базы, если она есть.
+Если необходима дополнительная информация, сначала выполни поиск в базе,
+затем сформируй ответ. Не выдавай вымышленных фактов.
+"""
+
+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="Добавляет новый документ в базу знаний."
     )
-    return agent
+]
+
+agent = create_agent(
+    llm=llm,
+    tools=tools,
+    system_prompt=system_prompt
+)
+
+executor = AgentExecutor(agent=agent, tools=tools, verbose=True)
 
 
 # --------------------------------------------------------------------------- #
-# Загрузка документов из директории
+# Инициализация: загрузка документов из директории
 # --------------------------------------------------------------------------- #
 
 def load_documents_from_dir(directory: Path) -> None:
@@ -150,95 +168,86 @@ def load_documents_from_dir(directory: Path) -> None:
     Загружает все текстовые файлы из указанной папки в базу знаний.
     """
     for file_path in directory.rglob("*"):
-        if file_path.suffix.lower() not in {".txt", ".md"}:
-            continue
-        content = file_path.read_text(encoding="utf-8")
-        title = file_path.stem
-        kb.add_document(content, title)
-        print(f"Загружено: {file_path}")
+        if file_path.is_file() and file_path.suffix.lower() in {".txt", ".md"}:
+            text = file_path.read_text(encoding="utf-8")
+            title = file_path.stem
+            docs = split_text_to_documents(text, title)
+            kb.add_documents(docs)
+            print(f"Загружен документ: {title}")
 
 
 # --------------------------------------------------------------------------- #
-# CLI клиент
+# Интерактивный клиент
 # --------------------------------------------------------------------------- #
 
-def interactive_cli(agent) -> None:
+def interactive_cli() -> None:
     """
-    Простая REPL‑интерфейс с командами /add, /search и /quit.
-    Любой другой ввод считается запросом к агенту.
+    Простая CLI с командами /add, /search и /quit.
     """
-    print("=== Агент RAG ===")
+    print("=== RAG‑Агент ===")
     print("Команды:")
-    print("/add <title> – добавить новый документ (будет запрошен контент)")
+    print("/add <title> – добавить новый документ (будет запрошен ввод текста)")
     print("/search <query> – поиск в базе знаний")
     print("/quit – выйти")
 
     while True:
         try:
-            inp = input("\n> ").strip()
+            user_input = input("\n> ").strip()
         except EOFError:
             break
-        if not inp:
+
+        if not user_input:
             continue
 
-        if inp.lower() == "/quit":
+        if user_input.lower() == "/quit":
+            print("До свидания!")
             break
 
-        elif inp.startswith("/add "):
-            title = inp[5:].strip()
-            print("Введите содержимое документа (завершите ввод Ctrl+D / Ctrl+Z):")
-            try:
-                content_lines = []
-                while True:
-                    line = input()
-                    content_lines.append(line)
-            except EOFError:
-                pass
-            content = "\n".join(content_lines)
-            result = add_to_knowledge_base.invoke({"content": content, "title": title})
+        if user_input.startswith("/add"):
+            parts = user_input.split(maxsplit=1)
+            title = parts[1] if len(parts) > 1 else "Untitled"
+            print(f"Введите текст для документа '{title}'. Завершите ввод пустой строкой.")
+            lines = []
+            while True:
+                line = input()
+                if line == "":
+                    break
+                lines.append(line)
+            content = "\n".join(lines)
+            result = add_to_knowledge_base(content, title)
             print(result)
 
-        elif inp.startswith("/search "):
-            query = inp[8:].strip()
-            results = kb.search(query, max_results=5)
+        elif user_input.startswith("/search"):
+            query = user_input[len("/search"):].strip()
+            if not query:
+                print("Введите запрос после /search")
+                continue
+            results = search_knowledge_base(query, max_results=3)
             if not results:
                 print("Ничего не найдено.")
-            else:
-                for r in results:
-                    print(f"\n**{r['title']}** ({r['source']}) – {r['score']:.2f}")
-                    print(r["content"][:400] + "...")
+                continue
+            for i, res in enumerate(results, 1):
+                print(f"\nРезультат {i}:")
+                print(f"Титул: {res['metadata'].get('title', 'Неизвестно')}")
+                print(res["content"][:500] + ("..." if len(res["content"]) > 500 else ""))
 
         else:
-            # Любой другой ввод – запрос к агенту
-            # Актуальный API: invoke принимает {"messages": [...]}
-            result = agent.invoke(
-                {"messages": [{"role": "user", "content": inp}]}
-            )
-            # Извлекаем текст последнего сообщения агента
-            messages = result.get("messages", [])
-            if messages:
-                print(messages[-1].content)
+            # Любой другой ввод считается запросом к агенту
+            response = executor.invoke({"input": user_input})
+            print(response.get("output", "Нет ответа."))
 
 
 # --------------------------------------------------------------------------- #
 # Точка входа
 # --------------------------------------------------------------------------- #
 
-def main() -> None:
-    parser = argparse.ArgumentParser(description="Агент с RAG‑памятью")
-    parser.add_argument(
-        "--load-dir",
-        type=str,
-        help="Путь к директории с документами для предварительной загрузки"
-    )
-    args = parser.parse_args()
-
-    if args.load_dir:
-        load_documents_from_dir(Path(args.load_dir))
-
-    agent = build_agent()
-    interactive_cli(agent)
-
-
 if __name__ == "__main__":
-    main()
+    if len(sys.argv) > 1 and sys.argv[1] == "--init":
+        # Инициализация: загрузка документов из папки docs/
+        docs_dir = Path("docs")
+        if not docs_dir.exists():
+            print(f"Папка {docs_dir} не найдена.")
+            sys.exit(1)
+        load_documents_from_dir(docs_dir)
+    else:
+        interactive_cli()