rag-jev 0.2.0__py3-none-any.whl

rag_jev/__init__.py

@@ -0,0 +1,15 @@
+from rag_jev.models import Document, Judgment, Policy, SelectionResult, SelectRequest
+from rag_jev.provider import Jev, ProviderError
+from rag_jev.selector import ContextSelector, SelectOptions
+
+__all__ = [
+    "ContextSelector",
+    "Document",
+    "Jev",
+    "Judgment",
+    "Policy",
+    "ProviderError",
+    "SelectRequest",
+    "SelectionResult",
+    "SelectOptions",
+]

rag_jev/calibration.py

@@ -0,0 +1,353 @@
+"""Development-only policy selection with a frozen, separate held-out assessment."""
+
+from __future__ import annotations
+
+import argparse
+import asyncio
+import hashlib
+import json
+import math
+import random
+import re
+from collections import defaultdict
+from collections.abc import Sequence
+from pathlib import Path
+from statistics import mean
+from typing import Any
+
+from rag_jev.evaluation import EvalCase, Trace, collect_trace, fingerprint, measure
+from rag_jev.models import Policy, ScoringStrategy, SelectRequest
+from rag_jev.provider import Jev
+from rag_jev.selector import ContextSelector, apply_policy
+
+
+def digest(value: Any) -> str:
+    return hashlib.sha256(json.dumps(value, sort_keys=True).encode()).hexdigest()
+
+
+def validate_splits(development: list[EvalCase], evaluation: list[EvalCase]) -> None:
+    if len(development) < 2 or len(evaluation) < 2:
+        raise ValueError("Each split needs at least two questions")
+    ids: set[str] = set()
+    questions: set[str] = set()
+    for case in development + evaluation:
+        question = re.sub(r"\W+", " ", case.query.casefold()).strip()
+        if not question or case.id in ids or question in questions:
+            raise ValueError("Duplicate case ID or normalized question across calibration inputs")
+        ids.add(case.id)
+        questions.add(question)
+    if {c.family_id for c in development if c.family_id} & {
+        c.family_id for c in evaluation if c.family_id
+    }:
+        raise ValueError("Related question families cross development/evaluation splits")
+    if not any(c.relevant_ids for c in development):
+        raise ValueError("Development requires labeled relevant evidence for recall calibration")
+
+
+def candidate_policies(
+    cutoffs: Sequence[float], top_ns: Sequence[int], budget: int | None
+) -> list[Policy]:
+    policies = [
+        Policy(mode=mode, min_relevance=cutoff, max_context_tokens=budget)
+        for mode in ("filter", "filter_and_rerank")
+        for cutoff in sorted(set(cutoffs))
+    ]
+    policies += [
+        Policy(mode="rerank", top_n=k, max_context_tokens=budget) for k in sorted(set(top_ns))
+    ]
+    if not policies:
+        raise ValueError("At least one candidate policy is required")
+    return policies
+
+
+def assess(
+    cases: list[EvalCase], traces: list[Trace], policy: Policy | None, strategy: ScoringStrategy
+) -> dict[str, Any]:
+    by_id = {t.case_id: t for t in traces}
+    if len(by_id) != len(traces) or set(by_id) != {c.id for c in cases}:
+        raise ValueError("Trace IDs must match this split exactly")
+    rows = []
+    for case in cases:
+        trace = by_id[case.id]
+        request = SelectRequest(
+            query=case.query,
+            retrieval_query=case.retrieval_query,
+            documents=case.documents,
+            scoring_strategy=strategy,
+            **(policy or Policy(mode="rerank")).model_dump(),
+        )
+        if (
+            trace.input_hash != fingerprint(case)
+            or trace.scoring_strategy != strategy
+            or trace.prompt_version != request.prompt_version
+            or trace.relevance_guidance is not None
+        ):
+            raise ValueError("Trace inputs, strategy, prompt or guidance do not match calibration")
+        if len(trace.judgments) != len(case.documents):
+            raise ValueError("Incomplete trace judgments")
+        result = apply_policy(request, trace.judgments, usage=trace.scoring_usage)
+        by_doc = {d.id: d for d in case.documents}
+        selected = (
+            result.documents
+            if policy
+            else [by_doc[i] for i in case.baseline_ids]
+            if case.baseline_ids is not None
+            else case.documents
+        )
+        row = measure(case, selected).model_dump()
+        row.update(
+            {
+                "case_id": case.id,
+                "selected_ids": [d.id for d in selected],
+                "budget_exceeded": result.context_token_budget_exceeded if policy else None,
+            }
+        )
+        rows.append(row)
+    keys = [
+        "evidence_recall",
+        "precision",
+        "ndcg",
+        "context_tokens",
+        "documents",
+        "correctly_empty",
+        "all_evidence_lost",
+    ]
+    summary = {}
+    for key in keys:
+        values = [float(r[key]) for r in rows if r[key] is not None]
+        summary[key] = mean(values) if values else None
+    return {"summary": summary, "rows": rows}
+
+
+def choose(
+    cases: list[EvalCase],
+    traces: list[Trace],
+    policies: list[Policy],
+    strategy: ScoringStrategy,
+    min_recall: float,
+) -> dict[str, Any]:
+    if not math.isfinite(min_recall) or not 0 <= min_recall <= 1:
+        raise ValueError("min_recall must be finite and between zero and one")
+    results = {p.version: assess(cases, traces, p, strategy) for p in policies}
+    eligible = [
+        p
+        for p in policies
+        if results[p.version]["summary"]["evidence_recall"] is not None
+        and results[p.version]["summary"]["evidence_recall"] >= min_recall
+    ]
+    if not eligible:
+        return {
+            "eligible": False,
+            "policy": None,
+            "development": results,
+            "reason": "No development policy meets the requested observed evidence recall",
+        }
+    selected = min(
+        eligible,
+        key=lambda p: (
+            results[p.version]["summary"]["context_tokens"],
+            -(results[p.version]["summary"]["ndcg"] or 0),
+            p.version,
+        ),
+    )
+    return {"eligible": True, "policy": selected.model_dump(), "development": results}
+
+
+def paired_interval(
+    baseline: list[float], selected: list[float], families: list[str] | None = None
+) -> dict[str, Any]:
+    differences = [b - a for a, b in zip(baseline, selected, strict=True)]
+    grouped: dict[str, list[float]] = defaultdict(list)
+    for family, difference in zip(
+        families or [str(i) for i in range(len(differences))], differences, strict=True
+    ):
+        grouped[family].append(difference)
+    groups = list(grouped.values())
+    if len(groups) < 2:
+        return {
+            "delta": mean(differences),
+            "ci95": None,
+            "questions": len(differences),
+            "independent_units": len(groups),
+            "scope": "Too few independent families for an interval",
+        }
+    rng = random.Random(20260919)
+    draws = sorted(
+        mean(v for group in rng.choices(groups, k=len(groups)) for v in group) for _ in range(2000)
+    )
+    return {
+        "delta": mean(differences),
+        "ci95": [draws[49], draws[1949]],
+        "questions": len(differences),
+        "independent_units": len(groups),
+        "scope": "Descriptive paired bootstrap of supplied families (else questions); "
+        "query-weighted mean, 2,000 draws",
+    }
+
+
+def held_out(
+    cases: list[EvalCase],
+    traces: list[Trace],
+    policies: list[Policy],
+    strategy: ScoringStrategy,
+    frozen: dict[str, Any],
+) -> dict[str, Any]:
+    baseline = assess(cases, traces, None, strategy)
+    results = {
+        p.version: {"policy": p.model_dump(), **assess(cases, traces, p, strategy)}
+        for p in policies
+    }
+    report: dict[str, Any] = {
+        "baseline": baseline,
+        "policies": results,
+        "frozen_policy": frozen["policy"],
+        "eligible": frozen["eligible"],
+        "answer_quality": "not_evaluated",
+        "selection_rule": "Minimize development context tokens subject to observed recall; "
+        "ties higher NDCG then policy ID",
+        "scope": "Evidence already in candidate sets. No semantic accuracy or total-cost "
+        "claim. Held-out policies are descriptive; never reselect from this table.",
+    }
+    if frozen["eligible"]:
+        policy = Policy.model_validate(frozen["policy"])
+        selected = results[policy.version]
+        report["frozen_policy_held_out"] = selected
+        report["context_token_difference"] = paired_interval(
+            [r["context_tokens"] for r in baseline["rows"]],
+            [r["context_tokens"] for r in selected["rows"]],
+            [
+                "family:" + c.family_id if c.family_id is not None else "question:" + c.id
+                for c in cases
+            ],
+        )
+    return report
+
+
+async def run_calibration(args: argparse.Namespace) -> None:
+    if not math.isfinite(args.min_recall) or not 0 <= args.min_recall <= 1:
+        raise ValueError("min_recall must be finite and between zero and one")
+    development = [
+        EvalCase.model_validate(r)
+        for r in json.loads(await asyncio.to_thread(Path(args.development).read_text))
+    ]
+    evaluation = [
+        EvalCase.model_validate(r)
+        for r in json.loads(await asyncio.to_thread(Path(args.evaluation).read_text))
+    ]
+    validate_splits(development, evaluation)
+    policies = candidate_policies(args.cutoffs, args.top_ns, args.max_context_tokens)
+    output = Path(args.output)
+    frozen_path = output.with_suffix(".frozen.json")
+    config_path = output.with_suffix(".policy.json")
+    traces_path = output.with_suffix(".traces.json")
+    if any(p.exists() for p in [output, frozen_path, config_path, traces_path]):
+        raise ValueError(
+            "Use a fresh output path; previous calibration artifacts are never overwritten"
+        )
+    output.parent.mkdir(parents=True, exist_ok=True)
+    cached = (
+        [
+            Trace.model_validate(r)
+            for r in json.loads(await asyncio.to_thread(Path(args.replay).read_text))
+        ]
+        if args.replay
+        else []
+    )
+    all_ids = {c.id for c in development + evaluation}
+    if args.replay and (
+        len({t.case_id for t in cached}) != len(cached) or {t.case_id for t in cached} != all_ids
+    ):
+        raise ValueError("Replay traces must cover both splits exactly")
+
+    def save(path: Path, value: Any) -> None:
+        path.write_text(json.dumps(value, indent=2, allow_nan=False) + "\n")
+
+    collected: list[Trace] = []
+
+    async def traces(cases: list[EvalCase], selector: ContextSelector | None) -> list[Trace]:
+        if args.replay:
+            ids = {c.id for c in cases}
+            return [t for t in cached if t.case_id in ids]
+        assert selector is not None
+        subset = []
+        for case in cases:
+            trace = await collect_trace(case, selector, scoring_strategy=args.scoring_strategy)
+            subset.append(trace)
+            collected.append(trace)
+            save(traces_path, [t.model_dump(mode="json") for t in collected])
+        return subset
+
+    async def execute(selector: ContextSelector | None) -> None:
+        dev_traces = await traces(development, selector)
+        frozen = choose(development, dev_traces, policies, args.scoring_strategy, args.min_recall)
+        frozen.update(
+            {
+                "scoring_strategy": args.scoring_strategy,
+                "min_recall": args.min_recall,
+                "development_hash": digest([c.model_dump() for c in development]),
+                "evaluation_hash": digest([c.model_dump() for c in evaluation]),
+                "candidate_policies": [p.model_dump() for p in policies],
+                "sources": sorted({t.source for t in dev_traces}),
+                "models": sorted({j.model for t in dev_traces for j in t.judgments}),
+            }
+        )
+        # Persist the choice before inspecting held-out metrics or making held-out calls.
+        save(frozen_path, frozen)
+        eval_traces = await traces(evaluation, selector)
+        report = held_out(evaluation, eval_traces, policies, args.scoring_strategy, frozen)
+        report.update(
+            {
+                "freeze_hash": digest(frozen),
+                "development_questions": len(development),
+                "evaluation_questions": len(evaluation),
+                "minimum_development_recall": args.min_recall,
+                "sources": sorted({t.source for t in dev_traces + eval_traces}),
+                "models": sorted({j.model for t in dev_traces + eval_traces for j in t.judgments}),
+                "scoring_usage": {
+                    "input_tokens": sum(
+                        t.scoring_usage.input_tokens
+                        if t.scoring_usage
+                        else sum(j.input_tokens for j in t.judgments)
+                        for t in dev_traces + eval_traces
+                    ),
+                    "output_tokens": sum(
+                        t.scoring_usage.output_tokens
+                        if t.scoring_usage
+                        else sum(j.output_tokens for j in t.judgments)
+                        for t in dev_traces + eval_traces
+                    ),
+                },
+                "data_warning": "Artifacts include query and passage text; "
+                "keep private inputs local.",
+            }
+        )
+        save(output, report)
+        if frozen["eligible"]:
+            save(
+                config_path,
+                {**frozen["policy"], "scoring_strategy": args.scoring_strategy, "shadow": True},
+            )
+        print(
+            json.dumps(
+                {
+                    "report": str(output),
+                    "frozen": str(frozen_path),
+                    "policy": str(config_path) if frozen["eligible"] else None,
+                    "eligible": frozen["eligible"],
+                    "answer_quality": "not_evaluated",
+                }
+            )
+        )
+
+    if args.replay:
+        await execute(None)
+    else:
+        async with Jev(model=args.model) as provider:
+            await execute(
+                ContextSelector(
+                    provider,
+                    timeout_ms=args.timeout_ms,
+                    max_concurrency=args.max_concurrency,
+                    on_error="raise",
+                )
+            )

