engines/live_request.py

from __future__ import annotations

from dataclasses import dataclass
import os
from pathlib import Path
from typing import Any

from engines.fixture_builder import COMPACT_TASK_ROWS
from engines.presets import MODEL_BY_KEY, PRESET_BY_KEY


DEFAULT_LIVE_DECODE_STEPS = 48
BENCHMARK_TOKENS_PER_PAGE = 16
BENCHMARK_RECENT_WINDOW_TOKENS = 128
BENCHMARK_SHORTLIST_POLICY = "benchmark"

MODE_TO_VARIANT = {
    "compact_task": {
        "dense": "exact",
        "m0": "quality",
        "m0_selective_k": "systems",
    },
    "longbench_mini": {
        "dense": "exact",
        "m0": "quality",
        "m0_selective_k": "systems",
    },
    "backend_truth": {
        "dense": "exact",
        "m0": "shortlist_base",
        "m0_selective_k": "learned_selector",
    },
}

VARIANT_CONFIG = {
    "exact": {
        "execution_recent_window_tokens": 0,
        "execution_sink_window_tokens": 0,
        "execution_relevance_top_k": 0,
        "selector_required": False,
        "learned_page_selector_profile": "quality",
        "learned_page_selector_prompt_family": None,
        "learned_page_selector_prompt_variant": None,
    },
    "quality": {
        "execution_recent_window_tokens": 0,
        "execution_sink_window_tokens": 0,
        "execution_relevance_top_k": 0,
        "selector_required": True,
        "learned_page_selector_profile": "quality",
        "learned_page_selector_prompt_family": "cache",
        "learned_page_selector_prompt_variant": "locality",
    },
    "systems": {
        "execution_recent_window_tokens": 0,
        "execution_sink_window_tokens": 0,
        "execution_relevance_top_k": 0,
        "selector_required": True,
        "learned_page_selector_profile": "systems",
        "learned_page_selector_prompt_family": "cache",
        "learned_page_selector_prompt_variant": "locality",
    },
    "shortlist_base": {
        "execution_recent_window_tokens": 1024,
        "execution_sink_window_tokens": 256,
        "execution_relevance_top_k": 4,
        "selector_required": False,
        "learned_page_selector_profile": "systems",
        "learned_page_selector_prompt_family": None,
        "learned_page_selector_prompt_variant": None,
    },
    "learned_selector": {
        "execution_recent_window_tokens": 0,
        "execution_sink_window_tokens": 0,
        "execution_relevance_top_k": 0,
        "selector_required": True,
        "learned_page_selector_profile": "systems",
        "learned_page_selector_prompt_family": "cache",
        "learned_page_selector_prompt_variant": "locality",
    },
}

SELECTOR_ARTIFACT_BY_MODEL = {
    "qwen35_4b_hf": "qwen35_4b/linear_selector_model.json",
    "qwen35_9b_hf": "qwen35_9b/linear_selector_model.json",
    "qwen35_27b_hf": "qwen35_27b/linear_selector_model.json",
}

COMPACT_TASK_DECODE_STEPS = {
    "retrieval_passkey": 64,
    "reasoning_arithmetic": 64,
    "instruction_constraints": 32,
}


@dataclass(frozen=True)
class LiveRuntimeSettings:
    model_key: str
    model_id: str
    model_family: str
    benchmark_suite: str
    benchmark_variant: str
    context_length: int
    decode_steps: int
    max_live_context: int
    bits_k: int
    bits_v: int
    tokens_per_page: int
    recent_window_tokens: int
    execution_recent_window_tokens: int
    execution_sink_window_tokens: int
    execution_relevance_top_k: int
    learned_page_selector_profile: str
    learned_page_selector_prompt_family: str | None
    learned_page_selector_prompt_variant: str | None
    learned_page_selector_path: str | None
    compact_task_name: str | None
    prompt_text: str
    use_exact_length_prompt: bool
    is_custom_prompt: bool


def canonicalize_benchmark_payload(request: dict[str, Any]) -> dict[str, Any]:
    payload = dict(request)
    model_key = str(payload.get("model") or "")
    model = MODEL_BY_KEY.get(model_key)
    if model is None or model.family != "qwen35":
        return payload
    preset_key = str(payload.get("preset") or "")
    if preset_key not in PRESET_BY_KEY:
        return payload
    payload["page_size"] = BENCHMARK_TOKENS_PER_PAGE
    payload["bits_k"] = 4
    payload["bits_v"] = 4
    payload["recent_window"] = BENCHMARK_RECENT_WINDOW_TOKENS // BENCHMARK_TOKENS_PER_PAGE
    payload["sink_window"] = 0
    payload["shortlist_policy"] = BENCHMARK_SHORTLIST_POLICY
    return payload


def _selector_artifact_path(model_key: str) -> str | None:
    relative = SELECTOR_ARTIFACT_BY_MODEL.get(model_key)
    if relative is None:
        return None
    path = Path(__file__).resolve().parents[1] / "data" / "selector_artifacts" / relative
    return str(path) if path.exists() else None


