backend/capabilities.py

from __future__ import annotations
import json
from dataclasses import dataclass
from typing import Optional
import config
STATIC_CAPABILITIES: dict[str, dict] = {'gate.webrtcvad': {'available': True, 'reason': None, 'provenance': 'real_checkpoint'}, 'gate.silero_vad': {'available': True, 'reason': None, 'provenance': 'real_checkpoint'}, 'gate.none': {'available': True, 'reason': None, 'provenance': 'rule'}, 'acoustic.smart_turn_onnx': {'available': True, 'reason': None, 'provenance': 'real_checkpoint'}, 'semantic.qwen_local': {'available': True, 'reason': None, 'provenance': 'real_checkpoint'}, 'semantic.livekit_eou': {'available': True, 'reason': None, 'provenance': 'real_checkpoint'}, 'semantic.qwen_local_streaming': {'available': True, 'reason': None, 'provenance': 'architecture_reimplemented'}, 'semantic.groq_api': {'available': False, 'reason': 'no API key configured - set GROQ_API_KEY to enable (see docs/decision-log.md #17)', 'provenance': 'unavailable'}, 'semantic.openrouter_api': {'available': False, 'reason': 'no API key configured - set OPENROUTER_API_KEY to enable (see docs/decision-log.md #17)', 'provenance': 'unavailable'}, 'fusion.weighted_vote': {'available': True, 'reason': None, 'provenance': 'rule'}, 'fusion.easy_turn': {'available': False, 'reason': "Easy Turn's linguistic branch needs its own ASR component we haven't wired yet - checkpoint downloaded for reference only (see docs/decision-log.md #19)", 'provenance': 'unavailable'}, 'mode.full_duplex_bypass.moshi': {'available': False, 'reason': 'needs GPU VRAM not available on this machine; API routing deferred (see docs/decision-log.md #6, #17)', 'provenance': 'unavailable'}, 'mode.full_duplex_bypass.human1': {'available': False, 'reason': 'needs GPU VRAM not available on this machine; API routing deferred (see docs/decision-log.md #6, #17)', 'provenance': 'unavailable'}}

@dataclass
class CapabilityInfo:
    key: str
    available: bool
    reason: Optional[str]
    provenance: Optional[str]

def _head_checkpoint_path(encoder: str, pooling: str, head: str) -> tuple:
    stem = f'{encoder}_{pooling}_{head}'
    ckpt = config.CHECKPOINTS_DIR / f'{stem}.pt'
    meta = config.CHECKPOINTS_DIR / f'{stem}.metadata.json'
    return (ckpt, meta)

def head_capability_key(encoder: str, pooling: str, head: str) -> str:
    return f'acoustic.head.{encoder}.{pooling}.{head}'

def _check_trained_head(encoder: str, pooling: str, head: str) -> CapabilityInfo:
    key = head_capability_key(encoder, pooling, head)
    ckpt, meta = _head_checkpoint_path(encoder, pooling, head)
    if not (ckpt.exists() and meta.exists()):
        return CapabilityInfo(key=key, available=False, reason=f'not trained yet - run experiments/train_head.py to produce database/checkpoints/{encoder}_{pooling}_{head}.pt', provenance='unavailable')
    try:
        metadata = json.loads(meta.read_text())
    except (OSError, json.JSONDecodeError) as exc:
        return CapabilityInfo(key=key, available=False, reason=f'checkpoint metadata unreadable ({exc}) - retrain via experiments/train_head.py', provenance='unavailable')
    if metadata.get('encoder') != encoder or metadata.get('pooling') != pooling or metadata.get('head') != head:
        return CapabilityInfo(key=key, available=False, reason='checkpoint metadata does not match the requested (encoder, pooling, head) combination', provenance='unavailable')
    return CapabilityInfo(key=key, available=True, reason=None, provenance='trained_by_us')

def get(key: str) -> CapabilityInfo:
    if key.startswith('acoustic.head.'):
        _, _, encoder, pooling, head = key.split('.')
        return _check_trained_head(encoder, pooling, head)
    if key in STATIC_CAPABILITIES:
        entry = STATIC_CAPABILITIES[key]
        return CapabilityInfo(key=key, **entry)
    raise KeyError(f'unknown capability key: {key!r}')

def is_available(key: str) -> bool:
    return get(key).available

def list_trained_heads() -> list[CapabilityInfo]:
    found = []
    if not config.CHECKPOINTS_DIR.exists():
        return found
    for meta_path in sorted(config.CHECKPOINTS_DIR.glob('*.metadata.json')):
        try:
            metadata = json.loads(meta_path.read_text())
        except (OSError, json.JSONDecodeError):
            continue
        encoder, pooling, head = (metadata.get('encoder'), metadata.get('pooling'), metadata.get('head'))
        if not all([encoder, pooling, head]):
            continue
        info = _check_trained_head(encoder, pooling, head)
        if info.available:
            found.append(info)
    return found