stonks-oracle/services/shared/inference/capabilities.py

"""Capability probing for inference endpoints.

Discovers actual endpoint capabilities by sending probe requests,
rather than relying solely on declared capabilities. Probe results
are cached with a configurable TTL and used to validate that
required capabilities work before routing real requests.

Requirements: 2.10, 2.11
"""
from __future__ import annotations

import json
import logging
import time
from datetime import datetime, timezone
from typing import Any
from uuid import UUID

import httpx
from pydantic import BaseModel, Field

from services.shared.inference.models import (
    ChatMessage,
    InferenceTarget,
)

logger = logging.getLogger(__name__)

# Default TTL for probe result cache (seconds)
DEFAULT_PROBE_TTL_SECONDS = 300  # 5 minutes


# ---------------------------------------------------------------------------
# Probe result models
# ---------------------------------------------------------------------------


class HealthProbeResult(BaseModel):
    """Result of a health probe against an endpoint."""

    success: bool
    detail: str = ""
    timestamp: datetime = Field(default_factory=lambda: datetime.now(timezone.utc))
    software_version: str | None = None
    latency_ms: int = 0


class ModelListingResult(BaseModel):
    """Result of probing the /v1/models endpoint."""

    success: bool
    detail: str = ""
    timestamp: datetime = Field(default_factory=lambda: datetime.now(timezone.utc))
    software_version: str | None = None
    models: list[str] = Field(default_factory=list)
    target_model_found: bool = False


class JsonSchemaProbeResult(BaseModel):
    """Result of probing strict JSON Schema output."""

    success: bool
    detail: str = ""
    timestamp: datetime = Field(default_factory=lambda: datetime.now(timezone.utc))
    software_version: str | None = None
    schema_valid: bool = False
    structured_mode_used: str = ""


class UsageProbeResult(BaseModel):
    """Result of probing usage metadata return."""

    success: bool
    detail: str = ""
    timestamp: datetime = Field(default_factory=lambda: datetime.now(timezone.utc))
    software_version: str | None = None
    has_prompt_tokens: bool = False
    has_completion_tokens: bool = False


class SeedProbeResult(BaseModel):
    """Result of probing seed/determinism behavior."""

    success: bool
    detail: str = ""
    timestamp: datetime = Field(default_factory=lambda: datetime.now(timezone.utc))
    software_version: str | None = None
    outputs_match: bool = False


class OutputTokenFieldResult(BaseModel):
    """Result of probing max_completion_tokens field support."""

    success: bool
    detail: str = ""
    timestamp: datetime = Field(default_factory=lambda: datetime.now(timezone.utc))
    software_version: str | None = None
    field_accepted: bool = False


class FullProbeResult(BaseModel):
    """Aggregated results from all capability probes."""

    endpoint_id: UUID
    timestamp: datetime = Field(default_factory=lambda: datetime.now(timezone.utc))
    health: HealthProbeResult | None = None
    model_listing: ModelListingResult | None = None
    json_schema: JsonSchemaProbeResult | None = None
    usage_metadata: UsageProbeResult | None = None
    seed_determinism: SeedProbeResult | None = None
    output_token_field: OutputTokenFieldResult | None = None
    software_version: str | None = None
    probe_duration_ms: int = 0


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

_MINIMAL_PROBE_SCHEMA = {
    "title": "probe_response",
    "type": "object",
    "properties": {
        "status": {"type": "string"},
    },
    "required": ["status"],
}

_PROBE_MESSAGES = [
    ChatMessage(role="user", content="Respond with JSON: {\"status\": \"ok\"}"),
]


def _resolve_auth_headers(target: InferenceTarget) -> dict[str, str]:
    """Build auth headers for probe requests from the target."""
    import os

    headers: dict[str, str] = {"Content-Type": "application/json"}
    headers.update(target.extra_headers)

    if target.auth_secret_ref:
        secret = os.environ.get(target.auth_secret_ref)
        if secret:
            scheme = target.auth_scheme.lower()
            if scheme == "bearer":
                headers["Authorization"] = f"Bearer {secret}"
            else:
                headers[target.auth_scheme] = secret

    return headers


