Einen privaten RAG-Bot bauen

export VENICE_API_KEY=<my-key>

pip install "openai>=1.0.0" "qdrant-client[fastembed]>=1.14.1"

openai>=1.0.0
qdrant-client[fastembed]>=1.14.1

import os
import textwrap
import uuid
from dataclasses import dataclass
from pathlib import Path

from fastembed.rerank.cross_encoder import TextCrossEncoder
from openai import OpenAI
from qdrant_client import QdrantClient, models

VENICE_BASE_URL = "https://api.venice.ai/api/v1"
CHAT_MODEL = "kimi-k2-6"
EMBEDDING_MODEL = "text-embedding-bge-m3"
RERANKER_MODEL = "Xenova/ms-marco-MiniLM-L-6-v2"
COLLECTION_NAME = "private_rag_bot"


@dataclass
class SourceDocument:
    content: str
    metadata: dict


@dataclass
class RankedChunk:
    content: str
    metadata: dict
    vector_score: float
    rerank_score: float

curl "https://api.venice.ai/api/v1/models?type=embedding" \
  -H "Authorization: Bearer $VENICE_API_KEY"

curl "https://api.venice.ai/api/v1/models?type=text" \
  -H "Authorization: Bearer $VENICE_API_KEY"

venice = OpenAI(
    api_key=os.environ["VENICE_API_KEY"],
    base_url=VENICE_BASE_URL,
)

qdrant = QdrantClient(path="./qdrant_data")

TEXT_EXTENSIONS = {".md", ".rst", ".txt"}

def expand_paths(paths: list[Path]) -> list[Path]:
    files = []
    for path in paths:
        if path.is_dir():
            files.extend(
                sorted(
                    file_path
                    for file_path in path.rglob("*")
                    if file_path.is_file()
                    and file_path.suffix.lower() in TEXT_EXTENSIONS
                )
            )
        elif path.is_file():
            files.append(path)
        else:
            raise FileNotFoundError(f"Document path does not exist: {path}")
    return files

def chunk_text(text: str, chunk_size: int, chunk_overlap: int) -> list[str]:
    clean_text = textwrap.dedent(text).strip()
    if not clean_text:
        return []
    if len(clean_text) <= chunk_size:
        return [clean_text]

    chunks = []
    start = 0
    while start < len(clean_text):
        end = min(start + chunk_size, len(clean_text))

        if end < len(clean_text):
            paragraph_break = clean_text.rfind("\n\n", start, end)
            sentence_break = clean_text.rfind(". ", start, end)
            split_at = max(paragraph_break, sentence_break)
            if split_at > start + chunk_size // 2:
                end = split_at + 1

        chunk = clean_text[start:end].strip()
        if chunk:
            chunks.append(chunk)

        if end >= len(clean_text):
            break

        start = max(end - chunk_overlap, start + 1)

    return chunks

def embed(texts: list[str]) -> list[list[float]]:
    embeddings = []
    for start in range(0, len(texts), 32):
        batch = texts[start : start + 32]
        response = venice.embeddings.create(
            model="text-embedding-bge-m3",
            input=batch,
        )
        embeddings.extend(
            item.embedding
            for item in sorted(response.data, key=lambda item: item.index)
        )
    return embeddings

qdrant.create_collection(
    collection_name=COLLECTION_NAME,
    vectors_config=models.VectorParams(
        size=len(embeddings[0]),
        distance=models.Distance.COSINE,
    ),
)

points.append(
    models.PointStruct(
        id=chunk_id,
        vector=embedding,
        payload={
            "text": chunk.content,
            "source": source,
            "chunk_index": chunk_index,
        },
    )
)

qdrant.upsert(collection_name=COLLECTION_NAME, points=points)

query_vector = embed([question])[0]
hits = qdrant.query_points(
    collection_name=COLLECTION_NAME,
    query=query_vector,
    with_payload=True,
    limit=8,
).points

from fastembed.rerank.cross_encoder import TextCrossEncoder

reranker = TextCrossEncoder(model_name="Xenova/ms-marco-MiniLM-L-6-v2")

candidate_texts = [str((hit.payload or {}).get("text", "")) for hit in hits]
rerank_scores = list(reranker.rerank(question, candidate_texts))
reranked = sorted(
    zip(hits, rerank_scores),
    key=lambda hit_and_score: hit_and_score[1],
    reverse=True,
)

def format_context(chunks: list[RankedChunk]) -> str:
    if not chunks:
        return "No relevant context was retrieved."

    context_parts = []
    for index, chunk in enumerate(chunks, start=1):
        source = chunk.metadata.get("source", "unknown")
        context_parts.append(
            f"[{index}] Source: {source} | "
            f"Vector score: {chunk.vector_score:.4f} | "
            f"Rerank score: {chunk.rerank_score:.4f}\n"
            f"{chunk.content}"
        )
    return "\n\n---\n\n".join(context_parts)

response = venice.chat.completions.create(
    model="kimi-k2-6",
    temperature=0.2,
    messages=[
        {
            "role": "system",
            "content": (
                "You are a helpful RAG assistant. Answer using only the supplied "
                "context. If the context does not answer the question, say that "
                "you do not have enough information."
            ),
        },
        {
            "role": "user",
            "content": (
                f"Retrieved context:\n{context}\n\n"
                f"Question: {question}\n\n"
                "Answer with citations like [1] when the context supports the answer:"
            ),
        },
    ],
)

python rag_bot.py \
  --question "What does reranking improve in a RAG pipeline?"

python rag_bot.py \
  --docs ./docs \
  --question "What does this project do?"

python rag_bot.py \
  --docs ./docs \
  --qdrant-path ./qdrant_data \
  --chat

Answer
============================================================
Reranking improves retrieval quality by rescoring the top
vector-search candidates with a cross-encoder model [1].

Sources
============================================================
1. sample-docs (vector=0.8123, rerank=0.7342)

--show-context

python rag_bot.py \
  --docs ./docs \
  --qdrant-path ./qdrant_data \
  --recreate-collection \
  --question "Summarize the most important setup steps."

python rag_bot.py \
  --qdrant-path ./qdrant_data \
  --skip-ingest \
  --question "Which file explains deployment?"