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

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

@@ -1,15 +1,21 @@
-# rag_agent.py
+"""
+rag_agent.py — AI-агент с RAG-памятью на базе Ollama + Qdrant.
+
+Стек: langchain-ollama, langchain-qdrant, langgraph, qdrant-client, rich
+Установка: pip install langchain langchain-core langchain-ollama langchain-qdrant
+           langchain-text-splitters langgraph qdrant-client rich
+"""
+
 import os
 from pathlib import Path
 from typing import List
 
-from langchain_ollama import OllamaEmbeddings, OllamaLLM
+from langchain_ollama import OllamaEmbeddings, ChatOllama
 from langchain_qdrant import QdrantVectorStore
 from langchain_text_splitters import RecursiveCharacterTextSplitter
-from langchain.tools import tool
+from langchain_core.tools import tool
-from langchain.agents import create_react_agent, AgentExecutor
+from langchain_core.documents import Document
-from langchain.prompts import PromptTemplate
+from langgraph.prebuilt import create_react_agent
-from langchain.schema import Document
 from qdrant_client import QdrantClient as _QdrantClient
 from qdrant_client.http.models import Distance, VectorParams
 
@@ -33,10 +39,8 @@ class RagVectorStore:
     def __init__(self, collection_name: str = COLLECTION_NAME):
         self.embeddings = OllamaEmbeddings(model=EMBEDDING_MODEL)
 
-        # Низкоуровневый клиент для управления коллекцией
         self._client = _QdrantClient(host=QDRANT_HOST, port=QDRANT_PORT)
 
