runtime/efficient.py

"""DeepSeek-inspired efficient inference patterns — interface only.

This module defines the SEAMS for V4-Pro / DSpark-style efficiency in
openclinical-ai. It does NOT implement fake MoE routing — it defines the
contract so real adapters can plug in.

Three architectural patterns mirrored here (from DeepSeek-V4-Pro + DSpark):

1. **Hybrid Attention (CSA + HCA)**
   - Compressed Sparse Attention: compresses long context into compact
     representations, retaining semantic anchors.
   - Heavily Compressed Attention: aggressive compression for million-
     token contexts (full patient history, longitudinal records).
   - Result: 27% FLOPs, 10% KV cache vs DeepSeek V3.2.

2. **Fine-grained MoE Expert Routing**
   - V4-Pro: 1.6T total params, 49B (3%) activated per inference.
   - V4-Flash: 284B total, 13B (4.5%) activated.
   - For healthcare: medical specialty experts (cardiology, oncology,
     geriatrics, mental health, pediatrics, pharmacy, etc.).

3. **FP8 Quantized Inference**
   - 8-bit floating point halves memory bandwidth vs FP16.
   - 4-bit KV cache for long-context efficiency.
   - Policy hook lives in runtime.affordability (per-model decision).

The interface here lets real adapters plug in. For heuristic adapters
(MVP), the routing is trivial (single expert) — but the contract is
in place so production adapters slot in without changing the substrate.

References:
- DeepSeek-V4-Pro:        https://huggingface.co/deepseek-ai/DeepSeek-V4-Pro
- DeepSeek-V4-Pro-DSpark: https://huggingface.co/deepseek-ai/DeepSeek-V4-Pro-DSpark
"""
from __future__ import annotations

import logging
from dataclasses import dataclass, field
from typing import Any

logger = logging.getLogger("openclinical.runtime.efficient")


# --- Expert routing (MoE seam) -------------------------------------------


@dataclass
class ExpertRoute:
    """A single expert in a MoE model — what it handles, when it's selected.

    In V4-Pro, there are thousands of these. For openclinical-ai's first
    real adapter, expect ~10 medical specialty experts. The router picks
    which subset handles a given input.
    """
    expert_id: str
    specialty: str  # clinical specialty or model behavior
    handles: list[str]  # input patterns / query types this expert handles
    estimated_params_b: float = 0.0  # parameter count in billions (for cost)


class ExpertRouter:
    """Selects which expert(s) handle a given input.

    V4-Pro style: routes among many fine-grained experts (learned router).
    Heuristic MVP: routes based on keyword/pattern matching against
    `ExpertRoute.handles`. Future: learned router trained alongside
    the experts.

    The MVP returns all experts (no real selection yet). The seam is in
    place; real selection lands when real adapters ship.
    """

    def __init__(self, experts: list[ExpertRoute]) -> None:
        self.experts = experts

    def route(self, inputs: dict[str, Any]) -> list[ExpertRoute]:
        """Return the experts that should handle this input.

        MVP heuristic: returns all experts. This is the seam — real
        routing logic plugs in here when production adapters ship.

        Production routing will:
        1. Score each expert's `handles` against the input content
        2. Select top-k experts (typically 1-2)
        3. Apply auxiliary-loss-free load balancing (DeepSeek-style)
        """
        # MVP: return all experts. Seams in place, real routing is
        # a v2 deliverable alongside real adapters.
        if not self.experts:
            return []
        return list(self.experts)

    def activated_params_b(self, selected: list[ExpertRoute]) -> float:
        """Sum activated parameters across the selected experts."""
        return sum(e.estimated_params_b for e in selected)


# Default medical-specialty expert set (matches what a real clinical MoE
# adapter would expose). The MVP heuristic adapter doesn't use these —
# they're here so real adapters slot in without changing the substrate.
DEFAULT_MEDICAL_EXPERTS: list[ExpertRoute] = [
    ExpertRoute(
        expert_id="clinical-reasoning",
        specialty="General clinical reasoning",
        handles=["assessment", "diagnosis support", "care plan"],
        estimated_params_b=8.0,
    ),
    ExpertRoute(
        expert_id="drug-interaction",
        specialty="Pharmacology",
        handles=["medication", "drug interaction", "polypharmacy", "rx"],
        estimated_params_b=4.0,
    ),
    ExpertRoute(
        expert_id="variant-impact",
        specialty="Clinical genomics",
        handles=["variant", "genomic", "rare disease", "vcf"],
        estimated_params_b=6.0,
    ),
    ExpertRoute(
        expert_id="geriatrics",
        specialty="Geriatric medicine",
        handles=["elderly", "dementia", "falls", "ltc", "psw"],
        estimated_params_b=5.0,
    ),
    ExpertRoute(
        expert_id="mental-health",
        specialty="Mental health",
        handles=["depression", "anxiety", "mood", "suicide"],
        estimated_params_b=4.0,
    ),
    ExpertRoute(
        expert_id="pediatrics",
        specialty="Pediatrics",
        handles=["child", "infant", "neonate", "vaccination"],
        estimated_params_b=4.0,
    ),
    ExpertRoute(
        expert_id="emergency",
        specialty="Emergency / triage",
        handles=["triage", "acute", "trauma", "code"],
        estimated_params_b=5.0,
    ),
    ExpertRoute(
        expert_id="cardiology",
        specialty="Cardiology",
        handles=["heart", "bp", "blood pressure", "arrhythmia"],
        estimated_params_b=4.0,
    ),
    ExpertRoute(
        expert_id="oncology",
        specialty="Oncology",
        handles=["cancer", "tumor", "chemo", "metastasis"],
        estimated_params_b=5.0,
    ),
    ExpertRoute(
        expert_id="infectious-disease",
        specialty="Infectious disease",
        handles=["infection", "antibiotic", "sepsis", "outbreak"],
        estimated_params_b=4.0,
    ),
]


