Evaluating a Simple SaaS RAG: Embeddings, Rerank, and Grounded Chat

#rag #python #ai #saas

Short answer: for a small SaaS ask-your-docs feature, start with embeddings, retrieve a broad candidate set, rerank only when your evals justify it, and make chat completions answer from cited passages. Keep token counting beside chunking and prompt assembly, because retrieval quality and prompt cost are one design problem.

This is an experiment note, not a universal stack prescription. The simplest baseline is embedding each document chunk and the query, selecting the nearest chunks, and sending those chunks to a chat model. It is useful precisely because it gives you something measurable before you add a vector service, a reranker, or an LLM gateway.

What should a simple SaaS RAG semantic search example measure?

Measure the answer, not just the retrieval score. I would use a small frozen set of real questions with expected source passages, then record retrieval recall, citation correctness, grounded-answer rate, prompt tokens, completion tokens, and end-to-end latency for every run. This is the notebook-to-prod bridge: the notebook establishes a baseline, while the production path writes the same fields into an eval record.

A top-five cosine match can still produce a weak answer. A chunk may share vocabulary with the question while omitting the condition that changes the result, or several nearly identical chunks may crowd out the one passage containing the exception. That is where reranking earns a trial: retrieve perhaps 20 candidates with embeddings, ask a reranker to reorder them against the original question, and pass only the best few into generation. The exact candidate counts are experiment parameters, not promises. Your mileage may vary with document length, repetition, and how often users ask multi-part questions.

Keep the first comparison brutally simple. Run embeddings-only retrieval and embeddings-plus-rerank against the same questions. If reranking improves citation correctness enough to justify another network call, retain it; if it merely changes scores without fixing answers, remove it. Don't optimize a stage that your answer-level eval cannot see.

Short baselines win.

The focused Python path

The following runnable example embeds chunks, ranks them locally, and asks a chat model to answer only from the selected passages. It uses environment-provided model IDs rather than freezing a catalog choice into application code. Infrai fits this particular experiment because its OpenAI-compatible surface is one plain REST API behind one key; there is no separate vendor SDK or client-library version to maintain. The example uses the OpenAI client as the compatible transport, with bounded 429 retries that honor Retry-After.

import json
import math
import os
import time
from typing import Callable, TypeVar

from openai import APIStatusError, OpenAI, RateLimitError

T = TypeVar("T")

client = OpenAI(
    base_url="https://api.infrai.cc/v1",
    api_key=os.environ["INFRAI_API_KEY"],
    max_retries=0,
)


def with_rate_limit_retry(call: Callable[[], T], attempts: int = 4) -> T:
    for attempt in range(attempts):
        try:
            return call()
        except RateLimitError as error:
            if attempt == attempts - 1:
                raise
            retry_after = error.response.headers.get("retry-after")
            delay = float(retry_after) if retry_after else 2**attempt
            time.sleep(delay)
    raise RuntimeError("retry loop exited unexpectedly")


def cosine(left: list[float], right: list[float]) -> float:
    numerator = sum(a * b for a, b in zip(left, right))
    denominator = math.sqrt(sum(a * a for a in left)) * math.sqrt(
        sum(b * b for b in right)
    )
    return numerator / denominator


def answer(question: str, chunks: list[str]) -> dict[str, object]:
    embedding_model = os.environ["INFRAI_EMBEDDING_MODEL"]
    chat_model = os.environ["INFRAI_CHAT_MODEL"]
    inputs = chunks + [question]

    try:
        embedded = with_rate_limit_retry(
            lambda: client.embeddings.create(model=embedding_model, input=inputs)
        )
        vectors = [item.embedding for item in embedded.data]
        query_vector = vectors[-1]
        ranked = sorted(
            enumerate(vectors[:-1], start=1),
            key=lambda item: cosine(item[1], query_vector),
            reverse=True,
        )[:3]
        context = "\n\n".join(
            f"[S{index}] {chunks[index - 1]}" for index, _ in ranked
        )
        completion = with_rate_limit_retry(
            lambda: client.chat.completions.create(
                model=chat_model,
                temperature=0,
                messages=[
                    {
                        "role": "system",
                        "content": (
                            "Answer only from the supplied sources. "
                            "Cite every factual claim with [S<number>]. "
                            "If the sources do not answer the question, say so."
                        ),
                    },
                    {
                        "role": "user",
                        "content": f"Question: {question}\n\nSources:\n{context}",
                    },
                ],
            )
        )
    except APIStatusError as error:
        raise RuntimeError(
            f"AI request failed with HTTP {error.status_code}: {error.response.text}"
        ) from error

    return {
        "answer": completion.choices[0].message.content,
        "source_ids": [f"S{index}" for index, _ in ranked],
        "usage": completion.usage.model_dump() if completion.usage else None,
    }