rag_jev/cli.py

@@ -0,0 +1,206 @@
+from __future__ import annotations
+
+import argparse
+import asyncio
+import json
+import os
+import sys
+from pathlib import Path
+from typing import Any
+
+from dotenv import load_dotenv
+from pydantic import ValidationError
+from typesafe_sdk import TypeSafeError
+
+from rag_jev.demo import CASES, FixtureScorer, fixture_traces
+from rag_jev.evaluation import EvalCase, Trace, collect_trace, evaluate
+from rag_jev.models import SelectRequest
+from rag_jev.provider import Jev, ProviderError
+from rag_jev.selector import ContextSelector
+
+
+def emit(value: Any, path: str | None = None) -> None:
+    output = json.dumps(value, indent=2, ensure_ascii=False, allow_nan=False) + "\n"
+    if path:
+        Path(path).write_text(output)
+    else:
+        sys.stdout.write(output)
+
+
+def live_selector(provider: Jev, args: argparse.Namespace) -> ContextSelector:
+    return ContextSelector(
+        provider,
+        timeout_ms=args.timeout_ms,
+        max_concurrency=args.max_concurrency,
+        on_error="raise",
+    )
+
+
+async def run(args: argparse.Namespace) -> None:
+    if args.command == "review-results":
+        from rag_jev.review import summarize_review
+
+        emit(
+            await asyncio.to_thread(summarize_review, args.packet, args.labels, args.key),
+            args.output,
+        )
+        return
+    if args.command == "calibrate":
+        from rag_jev.calibration import run_calibration
+
+        await run_calibration(args)
+        return
+    if args.command == "demo":
+        selector = ContextSelector(FixtureScorer())
+        case = CASES[0]
+        selected = await selector.select(
+            query=case.query, documents=case.documents, min_relevance=0.2
+        )
+        report = await evaluate(CASES, fixture_traces(), min_relevance=0.2)
+        emit(
+            {
+                "notice": "OFFLINE FIXTURE DEMO: scores are hand-authored, not Jev predictions",
+                "selection": selected.model_dump(),
+                "evaluation": report,
+            }
+        )
+        return
+    if args.command == "select":
+        raw = await asyncio.to_thread(
+            sys.stdin.read if args.input == "-" else Path(args.input).read_text
+        )
+        request = SelectRequest.model_validate_json(raw)
+        async with Jev(model=args.model) as provider:
+            result = await live_selector(provider, args).select_request(request)
+        emit(result.model_dump(), args.output)
+        return
+    if args.command == "eval":
+        raw = await asyncio.to_thread(Path(args.input).read_text)
+        cases = [EvalCase.model_validate(x) for x in json.loads(raw)]
+        if args.replay:
+            raw = await asyncio.to_thread(Path(args.replay).read_text)
+            traces = [Trace.model_validate(x) for x in json.loads(raw)]
+        else:
+            async with Jev(model=args.model) as provider:
+                selector = live_selector(provider, args)
+                traces = [
+                    await collect_trace(case, selector, scoring_strategy=args.scoring_strategy)
+                    for case in cases
+                ]
+            if args.save_traces:
+                emit([t.model_dump() for t in traces], args.save_traces)
+        if any(t.scoring_strategy != args.scoring_strategy for t in traces):
+            raise ValueError("Replay strategy differs; pass the matching --scoring-strategy")
+        report = await evaluate(cases, traces, min_relevance=args.min_relevance, top_n=args.top_n)
+        emit(report, args.output)
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description="Select RAG context with TypeSafe Jev")
+    parser.add_argument("--env-file", default=".env", help="Loaded without overriding environment")
+    commands = parser.add_subparsers(dest="command", required=True)
+    commands.add_parser("demo", help="Offline, hand-authored fixture demonstration")
+    commands.add_parser("schema", help="Print the service OpenAPI schema")
+    review = commands.add_parser(
+        "review-results", help="Validate supplied human labels and report completeness"
+    )
+    review.add_argument("packet")
+    review.add_argument("labels")
+    review.add_argument("--key", help="Unblind only after independent review")
+    review.add_argument("--output")
+    serve = commands.add_parser("serve", help="Start the real Jev selection service")
+    serve.add_argument("--host", default="127.0.0.1")
+    serve.add_argument("--port", type=int, default=8000)
+    serve.add_argument(
+        "--replay-only",
+        action="store_true",
+        help="Browse saved runs without keys or inference calls",
+    )
+    select = commands.add_parser("select", help="Select using a JSON request file, or - for stdin")
+    select.add_argument("input")
+    evaluation = commands.add_parser(
+        "eval", help="Compare policies on labeled retrieved candidates"
+    )
+    evaluation.add_argument("input", help="JSON array of evaluation cases")
+    evaluation.add_argument("--replay", help="Use saved traces; no API requests")
+    evaluation.add_argument("--save-traces", help="Save live judgments for offline policy tuning")
+    evaluation.add_argument("--min-relevance", type=float, required=True)
+    evaluation.add_argument("--top-n", type=int)
+    evaluation.add_argument(
+        "--scoring-strategy", choices=["independent", "contextual"], default="independent"
+    )
+    calibration = commands.add_parser(
+        "calibrate",
+        help="Choose a policy on development questions, then assess a separate held-out set",
+    )
+    calibration.add_argument("development")
+    calibration.add_argument("evaluation")
+    calibration.add_argument("--replay", help="Saved traces covering both splits; no API calls")
+    calibration.add_argument(
+        "--scoring-strategy", choices=["independent", "contextual"], default="contextual"
+    )
+    calibration.add_argument("--cutoffs", type=float, nargs="+", default=[0.2, 0.35, 0.5])
+    calibration.add_argument("--top-ns", type=int, nargs="+", default=[3, 5, 8])
+    calibration.add_argument("--max-context-tokens", type=int)
+    calibration.add_argument("--min-recall", type=float, default=0.98)
+    for command in (select, evaluation, calibration):
+        command.add_argument("--model")
+        command.add_argument("--timeout-ms", type=float)
+        command.add_argument("--max-concurrency", type=int)
+        command.add_argument("--output", required=command is calibration)
+    args = parser.parse_args()
+    load_dotenv(args.env_file, override=False)
+    try:
+        if args.command in ("select", "eval", "calibrate"):
+            args.model = args.model or os.getenv("RAG_JEV_MODEL", "jev-latest")
+            if args.timeout_ms is None:
+                args.timeout_ms = float(os.getenv("RAG_JEV_TIMEOUT_MS", "5000"))
+            if args.max_concurrency is None:
+                args.max_concurrency = int(os.getenv("RAG_JEV_MAX_CONCURRENCY", "8"))
+        if args.command == "serve":
+            import uvicorn
+
+            from rag_jev.server import create_app
+
+            if not args.replay_only and not os.getenv("TYPESAFE_API_KEY", "").strip():
+                parser.error("set TYPESAFE_API_KEY in the environment or .env before serving")
+            if args.host not in ("127.0.0.1", "localhost", "::1") and not os.getenv(
+                "RAG_JEV_API_TOKEN"
+            ):
+                parser.error("set RAG_JEV_API_TOKEN before binding outside localhost")
+            uvicorn.run(
+                create_app(replay_only=args.replay_only),
+                host=args.host,
+                port=args.port,
+                access_log=False,
+            )
+        elif args.command == "schema":
+            from rag_jev.server import create_app
+
+            emit(create_app().openapi())
+        else:
+            asyncio.run(run(args))
+    except ValidationError as exc:
+        emit(
+            {
+                "error": "invalid_input",
+                "issues": [
+                    {"loc": list(e["loc"]), "message": e["msg"]}
+                    for e in exc.errors(include_input=False)
+                ],
+            }
+        )
+        raise SystemExit(2) from None
+    except ProviderError as exc:
+        print(f"Jev request failed: {exc.code}", file=sys.stderr)
+        raise SystemExit(1) from None
+    except TypeSafeError:
+        print("TypeSafe client failed; check API key and configuration.", file=sys.stderr)
+        raise SystemExit(1) from None
+    except (ValueError, OSError) as exc:
+        print(f"Invalid input or configuration ({type(exc).__name__}).", file=sys.stderr)
+        raise SystemExit(2) from None
+
+
+if __name__ == "__main__":
+    main()

