first
This commit is contained in:
valenti
@@ -0,0 +1,185 @@
|
||||
"""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
|
||||
Reference in New Issue
Block a user