def _extract_software_version(response: httpx.Response) -> str | None:
    """Try to extract software/version from response headers."""
    # Common version headers across providers
    for header in ("x-vllm-version", "server", "x-server-version", "x-version"):
        value = response.headers.get(header)
        if value:
            return value
    return None


# ---------------------------------------------------------------------------
# EndpointProber
# ---------------------------------------------------------------------------


class EndpointProber:
    """Probes inference endpoints to discover actual capabilities.

    Each probe method sends a minimal request to verify that a declared
    capability actually works, returning structured results with timing
    and software version metadata.
    """

    def __init__(
        self,
        *,
        http_client: httpx.AsyncClient | None = None,
        probe_timeout: float = 15.0,
    ) -> None:
        self._owns_client = http_client is None
        self._http = http_client or httpx.AsyncClient(timeout=probe_timeout)
        self._probe_timeout = probe_timeout

    async def probe_health(self, target: InferenceTarget) -> HealthProbeResult:
        """Check if the endpoint is reachable and responsive.

        Tries the configured health_path or falls back to common paths:
        /health, /v1/models, or the base URL itself.
        """
        headers = _resolve_auth_headers(target)
        base = target.base_url.rstrip("/")

        # Try common health paths in order
        health_paths = ["/health", "/v1/models", "/"]
        start = time.monotonic()

        for path in health_paths:
            url = f"{base}{path}"
            try:
                response = await self._http.get(
                    url, headers=headers, timeout=self._probe_timeout
                )
                latency_ms = int((time.monotonic() - start) * 1000)
                sw_version = _extract_software_version(response)

                if response.status_code < 500:
                    return HealthProbeResult(
                        success=True,
                        detail=f"Endpoint reachable via {path} (HTTP {response.status_code})",
                        latency_ms=latency_ms,
                        software_version=sw_version,
                    )
            except (httpx.ConnectError, httpx.TimeoutException) as exc:
                latency_ms = int((time.monotonic() - start) * 1000)
                return HealthProbeResult(
                    success=False,
                    detail=f"Connection failed: {type(exc).__name__}: {exc}",
                    latency_ms=latency_ms,
                )

        latency_ms = int((time.monotonic() - start) * 1000)
        return HealthProbeResult(
            success=False,
            detail="All health paths returned 5xx",
            latency_ms=latency_ms,
        )

    async def probe_model_listing(
        self, target: InferenceTarget
    ) -> ModelListingResult:
        """Check if the /v1/models endpoint works and lists the target model."""
        headers = _resolve_auth_headers(target)
        base = target.base_url.rstrip("/")
        url = f"{base}/v1/models"

        try:
            response = await self._http.get(
                url, headers=headers, timeout=self._probe_timeout
            )
            sw_version = _extract_software_version(response)

            if response.status_code != 200:
                return ModelListingResult(
                    success=False,
                    detail=f"Model listing returned HTTP {response.status_code}",
                    software_version=sw_version,
                )

            data = response.json()
            models_data = data.get("data", [])
            model_ids = [m.get("id", "") for m in models_data if isinstance(m, dict)]
            target_found = target.model in model_ids

            return ModelListingResult(
                success=True,
                detail=f"Found {len(model_ids)} model(s)",
                software_version=sw_version,
                models=model_ids,
                target_model_found=target_found,
            )
        except (httpx.ConnectError, httpx.TimeoutException) as exc:
            return ModelListingResult(
                success=False,
                detail=f"Connection failed: {type(exc).__name__}: {exc}",
            )
        except (json.JSONDecodeError, ValueError) as exc:
            return ModelListingResult(
                success=False,
                detail=f"Invalid JSON in model listing response: {exc}",
            )

    async def probe_json_schema(
        self, target: InferenceTarget
    ) -> JsonSchemaProbeResult:
        """Send a minimal schema-constrained request and verify the response validates."""
        headers = _resolve_auth_headers(target)
        base = target.base_url.rstrip("/")
        url = f"{base}/v1/chat/completions"

        body: dict[str, Any] = {
            "model": target.model,
            "messages": [{"role": m.role, "content": m.content} for m in _PROBE_MESSAGES],
            "temperature": 0,
            "max_tokens": 64,
            "response_format": {
                "type": "json_schema",
                "json_schema": {
                    "name": "probe_response",
                    "strict": True,
                    "schema": _MINIMAL_PROBE_SCHEMA,
                },
            },
        }

        try:
            response = await self._http.post(
                url, json=body, headers=headers, timeout=self._probe_timeout
            )
            sw_version = _extract_software_version(response)

            if response.status_code != 200:
                return JsonSchemaProbeResult(
                    success=False,
                    detail=f"JSON schema probe returned HTTP {response.status_code}",
                    software_version=sw_version,
                    structured_mode_used="json_schema",
                )

            data = response.json()
            choices = data.get("choices", [])
            if not choices:
                return JsonSchemaProbeResult(
                    success=False,
                    detail="Empty choices in schema probe response",
                    software_version=sw_version,
                    structured_mode_used="json_schema",
                )

            content = choices[0].get("message", {}).get("content", "")
            try:
                parsed = json.loads(content)
                # Validate against the minimal schema
                is_valid = (
                    isinstance(parsed, dict)
                    and "status" in parsed
                    and isinstance(parsed["status"], str)
                )
                return JsonSchemaProbeResult(
                    success=is_valid,
                    detail="Schema probe response validates" if is_valid else "Response did not match expected schema",
                    software_version=sw_version,
                    schema_valid=is_valid,
                    structured_mode_used="json_schema",
                )
            except (json.JSONDecodeError, ValueError):
                return JsonSchemaProbeResult(
                    success=False,
                    detail=f"Response content is not valid JSON: {content[:100]}",
                    software_version=sw_version,
                    schema_valid=False,
                    structured_mode_used="json_schema",
                )
        except (httpx.ConnectError, httpx.TimeoutException) as exc:
            return JsonSchemaProbeResult(
                success=False,
                detail=f"Connection failed: {type(exc).__name__}: {exc}",
            )

    async def probe_usage_metadata(
        self, target: InferenceTarget
    ) -> UsageProbeResult:
        """Check if usage tokens (prompt_tokens, completion_tokens) are returned."""
        headers = _resolve_auth_headers(target)
        base = target.base_url.rstrip("/")
        url = f"{base}/v1/chat/completions"

        body: dict[str, Any] = {
            "model": target.model,
            "messages": [{"role": m.role, "content": m.content} for m in _PROBE_MESSAGES],
            "temperature": 0,
            "max_tokens": 32,
        }

        try:
            response = await self._http.post(
                url, json=body, headers=headers, timeout=self._probe_timeout
            )
            sw_version = _extract_software_version(response)

            if response.status_code != 200:
                return UsageProbeResult(
                    success=False,
                    detail=f"Usage probe returned HTTP {response.status_code}",
                    software_version=sw_version,
                )

            data = response.json()
            usage = data.get("usage", {})
            has_prompt = "prompt_tokens" in usage and usage["prompt_tokens"] is not None
            has_completion = "completion_tokens" in usage and usage["completion_tokens"] is not None

            return UsageProbeResult(
                success=has_prompt or has_completion,
                detail=f"Usage metadata: prompt_tokens={'yes' if has_prompt else 'no'}, completion_tokens={'yes' if has_completion else 'no'}",
                software_version=sw_version,
                has_prompt_tokens=has_prompt,
                has_completion_tokens=has_completion,
            )
        except (httpx.ConnectError, httpx.TimeoutException) as exc:
            return UsageProbeResult(
                success=False,
                detail=f"Connection failed: {type(exc).__name__}: {exc}",
            )

    async def probe_seed_determinism(
        self, target: InferenceTarget
    ) -> SeedProbeResult:
        """Send the same request twice with the same seed and check if outputs match."""
        headers = _resolve_auth_headers(target)
        base = target.base_url.rstrip("/")
        url = f"{base}/v1/chat/completions"

        body: dict[str, Any] = {
            "model": target.model,
            "messages": [{"role": "user", "content": "Say exactly: hello"}],
            "temperature": 0,
            "max_tokens": 16,
            "seed": 42,
        }

        outputs: list[str] = []
        sw_version: str | None = None

        try:
            for _ in range(2):
                response = await self._http.post(
                    url, json=body, headers=headers, timeout=self._probe_timeout
                )
                if not sw_version:
                    sw_version = _extract_software_version(response)

                if response.status_code != 200:
                    return SeedProbeResult(
                        success=False,
                        detail=f"Seed probe returned HTTP {response.status_code}",
                        software_version=sw_version,
                    )

                data = response.json()
                choices = data.get("choices", [])
                if not choices:
                    return SeedProbeResult(
                        success=False,
                        detail="Empty choices in seed probe response",
                        software_version=sw_version,
                    )
                content = choices[0].get("message", {}).get("content", "")
                outputs.append(content)

            match = outputs[0] == outputs[1]
            return SeedProbeResult(
                success=True,
                detail=f"Outputs {'match' if match else 'differ'} with same seed",
                software_version=sw_version,
                outputs_match=match,
            )
        except (httpx.ConnectError, httpx.TimeoutException) as exc:
            return SeedProbeResult(
                success=False,
                detail=f"Connection failed: {type(exc).__name__}: {exc}",
            )

    async def probe_output_token_field(
        self, target: InferenceTarget
    ) -> OutputTokenFieldResult:
        """Check if the endpoint accepts max_completion_tokens field.

        Some providers use max_completion_tokens instead of (or in addition to)
        max_tokens. This probe tests if the field is accepted without error.
        """
        headers = _resolve_auth_headers(target)
        base = target.base_url.rstrip("/")
        url = f"{base}/v1/chat/completions"

        body: dict[str, Any] = {
            "model": target.model,
            "messages": [{"role": "user", "content": "Say hi"}],
            "temperature": 0,
            "max_completion_tokens": 16,
        }

        try:
            response = await self._http.post(
                url, json=body, headers=headers, timeout=self._probe_timeout
            )
            sw_version = _extract_software_version(response)

            # If the endpoint accepts the field, it should respond 200
            # If it rejects it, it will typically return 400 or 422
            accepted = response.status_code == 200
            detail = (
                "max_completion_tokens field accepted"
                if accepted
                else f"max_completion_tokens field rejected (HTTP {response.status_code})"
            )

            return OutputTokenFieldResult(
                success=accepted,
                detail=detail,
                software_version=sw_version,
                field_accepted=accepted,
            )
        except (httpx.ConnectError, httpx.TimeoutException) as exc:
            return OutputTokenFieldResult(
                success=False,
                detail=f"Connection failed: {type(exc).__name__}: {exc}",
            )

    async def run_full_probe(self, target: InferenceTarget) -> FullProbeResult:
        """Run all probes against a target and return aggregated results."""
        start = time.monotonic()

        health = await self.probe_health(target)
        if not health.success:
            # If health fails, skip remaining probes
            duration_ms = int((time.monotonic() - start) * 1000)
            return FullProbeResult(
                endpoint_id=target.endpoint_id,
                health=health,
                software_version=health.software_version,
                probe_duration_ms=duration_ms,
            )

        model_listing = await self.probe_model_listing(target)
        json_schema_result = await self.probe_json_schema(target)
        usage_result = await self.probe_usage_metadata(target)
        seed_result = await self.probe_seed_determinism(target)
        output_token_result = await self.probe_output_token_field(target)

        duration_ms = int((time.monotonic() - start) * 1000)

        # Determine software version from whichever probe returned one
        sw_version = (
            health.software_version
            or model_listing.software_version
            or json_schema_result.software_version
            or usage_result.software_version
            or seed_result.software_version
            or output_token_result.software_version
        )

        return FullProbeResult(
            endpoint_id=target.endpoint_id,
            health=health,
            model_listing=model_listing,
            json_schema=json_schema_result,
            usage_metadata=usage_result,
            seed_determinism=seed_result,
            output_token_field=output_token_result,
            software_version=sw_version,
            probe_duration_ms=duration_ms,
        )

    async def close(self) -> None:
        """Close the HTTP client if we own it."""
        if self._owns_client:
            await self._http.aclose()


