app/benchmarks/rag_observability/runner.py

"""Repeatable app-functioning RAG retrieval workload.

Drives the real retrieval critical path -- `app.websockets.handlers._get_chunks`,
the exact helper `CHAT_TURN`/`LEARN_NODE` call in production -- against the
fixed corpus/query-set in `backend/benchmarks/rag/`. Never times or profiles
the product independently: every duration this module writes out is read back
from the canonical `observe_operation` boundary's own local-store rows, never
from a clock the runner reads itself. What this module *does* do is supply
experiment identity through context: it wraps each measured `_get_chunks()`
call in its own `observe_operation("rag.benchmark_query", ...)`, calls
`op.set_experiment(identity)` on that wrapper before the call, and everything
`_get_chunks()` opens underneath (`rag.retrieval`, `embedding.query`,
`chroma.*`) shares that wrapper's trace -- the OTel ambient-context mechanism,
not an app-level side channel. The wrapper's own row (in `operational.jsonl`)
then carries both the identity and copies of the canonical retrieval
duration/attributes it read back from the nested `rag.retrieval` row, so every
local row, span, and metric for one call join on one `trace_id`.

Ingestion (`ingest_text`) is one-time setup, not part of what's measured, and
reasonably uses the real `ingest_text()` directly per the task-4 brief --
that's the same function the app's own upload path calls internally.

Corpus isolation: ingested chunks are tagged with a dedicated
`project_id` (`BENCHMARK_PROJECT_ID`) into the *same* "library" ChromaDB
collection every real session/project uses (matching how the product actually
isolates content -- by `project_id` metadata filter, never a separate
collection; `_get_chunks()` always queries the "library" collection by name).
Re-running the harness deletes and re-ingests only that project_id's chunks
first, so repeated runs stay idempotent without ever touching another
project's data.

Cache-state caveat: "cold" here means "first touch of this query in this
process, before its own warm-up repetition" -- a same-process proxy, not a
fresh-process/cold-OS-cache measurement. True process-level cold isolation is
out of scope for a single harness process; this is documented rather than
silently assumed away.
"""

from __future__ import annotations

import argparse
import asyncio
import json
import os
import random
import time
import uuid
from dataclasses import asdict
from pathlib import Path
from typing import Any

os.environ.setdefault("EVALUATION_RUN", "true")
os.environ.setdefault("CEREBRAS_API_KEY", "test-key")  # never called: retrieval-only workload

import chromadb
from opentelemetry import trace as otel_trace

from app.benchmarks.rag_observability import manifest as manifest_mod
from app.benchmarks.rag_observability import summarize as summarize_mod
from app.benchmarks.rag_observability.corpus import (
    Corpus,
    QueryLabel,
    QuerySet,
    load_corpus,
    load_query_set,
)
from app.benchmarks.rag_observability.quality import score_outcome
from app.observability import bootstrap
from app.observability.config import ObservabilityConfig, get_observability_config
from app.observability.contracts import ExperimentIdentity, RetrievalOutcome
from app.observability.local_store import LocalObservationStore
from app.observability.logging import trace_context_ids
from app.observability.operation import observe_operation
from app.rag.ingestion import LIBRARY_COLLECTION, ingest_text
from app.rag.pipeline_version import get_pipeline_version

BENCHMARK_PROJECT_ID = "rag-observability-benchmark"
DEFAULT_TOP_K = 5
DEFAULT_WARM_REPETITIONS = 10
DEFAULT_COLD_COUNT = 3
SMOKE_QUERY_COUNT = 2
SMOKE_WARM_REPETITIONS = 2


# --- Corpus setup (idempotent) -------------------------------------------------


def reset_and_ingest_corpus(db, corpus: Corpus) -> None:
    """Delete any previously-ingested benchmark chunks for our project_id, then
    re-ingest the fixed corpus fresh, so repeated runs never accumulate
    duplicate chunks or drift from the current corpus text."""
    db.delete_where(LIBRARY_COLLECTION, {"project_id": BENCHMARK_PROJECT_ID})
    for doc in corpus.documents:
        ingest_text(
            doc.text,
            source_label=doc.document_id,
            collection=LIBRARY_COLLECTION,
            chunk_type="content",
            db=db,
            project_id=BENCHMARK_PROJECT_ID,
            document_id=doc.document_id,
        )