# --- Context compression (CSA + HCA seam) --------------------------------


@dataclass
class CompressedContext:
    """A compressed representation of long context — CSA + HCA inspired.

    In production, V4-Pro's hybrid attention produces these natively.
    For MVP heuristic adapters, we provide a simple truncation +
    key-passage extraction so the cost/affordability story holds even
    before real models plug in.
    """
    original_tokens: int
    compressed_tokens: int
    compression_ratio: float
    method: str  # csa | hca | mixed
    semantic_anchors: list[str]  # key passages preserved verbatim

    def to_dict(self) -> dict[str, Any]:
        return {
            "original_tokens": self.original_tokens,
            "compressed_tokens": self.compressed_tokens,
            "compression_ratio": self.compression_ratio,
            "method": self.method,
            "semantic_anchors": self.semantic_anchors,
        }


class ContextCompressor:
    """Compresses long context into a compact representation.

    In production: V4-Pro's CSA/HCA hybrid attention produces these.
    For MVP heuristic adapters: simple truncation + key-passage extraction
    (the contract is right, the implementation is a placeholder until
    real models plug in).

    Healthcare use case:
    - Patient's 5-year medical history = ~500K tokens raw
    - Compressed to 32K tokens preserves diagnoses, medications, allergies
    - Inference runs at ~27% the cost of full-context inference
    """

    def compress(
        self,
        context: str,
        max_tokens: int,
    ) -> CompressedContext:
        """Compress `context` to <= max_tokens, preserving semantic anchors.

        MVP implementation: tokenize roughly (split on whitespace), pick
        the last `max_tokens` tokens plus any line containing clinical
        keywords (diagnoses, medications, allergies, vitals).
        """
        if not context:
            return CompressedContext(
                original_tokens=0,
                compressed_tokens=0,
                compression_ratio=1.0,
                method="csa",
                semantic_anchors=[],
            )

        tokens = context.split()
        original_tokens = len(tokens)

        if original_tokens <= max_tokens:
            return CompressedContext(
                original_tokens=original_tokens,
                compressed_tokens=original_tokens,
                compression_ratio=1.0,
                method="csa",
                semantic_anchors=[],
            )

        # Anchor keywords — lines containing these are preserved verbatim
        anchor_keywords = [
            "diagnosis", "diagnosed", "allergy", "allergic",
            "medication", "medications", "prescribed",
            "adverse", "reaction",
            "code", "code blue", "code white",
            "fall", "fell",
            "incident",
            "emergency",
            "transfer", "transferred",
            "discharge",
            "death", "died", "deceased",
        ]

        semantic_anchors = []
        for line in context.split("\n"):
            line_lower = line.lower()
            if any(kw in line_lower for kw in anchor_keywords):
                semantic_anchors.append(line.strip())
                if len(semantic_anchors) >= 20:
                    break

        # Reserve space for anchors
        anchor_tokens = sum(len(a.split()) for a in semantic_anchors)
        remaining = max(0, max_tokens - anchor_tokens)

        # Tail-truncate to remaining budget (most recent context is most
        # relevant for clinical reasoning — last visit, recent vitals)
        tail = tokens[-remaining:] if remaining > 0 else []

        compressed_tokens = anchor_tokens + len(tail)

        # Method selection: CSA for modest compression, HCA for aggressive
        ratio = compressed_tokens / original_tokens if original_tokens else 1.0
        method = "hca" if ratio < 0.1 else "csa" if ratio < 0.5 else "mixed"

        return CompressedContext(
            original_tokens=original_tokens,
            compressed_tokens=compressed_tokens,
            compression_ratio=round(ratio, 4),
            method=method,
            semantic_anchors=semantic_anchors,
        )


# --- Tiered model adapter (the substrate seam for V4-Pro / V4-Flash) -----


@dataclass
class TieredModelSpec:
    """Specification for a model at a particular tier.

    One model can have multiple tiers — each tier maps to a different
    activated-parameter count, quantization, and pricing.
    """
    model_id: str
    tier_id: str
    activated_params_b: float  # V4-Pro = 49, V4-Flash = 13
    quantization: str  # fp8 | fp16 | bf16
    context_length: int
    cost_input_usd_per_m: float
    cost_output_usd_per_m: float


class TieredModelAdapter:
    """Base class for adapters that support multiple tiers.

    Real adapters (V4-Pro, V4-Flash) override this. The MVP heuristic
    PSW adapter doesn't use it — but the contract is in place so
    production adapters slot in cleanly.

    The seam:
    - Each adapter declares its tier specs
    - The router picks the right tier based on tenant policy + request
    - Cost is computed from the tier spec, not hardcoded
    """

    def __init__(self, specs: list[TieredModelSpec]) -> None:
        self.specs = specs
        self._by_tier: dict[str, TieredModelSpec] = {s.tier_id: s for s in specs}

    def spec_for(self, tier_id: str) -> TieredModelSpec | None:
        """Return the spec for a given tier, or None if not supported."""
        return self._by_tier.get(tier_id)

    async def run(self, inputs: dict[str, Any]) -> dict[str, Any]:
        """Run inference. Override in subclasses.

        The base class raises — real adapters must implement.
        """
        raise NotImplementedError


def default_router() -> ExpertRouter:
    """Return a router configured with the default medical-specialty set."""
    return ExpertRouter(DEFAULT_MEDICAL_EXPERTS)


def default_compressor() -> ContextCompressor:
    """Return a default context compressor (CSA+HCA seam)."""
    return ContextCompressor()