def _compact_task_name(model_key: str, context_length: int) -> str:
    row = COMPACT_TASK_ROWS.get(model_key, {}).get(context_length)
    if row is None:
        raise ValueError(
            f"Live compact-task replay is only bundled for the benchmark contexts on this model. "
            f"No row exists for {MODEL_BY_KEY[model_key].label} at {context_length} tokens."
        )
    return str(row["task_name"])


def resolve_live_runtime_settings(
    request: dict[str, Any],
    *,
    decode_steps: int,
    max_live_context: int,
) -> LiveRuntimeSettings:
    request = canonicalize_benchmark_payload(request)
    model_key = str(request.get("model") or "")
    if model_key not in MODEL_BY_KEY:
        raise ValueError(f"Unsupported live model key: {model_key}")
    model = MODEL_BY_KEY[model_key]
    if model.family not in {"qwen35", "llama"}:
        raise ValueError(
            f"Live execution for model family '{model.family}' is not wired in this v1 Space yet. "
            "Preset-backed compare still works for this model."
        )
    if model.family == "llama" and not any(
        str(os.getenv(env_name) or "").strip() for env_name in ("HF_TOKEN", "HUGGINGFACE_HUB_TOKEN")
    ):
        raise ValueError(
            "Live execution for this gated Llama model requires HF_TOKEN in the Space settings. "
            "Preset-backed compare still works without it."
        )

    context_length = int(request.get("context_length") or 0)
    if context_length <= 0:
        raise ValueError("context_length must be positive for live execution")
    if context_length > int(max_live_context):
        raise ValueError(
            f"Live execution is currently capped at {int(max_live_context)} tokens. "
            "Use preset-backed compare for larger contexts or raise DOTCACHE_SPACE_MAX_LIVE_CONTEXT."
        )

    if model.family != "qwen35":
        raise ValueError("Only the Qwen benchmark-backed live lane is supported in this Space.")

    custom_prompt = str(request.get("custom_prompt") or "").strip()
    preset_key = str(request.get("preset") or "")
    if preset_key not in PRESET_BY_KEY:
        raise ValueError("Live Qwen compare requires a benchmark-backed preset selection.")
    if not custom_prompt and preset_key == "longbench_mini":
        raise ValueError(
            "Live LongBench replay is not wired in this Space yet. Preset-backed compare remains the valid benchmark view."
        )

    variant = MODE_TO_VARIANT[preset_key][str(request.get("mode") or "dense")]
    variant_config = VARIANT_CONFIG[variant]
    selector_artifact_path = _selector_artifact_path(model_key)
    if variant_config["selector_required"] and selector_artifact_path is None:
        raise ValueError(
            f"Live benchmark replay for {MODEL_BY_KEY[model_key].label} requires its learned selector artifact, "
            "and that artifact is not packaged in this Space."
        )

    compact_task_name = None
    prompt_text = custom_prompt
    use_exact_length_prompt = False
    benchmark_decode_steps = int(decode_steps or DEFAULT_LIVE_DECODE_STEPS)
    if not custom_prompt and preset_key == "compact_task":
        compact_task_name = _compact_task_name(model_key, context_length)
        benchmark_decode_steps = COMPACT_TASK_DECODE_STEPS[compact_task_name]
    elif not custom_prompt and preset_key == "backend_truth":
        prompt_text = "Cache locality matters for fast decoding."
        use_exact_length_prompt = True

    return LiveRuntimeSettings(
        model_key=model.key,
        model_id=str(model.model_id),
        model_family=str(model.family),
        benchmark_suite=preset_key,
        benchmark_variant=variant,
        context_length=context_length,
        decode_steps=benchmark_decode_steps,
        max_live_context=int(max_live_context),
        bits_k=4,
        bits_v=4,
        tokens_per_page=BENCHMARK_TOKENS_PER_PAGE,
        recent_window_tokens=BENCHMARK_RECENT_WINDOW_TOKENS,
        execution_recent_window_tokens=int(variant_config["execution_recent_window_tokens"]),
        execution_sink_window_tokens=int(variant_config["execution_sink_window_tokens"]),
        execution_relevance_top_k=int(variant_config["execution_relevance_top_k"]),
        learned_page_selector_profile=str(variant_config["learned_page_selector_profile"]),
        learned_page_selector_prompt_family=variant_config["learned_page_selector_prompt_family"],
        learned_page_selector_prompt_variant=variant_config["learned_page_selector_prompt_variant"],
        learned_page_selector_path=selector_artifact_path,
        compact_task_name=compact_task_name,
        prompt_text=prompt_text,
        use_exact_length_prompt=use_exact_length_prompt,
        is_custom_prompt=bool(custom_prompt),
    )


def default_dense_runner_script(repo_root: Path) -> Path:
    return repo_root / "scripts" / "space_dense_runner.py"


def default_dotcache_runner_script(repo_root: Path) -> Path:
    return repo_root / "scripts" / "space_dotcache_runner.py"


def default_llama_runner_script(repo_root: Path) -> Path:
    return repo_root / "scripts" / "space_llama_runner.py"