rag_jev/demo.py

@@ -0,0 +1,76 @@
+"""Fixed examples with hand-authored scores. This is NOT a Jev benchmark or model."""
+
+from rag_jev.evaluation import EvalCase, Trace, fingerprint
+from rag_jev.models import Document, Judgment
+from rag_jev.provider import ProviderError
+
+CASES = [
+    EvalCase(
+        id="refund",
+        query="How many days do I have to request a refund?",
+        documents=[
+            Document(
+                id="shipping",
+                text="Orders ship within 5 business days.",
+                metadata={"source": "shipping.md"},
+            ),
+            Document(
+                id="refund",
+                text="You may request a refund within 30 days of purchase.",
+                metadata={"source": "refunds.md", "page": 2},
+            ),
+            Document(
+                id="receipt",
+                text="Keep your receipt; it is required for refund requests.",
+                metadata={"source": "refunds.md", "page": 3},
+            ),
+        ],
+        relevant_ids=["refund", "receipt"],
+        reference_answer="30 days",
+    ),
+    EvalCase(
+        id="correction",
+        query="How do I renew my subscription every 30 days?",
+        documents=[
+            Document(id="annual", text="Subscriptions renew annually, not every 30 days."),
+            Document(id="renew", text="Annual subscriptions renew automatically unless canceled."),
+            Document(id="returns", text="Return physical products within 30 days."),
+        ],
+        relevant_ids=["annual", "renew"],
+        reference_answer="annually",
+    ),
+    EvalCase(
+        id="no-answer",
+        query="What is the office Wi-Fi password?",
+        documents=[
+            Document(id="hours", text="The office opens at 9 AM."),
+            Document(id="wifi", text="Wi-Fi is available in the office lobby."),
+        ],
+        relevant_ids=[],
+        reference_answer="insufficient evidence",
+    ),
+]
+SCORES = [[0.03, 0.98, 0.60], [0.97, 0.85, 0.02], [0.01, 0.08]]
+
+
+def fixture_traces() -> list[Trace]:
+    return [
+        Trace(
+            case_id=case.id,
+            input_hash=fingerprint(case),
+            elapsed_ms=0,
+            source="hand_authored_fixture_NOT_a_model_benchmark",
+            judgments=[Judgment(relevance=p, model="fixture-not-jev") for p in scores],
+        )
+        for case, scores in zip(CASES, SCORES, strict=True)
+    ]
+
+
+class FixtureScorer:
+    async def score(self, query: str, document: Document, guidance: str | None) -> Judgment:
+        for case, trace in zip(CASES, fixture_traces(), strict=True):
+            if query == case.query:
+                for candidate, judgment in zip(case.documents, trace.judgments, strict=True):
+                    if candidate.id == document.id and candidate.text == document.text:
+                        return judgment
+        raise ProviderError("unknown_demo_fixture")