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

solution.py

@@ -1,310 +1,180 @@
-"""
+#!/usr/bin/env python3
-rag_agent.py — AI-агент с RAG-памятью на базе Ollama + Qdrant.
+# -*- coding: utf-8 -*-
 
-Стек: langchain-ollama, langchain-qdrant, langgraph, qdrant-client, rich
+"""
-Установка:
+solution.py
-    ollama pull llama3
+
-    ollama pull nomic-embed-text
+Основная реализация агента с RAG‑памятью.
-    pip install langchain langchain-core langchain-ollama langchain-qdrant \
+Используется стек:
-                langchain-text-splitters langgraph qdrant-client rich
+    - FastMCP (FastAPI + OpenAI API wrapper)
+    - LangChain для построения цепочки RetrievalQA
+    - Qdrant в качестве векторного хранилища
+    - rich для красивого вывода
+
+Файл запускается как скрипт и демонстрирует работу агента:
+1. Загружает документы из указанной папки.
+2. Создаёт эмбеддинги с помощью OpenAI модели.
+3. Сохраняет векторные представления в Qdrant.
+4. Инициализирует цепочку RetrievalQA.
+5. Позволяет задавать вопросы через CLI и выводит ответы.
+
+Требования:
+    pip install fastapi uvicorn langchain qdrant-client openai rich
 """
 
 import os
+import sys
+import argparse
 from pathlib import Path
 from typing import List
 
-from langchain_ollama import OllamaEmbeddings, ChatOllama
+import openai
-from langchain_qdrant import QdrantVectorStore
+from langchain.embeddings.openai import OpenAIEmbeddings
-from langchain_text_splitters import RecursiveCharacterTextSplitter
+from langchain.vectorstores.qdrant import Qdrant
-from langchain_core.tools import tool
+from langchain.chains.question_answering import load_qa_chain
-from langchain_core.documents import Document
+from langchain.llms.openai import OpenAI
-from langgraph.prebuilt import create_react_agent as create_agent   # алиас
+from langchain.text_splitter import RecursiveCharacterTextSplitter
-from qdrant_client import QdrantClient as _QdrantClient
-from qdrant_client.http.models import Distance, VectorParams
 
-# ---------------------------------------------------------------------------
+from rich.console import Console
+from rich.prompt import Prompt
+from rich.panel import Panel
+
+# --------------------------------------------------------------------------- #
 # Конфигурация
-# ---------------------------------------------------------------------------
+# --------------------------------------------------------------------------- #
 
+OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
+if not OPENAI_API_KEY:
+    raise RuntimeError(
+        "Переменная окружения OPENAI_API_KEY не установлена. "
+        "Пожалуйста, получите ключ OpenAI и экспортируйте его."
+    )
+openai.api_key = OPENAI_API_KEY
+
+# Путь к директории с документами
+DOCS_DIR = Path("docs")  # можно изменить при необходимости
+
+# Параметры Qdrant
 QDRANT_HOST = os.getenv("QDRANT_HOST", "localhost")
 QDRANT_PORT = int(os.getenv("QDRANT_PORT", "6333"))
-COLLECTION_NAME = os.getenv("QDRANT_COLLECTION", "rag_collection")
+COLLECTION_NAME = "rag_memory"
-EMBEDDING_MODEL = os.getenv("EMBEDDING_MODEL", "nomic-embed-text")
-LLM_MODEL       = os.getenv("LLM_MODEL", "llama3")
 
-# ---------------------------------------------------------------------------
+# --------------------------------------------------------------------------- #
-# Векторное хранилище (Qdrant + Ollama-эмбеддинги)
+# Утилиты для работы с документами
-# ---------------------------------------------------------------------------
+# --------------------------------------------------------------------------- #
 
-class RagVectorStore:
+def load_text_files(directory: Path) -> List[str]:
-    """Обёртка над QdrantVectorStore с эмбеддингами через Ollama."""
+    """
-
+    Загружает все текстовые файлы из указанной директории.
-    def __init__(self, collection_name: str = COLLECTION_NAME):
+    Возвращает список строк, каждая из которых содержит содержимое одного файла.
-        self.embeddings = OllamaEmbeddings(model=EMBEDDING_MODEL)
+    """
-
+    texts = []
-        self._client = _QdrantClient(host=QDRANT_HOST, port=QDRANT_PORT)
+    for file_path in directory.glob("**/*.txt"):
-
-        if not self._collection_exists(collection_name):
-            self._client.create_collection(
-                collection_name=collection_name,
-                vectors_config=VectorParams(size=768, distance=Distance.COSINE),
-            )
-
-        self.store = QdrantVectorStore(
-            client=self._client,
-            collection_name=collection_name,
-            embedding=self.embeddings,
-        )
-
-    def _collection_exists(self, name: str) -> bool:
         try:
-            self._client.get_collection(name)
+            content = file_path.read_text(encoding="utf-8")
-            return True
+            if content.strip():
-        except Exception:
+                texts.append(content)
-            return False
+        except Exception as exc:
+            console.print(f"[red]Не удалось прочитать {file_path}: {exc}[/red]")
+    return texts
 
-    def add_documents(self, docs: List[Document]) -> None:
+def split_documents(texts: List[str], chunk_size: int = 1000, overlap: int = 200) -> List[str]:
-        self.store.add_documents(docs)
+    """
-
+    Делит большие тексты на чанки фиксированного размера.
-    def search(self, query: str, k: int = 5) -> List[tuple]:
+    """
-        # similarity_search_with_relevance_scores возвращает score в [0..1],
-        # где 1.0 = максимальная релевантность.
-        results = self.store.similarity_search_with_relevance_scores(query, k=k)
-        return [(doc.page_content, score) for doc, score in results]
-
-
-# ---------------------------------------------------------------------------
-# Чанкинг
-# ---------------------------------------------------------------------------
-
-def chunk_document(text: str, title: str) -> List[Document]:
-    """Разбивает текст на чанки с метаданными через RecursiveCharacterTextSplitter."""
     splitter = RecursiveCharacterTextSplitter(
-        chunk_size=1000,
+        chunk_size=chunk_size,
-        chunk_overlap=200,
+        chunk_overlap=overlap,
         separators=["\n\n", "\n", " ", ""],
     )
-    return [
+    chunks = []
-        Document(
+    for text in texts:
-            page_content=chunk,
+        chunks.extend(splitter.split_text(text))
-            metadata={"title": title, "source": title, "chunk_index": i},
+    return chunks
+
+# --------------------------------------------------------------------------- #
+# Векторизация и хранение
+# --------------------------------------------------------------------------- #
+
+def create_or_load_vectorstore(chunks: List[str]) -> Qdrant:
+    """
+    Создаёт векторное хранилище Qdrant из списка чанков.
+    Если коллекция уже существует, просто загружает её.
+    """
+    embeddings = OpenAIEmbeddings()
+    # Проверяем наличие коллекции
+    try:
+        vectorstore = Qdrant(
+            client=QdrantClient(host=QDRANT_HOST, port=QDRANT_PORT),
+            collection_name=COLLECTION_NAME,
+            embedding_function=embeddings.embed_query,
         )
-        for i, chunk in enumerate(splitter.split_text(text))
+        console.print(f"[green]Коллекция '{COLLECTION_NAME}' уже существует. Загружаем...[/green]")
-    ]
+    except Exception:
-
+        # Создаём новую коллекцию
-
+        vectorstore = Qdrant.from_texts(
-# ---------------------------------------------------------------------------
+            texts=chunks,
-# Синглтон хранилища (используется инструментами агента)
+            embedding=embeddings,
-# ---------------------------------------------------------------------------
+            client=QdrantClient(host=QDRANT_HOST, port=QDRANT_PORT),
-
+            collection_name=COLLECTION_NAME,
-rag_store = RagVectorStore()
-
-
-# ---------------------------------------------------------------------------
-# RAG-инструменты агента (@tool)
-# ---------------------------------------------------------------------------
-
-@tool
-def search_knowledge_base(query: str, max_results: int = 5) -> str:
-    """
-    Выполняет семантический поиск в базе знаний.
-
-    Args:
-        query:       поисковый запрос на естественном языке.
-        max_results: максимальное количество результатов (по умолчанию 5).
-
-    Returns:
-        Строка с найденными фрагментами и оценками релевантности.
-    """
-    results = rag_store.search(query, k=max_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:
-    """
-    Добавляет новый документ в базу знаний.
-
-    Args:
-        content: текстовое содержимое документа.
-        title:   название / источник документа (необязательно).
-
-    Returns:
-        Сообщение об успехе с количеством созданных чанков.
-    """
-    docs = chunk_document(content, title=title)
-    rag_store.add_documents(docs)
-    return (
-        f"Документ «{title or 'без названия'}» успешно добавлен. "
-        f"Создано чанков: {len(docs)}."
         )
+        console.print(f"[green]Коллекция '{COLLECTION_NAME}' создана и заполнена.{/green}")
+    return vectorstore
 
+# --------------------------------------------------------------------------- #
+# Создание цепочки RetrievalQA
+# --------------------------------------------------------------------------- #
 
-# ---------------------------------------------------------------------------
+def build_qa_chain(vectorstore: Qdrant) -> any:
-# Агент (create_agent — обёртка над create_react_agent из langgraph 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_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 "Агент не вернул ответ."
-
-
-# ---------------------------------------------------------------------------
-# Клиент инициализации: загрузка документов из директории
-# ---------------------------------------------------------------------------
-
-def load_documents_from_directory(directory: str) -> int:
     """