# ---------------------------------------------------------------------------
# ProbeResultStore — TTL-based cache for probe results
# ---------------------------------------------------------------------------


class ProbeResultStore:
    """In-memory cache for endpoint probe results with configurable TTL.

    Uses a simple dict with timestamps. Expired entries are lazily evicted
    on access. The TTL can be configured globally or per-store instance.
    """

    def __init__(self, ttl_seconds: float = DEFAULT_PROBE_TTL_SECONDS) -> None:
        self._ttl_seconds = ttl_seconds
        self._store: dict[UUID, tuple[float, FullProbeResult]] = {}

    @property
    def ttl_seconds(self) -> float:
        """Current TTL configuration."""
        return self._ttl_seconds

    def store(self, endpoint_id: UUID, result: FullProbeResult) -> None:
        """Store a probe result with the current timestamp."""
        self._store[endpoint_id] = (time.monotonic(), result)

    def get(self, endpoint_id: UUID) -> FullProbeResult | None:
        """Retrieve a probe result if it exists and hasn't expired.

        Returns None if no result is stored or if the TTL has elapsed.
        """
        entry = self._store.get(endpoint_id)
        if entry is None:
            return None

        stored_at, result = entry
        elapsed = time.monotonic() - stored_at
        if elapsed > self._ttl_seconds:
            # Expired — evict lazily
            del self._store[endpoint_id]
            return None

        return result

    def invalidate(self, endpoint_id: UUID) -> None:
        """Force cache eviction for an endpoint."""
        self._store.pop(endpoint_id, None)

    def clear(self) -> None:
        """Clear all cached results."""
        self._store.clear()

    def __len__(self) -> int:
        """Number of entries (including possibly expired ones)."""
        return len(self._store)


