Answers from your own documents (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"))  # folder with .txt and .md files
MODEL = os.environ.get("ZUKKO_MODEL", "claude-haiku-5-5")
CHUNK_SIZE = 1200  # rough chunk size in characters
TOP_K = 5          # how many chunks to send to the model

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 = "How many vacation days does a new employee get?"

response = client.chat.completions.create(
    model="claude-sonnet-5",
    messages=[
        # documents go first: this part never changes and lands in the prompt cache
        {"role": "system", "content": "Answer only from the documents below and name the source file. "
                                      "If the answer isn't there, say so.\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")  # downloads once, handles 100 languages


class Index:
    def __init__(self, chunks: list[dict]):
        self.chunks = chunks
        # e5 models expect "passage: " before texts and "query: " before questions
        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  # higher means closer in meaning
        best = scores.argsort()[::-1][:k]
        return [self.chunks[i] for i in best]