# --- Execution order -------------------------------------------------


def seeded_order(queries: list[QueryLabel], seed: int) -> list[QueryLabel]:
    """A fixed, seed-reproducible query order -- reused identically across
    pipeline versions so v1/v2 measure the same sequence (protocol: "fixed
    seeded execution order and repeat it for v1 and v2")."""
    order = list(range(len(queries)))
    random.Random(seed).shuffle(order)
    return [queries[i] for i in order]


# --- Reading back the canonical retrieval row -------------------------------

_RETRIEVAL_STAGES = (
    "embedding.query",
    "chroma.collection_lookup",
    "chroma.collection_count",
    "chroma.vector_search",
    "rag.result_prepare",
)


def _rows_for_trace(store: LocalObservationStore, trace_id: str | None) -> dict[str, dict[str, Any]]:
    """The canonical local-store rows belonging to one trace, keyed by operation name.

    Reads back what `_get_chunks()`'s own instrumentation already wrote via the
    shared `observe_operation` boundary -- not an independent measurement.
    """
    if not trace_id:
        return {}
    return {row["operation"]: row for row in store.all() if row.get("trace_id") == trace_id}


def _outcome_from_rows(rows: dict[str, dict[str, Any]], candidates: list[dict[str, Any]]) -> RetrievalOutcome:
    retrieval_row = rows.get("rag.retrieval")
    if retrieval_row and retrieval_row.get("retrieval"):
        return RetrievalOutcome(**retrieval_row["retrieval"])
    # Degraded fallback (telemetry mode disabled / row not captured): cannot
    # distinguish success_empty from error_fallback here, which is exactly why
    # the primary path above -- reading the canonical row -- is preferred.
    return RetrievalOutcome.success(candidates)


# --- One measured invocation -------------------------------------------------


async def invoke_once(
    label: QueryLabel,
    *,
    handlers_module: Any,
    global_store: LocalObservationStore,
    operational_store: LocalObservationStore,
    identity: ExperimentIdentity,
    top_k: int,
) -> tuple[str | None, RetrievalOutcome]:
    """Run one real `_get_chunks()` call wrapped in its own experiment-tagged
    operation, and return `(trace_id, outcome)` for the caller to score and
    (if this is a measured, non-warm-up call) persist.

    `identity.pipeline_version` is threaded straight into `_get_chunks()`'s
    own `pipeline_version` parameter, so this is the one and only place a
    non-default pipeline version enters the real retrieval critical path --
    ordinary product call sites never pass it. `_get_chunks()` treats
    `"rag-naive-v1"` and `None` identically (see `ChromaDBClient.query_observed`),
    so passing it explicitly here changes nothing for v1 runs.
    """
    with observe_operation(
        "rag.benchmark_query",
        subsystem="benchmark",
        consumer="benchmark_runner",
        evaluation_run=True,
        store=operational_store,
    ) as op:
        op.set_experiment(identity)
        trace_id, _span_id = trace_context_ids(otel_trace.get_current_span())

        candidates = await handlers_module._get_chunks(
            BENCHMARK_PROJECT_ID, label.question, n=top_k, consumer="benchmark_runner",
            pipeline_version=identity.pipeline_version,
        )

        rows = _rows_for_trace(global_store, trace_id)
        outcome = _outcome_from_rows(rows, candidates)
        op.set_retrieval(outcome)
        op.set("retrieval_status", outcome.status)
        op.set("empty_result", outcome.empty_result)
        op.set("retrieval_error", outcome.retrieval_error)
        if outcome.retrieval_error:
            op.mark_error(outcome.error_type)
        elif outcome.status == "success_empty":
            op.mark_terminal("success_empty")

        retrieval_row = rows.get("rag.retrieval")
        if retrieval_row is not None:
            if retrieval_row.get("duration_ms") is not None:
                op.add_stage("retrieval_ms", retrieval_row["duration_ms"])
            for key, value in (retrieval_row.get("attributes") or {}).items():
                if key not in ("error.type", "error.message"):
                    op.set(key, value)
        for stage in _RETRIEVAL_STAGES:
            stage_row = rows.get(stage)
            if stage_row is not None and stage_row.get("duration_ms") is not None:
                op.add_stage(stage, stage_row["duration_ms"])
            if stage_row is not None:
                for key, value in (stage_row.get("attributes") or {}).items():
                    op.set(f"{stage}.{key}", value)

    return trace_id, outcome