# ---------------------------------------------------------------------------
# Required capability validation
# ---------------------------------------------------------------------------


def validate_required_capabilities(
    target: InferenceTarget,
    probe_result: FullProbeResult,
) -> list[str]:
    """Validate that declared required capabilities actually work.

    Returns a list of failure descriptions. If the list is non-empty,
    the endpoint should NOT be activated for routing.

    Checks:
    - Health must succeed
    - If chat_completions is declared, health must pass
    - If json_schema is declared, json_schema probe must succeed and validate
    - If usage is declared, usage probe must show token counts
    - If seed is declared, seed probe must show matching outputs
    - If max_completion_tokens is declared, output_token_field probe must pass
    - If model_listing is declared, model listing must succeed and find the model
    """
    failures: list[str] = []
    caps = target.capabilities

    # Health is always required
    if probe_result.health is None or not probe_result.health.success:
        detail = probe_result.health.detail if probe_result.health else "No health probe result"
        failures.append(f"Health check failed: {detail}")
        # If health fails, all other checks are meaningless
        return failures

    # Model listing capability
    if caps.model_listing:
        if probe_result.model_listing is None or not probe_result.model_listing.success:
            detail = probe_result.model_listing.detail if probe_result.model_listing else "Not probed"
            failures.append(f"Model listing failed: {detail}")
        elif not probe_result.model_listing.target_model_found:
            failures.append(
                f"Model '{target.model}' not found in endpoint model list"
            )

    # JSON Schema capability
    if caps.json_schema:
        if probe_result.json_schema is None or not probe_result.json_schema.success:
            detail = probe_result.json_schema.detail if probe_result.json_schema else "Not probed"
            failures.append(f"JSON Schema structured output failed: {detail}")
        elif not probe_result.json_schema.schema_valid:
            failures.append(
                "JSON Schema probe succeeded but response did not validate"
            )

    # Usage metadata capability
    if caps.usage:
        if probe_result.usage_metadata is None or not probe_result.usage_metadata.success:
            detail = probe_result.usage_metadata.detail if probe_result.usage_metadata else "Not probed"
            failures.append(f"Usage metadata not available: {detail}")

    # Seed determinism capability
    if caps.seed:
        if probe_result.seed_determinism is None or not probe_result.seed_determinism.success:
            detail = probe_result.seed_determinism.detail if probe_result.seed_determinism else "Not probed"
            failures.append(f"Seed determinism probe failed: {detail}")
        elif not probe_result.seed_determinism.outputs_match:
            failures.append(
                "Seed declared but outputs differ with same seed"
            )

    # max_completion_tokens capability
    if caps.max_completion_tokens:
        if probe_result.output_token_field is None or not probe_result.output_token_field.success:
            detail = probe_result.output_token_field.detail if probe_result.output_token_field else "Not probed"
            failures.append(f"max_completion_tokens field not supported: {detail}")

    return failures