src/deberta_ime/sidecar.py

from __future__ import annotations

import json
from typing import TextIO

from .deberta import MODEL_ID, MODEL_REVISION
from .domain import Candidate, CandidateScorer, RerankRequest
from .mozc import MOZC_REVISION
from .profiles import (
    MOZC_PROFILE_SELECTION_REVISION,
    MOZC_PROFILES,
    MOZC_SIDECAR_SOURCE_ID,
)
from .reranker import Reranker

SIDECAR_SCHEMA_VERSION = 1
MAX_PROTOCOL_LINE_BYTES = 65_536


def _string_list(value: object, name: str) -> tuple[str, ...]:
    if not isinstance(value, list) or not all(isinstance(item, str) for item in value):
        raise ValueError(f"{name} must be a list of strings")
    return tuple(value)


def _preserved_response(
    *, request_id: str, candidate_count: int, reason: str
) -> dict[str, object]:
    return {
        "ok": True,
        "schema_version": SIDECAR_SCHEMA_VERSION,
        "request_id": request_id,
        "model": {"id": MODEL_ID, "revision": MODEL_REVISION},
        "profile": None,
        "decision": {"changed": False, "reason": reason, "margin": None},
        "ranked_original_ranks": list(range(candidate_count)),
    }


def _rerank_response(payload: dict[str, object], scorer: CandidateScorer) -> dict[str, object]:
    request_id = payload.get("request_id")
    if payload.get("schema_version") != SIDECAR_SCHEMA_VERSION:
        raise ValueError("unsupported schema_version")
    if not isinstance(request_id, str) or not request_id:
        raise ValueError("request_id must be a non-empty string")
    if payload.get("op") != "rerank":
        raise ValueError("unsupported operation")
    source = payload.get("candidate_source")
    if not isinstance(source, dict):
        raise ValueError("candidate_source must be an object")
    mode = payload.get("mode")
    reading = payload.get("reading")
    if not isinstance(reading, str):
        raise ValueError("reading must be a string")
    raw_candidates = payload.get("candidates")
    if not isinstance(raw_candidates, list):
        raise ValueError("candidates must be a list")
    candidates: list[Candidate] = []
    for raw_candidate in raw_candidates:
        if not isinstance(raw_candidate, dict):
            raise ValueError("each candidate must be an object")
        surface = raw_candidate.get("surface")
        prior_score = raw_candidate.get("prior_score")
        if (
            not isinstance(surface, str)
            or not isinstance(prior_score, int | float)
            or isinstance(prior_score, bool)
        ):
            raise ValueError("each candidate needs string surface and numeric prior_score")
        candidates.append(Candidate(surface=surface, prior_score=float(prior_score)))
    if source.get("id") != MOZC_SIDECAR_SOURCE_ID or source.get("revision") != MOZC_REVISION:
        return _preserved_response(
            request_id=request_id,
            candidate_count=len(candidates),
            reason="uncalibrated_source",
        )
    if not isinstance(mode, str) or mode not in MOZC_PROFILES:
        return _preserved_response(
            request_id=request_id,
            candidate_count=len(candidates),
            reason="unsupported_mode",
        )
    left_context = _string_list(payload.get("left_context", []), "left_context")
    right_context = _string_list(payload.get("right_context", []), "right_context")
    if mode == "incremental" and right_context:
        return _preserved_response(
            request_id=request_id,
            candidate_count=len(candidates),
            reason="context_mode_mismatch",
        )
    request = RerankRequest(
        reading=reading,
        candidates=tuple(candidates),
        left_context=left_context,
        right_context=right_context,
    )
    selected_profile = MOZC_PROFILES[mode]
    result = Reranker(scorer, selected_profile).rerank(request)
    return {
        "ok": True,
        "schema_version": SIDECAR_SCHEMA_VERSION,
        "request_id": request_id,
        "model": {"id": MODEL_ID, "revision": MODEL_REVISION},
        "profile": {
            "candidate_source_id": MOZC_SIDECAR_SOURCE_ID,
            "candidate_source_revision": MOZC_REVISION,
            "mode": mode,
            "prior_weight": selected_profile.prior_weight,
            "min_margin": selected_profile.min_margin,
            "selected_on_revision": MOZC_PROFILE_SELECTION_REVISION,
        },
        "decision": {
            "changed": result.changed,
            "reason": result.reason,
            "margin": result.margin,
        },
        "ranked_original_ranks": [item.original_rank for item in result.ranked],
    }


def serve(input_stream: TextIO, output_stream: TextIO, scorer: CandidateScorer) -> int:
    model_loaded = False
    for raw_line in input_stream:
        if not raw_line.strip():
            continue
        request_id: str | None = None
        try:
            decoded = json.loads(raw_line)
            if not isinstance(decoded, dict):
                raise ValueError("request must be a JSON object")
            candidate_request_id = decoded.get("request_id")
            if isinstance(candidate_request_id, str):
                request_id = candidate_request_id
            if len(raw_line.encode("utf-8")) > MAX_PROTOCOL_LINE_BYTES:
                raise ValueError(f"request line exceeds {MAX_PROTOCOL_LINE_BYTES} UTF-8 bytes")
            if decoded.get("schema_version") != SIDECAR_SCHEMA_VERSION:
                raise ValueError("unsupported schema_version")
            if not request_id:
                raise ValueError("request_id must be a non-empty string")
            operation = decoded.get("op")
            if operation == "health":
                response = {
                    "ok": True,
                    "schema_version": SIDECAR_SCHEMA_VERSION,
                    "request_id": request_id,
                    "operation": "health",
                    "model_loaded": model_loaded,
                }
            elif operation == "warmup":
                loader = getattr(scorer, "load", None)
                if callable(loader):
                    loader()
                model_loaded = True
                response = {
                    "ok": True,
                    "schema_version": SIDECAR_SCHEMA_VERSION,
                    "request_id": request_id,
                    "operation": "warmup",
                    "model_loaded": True,
                }
            else:
                response = _rerank_response(decoded, scorer)
        except Exception as error:
            response = {
                "ok": False,
                "schema_version": SIDECAR_SCHEMA_VERSION,
                "request_id": request_id,
                "error": str(error),
            }
        output_stream.write(json.dumps(response, ensure_ascii=False, separators=(",", ":")))
        output_stream.write("\n")
        output_stream.flush()
    return 0