def _append_jsonl(path: Path, row: dict[str, Any]) -> None:
    with path.open("a", encoding="utf-8") as handle:
        handle.write(json.dumps(row, ensure_ascii=False) + "\n")


def _quality_row(trace_id: str | None, identity: ExperimentIdentity, label: QueryLabel, outcome: RetrievalOutcome) -> dict[str, Any]:
    quality = score_outcome(label, outcome)
    return {
        "trace_id": trace_id,
        "experiment_id": identity.experiment_id,
        "run_id": identity.run_id,
        "pipeline_version": identity.pipeline_version,
        "query_id": identity.query_id,
        "query_category": identity.query_category,
        "repetition": identity.repetition,
        "cache_state": identity.cache_state,
        **asdict(quality),
    }


# --- Full workload -------------------------------------------------


async def run_workload(
    *,
    corpus: Corpus,
    query_set: QuerySet,
    pipeline_version: str,
    run_id: str,
    experiment_id: str,
    seed: int,
    top_k: int,
    warm_repetitions: int,
    cold_count: int,
    smoke: bool,
    output_dir: Path,
) -> dict[str, Any]:
    import app.websockets.handlers as handlers_module

    db = handlers_module.get_db()
    reset_and_ingest_corpus(db, corpus)

    global_store = bootstrap.get_state().local_store
    if global_store is None:
        raise RuntimeError(
            "observability must be initialized (mode='local' or 'full') before running "
            "the benchmark -- disabled mode cannot capture the canonical rag.retrieval row."
        )
    operational_store = LocalObservationStore(output_dir, filename="operational.jsonl")
    warmup_store = LocalObservationStore(output_dir, filename="_warmup_discard.jsonl")
    quality_path = output_dir / "quality.jsonl"

    queries = seeded_order(list(query_set.queries), seed)
    if smoke:
        queries = queries[:SMOKE_QUERY_COUNT]
        warm_repetitions = SMOKE_WARM_REPETITIONS
        cold_count = 0

    cold_ids = {q.query_id for q in queries[:cold_count]}

    def identity_for(label: QueryLabel, *, repetition: int, cache_state: str) -> ExperimentIdentity:
        return ExperimentIdentity(
            experiment_id=experiment_id,
            run_id=run_id,
            pipeline_version=pipeline_version,
            query_set_version=query_set.version,
            corpus_version=corpus.version,
            query_id=label.query_id,
            query_category=label.category,
            repetition=repetition,
            cache_state=cache_state,
            git_commit=manifest_mod.git_commit(),
            embedding_model=manifest_mod.EMBEDDING_MODEL,
            cerebras_model=None,
            chroma_version=chromadb.__version__,
            cognee_version=None,
        )

    completed = 0

    # Cold pass: exactly one measured call per cold-eligible query, before that
    # query's own warm-up has run in this process.
    for label in queries:
        if label.query_id not in cold_ids:
            continue
        identity = identity_for(label, repetition=0, cache_state="cold")
        trace_id, outcome = await invoke_once(
            label,
            handlers_module=handlers_module,
            global_store=global_store,
            operational_store=operational_store,
            identity=identity,
            top_k=top_k,
        )
        _append_jsonl(quality_path, _quality_row(trace_id, identity, label, outcome))
        completed += 1

    # Warm pass: one unmeasured warm-up (discarded) + N measured warm repetitions.
    for label in queries:
        warmup_identity = identity_for(label, repetition=0, cache_state="warmup")
        await invoke_once(
            label,
            handlers_module=handlers_module,
            global_store=global_store,
            operational_store=warmup_store,
            identity=warmup_identity,
            top_k=top_k,
        )
        for repetition in range(1, warm_repetitions + 1):
            identity = identity_for(label, repetition=repetition, cache_state="warm")
            trace_id, outcome = await invoke_once(
                label,
                handlers_module=handlers_module,
                global_store=global_store,
                operational_store=operational_store,
                identity=identity,
                top_k=top_k,
            )
            _append_jsonl(quality_path, _quality_row(trace_id, identity, label, outcome))
            completed += 1

    return {"completed_measured_calls": completed, "queries_run": len(queries)}