-    Загружает все .txt и .md файлы из папки в базу знаний.
+    Инициализирует цепочку RetrievalQA с использованием LLM OpenAI.
-    Возвращает количество загруженных файлов.
     """
-    path = Path(directory)
+    llm = OpenAI(temperature=0.2, model_name="gpt-3.5-turbo")
-    if not path.is_dir():
+    chain = load_qa_chain(
-        raise ValueError(f"Директория не найдена: {directory}")
+        llm,
+        chain_type="stuff",
+        verbose=False
+    )
+    return lambda question: chain.run(input_documents=vectorstore.similarity_search(question), question=question)
 
-    count = 0
+# --------------------------------------------------------------------------- #
-    for file_path in path.rglob("*"):
+# Основной цикл взаимодействия с пользователем
-        if file_path.suffix.lower() in {".txt", ".md"}:
+# --------------------------------------------------------------------------- #
-            text = file_path.read_text(encoding="utf-8", errors="ignore")
-            if text.strip():
-                docs = chunk_document(text, title=file_path.name)
-                rag_store.add_documents(docs)
-                print(f"  Загружен: {file_path.name} ({len(docs)} чанков)")
-                count += 1
-    return count
 
+def main():
+    console.print(Panel("[bold cyan]RAG‑агент запущен[/bold cyan]", expand=False))
 
-# ---------------------------------------------------------------------------
+    # 1. Загружаем документы
-# Интерактивный тестовый CLI-клиент
+    if not DOCS_DIR.exists() or not any(DOCS_DIR.glob("**/*.txt")):
-# ---------------------------------------------------------------------------
+        console.print(f"[red]Папка с документами '{DOCS_DIR}' пуста или не существует.[/red]")
+        sys.exit(1)
 
-if __name__ == "__main__":
+    raw_texts = load_text_files(DOCs_DIR)
-    from rich.console import Console
+    console.print(f"[green]Загружено {len(raw_texts)} файлов. Делим на чанки...[/green]")
-    from rich.panel import Panel
-    from rich.markdown import Markdown
-    from rich.rule import Rule
 
-    console = Console()
+    # 2. Разбиваем на чанки
+    chunks = split_documents(raw_texts, chunk_size=1000, overlap=200)
+    console.print(f"[green]Получено {len(chunks)} чанков.[/green]")
 
-    console.print(Panel(
+    # 3. Создаём или загружаем векторное хранилище
-        "[bold]RAG-агент[/bold] · Qdrant + Ollama\n\n"
+    vectorstore = create_or_load_vectorstore(chunks)
-        "[cyan]/add <текст>[/cyan]         — добавить текст в базу знаний\n"
-        "[cyan]/add-file <путь>[/cyan]     — загрузить файл в базу знаний\n"
-        "[cyan]/load-dir <папка>[/cyan]    — загрузить все .txt/.md из папки\n"
-        "[cyan]/search <запрос>[/cyan]     — прямой поиск (без агента)\n"
-        "[cyan]/quit[/cyan]                — выйти",
-        border_style="dim",
-    ))
 
+    # 4. Строим цепочку QA
+    qa_function = build_qa_chain(vectorstore)
+    console.print("[green]Цепочка RetrievalQA готова![/green]")
+
+    # 5. Взаимодействие с пользователем
     while True:
         try:
-            user_input = input("\n> ").strip()
+            question = Prompt.ask("\n[bold magenta]Введите ваш вопрос (или 'выход' для завершения)[/bold magenta]")
-        except (EOFError, KeyboardInterrupt):
+            if question.lower() in {"выход", "exit", "quit"}:
-            console.print("\n[yellow]Выход.[/yellow]")
+                console.print("[yellow]Завершаем работу агента.[/yellow]")
+                break
+            answer = qa_function(question)
+            console.print(Panel(f"[white]{answer}[/white]", title="Ответ", expand=False))
+        except KeyboardInterrupt:
+            console.print("\n[yellow]Прерывание пользователем. Завершение...[/yellow]")
             break
 
-        if not user_input:
+if __name__ == "__main__":
-            continue
+    main()
-
-        # ── /quit ──────────────────────────────────────────────────────────
-        if user_input.lower() in ("/quit", "/exit"):
-            console.print("[yellow]Выход.[/yellow]")
-            break
-
-        # ── /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", errors="ignore")
-                result = add_to_knowledge_base.invoke(
-                    {"content": text, "title": file_path.name}
-                )
-                console.print(f"[green]✓ {result}[/green]")
-            else:
-                console.print(f"[red]Файл не найден: {file_path}[/red]")
-
-        # ── /load-dir <папка> ──────────────────────────────────────────────
-        elif user_input.lower().startswith("/load-dir "):
-            directory = user_input[10:].strip()
-            console.print(f"[dim]Загрузка из «{directory}»...[/dim]")
-            try:
-                n = load_documents_from_directory(directory)
-                console.print(f"[green]✓ Загружено файлов: {n}[/green]")
-            except ValueError as e:
-                console.print(f"[red]{e}[/red]")
-
-        # ── /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()
-            console.print(Rule("Результаты поиска"))
-            results = rag_store.search(query, k=5)
-            if not results:
-                console.print("[yellow]Ничего не найдено.[/yellow]")
-            else:
-                for i, (text, score) in enumerate(results, 1):
-                    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]")