180 lines
7.6 KiB
Python
180 lines
7.6 KiB
Python
#!/usr/bin/env python3
|
|
# -*- coding: utf-8 -*-
|
|
|
|
"""
|
|
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
|
|
|
|
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
|
|
|
|
# --------------------------------------------------------------------------- #
|
|
# Конфигурация
|
|
# --------------------------------------------------------------------------- #
|
|
|
|
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 = "rag_memory"
|
|
|
|
# --------------------------------------------------------------------------- #
|
|
# Утилиты для работы с документами
|
|
# --------------------------------------------------------------------------- #
|
|
|
|
def load_text_files(directory: Path) -> List[str]:
|
|
"""
|
|
Загружает все текстовые файлы из указанной директории.
|
|
Возвращает список строк, каждая из которых содержит содержимое одного файла.
|
|
"""
|
|
texts = []
|
|
for file_path in directory.glob("**/*.txt"):
|
|
try:
|
|
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 split_documents(texts: List[str], chunk_size: int = 1000, overlap: int = 200) -> List[str]:
|
|
"""
|
|
Делит большие тексты на чанки фиксированного размера.
|
|
"""
|
|
splitter = RecursiveCharacterTextSplitter(
|
|
chunk_size=chunk_size,
|
|
chunk_overlap=overlap,
|
|
separators=["\n\n", "\n", " ", ""],
|
|
)
|
|
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,
|
|
)
|
|
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
|
|
# --------------------------------------------------------------------------- #
|
|
|
|
def build_qa_chain(vectorstore: Qdrant) -> any:
|
|
"""
|
|
Инициализирует цепочку RetrievalQA с использованием LLM OpenAI.
|
|
"""
|
|
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)
|
|
|
|
# --------------------------------------------------------------------------- #
|
|
# Основной цикл взаимодействия с пользователем
|
|
# --------------------------------------------------------------------------- #
|
|
|
|
def main():
|
|
console.print(Panel("[bold cyan]RAG‑агент запущен[/bold cyan]", expand=False))
|
|
|
|
# 1. Загружаем документы
|
|
if not DOCS_DIR.exists() or not any(DOCS_DIR.glob("**/*.txt")):
|
|
console.print(f"[red]Папка с документами '{DOCS_DIR}' пуста или не существует.[/red]")
|
|
sys.exit(1)
|
|
|
|
raw_texts = load_text_files(DOCs_DIR)
|
|
console.print(f"[green]Загружено {len(raw_texts)} файлов. Делим на чанки...[/green]")
|
|
|
|
# 2. Разбиваем на чанки
|
|
chunks = split_documents(raw_texts, chunk_size=1000, overlap=200)
|
|
console.print(f"[green]Получено {len(chunks)} чанков.[/green]")
|
|
|
|
# 3. Создаём или загружаем векторное хранилище
|
|
vectorstore = create_or_load_vectorstore(chunks)
|
|
|
|
# 4. Строим цепочку QA
|
|
qa_function = build_qa_chain(vectorstore)
|
|
console.print("[green]Цепочка RetrievalQA готова![/green]")
|
|
|
|
# 5. Взаимодействие с пользователем
|
|
while True:
|
|
try:
|
|
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 __name__ == "__main__":
|
|
main() |