Ответы по своим документам (RAG)

import os
import re
import sys
from pathlib import Path

import snowballstemmer
from openai import APIConnectionError, APIStatusError, OpenAI
from rank_bm25 import BM25Plus

DOCS_DIR = Path(os.environ.get("DOCS_DIR", "docs"))  # папка с файлами .txt и .md
MODEL = os.environ.get("ZUKKO_MODEL", "claude-haiku-5-5")
CHUNK_SIZE = 1200  # примерный размер фрагмента в символах
TOP_K = 5          # сколько фрагментов отправлять модели

from pathlib import Path
import os
from openai import OpenAI

client = OpenAI(base_url="https://api.zukko.pro/v1", api_key=os.environ["ZUKKO_API_KEY"], timeout=300)

docs = "\n\n".join(
    f"=== {p.name} ===\n{p.read_text(encoding='utf-8')}" for p in sorted(Path("docs").glob("*.md"))
)
question = "Сколько дней отпуска положено новому сотруднику?"

response = client.chat.completions.create(
    model="claude-sonnet-5",
    messages=[
        # документы — в начало: эта часть одинаковая и попадает в кэш промпта
        {"role": "system", "content": "Отвечай только по документам ниже и называй файл-источник. "
                                      "Если ответа нет — скажи об этом.\n\n" + docs},
        {"role": "user", "content": question},
    ],
)
print(response.choices[0].message.content)

# pip install sentence-transformers
from sentence_transformers import SentenceTransformer

encoder = SentenceTransformer("intfloat/multilingual-e5-small")  # скачается один раз, понимает русский


class Index:
    def __init__(self, chunks: list[dict]):
        self.chunks = chunks
        # у моделей e5 тексты помечают «passage: », а вопросы — «query: »
        self.vectors = encoder.encode(
            ["passage: " + c["text"] for c in chunks], normalize_embeddings=True
        )

    def search(self, query: str, k: int = TOP_K) -> list[dict]:
        q = encoder.encode(["query: " + query], normalize_embeddings=True)[0]
        scores = self.vectors @ q  # чем больше число, тем ближе смысл
        best = scores.argsort()[::-1][:k]
        return [self.chunks[i] for i in best]