From 161a28fb55f86e8748567c20b0afccaa2ed0bcd8 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=D0=9C=D0=B0=D1=80=D0=B8=D1=8F=20=D0=91=D0=B5=D1=80=D0=B4?= =?UTF-8?q?=D0=BD=D0=B8=D0=BA=D0=BE=D0=B2=D0=B0?= Date: Thu, 28 May 2026 10:04:28 +0000 Subject: [PATCH] =?UTF-8?q?=D0=90=D0=B3=D0=B5=D0=BD=D1=82=20=D1=81=20RAG-?= =?UTF-8?q?=D0=BF=D0=B0=D0=BC=D1=8F=D1=82=D1=8C=D1=8E:=20solution.py?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- solution.py | 416 ++++++++++++++++++---------------------------------- 1 file changed, 143 insertions(+), 273 deletions(-) diff --git a/solution.py b/solution.py index acefe31..a3456cc 100644 --- a/solution.py +++ b/solution.py @@ -1,310 +1,180 @@ -""" -rag_agent.py — AI-агент с RAG-памятью на базе Ollama + Qdrant. +#!/usr/bin/env python3 +# -*- coding: utf-8 -*- -Стек: langchain-ollama, langchain-qdrant, langgraph, qdrant-client, rich -Установка: - ollama pull llama3 - ollama pull nomic-embed-text - pip install langchain langchain-core langchain-ollama langchain-qdrant \ - langchain-text-splitters langgraph qdrant-client rich +""" +solution.py + +Основная реализация агента с RAG‑памятью. +Используется стек: + - 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 -from langchain_qdrant import QdrantVectorStore -from langchain_text_splitters import RecursiveCharacterTextSplitter -from langchain_core.tools import tool -from langchain_core.documents import Document -from langgraph.prebuilt import create_react_agent as create_agent # алиас -from qdrant_client import QdrantClient as _QdrantClient -from qdrant_client.http.models import Distance, VectorParams +import openai +from langchain.embeddings.openai import OpenAIEmbeddings +from langchain.vectorstores.qdrant import Qdrant +from langchain.chains.question_answering import load_qa_chain +from langchain.llms.openai import OpenAI +from langchain.text_splitter import RecursiveCharacterTextSplitter -# --------------------------------------------------------------------------- +from rich.console import Console +from rich.prompt import Prompt +from rich.panel import Panel + +# --------------------------------------------------------------------------- # # Конфигурация -# --------------------------------------------------------------------------- +# --------------------------------------------------------------------------- # -QDRANT_HOST = os.getenv("QDRANT_HOST", "localhost") -QDRANT_PORT = int(os.getenv("QDRANT_PORT", "6333")) -COLLECTION_NAME = os.getenv("QDRANT_COLLECTION", "rag_collection") -EMBEDDING_MODEL = os.getenv("EMBEDDING_MODEL", "nomic-embed-text") -LLM_MODEL = os.getenv("LLM_MODEL", "llama3") +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 -# --------------------------------------------------------------------------- -# Векторное хранилище (Qdrant + Ollama-эмбеддинги) -# --------------------------------------------------------------------------- +# Путь к директории с документами +DOCS_DIR = Path("docs") # можно изменить при необходимости -class RagVectorStore: - """Обёртка над QdrantVectorStore с эмбеддингами через Ollama.""" +# Параметры Qdrant +QDRANT_HOST = os.getenv("QDRANT_HOST", "localhost") +QDRANT_PORT = int(os.getenv("QDRANT_PORT", "6333")) +COLLECTION_NAME = "rag_memory" - 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, - 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: +def load_text_files(directory: Path) -> List[str]: + """ + Загружает все текстовые файлы из указанной директории. + Возвращает список строк, каждая из которых содержит содержимое одного файла. + """ + texts = [] + for file_path in directory.glob("**/*.txt"): try: - self._client.get_collection(name) - return True - except Exception: - return False + content = file_path.read_text(encoding="utf-8") + if content.strip(): + texts.append(content) + except Exception as exc: + console.print(f"[red]Не удалось прочитать {file_path}: {exc}[/red]") + return texts - def add_documents(self, docs: List[Document]) -> None: - 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.""" +def split_documents(texts: List[str], chunk_size: int = 1000, overlap: int = 200) -> List[str]: + """ + Делит большие тексты на чанки фиксированного размера. + """ splitter = RecursiveCharacterTextSplitter( - chunk_size=1000, - chunk_overlap=200, + chunk_size=chunk_size, + chunk_overlap=overlap, separators=["\n\n", "\n", " ", ""], ) - return [ - Document( - page_content=chunk, - metadata={"title": title, "source": title, "chunk_index": i}, + chunks = [] + for text in texts: + chunks.extend(splitter.split_text(text)) + 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, + ) + console.print(f"[green]Коллекция '{COLLECTION_NAME}' создана и заполнена.{/green}") + return vectorstore +# --------------------------------------------------------------------------- # +# Создание цепочки RetrievalQA +# --------------------------------------------------------------------------- # -# --------------------------------------------------------------------------- -# Синглтон хранилища (используется инструментами агента) -# --------------------------------------------------------------------------- - -rag_store = RagVectorStore() - - -# --------------------------------------------------------------------------- -# RAG-инструменты агента (@tool) -# --------------------------------------------------------------------------- - -@tool -def search_knowledge_base(query: str, max_results: int = 5) -> str: +def build_qa_chain(vectorstore: Qdrant) -> any: """ - Выполняет семантический поиск в базе знаний. - - Args: - query: поисковый запрос на естественном языке. - max_results: максимальное количество результатов (по умолчанию 5). - - Returns: - Строка с найденными фрагментами и оценками релевантности. + Инициализирует цепочку RetrievalQA с использованием LLM OpenAI. """ - 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)}." + llm = OpenAI(temperature=0.2, model_name="gpt-3.5-turbo") + chain = load_qa_chain( + llm, + chain_type="stuff", + verbose=False ) + return lambda question: chain.run(input_documents=vectorstore.similarity_search(question), question=question) +# --------------------------------------------------------------------------- # +# Основной цикл взаимодействия с пользователем +# --------------------------------------------------------------------------- # -# --------------------------------------------------------------------------- -# Агент (create_agent — обёртка над create_react_agent из langgraph v0.3+) -# --------------------------------------------------------------------------- +def main(): + console.print(Panel("[bold cyan]RAG‑агент запущен[/bold cyan]", expand=False)) -llm = ChatOllama(model=LLM_MODEL, temperature=0.3) + # 1. Загружаем документы + if not DOCS_DIR.exists() or not any(DOCS_DIR.glob("**/*.txt")): + console.print(f"[red]Папка с документами '{DOCS_DIR}' пуста или не существует.[/red]") + sys.exit(1) -SYSTEM_PROMPT = """Ты — интеллектуальный AI-ассистент с доступом к локальной базе знаний. + raw_texts = load_text_files(DOCs_DIR) + console.print(f"[green]Загружено {len(raw_texts)} файлов. Делим на чанки...[/green]") -Правила работы: -1. ВСЕГДА начинай с поиска через search_knowledge_base перед ответом на вопрос. -2. Если найдена релевантная информация — используй её и ссылайся на источник. -3. Если пользователь просит сохранить информацию — используй add_to_knowledge_base. -4. Если база знаний пуста по теме — честно сообщи об этом и помоги по своим знаниям. -5. Отвечай на том же языке, на котором задан вопрос. -""" + # 2. Разбиваем на чанки + chunks = split_documents(raw_texts, chunk_size=1000, overlap=200) + console.print(f"[green]Получено {len(chunks)} чанков.[/green]") -agent = create_agent( - model=llm, - tools=[search_knowledge_base, add_to_knowledge_base], - prompt=SYSTEM_PROMPT, -) + # 3. Создаём или загружаем векторное хранилище + vectorstore = create_or_load_vectorstore(chunks) + # 4. Строим цепочку QA + qa_function = build_qa_chain(vectorstore) + console.print("[green]Цепочка RetrievalQA готова![/green]") -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 файлы из папки в базу знаний. - Возвращает количество загруженных файлов. - """ - path = Path(directory) - if not path.is_dir(): - raise ValueError(f"Директория не найдена: {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", 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 - - -# --------------------------------------------------------------------------- -# Интерактивный тестовый CLI-клиент -# --------------------------------------------------------------------------- - -if __name__ == "__main__": - from rich.console import Console - 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]/load-dir <папка>[/cyan] — загрузить все .txt/.md из папки\n" - "[cyan]/search <запрос>[/cyan] — прямой поиск (без агента)\n" - "[cyan]/quit[/cyan] — выйти", - border_style="dim", - )) - + # 5. Взаимодействие с пользователем while True: try: - user_input = input("\n> ").strip() - except (EOFError, KeyboardInterrupt): - console.print("\n[yellow]Выход.[/yellow]") + question = Prompt.ask("\n[bold magenta]Введите ваш вопрос (или 'выход' для завершения)[/bold magenta]") + if question.lower() in {"выход", "exit", "quit"}: + 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: - continue - - # ── /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]") \ No newline at end of file +if __name__ == "__main__": + main() \ No newline at end of file