-        # Создаём коллекцию, если её ещё нет
         if not self._collection_exists(collection_name):
             self._client.create_collection(
                 collection_name=collection_name,
@@ -60,7 +64,9 @@ class RagVectorStore:
         self.store.add_documents(docs)
 
     def search(self, query: str, k: int = 5) -> List[tuple]:
-        results = self.store.similarity_search_with_score(query, k=k)
+        # similarity_search_with_relevance_scores возвращает score в [0..1],
+        # где 1.0 = максимальная релевантность (исправлено с with_score)
+        results = self.store.similarity_search_with_relevance_scores(query, k=k)
         return [(doc.page_content, score) for doc, score in results]
 
 
@@ -69,135 +75,222 @@ class RagVectorStore:
 # ---------------------------------------------------------------------------
 
 def chunk_document(text: str, title: str) -> List[Document]:
-    """Разбивает текст на чанки с метаданными."""
+    """Разбивает текст на чанки с метаданными через RecursiveCharacterTextSplitter."""
     splitter = RecursiveCharacterTextSplitter(
         chunk_size=1000,
         chunk_overlap=200,
         separators=["\n\n", "\n", " ", ""],
     )
     return [
-        Document(page_content=chunk, metadata={"title": title})
+        Document(
-        for chunk in splitter.split_text(text)
+            page_content=chunk,
+            metadata={"title": title, "source": title, "chunk_index": i},
+        )
+        for i, chunk in enumerate(splitter.split_text(text))
     ]
 
 
 # ---------------------------------------------------------------------------
-# Инициализация хранилища (глобальный синглтон для инструментов)
+# Синглтон хранилища (используется инструментами агента)
 # ---------------------------------------------------------------------------
 
 rag_store = RagVectorStore()
 
 
 # ---------------------------------------------------------------------------
-# LangChain-инструменты агента (@tool)
+# RAG-инструменты агента (@tool)
 # ---------------------------------------------------------------------------
 
 @tool
-def search_knowledge_base(query: str, max_results: int = 5) -> List[str]:
+def search_knowledge_base(query: str, max_results: int = 5) -> str:
     """
-    Семантический поиск в базе знаний.
+    Выполняет семантический поиск в базе знаний.
 
     Args:
-        query: поисковый запрос.
+        query:       поисковый запрос на естественном языке.
-        max_results: максимальное количество результатов.
+        max_results: максимальное количество результатов (по умолчанию 5).
 
     Returns:
-        Список найденных фрагментов с оценкой релевантности.
+        Строка с найденными фрагментами и оценками релевантности.
     """
     results = rag_store.search(query, k=max_results)
-    return [f"[score={score:.4f}] {text}" for text, score in results]
+    if not results:
+        return "В базе знаний ничего не найдено по данному запросу."
+    lines = [
+        f"[{i}] (score={score:.3f}) {text}"
+        for i, (text, score) in enumerate(results, start=1)
+    ]
+    return "\n\n".join(lines)
 
 
 @tool
-def add_to_knowledge_base(content: str, title: str) -> str:
+def add_to_knowledge_base(content: str, title: str = "") -> str:
     """
     Добавляет новый документ в базу знаний.
 
     Args:
-        content: полный текст документа.
+        content: текстовое содержимое документа.
-        title: заголовок документа.
+        title:   название / источник документа (необязательно).
 
     Returns:
-        Сообщение об успехе.
+        Сообщение об успехе с количеством созданных чанков.
     """
-    docs = chunk_document(content, title)
+    docs = chunk_document(content, title=title)
     rag_store.add_documents(docs)
-    return f"Документ '{title}' успешно добавлен в базу знаний."
+    return (
-
+        f"Документ «{title or 'без названия'}» успешно добавлен. "
-
+        f"Создано чанков: {len(docs)}."
-# ---------------------------------------------------------------------------
-# Агент на базе LangChain (create_react_agent)
-# ---------------------------------------------------------------------------
-
-llm = OllamaLLM(model=LLM_MODEL)
-
-tools = [search_knowledge_base, add_to_knowledge_base]
-
-prompt = PromptTemplate.from_template(
-    "Вы — интеллектуальный агент с доступом к локальной базе знаний.\n"
-    "При необходимости используйте инструменты search_knowledge_base и add_to_knowledge_base.\n\n"
-    "Доступные инструменты:\n{tools}\n\n"
-    "Формат ответа (строго соблюдайте):\n"
-    "Question: вопрос пользователя\n"
-    "Thought: что нужно сделать\n"
-    "Action: инструмент из [{tool_names}]\n"
-    "Action Input: входные данные для инструмента\n"
-    "Observation: результат инструмента\n"
-    "... (можно повторять Thought/Action/Observation)\n"
-    "Thought: теперь я знаю ответ\n"
-    "Final Answer: итоговый ответ пользователю\n\n"
-    "Begin!\n\n"
-    "Question: {input}\n"
-    "Thought: {agent_scratchpad}"
     )
 
-agent = create_react_agent(llm=llm, tools=tools, prompt=prompt)
+
-agent_executor = AgentExecutor(agent=agent, tools=tools, verbose=True)
+# ---------------------------------------------------------------------------
+# Агент (create_react_agent из langgraph — актуальный API v0.3+)
+# ---------------------------------------------------------------------------
+
+llm = ChatOllama(model=LLM_MODEL, temperature=0.3)
+
+SYSTEM_PROMPT = """Ты — интеллектуальный AI-ассистент с доступом к локальной базе знаний.
+
+Правила работы:
+1. ВСЕГДА начинай с поиска через search_knowledge_base перед ответом на вопрос.
+2. Если найдена релевантная информация — используй её и ссылайся на источник.
+3. Если пользователь просит сохранить информацию — используй add_to_knowledge_base.
+4. Если база знаний пуста по теме — честно сообщи об этом и помоги по своим знаниям.
+5. Отвечай на том же языке, на котором задан вопрос.
+"""
+
+agent = create_react_agent(
+    model=llm,
+    tools=[search_knowledge_base, add_to_knowledge_base],
+    prompt=SYSTEM_PROMPT,
+)
+
+
+def run_agent(user_input: str) -> str:
+    """Запускает агента и возвращает финальный ответ."""
+    result = agent.invoke({"messages": [("user", user_input)]})
+    for msg in reversed(result["messages"]):
+        if msg.type == "ai" and msg.content:
+            return msg.content
+    return "Агент не вернул ответ."
 
 
 # ---------------------------------------------------------------------------
-# Загрузка документов из директории (для init_loader.py)
+# Загрузка документов из директории
 # ---------------------------------------------------------------------------
 
-def load_documents_from_directory(directory: str) -> None:
+def load_documents_from_directory(directory: str) -> int:
-    """Загружает все .txt и .md файлы из папки в базу знаний."""
+    """
+    Загружает все .txt и .md файлы из папки в базу знаний.
+    Возвращает количество загруженных файлов.
+    """
     path = Path(directory)
+    count = 0
     for file_path in path.rglob("*"):
         if file_path.suffix.lower() in {".txt", ".md"}:
-            text = file_path.read_text(encoding="utf-8")
+            text = file_path.read_text(encoding="utf-8", errors="ignore")
-            add_to_knowledge_base.invoke({"content": text, "title": file_path.stem})
+            if text.strip():
-            print(f"Загружен: {file_path.name}")
+                docs = chunk_document(text, title=file_path.name)
+                rag_store.add_documents(docs)
+                print(f"  Загружен: {file_path.name} ({len(docs)} чанков)")
+                count += 1
+    return count
 
 
 # ---------------------------------------------------------------------------
-# Запуск напрямую — интерактивный CLI
+# Интерактивный CLI-клиент
 # ---------------------------------------------------------------------------
 
 if __name__ == "__main__":
-    print("=== RAG Агент ===")
+    from rich.console import Console
-    print("Команды: /add <path>, /search <query>, /quit")
+    from rich.panel import Panel
+    from rich.markdown import Markdown
+    from rich.rule import Rule
+
+    console = Console()
+
+    console.print(Panel(
+        "[bold]RAG-агент[/bold] · Qdrant + Ollama\n\n"
+        "[cyan]/add <текст>[/cyan]       — добавить текст в базу знаний\n"
+        "[cyan]/add-file <путь>[/cyan]   — загрузить файл в базу знаний\n"
+        "[cyan]/search <запрос>[/cyan]   — прямой поиск (без агента)\n"
+        "[cyan]/quit[/cyan]              — выйти",
+        border_style="dim",
+    ))
+
     while True:
         try:
-            user_input = input("> ").strip()
+            user_input = input("\n> ").strip()
         except (EOFError, KeyboardInterrupt):
+            console.print("\n[yellow]Выход.[/yellow]")
             break
 
         if not user_input:
             continue
-        if user_input.lower() == "/quit":
+
-            print("Выход.")
+        # /quit
+        if user_input.lower() in ("/quit", "/exit"):
+            console.print("[yellow]Выход.[/yellow]")
             break
-        elif user_input.startswith("/add "):
+
-            file_path = Path(user_input[5:].strip())
+        # /add-file <путь>
+        elif user_input.lower().startswith("/add-file "):
+            file_path = Path(user_input[10:].strip())
             if file_path.is_file():
-                text = file_path.read_text(encoding="utf-8")
+                text = file_path.read_text(encoding="utf-8", errors="ignore")
-                print(add_to_knowledge_base.invoke({"content": text, "title": file_path.stem}))
+                result = add_to_knowledge_base.invoke(
+                    {"content": text, "title": file_path.name}
+                )
+                console.print(f"[green]✓ {result}[/green]")
             else:
-                print(f"Файл не найден: {file_path}")
+                console.print(f"[red]Файл не найден: {file_path}[/red]")
-        elif user_input.startswith("/search "):
+
+        # /add <текст>  ← исправлено: теперь принимает текст напрямую,
+        #                 а не только путь к файлу
+        elif user_input.lower().startswith("/add "):
+            raw = user_input[5:].strip()
+            path = Path(raw)
+            if path.is_file():
+                # если передан путь к существующему файлу — загружаем файл
+                text = path.read_text(encoding="utf-8", errors="ignore")
+                result = add_to_knowledge_base.invoke(
+                    {"content": text, "title": path.name}
+                )
+            else:
+                # иначе сохраняем как текст напрямую
+                result = add_to_knowledge_base.invoke(
+                    {"content": raw, "title": "manual"}
+                )
+            console.print(f"[green]✓ {result}[/green]")
+
+        # /search <запрос>
+        elif user_input.lower().startswith("/search "):
             query = user_input[8:].strip()
-            for r in search_knowledge_base.invoke({"query": query, "max_results": 3}):
+            console.print(Rule("Результаты поиска"))
-                print(r)
+            results = rag_store.search(query, k=5)
+            if not results:
+                console.print("[yellow]Ничего не найдено.[/yellow]")
             else:
-            response = agent_executor.invoke({"input": user_input})
+                for i, (text, score) in enumerate(results, 1):
-            print(response["output"])
+                    console.print(Panel(
+                        text,
+                        title=f"[cyan]#{i}  score={score:.3f}[/cyan]",
+                        border_style="dim",
+                    ))
+
+        # /search без аргумента
+        elif user_input.lower() == "/search":
+            console.print("[yellow]Укажите запрос: /search <текст>[/yellow]")
+
+        # запрос к агенту
+        else:
+            console.print("[dim]Агент думает...[/dim]")
+            try:
+                answer = run_agent(user_input)
+                console.print(Panel(
+                    Markdown(answer),
+                    title="[green]Агент[/green]",
+                    border_style="green",
+                ))
+            except Exception as e:
+                console.print(f"[red]Ошибка агента: {e}[/red]")