233 lines
9.5 KiB
Python
233 lines
9.5 KiB
Python
#!/usr/bin/env python3
|
|
"""Root-owned Model Governance configuration reconciler.
|
|
|
|
The config contains provider/model/policy metadata only. Raw credentials are rejected by
|
|
aios_model_governance and never stored in the governance database.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import hashlib
|
|
import json
|
|
from pathlib import Path
|
|
from typing import Any, Mapping
|
|
|
|
from aios_model_governance import (
|
|
ModelGovernanceError,
|
|
ModelGovernanceStore,
|
|
ModelRecord,
|
|
ProviderRecord,
|
|
ScopePolicy,
|
|
)
|
|
|
|
SCHEMA = "aios.model-governance.bootstrap.v1"
|
|
MAX_CONFIG_BYTES = 512 * 1024
|
|
|
|
|
|
def canonical(value: object) -> bytes:
|
|
return json.dumps(
|
|
value, ensure_ascii=False, sort_keys=True, separators=(",", ":")
|
|
).encode("utf-8")
|
|
|
|
|
|
def fail(code: str) -> None:
|
|
raise SystemExit(code)
|
|
|
|
|
|
def _load(path: Path) -> dict[str, Any]:
|
|
if path.is_symlink() or not path.is_file() or path.stat().st_size > MAX_CONFIG_BYTES:
|
|
fail("MODEL_GOVERNANCE_CONFIG_INVALID")
|
|
try:
|
|
value = json.loads(path.read_text(encoding="utf-8"))
|
|
except json.JSONDecodeError:
|
|
fail("MODEL_GOVERNANCE_CONFIG_INVALID")
|
|
if not isinstance(value, Mapping) or value.get("schema") != SCHEMA:
|
|
fail("MODEL_GOVERNANCE_CONFIG_SCHEMA_INVALID")
|
|
return dict(value)
|
|
|
|
|
|
def _provider_desired(value: Mapping[str, Any]) -> ProviderRecord:
|
|
return ProviderRecord(
|
|
provider_id=str(value.get("provider_id") or ""),
|
|
credential_ref=value.get("credential_ref"),
|
|
enabled=bool(value.get("enabled", True)),
|
|
failure_domain=str(value.get("failure_domain") or "default"),
|
|
max_failures=int(value.get("max_failures") or 3),
|
|
cooldown_seconds=int(value.get("cooldown_seconds") or 60),
|
|
)
|
|
|
|
|
|
def _model_desired(value: Mapping[str, Any]) -> ModelRecord:
|
|
return ModelRecord(
|
|
model_id=str(value.get("model_id") or ""),
|
|
provider_id=str(value.get("provider_id") or ""),
|
|
version=str(value.get("version") or ""),
|
|
capabilities=tuple(value.get("capabilities") or ()),
|
|
input_cost_micros_per_1k=int(value.get("input_cost_micros_per_1k") or 0),
|
|
output_cost_micros_per_1k=int(value.get("output_cost_micros_per_1k") or 0),
|
|
max_input_tokens=int(value.get("max_input_tokens") or 0),
|
|
max_output_tokens=int(value.get("max_output_tokens") or 0),
|
|
evaluation_state=str(value.get("evaluation_state") or "APPROVED"),
|
|
evaluation_sha256=str(value.get("evaluation_sha256") or ""),
|
|
enabled=bool(value.get("enabled", True)),
|
|
priority=int(value.get("priority") or 100),
|
|
)
|
|
|
|
|
|
def _policy_desired(value: Mapping[str, Any]) -> ScopePolicy:
|
|
return ScopePolicy(
|
|
scope_type=str(value.get("scope_type") or ""),
|
|
tenant_id=value.get("tenant_id"),
|
|
project_id=value.get("project_id"),
|
|
allowed_models=tuple(value.get("allowed_models") or ()),
|
|
denied_providers=tuple(value.get("denied_providers") or ()),
|
|
required_capabilities=tuple(value.get("required_capabilities") or ()),
|
|
pinned_model_id=value.get("pinned_model_id"),
|
|
pinned_model_version=value.get("pinned_model_version"),
|
|
max_input_tokens_per_request=int(value.get("max_input_tokens_per_request") or 0),
|
|
max_output_tokens_per_request=int(value.get("max_output_tokens_per_request") or 0),
|
|
daily_token_budget=int(value.get("daily_token_budget") or 0),
|
|
daily_cost_budget_micros=int(value.get("daily_cost_budget_micros") or 0),
|
|
fallback_limit=int(value.get("fallback_limit") or 0),
|
|
)
|
|
|
|
|
|
def _row_dict(row: object) -> dict[str, Any] | None:
|
|
return dict(row) if row is not None else None
|
|
|
|
|
|
def _provider_matches(row: Mapping[str, Any], desired: ProviderRecord) -> bool:
|
|
return (
|
|
row["credential_ref"] == desired.credential_ref
|
|
and bool(row["enabled"]) == desired.enabled
|
|
and row["failure_domain"] == desired.failure_domain
|
|
and int(row["max_failures"]) == desired.max_failures
|
|
and int(row["cooldown_seconds"]) == desired.cooldown_seconds
|
|
)
|
|
|
|
|
|
def _model_matches(row: Mapping[str, Any], desired: ModelRecord) -> bool:
|
|
return (
|
|
row["provider_id"] == desired.provider_id
|
|
and row["version"] == desired.version
|
|
and tuple(json.loads(row["capabilities_json"])) == tuple(sorted(set(desired.capabilities)))
|
|
and int(row["input_cost_micros_per_1k"]) == desired.input_cost_micros_per_1k
|
|
and int(row["output_cost_micros_per_1k"]) == desired.output_cost_micros_per_1k
|
|
and int(row["max_input_tokens"]) == desired.max_input_tokens
|
|
and int(row["max_output_tokens"]) == desired.max_output_tokens
|
|
and row["evaluation_state"] == desired.evaluation_state
|
|
and row["evaluation_sha256"] == desired.evaluation_sha256
|
|
and bool(row["enabled"]) == desired.enabled
|
|
and int(row["priority"]) == desired.priority
|
|
)
|
|
|
|
|
|
def _policy_matches(row: Mapping[str, Any], desired: ScopePolicy) -> bool:
|
|
return (
|
|
row["scope_type"] == desired.scope_type
|
|
and row["tenant_id"] == desired.tenant_id
|
|
and row["project_id"] == desired.project_id
|
|
and tuple(json.loads(row["allowed_models_json"])) == tuple(sorted(set(desired.allowed_models)))
|
|
and tuple(json.loads(row["denied_providers_json"])) == tuple(sorted(set(desired.denied_providers)))
|
|
and tuple(json.loads(row["required_capabilities_json"])) == tuple(sorted(set(desired.required_capabilities)))
|
|
and row["pinned_model_id"] == desired.pinned_model_id
|
|
and row["pinned_model_version"] == desired.pinned_model_version
|
|
and int(row["max_input_tokens_per_request"]) == desired.max_input_tokens_per_request
|
|
and int(row["max_output_tokens_per_request"]) == desired.max_output_tokens_per_request
|
|
and int(row["daily_token_budget"]) == desired.daily_token_budget
|
|
and int(row["daily_cost_budget_micros"]) == desired.daily_cost_budget_micros
|
|
and int(row["fallback_limit"]) == desired.fallback_limit
|
|
)
|
|
|
|
|
|
def apply_config(db_path: Path, config_path: Path) -> dict[str, Any]:
|
|
config = _load(config_path)
|
|
config_sha = hashlib.sha256(canonical(config)).hexdigest()
|
|
providers = config.get("providers") or []
|
|
models = config.get("models") or []
|
|
policies = config.get("policies") or []
|
|
if not all(isinstance(value, list) for value in (providers, models, policies)):
|
|
fail("MODEL_GOVERNANCE_CONFIG_COLLECTION_INVALID")
|
|
|
|
store = ModelGovernanceStore(db_path)
|
|
changes = {"providers": 0, "models": 0, "policies": 0, "unchanged": 0}
|
|
try:
|
|
for item in providers:
|
|
if not isinstance(item, Mapping):
|
|
fail("MODEL_GOVERNANCE_PROVIDER_CONFIG_INVALID")
|
|
desired = _provider_desired(item)
|
|
row = store.conn.execute(
|
|
"SELECT * FROM model_governance_providers WHERE provider_id=?",
|
|
(desired.provider_id,),
|
|
).fetchone()
|
|
if row is not None and _provider_matches(row, desired):
|
|
changes["unchanged"] += 1
|
|
continue
|
|
expected = int(row["revision"]) if row is not None else 0
|
|
store.register_provider(desired, expected_revision=expected)
|
|
changes["providers"] += 1
|
|
|
|
for item in models:
|
|
if not isinstance(item, Mapping):
|
|
fail("MODEL_GOVERNANCE_MODEL_CONFIG_INVALID")
|
|
desired = _model_desired(item)
|
|
row = store.conn.execute(
|
|
"SELECT * FROM model_governance_models WHERE model_id=?",
|
|
(desired.model_id,),
|
|
).fetchone()
|
|
if row is not None and _model_matches(row, desired):
|
|
changes["unchanged"] += 1
|
|
continue
|
|
expected = int(row["revision"]) if row is not None else 0
|
|
store.register_model(desired, expected_revision=expected)
|
|
changes["models"] += 1
|
|
|
|
for item in policies:
|
|
if not isinstance(item, Mapping):
|
|
fail("MODEL_GOVERNANCE_POLICY_CONFIG_INVALID")
|
|
desired = _policy_desired(item)
|
|
row = store.conn.execute(
|
|
"SELECT * FROM model_governance_scope_policies WHERE scope_key=?",
|
|
(desired.scope_key,),
|
|
).fetchone()
|
|
if row is not None and _policy_matches(row, desired):
|
|
changes["unchanged"] += 1
|
|
continue
|
|
expected = int(row["revision"]) if row is not None else 0
|
|
store.set_scope_policy(desired, expected_revision=expected)
|
|
changes["policies"] += 1
|
|
|
|
if not store.verify_audit_chain():
|
|
fail("MODEL_GOVERNANCE_AUDIT_CHAIN_INVALID")
|
|
return {
|
|
"schema": "aios.model-governance.bootstrap-result.v1",
|
|
"database": str(db_path),
|
|
"config_sha256": config_sha,
|
|
"changes": changes,
|
|
"audit_chain_valid": True,
|
|
}
|
|
finally:
|
|
store.close()
|
|
|
|
|
|
def main() -> int:
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--db", required=True)
|
|
parser.add_argument("--config", required=True)
|
|
parser.add_argument("--output")
|
|
args = parser.parse_args()
|
|
try:
|
|
result = apply_config(Path(args.db), Path(args.config))
|
|
except ModelGovernanceError as exc:
|
|
print(json.dumps({"status": "FAILED", "error": exc.code}, sort_keys=True))
|
|
return 20
|
|
data = json.dumps(result, indent=2, sort_keys=True) + "\n"
|
|
if args.output:
|
|
Path(args.output).write_text(data, encoding="utf-8")
|
|
print(data, end="")
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|