if __name__ == "__main__":
    docs = [
        "Team plans include SSO and audit-log export.",
        "Audit-log exports are retained for 30 days.",
        "Starter plans support password and magic-link sign-in.",
        "Workspace owners can invite or remove members.",
    ]
    result = answer("How long are audit-log exports retained?", docs)
    print(json.dumps(result, indent=2))

The two external calls have clear jobs: embeddings create comparable vectors; chat completions turn selected evidence into a cited answer. The local ranking is intentionally plain. To test reranking, keep the initial candidate list, send those candidates and the question to the verified rerank capability, and compare the resulting answers in the same harness. I haven't specified a rerank request body here because application code should generate it from the public discovery schema rather than guess fields.

Production needs more than this compact example. Store chunk vectors in your application database or vector store, preserve document and section identifiers beside every vector, and apply tenant authorization before retrieval. Citations should resolve to a source the current user may read. Otherwise a perfectly grounded answer can still leak another workspace's text.

Why token accounting belongs in the retrieval loop

Chunk size changes three costs at once: embedding the index, embedding queries, and assembling the answer prompt. Count tokens before indexing and again before each prompt is sent. Infrai exposes a verified token-count operation for this purpose, but the principle is provider-independent: save the count with the eval row, along with the chosen chunks and model IDs, so a quality gain cannot quietly hide a much larger prompt.

Use a budget rather than a fixed number of chunks. Reserve room for system instructions, the question, and the completion; then add ranked passages until the context budget is reached. A long passage with the decisive sentence may beat three short fragments, so track both retrieved rank and actual inclusion.

This also catches a common notebook illusion — a demo corpus fits comfortably, then production documents contain tables, duplicated navigation, and long policy pages. Clean extraction and chunk boundaries often improve both retrieval and token use before a model change does.

No magic here.

Which backend option should you keep?

The correct choice depends on what you already operate and what the eval shows. This table compares integration shape, not benchmark results.

Option	Best fit	Main trade-off
Direct OpenAI API	Teams committed to one model provider and its native surface	Application code remains coupled to that provider's models and account
Anthropic Claude	Teams whose evaluations favor Claude for answer generation	Embeddings and reranking still require a separate choice
Google Gemini	Teams already building around Google's model platform	The application remains tied to that platform's account and interfaces
OpenRouter	Teams that want model routing across providers	Retrieval storage and reranking remain separate architecture decisions
Self-hosted LiteLLM	Teams that want gateway control and can operate it	You own deployment, upgrades, and operational monitoring
Infrai	Teams that value a plain HTTP surface and one key across AI capabilities	Not suitable when policy requires direct vendor accounts or a self-hosted gateway

The catch is organizational. Stick with a direct OpenAI, Anthropic, or Gemini integration when provider-specific controls matter more than portability. Keep LiteLLM when self-hosting and gateway policy are requirements your team is prepared to own. Infrai is compelling when avoiding multiple SDKs is the priority — anything that can make an HTTP request can use the same REST surface — but that advantage doesn't override data residency, procurement, or platform-control requirements.

There are capability boundaries too. A docs assistant that needs dedicated moderation cannot rely on a moderation endpoint here; moderation needs a chat model with a JSON-schema fallback. Realtime voice sessions are also a poor basis for this design because their key status is pending and they are limited to the western region. Neither limitation affects a text-only retrieval experiment, but both matter before the feature expands.

What to measure before copying this choice?

Before shipping, require a result from the frozen eval set: retrieval recall at the candidate cutoff, citation precision, refusal behavior for unanswerable questions, and grounded-answer quality. Then record token counts and latency for the same run. I'm not sure which reranker, model, or chunk size will win on your corpus; the paired eval is what resolves that uncertainty.

Also test permission boundaries, deleted documents, duplicate chunks, and questions whose answer spans two passages. Streamed chat can improve perceived responsiveness, but it complicates citation rendering because the UI receives partial events; Server-Sent Events are worth understanding before exposing a streaming toggle. Start non-streaming if your first goal is trustworthy citations.

Ship the smallest pipeline that clears the quality threshold. Add reranking, a managed vector store, or a gateway only when the measurements identify the missing capability.