RAG на практике: поиск по документам с LangChain и pgvector

pip install langchain langchain-postgres langchain-openai \
            psycopg[binary] python-dotenv pypdf

docker run -d \
  --name pgvector-db \
  -e POSTGRES_USER=rag_user \
  -e POSTGRES_PASSWORD=secret \
  -e POSTGRES_DB=rag_db \
  -p 5432:5432 \
  pgvector/pgvector:pg16

CREATE EXTENSION IF NOT EXISTS vector;

rag-system/
├── .env
├── ingest.py        # Загрузка и индексация документов
├── query.py         # RAG-запросы
└── documents/       # PDF, TXT, MD-файлы

OPENAI_API_KEY=sk-...
DATABASE_URL=postgresql+psycopg://rag_user:secret@localhost:5432/rag_db

# ingest.py
import os
from pathlib import Path
from dotenv import load_dotenv

from langchain_community.document_loaders import PyPDFLoader, TextLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_openai import OpenAIEmbeddings
from langchain_postgres import PGEngine, PGVectorStore

load_dotenv()

# Настройка движка БД
engine = PGEngine.from_connection_string(url=os.environ["DATABASE_URL"])

# Инициализация таблицы (создаётся один раз)
engine.init_vectorstore_table(
    table_name="documents",
    vector_size=1536,  # размер для text-embedding-3-small
)

# Модель эмбеддингов
embeddings = OpenAIEmbeddings(model="text-embedding-3-small")

# Создание векторного хранилища
store = PGVectorStore.create_sync(
    engine=engine,
    table_name="documents",
    embedding_service=embeddings,
)

def load_documents(folder: str):
    docs = []
    for path in Path(folder).rglob("*"):
        if path.suffix == ".pdf":
            loader = PyPDFLoader(str(path))
        elif path.suffix in (".txt", ".md"):
            loader = TextLoader(str(path), encoding="utf-8")
        else:
            continue
        docs.extend(loader.load())
    return docs

def ingest(folder: str = "documents"):
    raw_docs = load_documents(folder)
    print(f"Загружено документов: {len(raw_docs)}")

    splitter = RecursiveCharacterTextSplitter(
        chunk_size=1000,
        chunk_overlap=200,
        separators=["\n\n", "\n", ".", " "],
    )
    chunks = splitter.split_documents(raw_docs)
    print(f"Чанков для индексации: {len(chunks)}")

    store.add_documents(chunks)
    print("Индексация завершена.")

if __name__ == "__main__":
    ingest()

-- HNSW-индекс для косинусного сходства
CREATE INDEX ON documents
USING hnsw (embedding vector_cosine_ops)
WITH (m = 16, ef_construction = 64);

# query.py
import os
from dotenv import load_dotenv
from langchain_openai import OpenAIEmbeddings, ChatOpenAI
from langchain_postgres import PGEngine, PGVectorStore
from langchain.prompts import ChatPromptTemplate
from langchain_core.output_parsers import StrOutputParser
from langchain_core.runnables import RunnablePassthrough

load_dotenv()

engine = PGEngine.from_connection_string(url=os.environ["DATABASE_URL"])
embeddings = OpenAIEmbeddings(model="text-embedding-3-small")

store = PGVectorStore.create_sync(
    engine=engine,
    table_name="documents",
    embedding_service=embeddings,
)

# Ретривер: топ-5 релевантных чанков
retriever = store.as_retriever(
    search_type="similarity",
    search_kwargs={"k": 5},
)

# Промпт с контекстом
PROMPT_TEMPLATE = """Ты — помощник по работе с документами компании.
Отвечай на вопрос ТОЛЬКО на основе предоставленного контекста.
Если информации недостаточно — так и скажи, не придумывай.

Контекст:
{context}

Вопрос: {question}

Ответ:"""

prompt = ChatPromptTemplate.from_template(PROMPT_TEMPLATE)
llm = ChatOpenAI(model="gpt-4o-mini", temperature=0)

def format_docs(docs):
    return "\n\n---\n\n".join(
        f"[Источник: {d.metadata.get('source', 'неизвестно')}]\n{d.page_content}"
        for d in docs
    )

# LCEL-цепочка
rag_chain = (
    {"context": retriever | format_docs, "question": RunnablePassthrough()}
    | prompt
    | llm
    | StrOutputParser()
)

def ask(question: str) -> str:
    return rag_chain.invoke(question)

if __name__ == "__main__":
    answer = ask("Как оформить заявление на отпуск?")
    print(answer)

-- Добавляем полнотекстовый поиск к таблице
ALTER TABLE documents ADD COLUMN fts_content tsvector
    GENERATED ALWAYS AS (to_tsvector('russian', cmetadata->>'content')) STORED;

CREATE INDEX ON documents USING gin(fts_content);

# Искать только в документах отдела HR
results = store.similarity_search(
    query="отпуск",
    k=5,
    filter={"department": "HR", "year": 2024},
)

pip install ragas

from ragas import evaluate
from ragas.metrics import faithfulness, answer_relevancy, context_recall

# Запускаем оценку на тестовом датасете
results = evaluate(dataset, metrics=[faithfulness, answer_relevancy, context_recall])
print(results)