asf/services/asf-ai/app/gateway.py

"""AIGateway — the high-level capability surface called by business code.

It handles:
  1. Routing: picks which provider + model to use given the request.
  2. Policy enforcement: tenant-level AI policy (which providers/models are
     allowed; budget caps will plug in here cleanly later).
  3. Usage recording: emits UsageRecords for metering + future billing.
  4. Error translation (provider-specific → stable ASF exceptions).
"""

from __future__ import annotations

from dataclasses import dataclass, field
from datetime import datetime, timezone
from typing import Callable, Optional

import structlog

from .capabilities import (
    ChatRequest,
    ChatResponse,
    CodeAnalysisRequest,
    CodeAnalysisResult,
    TestGenerationRequest,
    GeneratedTests,
    RequirementDecomposeRequest,
    RequirementDecomposition,
)
from .providers import (
    AIProvider,
    ProviderRegistry,
    default_registry,
)
from .usage import UsageRecord

log = structlog.get_logger(__name__)


@dataclass
class AIGatewayConfig:
    default_provider: str = "mock"
    default_model_by_capability: dict[str, str] = field(default_factory=dict)
    # Hook for future tenant policy layer: given (tenant_id, capability),
    # return the provider name to route to. None → use default_provider.
    policy_selector: Optional[Callable[..., Optional[str]]] = None


class AIGateway:
    """Entry point for all ASF AI operations."""

    def __init__(
        self,
        *,
        config: Optional[AIGatewayConfig] = None,
        registry: Optional[ProviderRegistry] = None,
    ) -> None:
        self._config = config or AIGatewayConfig()
        self._registry = registry or default_registry()
        self._usage_callbacks: list[Callable[[UsageRecord], None]] = []

    # --- callback hooks (future usage metering / billing) ---
    def on_usage(self, cb: Callable[[UsageRecord], None]) -> None:
        self._usage_callbacks.append(cb)

    # --- routing helper ---
    def _route(self, req, capability: str) -> AIProvider:
        tenant_id = str(req.trace.tenant_id) if req.trace else None
        chosen: Optional[str] = None
        if self._config.policy_selector and tenant_id:
            try:
                chosen = self._config.policy_selector(
                    tenant_id=req.trace.tenant_id,
                    capability=capability,
                    preferred=req.preferred_model,
                )
            except Exception:  # noqa: BLE001
                log.exception("ai_policy_selector_failed")

        chosen = chosen or self._config.default_provider
        return self._registry.get(chosen)

    def _emit_usage(self, record: UsageRecord) -> None:
        for cb in self._usage_callbacks:
            try:
                cb(record)
            except Exception:  # noqa: BLE001
                log.exception("ai_usage_callback_failed", provider=record.provider)

    # ------------------------------------------------------------------
    # Capability methods
    # ------------------------------------------------------------------

    async def chat(self, req: ChatRequest) -> ChatResponse:
        provider = self._route(req, "chat")
        log.debug("ai_chat_start", provider=provider.name, tenant=str(req.trace.tenant_id))
        resp = await provider.chat(req)
        self._emit_usage(
            UsageRecord(
                tenant_id=req.trace.tenant_id,
                project_id=req.trace.project_id,
                job_id=req.trace.job_id,
                execution_id=resp.execution_id,
                provider=resp.provider,
                model=resp.model_used,
                capability="chat",
                input_tokens=resp.input_tokens,
                output_tokens=resp.output_tokens,
                estimated_cost=resp.estimated_cost,
                duration_ms=resp.duration_ms,
                recorded_at=datetime.now(tz=timezone.utc),
            )
        )
        return resp

    async def analyze_code(self, req: CodeAnalysisRequest) -> CodeAnalysisResult:
        provider = self._route(req, "analyze_code")
        resp = await provider.analyze_code(req)
        self._emit_usage(
            UsageRecord(
                tenant_id=req.trace.tenant_id,
                project_id=req.trace.project_id,
                job_id=req.trace.job_id,
                execution_id=resp.execution_id,
                provider=resp.provider,
                model=resp.model_used,
                capability="analyze_code",
                input_tokens=resp.input_tokens,
                output_tokens=resp.output_tokens,
                estimated_cost=resp.estimated_cost,
                duration_ms=resp.duration_ms,
                recorded_at=datetime.now(tz=timezone.utc),
            )
        )
        return resp

    async def generate_tests(self, req: TestGenerationRequest) -> GeneratedTests:
        provider = self._route(req, "generate_tests")
        resp = await provider.generate_tests(req)
        self._emit_usage(
            UsageRecord(
                tenant_id=req.trace.tenant_id,
                project_id=req.trace.project_id,
                job_id=req.trace.job_id,
                execution_id=resp.execution_id,
                provider=resp.provider,
                model=resp.model_used,
                capability="generate_tests",
                input_tokens=resp.input_tokens,
                output_tokens=resp.output_tokens,
                estimated_cost=resp.estimated_cost,
                duration_ms=resp.duration_ms,
                recorded_at=datetime.now(tz=timezone.utc),
            )
        )
        return resp

    async def decompose_requirement(
        self, req: RequirementDecomposeRequest
    ) -> RequirementDecomposition:
        provider = self._route(req, "decompose_requirement")
        resp = await provider.decompose_requirement(req)
        self._emit_usage(
            UsageRecord(
                tenant_id=req.trace.tenant_id,
                project_id=req.trace.project_id,
                job_id=req.trace.job_id,
                execution_id=resp.execution_id,
                provider=resp.provider,
                model=resp.model_used,
                capability="decompose_requirement",
                input_tokens=resp.input_tokens,
                output_tokens=resp.output_tokens,
                estimated_cost=resp.estimated_cost,
                duration_ms=resp.duration_ms,
                recorded_at=datetime.now(tz=timezone.utc),
            )
        )
        return resp


def build_default_gateway() -> AIGateway:
    """Default bootstrap-configured gateway — uses MockProvider, stdout usage logging."""
    gw = AIGateway()
    gw.on_usage(lambda record: log.debug("ai_usage_record", **record.model_dump(mode="json")))
    return gw