1438 lines
58 KiB
Python
1438 lines
58 KiB
Python
"""Durable Model Governance for AIOS external Agent runtimes.
|
|
|
|
Single-process/single-SQLite governance layer:
|
|
- provider/model registry without raw secrets;
|
|
- tenant/project/system scope policies;
|
|
- evaluated + version-pinned deterministic routing;
|
|
- atomic budget reservations and usage settlement;
|
|
- provider circuit/fallback;
|
|
- idempotent request decisions;
|
|
- tamper-evident audit chain.
|
|
|
|
The database stores references and accounting only. Raw provider credentials are forbidden.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
from dataclasses import asdict, dataclass
|
|
from datetime import datetime, timezone
|
|
import hashlib
|
|
import json
|
|
import re
|
|
import sqlite3
|
|
import threading
|
|
import time
|
|
from pathlib import Path
|
|
from typing import Any, Callable, Mapping, Sequence
|
|
|
|
MODEL_ID_RE = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._:/-]{1,159}$")
|
|
PROVIDER_ID_RE = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._-]{0,79}$")
|
|
TASK_ID_RE = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._:-]{2,179}$")
|
|
REQUEST_ID_RE = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._:-]{7,254}$")
|
|
TRACE_ID_RE = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._:-]{2,254}$")
|
|
SHA256_RE = re.compile(r"^[0-9a-f]{64}$")
|
|
SECRET_RE = re.compile(
|
|
r"(?i)(authorization\s*[:=]\s*bearer\s+\S+|"
|
|
r"(?:password|passwd|secret|token|api[_-]?key)\s*[:=]\s*\S+)"
|
|
)
|
|
CREDENTIAL_REF_PREFIXES = (
|
|
"systemd-credential:",
|
|
"runtime-profile:",
|
|
"credential-broker:",
|
|
)
|
|
APPROVED = "APPROVED"
|
|
DISABLED = "DISABLED"
|
|
|
|
|
|
class ModelGovernanceError(RuntimeError):
|
|
def __init__(self, code: str, *, details: Mapping[str, Any] | None = None) -> None:
|
|
super().__init__(code)
|
|
self.code = code
|
|
self.details = dict(details or {})
|
|
|
|
|
|
def _canonical(value: object) -> bytes:
|
|
return json.dumps(
|
|
value, ensure_ascii=False, sort_keys=True, separators=(",", ":")
|
|
).encode("utf-8")
|
|
|
|
|
|
def _sha(value: object) -> str:
|
|
return hashlib.sha256(_canonical(value)).hexdigest()
|
|
|
|
|
|
def _now_iso(now: float | None = None) -> str:
|
|
return datetime.fromtimestamp(
|
|
time.time() if now is None else now,
|
|
timezone.utc,
|
|
).isoformat()
|
|
|
|
|
|
def _day_key(now: float) -> str:
|
|
return datetime.fromtimestamp(now, timezone.utc).strftime("%Y-%m-%d")
|
|
|
|
|
|
def _text(value: object, code: str, pattern: re.Pattern[str] | None = None) -> str:
|
|
if not isinstance(value, str) or not value.strip():
|
|
raise ModelGovernanceError(code)
|
|
value = value.strip()
|
|
if pattern is not None and pattern.fullmatch(value) is None:
|
|
raise ModelGovernanceError(code)
|
|
return value
|
|
|
|
|
|
def _json_text(value: object) -> str:
|
|
return json.dumps(value, ensure_ascii=False, sort_keys=True, separators=(",", ":"))
|
|
|
|
|
|
def _json_list(value: object, code: str) -> tuple[str, ...]:
|
|
if value is None:
|
|
return ()
|
|
if not isinstance(value, Sequence) or isinstance(value, (str, bytes)):
|
|
raise ModelGovernanceError(code)
|
|
result = tuple(sorted({_text(item, code) for item in value}))
|
|
return result
|
|
|
|
|
|
def _scope_key(
|
|
scope_type: str,
|
|
tenant_id: str | None,
|
|
project_id: str | None,
|
|
) -> str:
|
|
scope_type = _text(scope_type, "MODEL_SCOPE_TYPE_INVALID")
|
|
tenant = (tenant_id or "").strip()
|
|
project = (project_id or "").strip()
|
|
if scope_type == "system":
|
|
if tenant or project:
|
|
raise ModelGovernanceError("MODEL_SCOPE_SYSTEM_FIELDS_DENY")
|
|
return "system"
|
|
if scope_type == "legacy":
|
|
if tenant or project:
|
|
raise ModelGovernanceError("MODEL_SCOPE_LEGACY_FIELDS_DENY")
|
|
return "legacy"
|
|
if scope_type != "tenant":
|
|
raise ModelGovernanceError("MODEL_SCOPE_TYPE_INVALID")
|
|
if not tenant:
|
|
raise ModelGovernanceError("MODEL_SCOPE_TENANT_REQUIRED")
|
|
if "/" in tenant or "/" in project:
|
|
raise ModelGovernanceError("MODEL_SCOPE_KEY_INVALID")
|
|
return "tenant:" + tenant + ("/project:" + project if project else "")
|
|
|
|
|
|
def _credential_ref(value: str | None) -> str | None:
|
|
if value is None:
|
|
return None
|
|
value = _text(value, "MODEL_CREDENTIAL_REF_INVALID")
|
|
if SECRET_RE.search(value):
|
|
raise ModelGovernanceError("MODEL_CREDENTIAL_SECRET_EXPOSURE_DENY")
|
|
if not value.startswith(CREDENTIAL_REF_PREFIXES):
|
|
raise ModelGovernanceError("MODEL_CREDENTIAL_REF_INVALID")
|
|
return value
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class ProviderRecord:
|
|
provider_id: str
|
|
credential_ref: str | None = None
|
|
enabled: bool = True
|
|
failure_domain: str = "default"
|
|
max_failures: int = 3
|
|
cooldown_seconds: int = 60
|
|
revision: int = 1
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class ModelRecord:
|
|
model_id: str
|
|
provider_id: str
|
|
version: str
|
|
capabilities: tuple[str, ...]
|
|
input_cost_micros_per_1k: int = 0
|
|
output_cost_micros_per_1k: int = 0
|
|
max_input_tokens: int = 0
|
|
max_output_tokens: int = 0
|
|
evaluation_state: str = APPROVED
|
|
evaluation_sha256: str = "0" * 64
|
|
enabled: bool = True
|
|
priority: int = 100
|
|
revision: int = 1
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class ScopePolicy:
|
|
scope_type: str
|
|
tenant_id: str | None
|
|
project_id: str | None
|
|
allowed_models: tuple[str, ...] = ()
|
|
denied_providers: tuple[str, ...] = ()
|
|
required_capabilities: tuple[str, ...] = ()
|
|
pinned_model_id: str | None = None
|
|
pinned_model_version: str | None = None
|
|
max_input_tokens_per_request: int = 0
|
|
max_output_tokens_per_request: int = 0
|
|
daily_token_budget: int = 0
|
|
daily_cost_budget_micros: int = 0
|
|
fallback_limit: int = 2
|
|
revision: int = 1
|
|
|
|
@property
|
|
def scope_key(self) -> str:
|
|
return _scope_key(self.scope_type, self.tenant_id, self.project_id)
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class RouteRequest:
|
|
request_id: str
|
|
task_id: str
|
|
trace_id: str
|
|
scope_type: str
|
|
tenant_id: str | None
|
|
project_id: str | None
|
|
required_capabilities: tuple[str, ...]
|
|
estimated_input_tokens: int
|
|
max_output_tokens: int
|
|
preferred_models: tuple[str, ...] = ()
|
|
exclude_providers: tuple[str, ...] = ()
|
|
|
|
@property
|
|
def scope_key(self) -> str:
|
|
return _scope_key(self.scope_type, self.tenant_id, self.project_id)
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class RouteDecision:
|
|
decision_id: str
|
|
request_id: str
|
|
task_id: str
|
|
trace_id: str
|
|
scope_key: str
|
|
provider_id: str
|
|
model_id: str
|
|
model_version: str
|
|
provider_credential_ref: str | None
|
|
reserved_input_tokens: int
|
|
reserved_output_tokens: int
|
|
reserved_cost_micros: int
|
|
policy_revision: int
|
|
model_revision: int
|
|
fallback_index: int
|
|
created_at: str
|
|
|
|
|
|
class ModelGovernanceStore:
|
|
def __init__(
|
|
self,
|
|
path: str | Path,
|
|
*,
|
|
now_fn: Callable[[], float] = time.time,
|
|
) -> None:
|
|
self.path = Path(path)
|
|
self.now_fn = now_fn
|
|
self.lock = threading.RLock()
|
|
self.path.parent.mkdir(parents=True, exist_ok=True)
|
|
self.conn = sqlite3.connect(
|
|
self.path,
|
|
timeout=30.0,
|
|
check_same_thread=False,
|
|
isolation_level=None,
|
|
)
|
|
self.conn.row_factory = sqlite3.Row
|
|
self.conn.execute("PRAGMA foreign_keys=ON")
|
|
self.conn.execute("PRAGMA busy_timeout=30000")
|
|
self.conn.execute("PRAGMA journal_mode=WAL")
|
|
self._migrate()
|
|
|
|
def close(self) -> None:
|
|
self.conn.close()
|
|
|
|
def _migrate(self) -> None:
|
|
with self.lock:
|
|
self.conn.executescript(
|
|
"""
|
|
CREATE TABLE IF NOT EXISTS model_governance_providers(
|
|
provider_id TEXT PRIMARY KEY,
|
|
credential_ref TEXT,
|
|
enabled INTEGER NOT NULL,
|
|
failure_domain TEXT NOT NULL,
|
|
max_failures INTEGER NOT NULL,
|
|
cooldown_seconds INTEGER NOT NULL,
|
|
revision INTEGER NOT NULL,
|
|
updated_at TEXT NOT NULL
|
|
);
|
|
CREATE TABLE IF NOT EXISTS model_governance_models(
|
|
model_id TEXT PRIMARY KEY,
|
|
provider_id TEXT NOT NULL,
|
|
version TEXT NOT NULL,
|
|
capabilities_json TEXT NOT NULL,
|
|
input_cost_micros_per_1k INTEGER NOT NULL,
|
|
output_cost_micros_per_1k INTEGER NOT NULL,
|
|
max_input_tokens INTEGER NOT NULL,
|
|
max_output_tokens INTEGER NOT NULL,
|
|
evaluation_state TEXT NOT NULL,
|
|
evaluation_sha256 TEXT NOT NULL,
|
|
enabled INTEGER NOT NULL,
|
|
priority INTEGER NOT NULL,
|
|
revision INTEGER NOT NULL,
|
|
updated_at TEXT NOT NULL,
|
|
FOREIGN KEY(provider_id) REFERENCES model_governance_providers(provider_id)
|
|
);
|
|
CREATE TABLE IF NOT EXISTS model_governance_scope_policies(
|
|
scope_key TEXT PRIMARY KEY,
|
|
scope_type TEXT NOT NULL,
|
|
tenant_id TEXT,
|
|
project_id TEXT,
|
|
allowed_models_json TEXT NOT NULL,
|
|
denied_providers_json TEXT NOT NULL,
|
|
required_capabilities_json TEXT NOT NULL,
|
|
pinned_model_id TEXT,
|
|
pinned_model_version TEXT,
|
|
max_input_tokens_per_request INTEGER NOT NULL,
|
|
max_output_tokens_per_request INTEGER NOT NULL,
|
|
daily_token_budget INTEGER NOT NULL,
|
|
daily_cost_budget_micros INTEGER NOT NULL,
|
|
fallback_limit INTEGER NOT NULL,
|
|
revision INTEGER NOT NULL,
|
|
updated_at TEXT NOT NULL
|
|
);
|
|
CREATE TABLE IF NOT EXISTS model_governance_provider_health(
|
|
provider_id TEXT PRIMARY KEY,
|
|
consecutive_failures INTEGER NOT NULL DEFAULT 0,
|
|
open_until REAL NOT NULL DEFAULT 0,
|
|
last_failure_class TEXT,
|
|
revision INTEGER NOT NULL DEFAULT 1,
|
|
updated_at TEXT NOT NULL,
|
|
FOREIGN KEY(provider_id) REFERENCES model_governance_providers(provider_id)
|
|
);
|
|
CREATE TABLE IF NOT EXISTS model_governance_usage(
|
|
scope_key TEXT NOT NULL,
|
|
day_key TEXT NOT NULL,
|
|
input_tokens INTEGER NOT NULL DEFAULT 0,
|
|
output_tokens INTEGER NOT NULL DEFAULT 0,
|
|
cost_micros INTEGER NOT NULL DEFAULT 0,
|
|
revision INTEGER NOT NULL DEFAULT 1,
|
|
updated_at TEXT NOT NULL,
|
|
PRIMARY KEY(scope_key,day_key)
|
|
);
|
|
CREATE TABLE IF NOT EXISTS model_governance_reservations(
|
|
request_id TEXT PRIMARY KEY,
|
|
task_id TEXT NOT NULL,
|
|
trace_id TEXT NOT NULL,
|
|
scope_key TEXT NOT NULL,
|
|
provider_id TEXT NOT NULL,
|
|
model_id TEXT NOT NULL,
|
|
model_version TEXT NOT NULL,
|
|
reserved_input_tokens INTEGER NOT NULL,
|
|
reserved_output_tokens INTEGER NOT NULL,
|
|
reserved_cost_micros INTEGER NOT NULL,
|
|
fallback_index INTEGER NOT NULL,
|
|
state TEXT NOT NULL,
|
|
actual_input_tokens INTEGER,
|
|
actual_output_tokens INTEGER,
|
|
actual_cost_micros INTEGER,
|
|
usage_source TEXT,
|
|
created_at TEXT NOT NULL,
|
|
updated_at TEXT NOT NULL
|
|
);
|
|
CREATE TABLE IF NOT EXISTS model_governance_route_decisions(
|
|
decision_id TEXT PRIMARY KEY,
|
|
request_id TEXT NOT NULL UNIQUE,
|
|
task_id TEXT NOT NULL,
|
|
trace_id TEXT NOT NULL,
|
|
scope_key TEXT NOT NULL,
|
|
provider_id TEXT NOT NULL,
|
|
model_id TEXT NOT NULL,
|
|
model_version TEXT NOT NULL,
|
|
provider_credential_ref TEXT,
|
|
policy_revision INTEGER NOT NULL,
|
|
model_revision INTEGER NOT NULL,
|
|
fallback_index INTEGER NOT NULL,
|
|
request_sha256 TEXT NOT NULL,
|
|
created_at TEXT NOT NULL
|
|
);
|
|
CREATE TABLE IF NOT EXISTS model_governance_audit(
|
|
sequence INTEGER PRIMARY KEY AUTOINCREMENT,
|
|
event_type TEXT NOT NULL,
|
|
task_id TEXT,
|
|
trace_id TEXT,
|
|
scope_key TEXT,
|
|
payload_json TEXT NOT NULL,
|
|
previous_sha256 TEXT NOT NULL,
|
|
event_sha256 TEXT NOT NULL UNIQUE,
|
|
created_at TEXT NOT NULL
|
|
);
|
|
"""
|
|
)
|
|
|
|
def _audit(
|
|
self,
|
|
event_type: str,
|
|
*,
|
|
task_id: str | None,
|
|
trace_id: str | None,
|
|
scope_key: str | None,
|
|
payload: Mapping[str, Any],
|
|
now: float,
|
|
) -> str:
|
|
if SECRET_RE.search(_json_text(payload)):
|
|
raise ModelGovernanceError("MODEL_AUDIT_SECRET_EXPOSURE_DENY")
|
|
row = self.conn.execute(
|
|
"SELECT event_sha256 FROM model_governance_audit ORDER BY sequence DESC LIMIT 1"
|
|
).fetchone()
|
|
previous = str(row[0]) if row else "0" * 64
|
|
material = {
|
|
"event_type": event_type,
|
|
"task_id": task_id,
|
|
"trace_id": trace_id,
|
|
"scope_key": scope_key,
|
|
"payload": dict(payload),
|
|
"previous_sha256": previous,
|
|
"created_at": _now_iso(now),
|
|
}
|
|
digest = _sha(material)
|
|
self.conn.execute(
|
|
"""
|
|
INSERT INTO model_governance_audit(
|
|
event_type,task_id,trace_id,scope_key,payload_json,
|
|
previous_sha256,event_sha256,created_at
|
|
) VALUES(?,?,?,?,?,?,?,?)
|
|
""",
|
|
(
|
|
event_type,
|
|
task_id,
|
|
trace_id,
|
|
scope_key,
|
|
_json_text(payload),
|
|
previous,
|
|
digest,
|
|
material["created_at"],
|
|
),
|
|
)
|
|
return digest
|
|
|
|
def verify_audit_chain(self) -> bool:
|
|
previous = "0" * 64
|
|
rows = self.conn.execute(
|
|
"SELECT * FROM model_governance_audit ORDER BY sequence"
|
|
).fetchall()
|
|
for row in rows:
|
|
payload = json.loads(row["payload_json"])
|
|
material = {
|
|
"event_type": row["event_type"],
|
|
"task_id": row["task_id"],
|
|
"trace_id": row["trace_id"],
|
|
"scope_key": row["scope_key"],
|
|
"payload": payload,
|
|
"previous_sha256": previous,
|
|
"created_at": row["created_at"],
|
|
}
|
|
if row["previous_sha256"] != previous or row["event_sha256"] != _sha(material):
|
|
return False
|
|
previous = row["event_sha256"]
|
|
return True
|
|
|
|
def register_provider(
|
|
self,
|
|
provider: ProviderRecord,
|
|
*,
|
|
expected_revision: int | None = None,
|
|
) -> dict[str, Any]:
|
|
provider_id = _text(
|
|
provider.provider_id,
|
|
"MODEL_PROVIDER_ID_INVALID",
|
|
PROVIDER_ID_RE,
|
|
)
|
|
credential_ref = _credential_ref(provider.credential_ref)
|
|
if provider.max_failures < 1 or provider.cooldown_seconds < 1:
|
|
raise ModelGovernanceError("MODEL_PROVIDER_POLICY_INVALID")
|
|
now = self.now_fn()
|
|
with self.lock:
|
|
self.conn.execute("BEGIN IMMEDIATE")
|
|
try:
|
|
current = self.conn.execute(
|
|
"SELECT revision FROM model_governance_providers WHERE provider_id=?",
|
|
(provider_id,),
|
|
).fetchone()
|
|
if current:
|
|
revision = int(current["revision"])
|
|
if expected_revision is None or revision != expected_revision:
|
|
raise ModelGovernanceError("MODEL_PROVIDER_CAS_MISMATCH")
|
|
next_revision = revision + 1
|
|
self.conn.execute(
|
|
"""
|
|
UPDATE model_governance_providers
|
|
SET credential_ref=?,enabled=?,failure_domain=?,
|
|
max_failures=?,cooldown_seconds=?,revision=?,updated_at=?
|
|
WHERE provider_id=? AND revision=?
|
|
""",
|
|
(
|
|
credential_ref,
|
|
int(provider.enabled),
|
|
provider.failure_domain,
|
|
provider.max_failures,
|
|
provider.cooldown_seconds,
|
|
next_revision,
|
|
_now_iso(now),
|
|
provider_id,
|
|
revision,
|
|
),
|
|
)
|
|
else:
|
|
if expected_revision not in (None, 0):
|
|
raise ModelGovernanceError("MODEL_PROVIDER_CAS_MISMATCH")
|
|
next_revision = 1
|
|
self.conn.execute(
|
|
"""
|
|
INSERT INTO model_governance_providers(
|
|
provider_id,credential_ref,enabled,failure_domain,
|
|
max_failures,cooldown_seconds,revision,updated_at
|
|
) VALUES(?,?,?,?,?,?,?,?)
|
|
""",
|
|
(
|
|
provider_id,
|
|
credential_ref,
|
|
int(provider.enabled),
|
|
provider.failure_domain,
|
|
provider.max_failures,
|
|
provider.cooldown_seconds,
|
|
next_revision,
|
|
_now_iso(now),
|
|
),
|
|
)
|
|
self.conn.execute(
|
|
"""
|
|
INSERT OR IGNORE INTO model_governance_provider_health(
|
|
provider_id,consecutive_failures,open_until,last_failure_class,
|
|
revision,updated_at
|
|
) VALUES(?,0,0,NULL,1,?)
|
|
""",
|
|
(provider_id, _now_iso(now)),
|
|
)
|
|
self._audit(
|
|
"MODEL_PROVIDER_REGISTERED",
|
|
task_id=None,
|
|
trace_id=None,
|
|
scope_key=None,
|
|
payload={
|
|
"provider_id": provider_id,
|
|
"credential_ref": credential_ref,
|
|
"enabled": bool(provider.enabled),
|
|
"revision": next_revision,
|
|
},
|
|
now=now,
|
|
)
|
|
self.conn.commit()
|
|
return {
|
|
"provider_id": provider_id,
|
|
"revision": next_revision,
|
|
"credential_ref": credential_ref,
|
|
}
|
|
except Exception:
|
|
self.conn.rollback()
|
|
raise
|
|
|
|
def register_model(
|
|
self,
|
|
model: ModelRecord,
|
|
*,
|
|
expected_revision: int | None = None,
|
|
) -> dict[str, Any]:
|
|
model_id = _text(model.model_id, "MODEL_ID_INVALID", MODEL_ID_RE)
|
|
provider_id = _text(
|
|
model.provider_id,
|
|
"MODEL_PROVIDER_ID_INVALID",
|
|
PROVIDER_ID_RE,
|
|
)
|
|
version = _text(model.version, "MODEL_VERSION_INVALID")
|
|
capabilities = _json_list(model.capabilities, "MODEL_CAPABILITY_INVALID")
|
|
if (
|
|
model.input_cost_micros_per_1k < 0
|
|
or model.output_cost_micros_per_1k < 0
|
|
or model.max_input_tokens < 0
|
|
or model.max_output_tokens < 0
|
|
or model.priority < 0
|
|
):
|
|
raise ModelGovernanceError("MODEL_RECORD_INVALID")
|
|
if model.evaluation_state != APPROVED:
|
|
raise ModelGovernanceError("MODEL_EVALUATION_GATE_DENY")
|
|
if SHA256_RE.fullmatch(model.evaluation_sha256) is None:
|
|
raise ModelGovernanceError("MODEL_EVALUATION_SHA256_INVALID")
|
|
|
|
now = self.now_fn()
|
|
with self.lock:
|
|
self.conn.execute("BEGIN IMMEDIATE")
|
|
try:
|
|
provider = self.conn.execute(
|
|
"SELECT enabled FROM model_governance_providers WHERE provider_id=?",
|
|
(provider_id,),
|
|
).fetchone()
|
|
if not provider or not int(provider["enabled"]):
|
|
raise ModelGovernanceError("MODEL_PROVIDER_UNAVAILABLE")
|
|
current = self.conn.execute(
|
|
"SELECT revision FROM model_governance_models WHERE model_id=?",
|
|
(model_id,),
|
|
).fetchone()
|
|
params = (
|
|
provider_id,
|
|
version,
|
|
_json_text(capabilities),
|
|
model.input_cost_micros_per_1k,
|
|
model.output_cost_micros_per_1k,
|
|
model.max_input_tokens,
|
|
model.max_output_tokens,
|
|
model.evaluation_state,
|
|
model.evaluation_sha256,
|
|
int(model.enabled),
|
|
model.priority,
|
|
)
|
|
if current:
|
|
revision = int(current["revision"])
|
|
if expected_revision is None or expected_revision != revision:
|
|
raise ModelGovernanceError("MODEL_REGISTRY_CAS_MISMATCH")
|
|
next_revision = revision + 1
|
|
self.conn.execute(
|
|
"""
|
|
UPDATE model_governance_models
|
|
SET provider_id=?,version=?,capabilities_json=?,
|
|
input_cost_micros_per_1k=?,output_cost_micros_per_1k=?,
|
|
max_input_tokens=?,max_output_tokens=?,
|
|
evaluation_state=?,evaluation_sha256=?,enabled=?,priority=?,
|
|
revision=?,updated_at=?
|
|
WHERE model_id=? AND revision=?
|
|
""",
|
|
(*params, next_revision, _now_iso(now), model_id, revision),
|
|
)
|
|
else:
|
|
if expected_revision not in (None, 0):
|
|
raise ModelGovernanceError("MODEL_REGISTRY_CAS_MISMATCH")
|
|
next_revision = 1
|
|
self.conn.execute(
|
|
"""
|
|
INSERT INTO model_governance_models(
|
|
model_id,provider_id,version,capabilities_json,
|
|
input_cost_micros_per_1k,output_cost_micros_per_1k,
|
|
max_input_tokens,max_output_tokens,evaluation_state,
|
|
evaluation_sha256,enabled,priority,revision,updated_at
|
|
) VALUES(?,?,?,?,?,?,?,?,?,?,?,?,?,?)
|
|
""",
|
|
(
|
|
model_id,
|
|
*params,
|
|
next_revision,
|
|
_now_iso(now),
|
|
),
|
|
)
|
|
self._audit(
|
|
"MODEL_REGISTERED",
|
|
task_id=None,
|
|
trace_id=None,
|
|
scope_key=None,
|
|
payload={
|
|
"model_id": model_id,
|
|
"provider_id": provider_id,
|
|
"version": version,
|
|
"evaluation_sha256": model.evaluation_sha256,
|
|
"revision": next_revision,
|
|
},
|
|
now=now,
|
|
)
|
|
self.conn.commit()
|
|
return {
|
|
"model_id": model_id,
|
|
"provider_id": provider_id,
|
|
"version": version,
|
|
"revision": next_revision,
|
|
}
|
|
except Exception:
|
|
self.conn.rollback()
|
|
raise
|
|
|
|
def set_scope_policy(
|
|
self,
|
|
policy: ScopePolicy,
|
|
*,
|
|
expected_revision: int | None = None,
|
|
) -> dict[str, Any]:
|
|
scope_key = policy.scope_key
|
|
allowed = _json_list(policy.allowed_models, "MODEL_POLICY_ALLOWED_MODELS_INVALID")
|
|
denied = _json_list(policy.denied_providers, "MODEL_POLICY_DENIED_PROVIDERS_INVALID")
|
|
required = _json_list(
|
|
policy.required_capabilities,
|
|
"MODEL_POLICY_CAPABILITIES_INVALID",
|
|
)
|
|
if (
|
|
policy.max_input_tokens_per_request < 0
|
|
or policy.max_output_tokens_per_request < 0
|
|
or policy.daily_token_budget < 0
|
|
or policy.daily_cost_budget_micros < 0
|
|
or not 0 <= policy.fallback_limit <= 8
|
|
):
|
|
raise ModelGovernanceError("MODEL_POLICY_LIMIT_INVALID")
|
|
if policy.pinned_model_id:
|
|
_text(policy.pinned_model_id, "MODEL_POLICY_PIN_INVALID", MODEL_ID_RE)
|
|
if not policy.pinned_model_version:
|
|
raise ModelGovernanceError("MODEL_POLICY_PIN_VERSION_REQUIRED")
|
|
|
|
now = self.now_fn()
|
|
with self.lock:
|
|
self.conn.execute("BEGIN IMMEDIATE")
|
|
try:
|
|
current = self.conn.execute(
|
|
"SELECT revision FROM model_governance_scope_policies WHERE scope_key=?",
|
|
(scope_key,),
|
|
).fetchone()
|
|
fields = (
|
|
policy.scope_type,
|
|
policy.tenant_id,
|
|
policy.project_id,
|
|
_json_text(allowed),
|
|
_json_text(denied),
|
|
_json_text(required),
|
|
policy.pinned_model_id,
|
|
policy.pinned_model_version,
|
|
policy.max_input_tokens_per_request,
|
|
policy.max_output_tokens_per_request,
|
|
policy.daily_token_budget,
|
|
policy.daily_cost_budget_micros,
|
|
policy.fallback_limit,
|
|
)
|
|
if current:
|
|
revision = int(current["revision"])
|
|
if expected_revision is None or expected_revision != revision:
|
|
raise ModelGovernanceError("MODEL_POLICY_CAS_MISMATCH")
|
|
next_revision = revision + 1
|
|
self.conn.execute(
|
|
"""
|
|
UPDATE model_governance_scope_policies
|
|
SET scope_type=?,tenant_id=?,project_id=?,
|
|
allowed_models_json=?,denied_providers_json=?,
|
|
required_capabilities_json=?,pinned_model_id=?,
|
|
pinned_model_version=?,max_input_tokens_per_request=?,
|
|
max_output_tokens_per_request=?,daily_token_budget=?,
|
|
daily_cost_budget_micros=?,fallback_limit=?,
|
|
revision=?,updated_at=?
|
|
WHERE scope_key=? AND revision=?
|
|
""",
|
|
(*fields, next_revision, _now_iso(now), scope_key, revision),
|
|
)
|
|
else:
|
|
if expected_revision not in (None, 0):
|
|
raise ModelGovernanceError("MODEL_POLICY_CAS_MISMATCH")
|
|
next_revision = 1
|
|
self.conn.execute(
|
|
"""
|
|
INSERT INTO model_governance_scope_policies(
|
|
scope_key,scope_type,tenant_id,project_id,
|
|
allowed_models_json,denied_providers_json,
|
|
required_capabilities_json,pinned_model_id,pinned_model_version,
|
|
max_input_tokens_per_request,max_output_tokens_per_request,
|
|
daily_token_budget,daily_cost_budget_micros,fallback_limit,
|
|
revision,updated_at
|
|
) VALUES(?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?)
|
|
""",
|
|
(scope_key, *fields, next_revision, _now_iso(now)),
|
|
)
|
|
self._audit(
|
|
"MODEL_SCOPE_POLICY_SET",
|
|
task_id=None,
|
|
trace_id=None,
|
|
scope_key=scope_key,
|
|
payload={
|
|
"scope_key": scope_key,
|
|
"revision": next_revision,
|
|
"pinned_model_id": policy.pinned_model_id,
|
|
"pinned_model_version": policy.pinned_model_version,
|
|
},
|
|
now=now,
|
|
)
|
|
self.conn.commit()
|
|
return {"scope_key": scope_key, "revision": next_revision}
|
|
except Exception:
|
|
self.conn.rollback()
|
|
raise
|
|
|
|
@staticmethod
|
|
def _estimated_cost(
|
|
model: sqlite3.Row,
|
|
input_tokens: int,
|
|
output_tokens: int,
|
|
) -> int:
|
|
# ceil(tokens / 1000 * micros_per_1k) without floating point.
|
|
return (
|
|
(input_tokens * int(model["input_cost_micros_per_1k"]) + 999) // 1000
|
|
+ (output_tokens * int(model["output_cost_micros_per_1k"]) + 999) // 1000
|
|
)
|
|
|
|
def _policy(self, scope_key: str) -> sqlite3.Row:
|
|
row = self.conn.execute(
|
|
"SELECT * FROM model_governance_scope_policies WHERE scope_key=?",
|
|
(scope_key,),
|
|
).fetchone()
|
|
if row:
|
|
return row
|
|
if scope_key.startswith("tenant:") and "/project:" in scope_key:
|
|
tenant_key = scope_key.split("/project:", 1)[0]
|
|
row = self.conn.execute(
|
|
"SELECT * FROM model_governance_scope_policies WHERE scope_key=?",
|
|
(tenant_key,),
|
|
).fetchone()
|
|
if row:
|
|
return row
|
|
# Tenant/system/legacy scopes never cross-fallback. A global/system policy
|
|
# cannot silently authorize a tenant that lacks its own policy.
|
|
raise ModelGovernanceError("MODEL_SCOPE_POLICY_MISSING")
|
|
|
|
def _usage_and_reserved(
|
|
self,
|
|
scope_key: str,
|
|
day_key: str,
|
|
now: float,
|
|
) -> tuple[int, int]:
|
|
usage = self.conn.execute(
|
|
"""
|
|
SELECT input_tokens+output_tokens AS tokens,cost_micros
|
|
FROM model_governance_usage
|
|
WHERE scope_key=? AND day_key=?
|
|
""",
|
|
(scope_key, day_key),
|
|
).fetchone()
|
|
tokens = int(usage["tokens"]) if usage else 0
|
|
cost = int(usage["cost_micros"]) if usage else 0
|
|
reservations = self.conn.execute(
|
|
"""
|
|
SELECT COALESCE(SUM(reserved_input_tokens+reserved_output_tokens),0) AS tokens,
|
|
COALESCE(SUM(reserved_cost_micros),0) AS cost
|
|
FROM model_governance_reservations
|
|
WHERE scope_key=? AND state='RESERVED'
|
|
""",
|
|
(scope_key,),
|
|
).fetchone()
|
|
return (
|
|
tokens + int(reservations["tokens"]),
|
|
cost + int(reservations["cost"]),
|
|
)
|
|
|
|
def _candidate_models(
|
|
self,
|
|
request: RouteRequest,
|
|
policy: sqlite3.Row,
|
|
now: float,
|
|
) -> list[sqlite3.Row]:
|
|
required = set(request.required_capabilities)
|
|
required.update(json.loads(policy["required_capabilities_json"]))
|
|
allowed = set(json.loads(policy["allowed_models_json"]))
|
|
denied_providers = set(json.loads(policy["denied_providers_json"]))
|
|
denied_providers.update(request.exclude_providers)
|
|
preferred = list(request.preferred_models)
|
|
|
|
rows = self.conn.execute(
|
|
"""
|
|
SELECT m.*,p.credential_ref,p.enabled AS provider_enabled,
|
|
p.max_failures,p.cooldown_seconds,
|
|
COALESCE(h.consecutive_failures,0) AS consecutive_failures,
|
|
COALESCE(h.open_until,0) AS open_until
|
|
FROM model_governance_models m
|
|
JOIN model_governance_providers p ON p.provider_id=m.provider_id
|
|
LEFT JOIN model_governance_provider_health h ON h.provider_id=p.provider_id
|
|
WHERE m.enabled=1 AND p.enabled=1 AND m.evaluation_state=?
|
|
""",
|
|
(APPROVED,),
|
|
).fetchall()
|
|
|
|
candidates: list[sqlite3.Row] = []
|
|
for row in rows:
|
|
if float(row["open_until"]) > now:
|
|
continue
|
|
if row["provider_id"] in denied_providers:
|
|
continue
|
|
if allowed and row["model_id"] not in allowed:
|
|
continue
|
|
if policy["pinned_model_id"] and row["model_id"] != policy["pinned_model_id"]:
|
|
continue
|
|
if (
|
|
policy["pinned_model_version"]
|
|
and row["version"] != policy["pinned_model_version"]
|
|
):
|
|
continue
|
|
capabilities = set(json.loads(row["capabilities_json"]))
|
|
if not required.issubset(capabilities):
|
|
continue
|
|
if int(row["max_input_tokens"]) and request.estimated_input_tokens > int(row["max_input_tokens"]):
|
|
continue
|
|
if int(row["max_output_tokens"]) and request.max_output_tokens > int(row["max_output_tokens"]):
|
|
continue
|
|
candidates.append(row)
|
|
|
|
def key(row: sqlite3.Row) -> tuple[int, int, str, str]:
|
|
pref_rank = preferred.index(row["model_id"]) if row["model_id"] in preferred else 10_000
|
|
return (
|
|
pref_rank,
|
|
int(row["priority"]),
|
|
str(row["provider_id"]),
|
|
str(row["model_id"]),
|
|
)
|
|
|
|
return sorted(candidates, key=key)
|
|
|
|
def route(
|
|
self,
|
|
request: RouteRequest,
|
|
*,
|
|
fallback_index: int = 0,
|
|
) -> RouteDecision:
|
|
request_id = _text(request.request_id, "MODEL_REQUEST_ID_INVALID", REQUEST_ID_RE)
|
|
task_id = _text(request.task_id, "MODEL_TASK_ID_INVALID", TASK_ID_RE)
|
|
trace_id = _text(request.trace_id, "MODEL_TRACE_ID_INVALID", TRACE_ID_RE)
|
|
if request.estimated_input_tokens < 0 or request.max_output_tokens < 0:
|
|
raise ModelGovernanceError("MODEL_TOKEN_ESTIMATE_INVALID")
|
|
scope_key = request.scope_key
|
|
now = self.now_fn()
|
|
request_material = {
|
|
**asdict(request),
|
|
"scope_key": scope_key,
|
|
"fallback_index": fallback_index,
|
|
}
|
|
request_sha = _sha(request_material)
|
|
|
|
with self.lock:
|
|
self.conn.execute("BEGIN IMMEDIATE")
|
|
try:
|
|
existing = self.conn.execute(
|
|
"""
|
|
SELECT d.*,r.reserved_input_tokens,r.reserved_output_tokens,
|
|
r.reserved_cost_micros
|
|
FROM model_governance_route_decisions d
|
|
JOIN model_governance_reservations r USING(request_id)
|
|
WHERE d.request_id=?
|
|
""",
|
|
(request_id,),
|
|
).fetchone()
|
|
if existing:
|
|
if existing["request_sha256"] != request_sha:
|
|
raise ModelGovernanceError("MODEL_ROUTE_IDEMPOTENCY_CONFLICT")
|
|
decision = RouteDecision(
|
|
existing["decision_id"],
|
|
request_id,
|
|
existing["task_id"],
|
|
existing["trace_id"],
|
|
existing["scope_key"],
|
|
existing["provider_id"],
|
|
existing["model_id"],
|
|
existing["model_version"],
|
|
existing["provider_credential_ref"],
|
|
int(existing["reserved_input_tokens"]),
|
|
int(existing["reserved_output_tokens"]),
|
|
int(existing["reserved_cost_micros"]),
|
|
int(existing["policy_revision"]),
|
|
int(existing["model_revision"]),
|
|
int(existing["fallback_index"]),
|
|
existing["created_at"],
|
|
)
|
|
self.conn.commit()
|
|
return decision
|
|
|
|
policy = self._policy(scope_key)
|
|
if fallback_index > int(policy["fallback_limit"]):
|
|
raise ModelGovernanceError("MODEL_FALLBACK_LIMIT_EXCEEDED")
|
|
if (
|
|
int(policy["max_input_tokens_per_request"])
|
|
and request.estimated_input_tokens > int(policy["max_input_tokens_per_request"])
|
|
):
|
|
raise ModelGovernanceError("MODEL_INPUT_TOKEN_LIMIT_DENY")
|
|
if (
|
|
int(policy["max_output_tokens_per_request"])
|
|
and request.max_output_tokens > int(policy["max_output_tokens_per_request"])
|
|
):
|
|
raise ModelGovernanceError("MODEL_OUTPUT_TOKEN_LIMIT_DENY")
|
|
|
|
candidates = self._candidate_models(request, policy, now)
|
|
if not candidates:
|
|
raise ModelGovernanceError("MODEL_ROUTE_NO_ELIGIBLE_MODEL")
|
|
|
|
used_tokens, used_cost = self._usage_and_reserved(
|
|
scope_key,
|
|
_day_key(now),
|
|
now,
|
|
)
|
|
chosen: sqlite3.Row | None = None
|
|
chosen_cost = 0
|
|
for row in candidates:
|
|
cost = self._estimated_cost(
|
|
row,
|
|
request.estimated_input_tokens,
|
|
request.max_output_tokens,
|
|
)
|
|
token_reservation = request.estimated_input_tokens + request.max_output_tokens
|
|
if (
|
|
int(policy["daily_token_budget"])
|
|
and used_tokens + token_reservation > int(policy["daily_token_budget"])
|
|
):
|
|
continue
|
|
if (
|
|
int(policy["daily_cost_budget_micros"])
|
|
and used_cost + cost > int(policy["daily_cost_budget_micros"])
|
|
):
|
|
continue
|
|
chosen = row
|
|
chosen_cost = cost
|
|
break
|
|
|
|
if chosen is None:
|
|
raise ModelGovernanceError("MODEL_BUDGET_DENY")
|
|
|
|
created_at = _now_iso(now)
|
|
decision_id = "model-route-" + hashlib.sha256(
|
|
f"{request_id}:{chosen['model_id']}:{chosen['version']}:{fallback_index}".encode()
|
|
).hexdigest()[:32]
|
|
credential_ref = _credential_ref(chosen["credential_ref"])
|
|
|
|
self.conn.execute(
|
|
"""
|
|
INSERT INTO model_governance_reservations(
|
|
request_id,task_id,trace_id,scope_key,provider_id,model_id,
|
|
model_version,reserved_input_tokens,reserved_output_tokens,
|
|
reserved_cost_micros,fallback_index,state,created_at,updated_at
|
|
) VALUES(?,?,?,?,?,?,?,?,?,?,?,'RESERVED',?,?)
|
|
""",
|
|
(
|
|
request_id,
|
|
task_id,
|
|
trace_id,
|
|
scope_key,
|
|
chosen["provider_id"],
|
|
chosen["model_id"],
|
|
chosen["version"],
|
|
request.estimated_input_tokens,
|
|
request.max_output_tokens,
|
|
chosen_cost,
|
|
fallback_index,
|
|
created_at,
|
|
created_at,
|
|
),
|
|
)
|
|
self.conn.execute(
|
|
"""
|
|
INSERT INTO model_governance_route_decisions(
|
|
decision_id,request_id,task_id,trace_id,scope_key,provider_id,
|
|
model_id,model_version,provider_credential_ref,policy_revision,
|
|
model_revision,fallback_index,request_sha256,created_at
|
|
) VALUES(?,?,?,?,?,?,?,?,?,?,?,?,?,?)
|
|
""",
|
|
(
|
|
decision_id,
|
|
request_id,
|
|
task_id,
|
|
trace_id,
|
|
scope_key,
|
|
chosen["provider_id"],
|
|
chosen["model_id"],
|
|
chosen["version"],
|
|
credential_ref,
|
|
int(policy["revision"]),
|
|
int(chosen["revision"]),
|
|
fallback_index,
|
|
request_sha,
|
|
created_at,
|
|
),
|
|
)
|
|
self._audit(
|
|
"MODEL_ROUTE_RESERVED",
|
|
task_id=task_id,
|
|
trace_id=trace_id,
|
|
scope_key=scope_key,
|
|
payload={
|
|
"decision_id": decision_id,
|
|
"request_id": request_id,
|
|
"provider_id": chosen["provider_id"],
|
|
"model_id": chosen["model_id"],
|
|
"model_version": chosen["version"],
|
|
"reserved_input_tokens": request.estimated_input_tokens,
|
|
"reserved_output_tokens": request.max_output_tokens,
|
|
"reserved_cost_micros": chosen_cost,
|
|
"policy_revision": int(policy["revision"]),
|
|
"model_revision": int(chosen["revision"]),
|
|
"fallback_index": fallback_index,
|
|
},
|
|
now=now,
|
|
)
|
|
self.conn.commit()
|
|
return RouteDecision(
|
|
decision_id,
|
|
request_id,
|
|
task_id,
|
|
trace_id,
|
|
scope_key,
|
|
chosen["provider_id"],
|
|
chosen["model_id"],
|
|
chosen["version"],
|
|
credential_ref,
|
|
request.estimated_input_tokens,
|
|
request.max_output_tokens,
|
|
chosen_cost,
|
|
int(policy["revision"]),
|
|
int(chosen["revision"]),
|
|
fallback_index,
|
|
created_at,
|
|
)
|
|
except Exception:
|
|
self.conn.rollback()
|
|
raise
|
|
|
|
def record_success(
|
|
self,
|
|
request_id: str,
|
|
*,
|
|
input_tokens: int | None,
|
|
output_tokens: int | None,
|
|
actual_cost_micros: int | None = None,
|
|
usage_source: str = "runtime",
|
|
) -> dict[str, Any]:
|
|
request_id = _text(request_id, "MODEL_REQUEST_ID_INVALID", REQUEST_ID_RE)
|
|
if input_tokens is not None and input_tokens < 0:
|
|
raise ModelGovernanceError("MODEL_USAGE_INVALID")
|
|
if output_tokens is not None and output_tokens < 0:
|
|
raise ModelGovernanceError("MODEL_USAGE_INVALID")
|
|
if actual_cost_micros is not None and actual_cost_micros < 0:
|
|
raise ModelGovernanceError("MODEL_USAGE_INVALID")
|
|
now = self.now_fn()
|
|
|
|
with self.lock:
|
|
self.conn.execute("BEGIN IMMEDIATE")
|
|
try:
|
|
reservation = self.conn.execute(
|
|
"SELECT * FROM model_governance_reservations WHERE request_id=?",
|
|
(request_id,),
|
|
).fetchone()
|
|
if not reservation:
|
|
raise ModelGovernanceError("MODEL_RESERVATION_NOT_FOUND")
|
|
if reservation["state"] == "SETTLED":
|
|
result = {
|
|
"request_id": request_id,
|
|
"state": "SETTLED",
|
|
"input_tokens": int(reservation["actual_input_tokens"]),
|
|
"output_tokens": int(reservation["actual_output_tokens"]),
|
|
"cost_micros": int(reservation["actual_cost_micros"]),
|
|
"usage_source": reservation["usage_source"],
|
|
"idempotent": True,
|
|
}
|
|
self.conn.commit()
|
|
return result
|
|
if reservation["state"] != "RESERVED":
|
|
raise ModelGovernanceError("MODEL_RESERVATION_STATE_INVALID")
|
|
|
|
if input_tokens is None or output_tokens is None:
|
|
actual_input = int(reservation["reserved_input_tokens"])
|
|
actual_output = int(reservation["reserved_output_tokens"])
|
|
source = "reserved_upper_bound"
|
|
else:
|
|
actual_input = input_tokens
|
|
actual_output = output_tokens
|
|
source = usage_source
|
|
if (
|
|
actual_input > int(reservation["reserved_input_tokens"])
|
|
or actual_output > int(reservation["reserved_output_tokens"])
|
|
):
|
|
raise ModelGovernanceError("MODEL_USAGE_EXCEEDS_RESERVATION")
|
|
|
|
if actual_cost_micros is None:
|
|
reserved_total = (
|
|
int(reservation["reserved_input_tokens"])
|
|
+ int(reservation["reserved_output_tokens"])
|
|
)
|
|
actual_total = actual_input + actual_output
|
|
reserved_cost = int(reservation["reserved_cost_micros"])
|
|
cost = (
|
|
reserved_cost
|
|
if reserved_total == 0
|
|
else (reserved_cost * actual_total + reserved_total - 1) // reserved_total
|
|
)
|
|
else:
|
|
cost = actual_cost_micros
|
|
if cost > int(reservation["reserved_cost_micros"]):
|
|
raise ModelGovernanceError("MODEL_COST_EXCEEDS_RESERVATION")
|
|
|
|
day_key = _day_key(now)
|
|
row = self.conn.execute(
|
|
"SELECT revision FROM model_governance_usage WHERE scope_key=? AND day_key=?",
|
|
(reservation["scope_key"], day_key),
|
|
).fetchone()
|
|
if row:
|
|
revision = int(row["revision"])
|
|
cursor = self.conn.execute(
|
|
"""
|
|
UPDATE model_governance_usage
|
|
SET input_tokens=input_tokens+?,
|
|
output_tokens=output_tokens+?,
|
|
cost_micros=cost_micros+?,
|
|
revision=revision+1,updated_at=?
|
|
WHERE scope_key=? AND day_key=? AND revision=?
|
|
""",
|
|
(
|
|
actual_input,
|
|
actual_output,
|
|
cost,
|
|
_now_iso(now),
|
|
reservation["scope_key"],
|
|
day_key,
|
|
revision,
|
|
),
|
|
)
|
|
if cursor.rowcount != 1:
|
|
raise ModelGovernanceError("MODEL_USAGE_CAS_MISMATCH")
|
|
else:
|
|
self.conn.execute(
|
|
"""
|
|
INSERT INTO model_governance_usage(
|
|
scope_key,day_key,input_tokens,output_tokens,cost_micros,
|
|
revision,updated_at
|
|
) VALUES(?,?,?,?,?,1,?)
|
|
""",
|
|
(
|
|
reservation["scope_key"],
|
|
day_key,
|
|
actual_input,
|
|
actual_output,
|
|
cost,
|
|
_now_iso(now),
|
|
),
|
|
)
|
|
|
|
self.conn.execute(
|
|
"""
|
|
UPDATE model_governance_reservations
|
|
SET state='SETTLED',actual_input_tokens=?,actual_output_tokens=?,
|
|
actual_cost_micros=?,usage_source=?,updated_at=?
|
|
WHERE request_id=? AND state='RESERVED'
|
|
""",
|
|
(
|
|
actual_input,
|
|
actual_output,
|
|
cost,
|
|
source,
|
|
_now_iso(now),
|
|
request_id,
|
|
),
|
|
)
|
|
self.conn.execute(
|
|
"""
|
|
UPDATE model_governance_provider_health
|
|
SET consecutive_failures=0,open_until=0,last_failure_class=NULL,
|
|
revision=revision+1,updated_at=?
|
|
WHERE provider_id=?
|
|
""",
|
|
(_now_iso(now), reservation["provider_id"]),
|
|
)
|
|
self._audit(
|
|
"MODEL_USAGE_SETTLED",
|
|
task_id=reservation["task_id"],
|
|
trace_id=reservation["trace_id"],
|
|
scope_key=reservation["scope_key"],
|
|
payload={
|
|
"request_id": request_id,
|
|
"provider_id": reservation["provider_id"],
|
|
"model_id": reservation["model_id"],
|
|
"model_version": reservation["model_version"],
|
|
"input_tokens": actual_input,
|
|
"output_tokens": actual_output,
|
|
"cost_micros": cost,
|
|
"usage_source": source,
|
|
},
|
|
now=now,
|
|
)
|
|
self.conn.commit()
|
|
return {
|
|
"request_id": request_id,
|
|
"state": "SETTLED",
|
|
"input_tokens": actual_input,
|
|
"output_tokens": actual_output,
|
|
"cost_micros": cost,
|
|
"usage_source": source,
|
|
"idempotent": False,
|
|
}
|
|
except Exception:
|
|
self.conn.rollback()
|
|
raise
|
|
|
|
def record_failure(
|
|
self,
|
|
request_id: str,
|
|
failure_class: str,
|
|
) -> dict[str, Any]:
|
|
request_id = _text(request_id, "MODEL_REQUEST_ID_INVALID", REQUEST_ID_RE)
|
|
failure_class = _text(failure_class, "MODEL_FAILURE_CLASS_INVALID")
|
|
now = self.now_fn()
|
|
with self.lock:
|
|
self.conn.execute("BEGIN IMMEDIATE")
|
|
try:
|
|
reservation = self.conn.execute(
|
|
"SELECT * FROM model_governance_reservations WHERE request_id=?",
|
|
(request_id,),
|
|
).fetchone()
|
|
if not reservation:
|
|
raise ModelGovernanceError("MODEL_RESERVATION_NOT_FOUND")
|
|
if reservation["state"] == "FAILED":
|
|
health = self.conn.execute(
|
|
"SELECT * FROM model_governance_provider_health WHERE provider_id=?",
|
|
(reservation["provider_id"],),
|
|
).fetchone()
|
|
self.conn.commit()
|
|
return {
|
|
"request_id": request_id,
|
|
"state": "FAILED",
|
|
"provider_id": reservation["provider_id"],
|
|
"open_until": float(health["open_until"] if health else 0),
|
|
"idempotent": True,
|
|
}
|
|
if reservation["state"] != "RESERVED":
|
|
raise ModelGovernanceError("MODEL_RESERVATION_STATE_INVALID")
|
|
|
|
provider = self.conn.execute(
|
|
"SELECT * FROM model_governance_providers WHERE provider_id=?",
|
|
(reservation["provider_id"],),
|
|
).fetchone()
|
|
health = self.conn.execute(
|
|
"SELECT * FROM model_governance_provider_health WHERE provider_id=?",
|
|
(reservation["provider_id"],),
|
|
).fetchone()
|
|
failures = int(health["consecutive_failures"] if health else 0) + 1
|
|
threshold = int(provider["max_failures"])
|
|
open_until = (
|
|
now + int(provider["cooldown_seconds"])
|
|
if failures >= threshold
|
|
else 0.0
|
|
)
|
|
if health:
|
|
revision = int(health["revision"])
|
|
cursor = self.conn.execute(
|
|
"""
|
|
UPDATE model_governance_provider_health
|
|
SET consecutive_failures=?,open_until=?,last_failure_class=?,
|
|
revision=revision+1,updated_at=?
|
|
WHERE provider_id=? AND revision=?
|
|
""",
|
|
(
|
|
failures,
|
|
open_until,
|
|
failure_class,
|
|
_now_iso(now),
|
|
reservation["provider_id"],
|
|
revision,
|
|
),
|
|
)
|
|
if cursor.rowcount != 1:
|
|
raise ModelGovernanceError("MODEL_PROVIDER_HEALTH_CAS_MISMATCH")
|
|
else:
|
|
self.conn.execute(
|
|
"""
|
|
INSERT INTO model_governance_provider_health(
|
|
provider_id,consecutive_failures,open_until,last_failure_class,
|
|
revision,updated_at
|
|
) VALUES(?,?,?,?,1,?)
|
|
""",
|
|
(
|
|
reservation["provider_id"],
|
|
failures,
|
|
open_until,
|
|
failure_class,
|
|
_now_iso(now),
|
|
),
|
|
)
|
|
self.conn.execute(
|
|
"""
|
|
UPDATE model_governance_reservations
|
|
SET state='FAILED',updated_at=?
|
|
WHERE request_id=? AND state='RESERVED'
|
|
""",
|
|
(_now_iso(now), request_id),
|
|
)
|
|
self._audit(
|
|
"MODEL_ROUTE_FAILED",
|
|
task_id=reservation["task_id"],
|
|
trace_id=reservation["trace_id"],
|
|
scope_key=reservation["scope_key"],
|
|
payload={
|
|
"request_id": request_id,
|
|
"provider_id": reservation["provider_id"],
|
|
"model_id": reservation["model_id"],
|
|
"failure_class": failure_class,
|
|
"consecutive_failures": failures,
|
|
"open_until": open_until,
|
|
},
|
|
now=now,
|
|
)
|
|
self.conn.commit()
|
|
return {
|
|
"request_id": request_id,
|
|
"state": "FAILED",
|
|
"provider_id": reservation["provider_id"],
|
|
"open_until": open_until,
|
|
"idempotent": False,
|
|
}
|
|
except Exception:
|
|
self.conn.rollback()
|
|
raise
|
|
|
|
def fallback(
|
|
self,
|
|
original: RouteRequest,
|
|
failed_request_id: str,
|
|
) -> RouteDecision:
|
|
with self.lock:
|
|
failed = self.conn.execute(
|
|
"SELECT * FROM model_governance_reservations WHERE request_id=?",
|
|
(failed_request_id,),
|
|
).fetchone()
|
|
if not failed or failed["state"] != "FAILED":
|
|
raise ModelGovernanceError("MODEL_FALLBACK_SOURCE_INVALID")
|
|
next_index = int(failed["fallback_index"]) + 1
|
|
next_request = RouteRequest(
|
|
request_id=f"{original.request_id}:fallback:{next_index}",
|
|
task_id=original.task_id,
|
|
trace_id=original.trace_id,
|
|
scope_type=original.scope_type,
|
|
tenant_id=original.tenant_id,
|
|
project_id=original.project_id,
|
|
required_capabilities=original.required_capabilities,
|
|
estimated_input_tokens=original.estimated_input_tokens,
|
|
max_output_tokens=original.max_output_tokens,
|
|
preferred_models=original.preferred_models,
|
|
exclude_providers=tuple(
|
|
sorted(set((*original.exclude_providers, str(failed["provider_id"]))))
|
|
),
|
|
)
|
|
return self.route(next_request, fallback_index=next_index)
|
|
|
|
def usage_snapshot(
|
|
self,
|
|
scope_type: str,
|
|
tenant_id: str | None,
|
|
project_id: str | None,
|
|
*,
|
|
now: float | None = None,
|
|
) -> dict[str, Any]:
|
|
current = self.now_fn() if now is None else now
|
|
scope_key = _scope_key(scope_type, tenant_id, project_id)
|
|
day_key = _day_key(current)
|
|
row = self.conn.execute(
|
|
"SELECT * FROM model_governance_usage WHERE scope_key=? AND day_key=?",
|
|
(scope_key, day_key),
|
|
).fetchone()
|
|
reserved = self.conn.execute(
|
|
"""
|
|
SELECT COALESCE(SUM(reserved_input_tokens),0) AS input_tokens,
|
|
COALESCE(SUM(reserved_output_tokens),0) AS output_tokens,
|
|
COALESCE(SUM(reserved_cost_micros),0) AS cost_micros
|
|
FROM model_governance_reservations
|
|
WHERE scope_key=? AND state='RESERVED'
|
|
""",
|
|
(scope_key,),
|
|
).fetchone()
|
|
return {
|
|
"scope_key": scope_key,
|
|
"day_key": day_key,
|
|
"settled_input_tokens": int(row["input_tokens"] if row else 0),
|
|
"settled_output_tokens": int(row["output_tokens"] if row else 0),
|
|
"settled_cost_micros": int(row["cost_micros"] if row else 0),
|
|
"reserved_input_tokens": int(reserved["input_tokens"]),
|
|
"reserved_output_tokens": int(reserved["output_tokens"]),
|
|
"reserved_cost_micros": int(reserved["cost_micros"]),
|
|
}
|