# --- CLI -------------------------------------------------


def _generate_run_id(pipeline_version: str) -> str:
    stamp = time.strftime("%Y%m%dT%H%M%SZ", time.gmtime())
    return f"{pipeline_version}-{stamp}-{uuid.uuid4().hex[:8]}"


def parse_args(argv: list[str] | None = None) -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="RAG observability benchmark runner")
    parser.add_argument("--pipeline", default="rag-naive-v1", help="pipeline_version tag, e.g. rag-naive-v1")
    parser.add_argument("--smoke", action="store_true", help="run a fast 2-query smoke workload, marked smoke=true")
    parser.add_argument("--seed", type=int, default=42)
    parser.add_argument("--top-k", type=int, default=DEFAULT_TOP_K)
    parser.add_argument("--warm-repetitions", type=int, default=DEFAULT_WARM_REPETITIONS)
    parser.add_argument("--cold-count", type=int, default=DEFAULT_COLD_COUNT)
    parser.add_argument("--otel-mode", default="local", choices=["disabled", "local", "full"])
    parser.add_argument("--run-id", default=None)
    parser.add_argument("--output-dir", default=None, type=Path)
    parser.add_argument(
        "--corpus", default=None, type=Path, help="override path to corpus.v1.json"
    )
    parser.add_argument(
        "--query-set", default=None, type=Path, help="override path to query_set.v1.json"
    )
    return parser.parse_args(argv)


def main(argv: list[str] | None = None) -> dict[str, Any]:
    args = parse_args(argv)

    # Reject an unknown --pipeline before writing any run artifacts (manifest,
    # local store, OTel export) -- an evaluation run for a mistyped/unregistered
    # pipeline version must fail loudly at startup, never silently measure
    # rag-naive-v1 under the wrong label.
    get_pipeline_version(args.pipeline)

    corpus = load_corpus(args.corpus) if args.corpus else load_corpus()
    query_set = load_query_set(args.query_set) if args.query_set else load_query_set()

    run_id = args.run_id or _generate_run_id(args.pipeline)
    experiment_id = f"{args.pipeline}-{run_id}"

    base_cfg = get_observability_config()
    output_dir = args.output_dir or (base_cfg.artifact_root / "benchmarks" / run_id)
    output_dir = Path(output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)

    cfg = ObservabilityConfig(
        enabled=True,
        mode=args.otel_mode,
        otlp_endpoint=base_cfg.otlp_endpoint,
        signoz_ui_url=base_cfg.signoz_ui_url,
        artifact_root=output_dir,
    )
    bootstrap.initialize_observability(config=cfg)

    created_at = time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime())
    run_manifest = manifest_mod.build_manifest(
        run_id=run_id,
        experiment_id=experiment_id,
        pipeline_version=args.pipeline,
        smoke=args.smoke,
        seed=args.seed,
        warm_repetitions=(SMOKE_WARM_REPETITIONS if args.smoke else args.warm_repetitions),
        cold_query_count=(0 if args.smoke else args.cold_count),
        created_at=created_at,
        corpus=corpus,
        query_set=query_set,
    )
    manifest_mod.write_manifest(run_manifest, output_dir / "manifest.json")

    try:
        result = asyncio.run(
            run_workload(
                corpus=corpus,
                query_set=query_set,
                pipeline_version=args.pipeline,
                run_id=run_id,
                experiment_id=experiment_id,
                seed=args.seed,
                top_k=args.top_k,
                warm_repetitions=args.warm_repetitions,
                cold_count=args.cold_count,
                smoke=args.smoke,
                output_dir=output_dir,
            )
        )
    finally:
        bootstrap.shutdown_observability()

    summary = summarize_mod.summarize_run_dir(output_dir)
    summarize_mod.write_summary(summary, output_dir / "summary.json")

    print(f"run_id={run_id}")
    print(f"output_dir={output_dir}")
    print(f"smoke={args.smoke}")
    print(f"completed_measured_calls={result['completed_measured_calls']}")
    print(f"queries_run={result['queries_run']}")
    return {"run_id": run_id, "output_dir": str(output_dir), "summary": summary, **result}


if __name__ == "__main__":
    main()