diff --git a/.github/skills/backend-standards/SKILL.md b/.github/skills/backend-standards/SKILL.md index 7c8d3bc35d..a36b22fb68 100644 --- a/.github/skills/backend-standards/SKILL.md +++ b/.github/skills/backend-standards/SKILL.md @@ -6,7 +6,7 @@ description: Backend conventions for the otari gateway (`src/gateway/`), async S # Backend Standards: otari gateway (`src/gateway/`) The gateway is an async FastAPI service: request handlers in `api/routes/`, business logic in -`services/`, ORM in `models/` (`entities.py` plus `tenancy.py`), migrations in +`services/`, ORM in `models/` (one module per domain), migrations in `alembic/versions/`. This guide is the backend counterpart to the frontend skill and to the path-scoped review instructions in `.github/instructions/` (performance and security). `AGENTS.md` is the source of truth for @@ -40,7 +40,7 @@ count = (await db.execute(select(func.count()).select_from(ModelPricing))).scala ## The SQLModel half: the reconciled control plane's tables `models/tenancy.py` (organizations, workspaces, identities, memberships) is SQLModel rather -than `entities.py`'s declarative style, because its `Create`/`Update`/`Public` schemas are the +than the declarative `Base` style of the other domain modules, because its `Create`/`Update`/`Public` schemas are the endpoint contracts the generated dashboard client is built from. Same session, same chain, three extra rules: diff --git a/ARCHITECTURE.md b/ARCHITECTURE.md index a48d6da702..6086df0002 100644 --- a/ARCHITECTURE.md +++ b/ARCHITECTURE.md @@ -39,7 +39,7 @@ Provider resolution is the seam on the request hot path, but it is not the only Today the control plane resolves two ways, selected by [mode](docs/modes.md): -- **Standalone** (default): the gateway resolves against its own local database (users, keys, budgets, usage in `src/gateway/models/entities.py`). This is the open-source control plane in its simplest form. +- **Standalone** (default): the gateway resolves against its own local database (users, keys, budgets, usage in `src/gateway/models/`). This is the open-source control plane in its simplest form. - **Hybrid** (`OTARI_AI_TOKEN` set): the gateway delegates resolution to a peer over HTTP (`src/gateway/api/routes/_platform.py`). Any service that implements the protocol can answer; otari.ai is the reference peer. Hybrid mode is a *network* form of this seam, and it is worth not conflating it with an overlay. In hybrid mode a remote control plane answers the resolve protocol over HTTP, out of the gateway's process; the peer can be any service that implements the protocol. An overlay, by contrast, is an *in-process* build that binds its own adapters into the composition container (see [How a port is resolved](#how-a-port-is-resolved)) and runs in the same process as the core. Both let something other than the plain local logic answer; the difference is whether that something runs over the network or in the same process. So a hybrid peer and an overlay are two ways to reach the seam, not the same thing. diff --git a/alembic/env.py b/alembic/env.py index 1bef59d6e8..cf57004486 100644 --- a/alembic/env.py +++ b/alembic/env.py @@ -11,10 +11,9 @@ from gateway.core.database import to_sync_url -# Importing anything from gateway.models registers every model module on this -# metadata (see gateway/models/__init__.py), which is what makes the comparison -# below cover the whole schema rather than the half this file names. -from gateway.models.entities import Base +# Importing any gateway.models module registers every table, so autogenerate +# compares against the whole schema. +from gateway.models.base import Base logger = logging.getLogger("alembic") diff --git a/docs/public/openapi.json b/docs/public/openapi.json index 1679dacfd2..5ea2835c4f 100644 --- a/docs/public/openapi.json +++ b/docs/public/openapi.json @@ -7740,7 +7740,7 @@ "type": "object" }, "OrgProviderKeyCreateRequest": { - "description": "What a caller sends to create a key.\n\nThe plaintext key is never stored as sent: the service encrypts it\n(`services/secret_box.py`) and keeps only the ciphertext and ``last4``,\nthe same convention `entities.ProviderCredential` already uses.", + "description": "What a caller sends to create a key.\n\nThe plaintext key is never stored as sent: the service encrypts it\n(`services/secret_box.py`) and keeps only the ciphertext and ``last4``,\nthe same convention `providers.ProviderCredential` already uses.", "properties": { "api_base": { "anyOf": [ @@ -8472,7 +8472,7 @@ "type": "object" }, "OrganizationGuardrailCreate": { - "description": "Request body for mandating a guardrail across an organization.\n\n``credential`` is never stored as sent: it is encrypted with\n``OTARI_SECRET_KEY`` and only the ciphertext is kept, the same convention\n`entities.WorkspaceMcpServer` and `entities.ProviderCredential` use. It is\nsent to the endpoint as ``Authorization: Bearer`` when the guardrail runs,\nso it authenticates this gateway to the guardrails service the entry names.\nA guardrail *vendor's* own key is not this: the guardrails service builds\nits guardrails from the operator's YAML and holds those itself.\n\nA credential therefore requires ``url``. The deployment's ``guardrails_url``\nis not necessarily encrypted (the shipped compose file makes it a same-host\n``http://`` sidecar) and this row cannot see what it is set to, so an entry\nthat fell back to it could not promise the bearer was sent over https. See\n`_require_url_for_credential`.\n\n``on`` is not offered. This plane mandates input-direction checks, which is\nthe only direction the request path enforces\n(`services.guardrails.run_input_guardrails`); an organization that could\nstore an output-direction mandate would be storing one nothing runs.", + "description": "Request body for mandating a guardrail across an organization.\n\n``credential`` is never stored as sent: it is encrypted with\n``OTARI_SECRET_KEY`` and only the ciphertext is kept, the same convention\n`tools.WorkspaceMcpServer` and `providers.ProviderCredential` use. It is\nsent to the endpoint as ``Authorization: Bearer`` when the guardrail runs,\nso it authenticates this gateway to the guardrails service the entry names.\nA guardrail *vendor's* own key is not this: the guardrails service builds\nits guardrails from the operator's YAML and holds those itself.\n\nA credential therefore requires ``url``. The deployment's ``guardrails_url``\nis not necessarily encrypted (the shipped compose file makes it a same-host\n``http://`` sidecar) and this row cannot see what it is set to, so an entry\nthat fell back to it could not promise the bearer was sent over https. See\n`_require_url_for_credential`.\n\n``on`` is not offered. This plane mandates input-direction checks, which is\nthe only direction the request path enforces\n(`services.guardrails.run_input_guardrails`); an organization that could\nstore an output-direction mandate would be storing one nothing runs.", "example": { "applies_to_all_workspaces": true, "credential": "sk-guardrails-...", @@ -16484,7 +16484,7 @@ "type": "object" }, "WorkspaceMcpServerCreate": { - "description": "Request body for registering a server.\n\n``authorization_token`` is never stored as sent: it is encrypted with\n``OTARI_SECRET_KEY`` and only the ciphertext is kept, the same convention\n`entities.ProviderCredential` and `OrgProviderKey` already use.", + "description": "Request body for registering a server.\n\n``authorization_token`` is never stored as sent: it is encrypted with\n``OTARI_SECRET_KEY`` and only the ciphertext is kept, the same convention\n`providers.ProviderCredential` and `OrgProviderKey` already use.", "properties": { "allowed_tools": { "anyOf": [ diff --git a/scripts/demo_gif/seed.py b/scripts/demo_gif/seed.py index 25610822d6..b1d8b3fe77 100644 --- a/scripts/demo_gif/seed.py +++ b/scripts/demo_gif/seed.py @@ -27,16 +27,14 @@ from sqlalchemy.orm import sessionmaker from sqlmodel import col -from gateway.models.entities import ( - APIKey, - Budget, - ModelAlias, - ModelPricing, - UsageLog, - User, -) +from gateway.models.api_keys import APIKey +from gateway.models.budgets import Budget from gateway.models.money import to_usd +from gateway.models.pricing import ModelPricing +from gateway.models.providers import ModelAlias from gateway.models.tenancy import Organization, Workspace +from gateway.models.usage import UsageLog +from gateway.models.users import User URL = sys.argv[1] if len(sys.argv) > 1 else "sqlite:///./scripts/demo_gif/demo.db" rng = random.Random(4242) # deterministic values across runs diff --git a/scripts/seed_usage_smoke.py b/scripts/seed_usage_smoke.py index cc5a7acb4b..f4763541ac 100644 --- a/scripts/seed_usage_smoke.py +++ b/scripts/seed_usage_smoke.py @@ -19,8 +19,11 @@ from sqlalchemy.orm import sessionmaker from sqlmodel import col -from gateway.models.entities import APIKey, ModelPricing, UsageLog, User +from gateway.models.api_keys import APIKey +from gateway.models.pricing import ModelPricing from gateway.models.tenancy import Organization, Workspace +from gateway.models.usage import UsageLog +from gateway.models.users import User from gateway.services.pricing_service import gateway_tool_pricing_key from gateway.services.tool_usage import TOOL_METER_NAMESPACE diff --git a/src/gateway/AGENTS.md b/src/gateway/AGENTS.md index 6e27f5c756..6f7120a7c7 100644 --- a/src/gateway/AGENTS.md +++ b/src/gateway/AGENTS.md @@ -209,12 +209,16 @@ endpoints. ## Data and migrations -Gateway ORM entities live in `models/entities.py`. Reconciled control-plane -SQLModel tables live in `models/tenancy.py`, and the newer tenancy-scoped -gateway tables whose `Public` schemas are endpoint contracts follow its style in -their own modules (`models/provider_keys.py`, `models/playground.py`). All of -them share `SQLModel.metadata`; `models/__init__.py` imports every table module -before Alembic uses it. +Put a table in its domain's model module (`models/budgets.py`, +`models/tenancy.py`, and so on). `models/base.py` holds `Base` and the shared +column types and mixins. Tables use the declarative `Base`, except those whose +`Public` schemas are endpoint contracts, which use SQLModel (`models/tenancy.py`, +`models/provider_keys.py`, `models/playground.py`). A new table module must join +the import list in `models/__init__.py`, or Alembic proposes dropping its tables. + +Two classes are named `User`: `models/users.py` is the billing identity that +keys, budgets, and usage attach to; `models/tenancy.py` is the dashboard sign-in +identity. Request code gets a session through `get_db`; non-request code uses `create_session()`; the usage-log writer uses `create_log_session()`, which diff --git a/src/gateway/adapters/telemetry_storage_adapter.py b/src/gateway/adapters/telemetry_storage_adapter.py index 2d0faff00b..086b3a87d1 100644 --- a/src/gateway/adapters/telemetry_storage_adapter.py +++ b/src/gateway/adapters/telemetry_storage_adapter.py @@ -24,7 +24,7 @@ from gateway.core.sql import bucket_expr, canonical_bucket, dialect_name, match_any, utc_bound from gateway.log_config import logger -from gateway.models.entities import AgentTelemetry +from gateway.models.usage import AgentTelemetry from gateway.ports.telemetry_storage_port import ( BehaviorCounts, BehaviorGroup, diff --git a/src/gateway/api/deps.py b/src/gateway/api/deps.py index 381b851231..94b71abdc3 100644 --- a/src/gateway/api/deps.py +++ b/src/gateway/api/deps.py @@ -16,7 +16,7 @@ from gateway.core.feature import CoreFeature from gateway.log_config import logger from gateway.metrics import REGISTRY, Counter -from gateway.models.entities import APIKey +from gateway.models.api_keys import APIKey from gateway.models.tenancy import User as TenancyUser from gateway.ports.billing_port import BillingPort from gateway.ports.entitlement_port import EntitlementPort diff --git a/src/gateway/api/routes/_passthrough.py b/src/gateway/api/routes/_passthrough.py index 8b722df18e..63d391fb27 100644 --- a/src/gateway/api/routes/_passthrough.py +++ b/src/gateway/api/routes/_passthrough.py @@ -51,7 +51,9 @@ from gateway.inflight import track_request from gateway.log_config import logger from gateway.model_labeling import relabel_model -from gateway.models.entities import APIKey, ModelPricing, UsageLog +from gateway.models.api_keys import APIKey +from gateway.models.pricing import ModelPricing +from gateway.models.usage import UsageLog from gateway.rate_limit import check_rate_limit from gateway.services.budget_service import ( ZERO, diff --git a/src/gateway/api/routes/_pipeline.py b/src/gateway/api/routes/_pipeline.py index 44761c29d8..c891b6fc96 100644 --- a/src/gateway/api/routes/_pipeline.py +++ b/src/gateway/api/routes/_pipeline.py @@ -121,10 +121,12 @@ from gateway.metrics import REGISTRY, Histogram from gateway.metrics import Counter as PrometheusCounter from gateway.model_labeling import relabel_model -from gateway.models.entities import APIKey, ModelPricing, UsageLog +from gateway.models.api_keys import APIKey from gateway.models.guardrails import GuardrailConfig from gateway.models.mcp import McpServerConfig from gateway.models.money import to_usd +from gateway.models.pricing import ModelPricing +from gateway.models.usage import UsageLog from gateway.ports.model_provider_port import HostedAccessDeniedError, ModelProviderPort from gateway.rate_limit import RateLimitInfo, check_rate_limit from gateway.services.budget_service import ( diff --git a/src/gateway/api/routes/agent_telemetry.py b/src/gateway/api/routes/agent_telemetry.py index b4c356b8a8..493940a3f9 100644 --- a/src/gateway/api/routes/agent_telemetry.py +++ b/src/gateway/api/routes/agent_telemetry.py @@ -38,7 +38,7 @@ _resolve_window, ) from gateway.core.sql import MAX_FILTER_VALUES, bucket_expr, canonical_bucket, dialect_name, match_any -from gateway.models.entities import UsageLog +from gateway.models.usage import UsageLog from gateway.ports.telemetry_storage_port import ( BehaviorCounts, TelemetryFilter, diff --git a/src/gateway/api/routes/aliases.py b/src/gateway/api/routes/aliases.py index b1209d85a6..fe8a162fa1 100644 --- a/src/gateway/api/routes/aliases.py +++ b/src/gateway/api/routes/aliases.py @@ -31,7 +31,7 @@ from gateway.api.routes._helpers import resolve_managed_workspace_id from gateway.core.config import GatewayConfig from gateway.log_config import logger -from gateway.models.entities import ModelAlias +from gateway.models.providers import ModelAlias from gateway.repositories.users_repository import get_active_user from gateway.services.alias_service import all_alias_names, refresh_alias_cache from gateway.services.policy_store import all_policy_names diff --git a/src/gateway/api/routes/audio.py b/src/gateway/api/routes/audio.py index 20e3c8e243..d60f3ccb90 100644 --- a/src/gateway/api/routes/audio.py +++ b/src/gateway/api/routes/audio.py @@ -14,7 +14,8 @@ from gateway.api.routes._schema_derive import derive_request_base from gateway.api.routes._tools import _strip_gateway_fields from gateway.core.config import GatewayConfig -from gateway.models.entities import APIKey, ModelPricing +from gateway.models.api_keys import APIKey +from gateway.models.pricing import ModelPricing from gateway.services.log_writer import LogWriter from gateway.services.pricing_service import flat_request_cost, per_request_meters from gateway.services.provider_kwargs import ResolvedProvider diff --git a/src/gateway/api/routes/batches.py b/src/gateway/api/routes/batches.py index 7843e1ed56..e4797bdd1d 100644 --- a/src/gateway/api/routes/batches.py +++ b/src/gateway/api/routes/batches.py @@ -24,7 +24,9 @@ from gateway.core.metered_pricing import calculate_token_cost, quantize_cost from gateway.core.usage import cache_read_tokens_of from gateway.log_config import logger -from gateway.models.entities import APIKey, BatchRecord, UsageLog +from gateway.models.api_keys import APIKey +from gateway.models.inference import BatchRecord +from gateway.models.usage import UsageLog from gateway.rate_limit import check_rate_limit from gateway.services.batch_service import ( claim_batch_accounting, diff --git a/src/gateway/api/routes/budgets.py b/src/gateway/api/routes/budgets.py index 96130c610a..3cc677781c 100644 --- a/src/gateway/api/routes/budgets.py +++ b/src/gateway/api/routes/budgets.py @@ -10,16 +10,10 @@ from sqlmodel import col from gateway.api.deps import get_db, require_deployment_operator -from gateway.models.entities import ( - MAX_COUNT_LIMIT, - Budget, - BudgetResetLog, - ScopedBudget, - User, - WorkspaceBudgetDefault, -) +from gateway.models.budgets import MAX_COUNT_LIMIT, Budget, BudgetResetLog, ScopedBudget, WorkspaceBudgetDefault from gateway.models.money import MAX_USD_LIMIT, as_float, to_usd, to_usd_or_none from gateway.models.tenancy import Workspace +from gateway.models.users import User from gateway.services.budget_retiming import cadence_of, retime_ceilings_for_budget from gateway.services.scoped_budget_service import ResetAlignment diff --git a/src/gateway/api/routes/catalog.py b/src/gateway/api/routes/catalog.py index fb8e30ddc2..39a34a1c94 100644 --- a/src/gateway/api/routes/catalog.py +++ b/src/gateway/api/routes/catalog.py @@ -44,9 +44,11 @@ ) from gateway.core.config import HOSTED_OFFERING_INSTANCE, GatewayConfig from gateway.core.metered_pricing import effective_rates -from gateway.models.entities import APIKey, PricingSnapshot, UsageLog +from gateway.models.api_keys import APIKey +from gateway.models.pricing import PricingSnapshot from gateway.models.tenancy import User as TenancyUser from gateway.models.tenancy import Workspace +from gateway.models.usage import UsageLog from gateway.services.catalog_selectors import ( current_selector_index, model_selector_for_slug, diff --git a/src/gateway/api/routes/embeddings.py b/src/gateway/api/routes/embeddings.py index a2c08c4ca5..701ac397a7 100644 --- a/src/gateway/api/routes/embeddings.py +++ b/src/gateway/api/routes/embeddings.py @@ -12,7 +12,8 @@ from gateway.api.deps import get_config, get_db, get_log_writer, verify_api_key_or_master_key from gateway.api.routes._passthrough import BillingMeters, run_passthrough from gateway.core.config import GatewayConfig -from gateway.models.entities import APIKey, ModelPricing +from gateway.models.api_keys import APIKey +from gateway.models.pricing import ModelPricing from gateway.services.budget_service import estimate_cost from gateway.services.log_writer import LogWriter from gateway.services.pricing_service import input_token_cost diff --git a/src/gateway/api/routes/files.py b/src/gateway/api/routes/files.py index a046d623db..9f51ebb707 100644 --- a/src/gateway/api/routes/files.py +++ b/src/gateway/api/routes/files.py @@ -31,7 +31,8 @@ from gateway.api.routes._helpers import resolve_user_id from gateway.core.config import GatewayConfig from gateway.log_config import logger -from gateway.models.entities import APIKey, FileObject +from gateway.models.api_keys import APIKey +from gateway.models.tools import FileObject from gateway.services.file_service import fetch_file from gateway.services.file_store import FileStore from gateway.services.workspace_scope import default_workspace_id diff --git a/src/gateway/api/routes/images.py b/src/gateway/api/routes/images.py index 1035507cba..fae9d4f671 100644 --- a/src/gateway/api/routes/images.py +++ b/src/gateway/api/routes/images.py @@ -13,7 +13,8 @@ from gateway.api.routes._schema_derive import derive_request_base from gateway.api.routes._tools import _strip_gateway_fields from gateway.core.config import GatewayConfig -from gateway.models.entities import APIKey, ModelPricing +from gateway.models.api_keys import APIKey +from gateway.models.pricing import ModelPricing from gateway.services.log_writer import LogWriter from gateway.services.pricing_service import per_image_cost from gateway.services.provider_kwargs import ResolvedProvider diff --git a/src/gateway/api/routes/keys.py b/src/gateway/api/routes/keys.py index 4e37a09fcf..b742500e0d 100644 --- a/src/gateway/api/routes/keys.py +++ b/src/gateway/api/routes/keys.py @@ -12,8 +12,9 @@ from gateway.api.deps import CallerOrganization, get_config, get_db, require_deployment_operator from gateway.auth.models import generate_api_key, hash_key, key_prefix, key_suffix from gateway.core.config import GatewayConfig -from gateway.models.entities import APIKey, User +from gateway.models.api_keys import APIKey from gateway.models.tenancy import Workspace +from gateway.models.users import User from gateway.repositories.users_repository import get_or_create_default_user, owned_by_organization from gateway.services.model_access import is_allowlist_subset, validate_allowed_models from gateway.services.workspace_scope import organization_default_workspace_id diff --git a/src/gateway/api/routes/mcp.py b/src/gateway/api/routes/mcp.py index 06f3b5eb68..e2f4ae8a32 100644 --- a/src/gateway/api/routes/mcp.py +++ b/src/gateway/api/routes/mcp.py @@ -54,7 +54,7 @@ from gateway.core.database import release_session from gateway.inflight import track_request from gateway.log_config import logger -from gateway.models.entities import APIKey +from gateway.models.api_keys import APIKey from gateway.rate_limit import check_rate_limit from gateway.repositories.users_repository import get_active_user diff --git a/src/gateway/api/routes/models.py b/src/gateway/api/routes/models.py index b3397c77cf..6903827d87 100644 --- a/src/gateway/api/routes/models.py +++ b/src/gateway/api/routes/models.py @@ -15,7 +15,8 @@ verify_catalog_reader, ) from gateway.core.config import GatewayConfig -from gateway.models.entities import APIKey, ModelPricing +from gateway.models.api_keys import APIKey +from gateway.models.pricing import ModelPricing from gateway.models.tenancy import User as TenancyUser from gateway.services.merged_catalog_service import ( ModelObject, diff --git a/src/gateway/api/routes/moderations.py b/src/gateway/api/routes/moderations.py index e5716b1e26..c9bd59a01d 100644 --- a/src/gateway/api/routes/moderations.py +++ b/src/gateway/api/routes/moderations.py @@ -11,7 +11,8 @@ from gateway.api.deps import get_config, get_db, get_log_writer, verify_api_key_or_master_key from gateway.api.routes._passthrough import BillingMeters, run_passthrough from gateway.core.config import GatewayConfig -from gateway.models.entities import APIKey, ModelPricing +from gateway.models.api_keys import APIKey +from gateway.models.pricing import ModelPricing from gateway.services.log_writer import LogWriter from gateway.services.pricing_service import flat_request_cost, per_request_meters from gateway.services.provider_kwargs import ResolvedProvider diff --git a/src/gateway/api/routes/organization_keys.py b/src/gateway/api/routes/organization_keys.py index 5f74290e5d..a6ad657a57 100644 --- a/src/gateway/api/routes/organization_keys.py +++ b/src/gateway/api/routes/organization_keys.py @@ -58,9 +58,10 @@ ) from gateway.auth.models import generate_api_key, hash_key, key_prefix, key_suffix from gateway.core.config import GatewayConfig -from gateway.models.entities import APIKey, User +from gateway.models.api_keys import APIKey from gateway.models.tenancy import User as TenancyUser from gateway.models.tenancy import Workspace +from gateway.models.users import User from gateway.ports.growth_signal_port import GrowthActivationEvent from gateway.repositories.users_repository import get_or_create_attribution_user from gateway.services.model_access import is_allowlist_subset, validate_allowed_models diff --git a/src/gateway/api/routes/organization_pricing.py b/src/gateway/api/routes/organization_pricing.py index 847792a75c..ada3788e10 100644 --- a/src/gateway/api/routes/organization_pricing.py +++ b/src/gateway/api/routes/organization_pricing.py @@ -29,8 +29,8 @@ from gateway.api.deps import CurrentIdentity, ModelProviderPortDep, get_config, get_db, verify_master_key from gateway.core.config import GatewayConfig -from gateway.models.entities import OrganizationModelPricing from gateway.models.money import as_float +from gateway.models.pricing import OrganizationModelPricing # The tier shape comes from the deployment pricing route rather than a second # copy here. An override resolves into a transient ``ModelPricing`` and is read by diff --git a/src/gateway/api/routes/organization_routing.py b/src/gateway/api/routes/organization_routing.py index 4aacb26bc9..12b8426a3a 100644 --- a/src/gateway/api/routes/organization_routing.py +++ b/src/gateway/api/routes/organization_routing.py @@ -77,8 +77,8 @@ ) from gateway.core.config import GatewayConfig from gateway.log_config import logger -from gateway.models.entities import ModelAlias, RoutingPolicy -from gateway.models.routing import PolicySpec +from gateway.models.providers import ModelAlias +from gateway.models.routing import PolicySpec, RoutingPolicy from gateway.models.tenancy import User as TenancyUser from gateway.models.tenancy import Workspace from gateway.services.alias_service import all_alias_names diff --git a/src/gateway/api/routes/organization_usage.py b/src/gateway/api/routes/organization_usage.py index 98a612f612..666c17c6cf 100644 --- a/src/gateway/api/routes/organization_usage.py +++ b/src/gateway/api/routes/organization_usage.py @@ -82,9 +82,11 @@ _usage_filters, ) from gateway.core.sql import MAX_FILTER_VALUES -from gateway.models.entities import APIKey, UsageLog, User +from gateway.models.api_keys import APIKey from gateway.models.tenancy import User as TenancyUser from gateway.models.tenancy import Workspace +from gateway.models.usage import UsageLog +from gateway.models.users import User from gateway.services.tenancy import OrganizationService from gateway.services.tenancy.authorization import ( resolve_visible_workspace_scope, diff --git a/src/gateway/api/routes/otlp.py b/src/gateway/api/routes/otlp.py index 6c102f3653..609e481477 100644 --- a/src/gateway/api/routes/otlp.py +++ b/src/gateway/api/routes/otlp.py @@ -64,7 +64,7 @@ from gateway.api.deps import TelemetryStoragePortDep, get_config, get_db, verify_api_key_or_master_key from gateway.core.config import GatewayConfig from gateway.log_config import logger -from gateway.models.entities import APIKey +from gateway.models.api_keys import APIKey from gateway.ports.telemetry_storage_port import TelemetryRecord from gateway.services.agent_telemetry_service import ( CUMULATIVE, diff --git a/src/gateway/api/routes/pricing.py b/src/gateway/api/routes/pricing.py index c0dee32043..db937b21ed 100644 --- a/src/gateway/api/routes/pricing.py +++ b/src/gateway/api/routes/pricing.py @@ -11,8 +11,8 @@ from gateway.api.deps import get_config, get_db, require_deployment_operator, verify_catalog_reader from gateway.core.config import GatewayConfig -from gateway.models.entities import ModelPricing from gateway.models.money import as_float, to_usd, to_usd_or_none +from gateway.models.pricing import ModelPricing from gateway.models.pricing_schemas import PricingTier from gateway.services.alias_service import all_alias_names, resolve_effective_alias from gateway.services.policy_store import all_policy_names, resolve_effective_policy diff --git a/src/gateway/api/routes/providers.py b/src/gateway/api/routes/providers.py index b954e4e8ed..d4badec5b2 100644 --- a/src/gateway/api/routes/providers.py +++ b/src/gateway/api/routes/providers.py @@ -25,7 +25,7 @@ from gateway.api.deps import get_config, get_db, require_deployment_operator, verify_catalog_reader from gateway.core.config import PROVIDER_TYPE_ALIASES, RESERVED_PROVIDER_INSTANCE_NAMES, GatewayConfig from gateway.log_config import logger -from gateway.models.entities import ProviderCredential +from gateway.models.providers import ProviderCredential from gateway.services.model_discovery_service import ( background_discovery_enabled, discover_provider_models, diff --git a/src/gateway/api/routes/rerank.py b/src/gateway/api/routes/rerank.py index dd9c8fe6dc..841a6c75b7 100644 --- a/src/gateway/api/routes/rerank.py +++ b/src/gateway/api/routes/rerank.py @@ -12,7 +12,8 @@ from gateway.api.deps import get_config, get_db, get_log_writer, verify_api_key_or_master_key from gateway.api.routes._passthrough import BillingMeters, run_passthrough from gateway.core.config import GatewayConfig -from gateway.models.entities import APIKey, ModelPricing +from gateway.models.api_keys import APIKey +from gateway.models.pricing import ModelPricing from gateway.services.budget_service import estimate_cost from gateway.services.log_writer import LogWriter from gateway.services.pricing_service import input_token_cost diff --git a/src/gateway/api/routes/routing.py b/src/gateway/api/routes/routing.py index 80659ee65d..0e05c8998d 100644 --- a/src/gateway/api/routes/routing.py +++ b/src/gateway/api/routes/routing.py @@ -30,8 +30,7 @@ from gateway.api.routes._helpers import resolve_managed_workspace_id from gateway.core.config import GatewayConfig from gateway.log_config import logger -from gateway.models.entities import RoutingPolicy -from gateway.models.routing import PolicySpec +from gateway.models.routing import PolicySpec, RoutingPolicy from gateway.repositories.users_repository import get_active_user from gateway.services.alias_service import all_alias_names from gateway.services.policy_store import ( diff --git a/src/gateway/api/routes/routing_memory.py b/src/gateway/api/routes/routing_memory.py index 11438702ee..ff4051f16e 100644 --- a/src/gateway/api/routes/routing_memory.py +++ b/src/gateway/api/routes/routing_memory.py @@ -52,7 +52,7 @@ from gateway.api.routes._helpers import resolve_managed_workspace_id from gateway.core.config import GatewayConfig from gateway.log_config import logger -from gateway.models.entities import RouterPreference, RoutingMemory +from gateway.models.routing import RouterPreference, RoutingMemory from gateway.repositories.users_repository import get_active_user from gateway.services.policy_store import effective_policies from gateway.services.provider_kwargs import resolve_provider_selector diff --git a/src/gateway/api/routes/scoped_budgets.py b/src/gateway/api/routes/scoped_budgets.py index 2f0c6b6a5a..6902c09dfe 100644 --- a/src/gateway/api/routes/scoped_budgets.py +++ b/src/gateway/api/routes/scoped_budgets.py @@ -17,7 +17,8 @@ from sqlalchemy.ext.asyncio import AsyncSession from gateway.api.deps import get_db, require_deployment_operator -from gateway.models.entities import APIKey, Budget, ScopedBudget +from gateway.models.api_keys import APIKey +from gateway.models.budgets import Budget, ScopedBudget from gateway.models.money import as_float from gateway.models.tenancy import Organization, OrganizationMember, Workspace, WorkspaceMember diff --git a/src/gateway/api/routes/search.py b/src/gateway/api/routes/search.py index 11bfa9a9ac..e313efc4f5 100644 --- a/src/gateway/api/routes/search.py +++ b/src/gateway/api/routes/search.py @@ -66,7 +66,8 @@ from gateway.core.metered_pricing import quantize_cost from gateway.inflight import track_request from gateway.log_config import logger -from gateway.models.entities import APIKey, UsageLog +from gateway.models.api_keys import APIKey +from gateway.models.usage import UsageLog from gateway.rate_limit import check_rate_limit from gateway.services.budget_service import reconcile_reservation, refund_reservation, reserve_budget from gateway.services.log_writer import LogWriter diff --git a/src/gateway/api/routes/search_tools.py b/src/gateway/api/routes/search_tools.py index ae44acf249..a24d972101 100644 --- a/src/gateway/api/routes/search_tools.py +++ b/src/gateway/api/routes/search_tools.py @@ -35,7 +35,7 @@ validate_search_tool_transport, ) from gateway.log_config import logger -from gateway.models.entities import SearchToolCredential +from gateway.models.tools import SearchToolCredential from gateway.services.search_backend import default_api_base from gateway.services.search_tool_store_service import ( UNSET, diff --git a/src/gateway/api/routes/usage.py b/src/gateway/api/routes/usage.py index caff53fe63..4aab006816 100644 --- a/src/gateway/api/routes/usage.py +++ b/src/gateway/api/routes/usage.py @@ -29,8 +29,10 @@ ) from gateway.core.usage_source import is_served_here, not_served_here from gateway.inflight import get_registry -from gateway.models.entities import APIKey, UsageLog, User +from gateway.models.api_keys import APIKey from gateway.models.money import as_float +from gateway.models.usage import UsageLog +from gateway.models.users import User from gateway.services.external_usage_service import ( ExternalEventsRequest, ExternalIngestResult, diff --git a/src/gateway/api/routes/users.py b/src/gateway/api/routes/users.py index ba44a9fbf2..3b263cb6f9 100644 --- a/src/gateway/api/routes/users.py +++ b/src/gateway/api/routes/users.py @@ -18,8 +18,11 @@ ) from gateway.core.config import GatewayConfig from gateway.log_config import logger -from gateway.models.entities import APIKey, Budget, UsageLog, User +from gateway.models.api_keys import APIKey +from gateway.models.budgets import Budget from gateway.models.money import as_float +from gateway.models.usage import UsageLog +from gateway.models.users import User from gateway.repositories.users_repository import in_organization from gateway.services.budget_periods import budget_window from gateway.services.model_access import validate_allowed_models diff --git a/src/gateway/db/__init__.py b/src/gateway/db/__init__.py index 23d82e3481..c4b02252f0 100644 --- a/src/gateway/db/__init__.py +++ b/src/gateway/db/__init__.py @@ -1,14 +1,10 @@ from gateway.core.database import create_session, get_db, init_db, reset_db -from gateway.models.entities import ( - APIKey, - Base, - Budget, - BudgetResetLog, - ModelPricing, - PricingSnapshot, - UsageLog, - User, -) +from gateway.models.api_keys import APIKey +from gateway.models.base import Base +from gateway.models.budgets import Budget, BudgetResetLog +from gateway.models.pricing import ModelPricing, PricingSnapshot +from gateway.models.usage import UsageLog +from gateway.models.users import User from gateway.repositories.users_repository import get_active_user __all__ = [ diff --git a/src/gateway/models/__init__.py b/src/gateway/models/__init__.py index 5a166d433f..24b38582ef 100644 --- a/src/gateway/models/__init__.py +++ b/src/gateway/models/__init__.py @@ -1,18 +1,26 @@ -"""ORM models, and the one place that guarantees the schema is whole. +"""ORM models. Importing any module in this package registers every table. -Importing any module in this package runs this file first, which imports every -model module in turn. That is what keeps ``Base.metadata`` (shared with -``SQLModel.metadata``, see `entities`) complete for the three operations that -are only correct against the entire schema: Alembic's autogenerate comparison, -``create_all``, and ``drop_all``. Without it, whether a table exists in the -metadata would depend on which model modules the caller happened to import, so -Alembic would propose dropping the tables it could not see, and a test's -``drop_all`` teardown would leave the rest behind for the next test to collide -with. +Alembic autogenerate and ``create_all`` need the whole schema in ``Base.metadata``. +A table module missing from the import list below looks deleted to Alembic, which +then proposes dropping its tables. -A new model module that declares tables belongs in the import list below; the -schema-less request/response modules beside them (`guardrails`, `mcp`, -`routing`) contribute nothing to the metadata and stay out of it. +Put a table in its domain's model module. A module that declares tables must be +in the list; one that declares none (`base`, `mcp`) stays out. """ -from gateway.models import entities, playground, provider_keys, tenancy # noqa: F401 +from gateway.models import ( # noqa: F401 + api_keys, + budgets, + guardrails, + inference, + platform, + playground, + pricing, + provider_keys, + providers, + routing, + tenancy, + tools, + usage, + users, +) diff --git a/src/gateway/models/api_keys.py b/src/gateway/models/api_keys.py new file mode 100644 index 0000000000..b8aa9fe888 --- /dev/null +++ b/src/gateway/models/api_keys.py @@ -0,0 +1,88 @@ +"""ORM table for API keys.""" + +import uuid +from datetime import UTC, datetime +from typing import Any + +from sqlalchemy import JSON, DateTime, ForeignKey, Uuid +from sqlalchemy.orm import Mapped, mapped_column, relationship + +from gateway.models.base import Base + + +class APIKey(Base): + """API Key model for authentication and authorization.""" + + __tablename__ = "api_keys" + + id: Mapped[str] = mapped_column(primary_key=True) + key_hash: Mapped[str] = mapped_column(unique=True, index=True) + # NOT NULL: every row belongs to a workspace. RESTRICT: deleting a workspace must + # not silently delete its keys, usage, aliases, and policies. + workspace_id: Mapped[uuid.UUID] = mapped_column( + Uuid, ForeignKey("workspace.id", ondelete="RESTRICT"), nullable=False, index=True + ) + # Display-only leading characters of the plaintext key, kept so the dashboard can + # recognize a key after its one-time reveal. Nullable: keys minted before this + # column existed cannot be back-filled (the plaintext is unrecoverable). + key_prefix: Mapped[str | None] = mapped_column() + # Display-only trailing characters, stored so the dashboard can tell two keys + # apart when they share a prefix. Nullable for the same reason as ``key_prefix`` + # and for one more: every key minted before this column will show prefix-only + # forever, because the plaintext is unrecoverable. Named ``key_suffix`` rather + # than the ``last4`` its provider-credential counterpart uses, so that the pair + # on this table reads as a pair. + key_suffix: Mapped[str | None] = mapped_column() + key_name: Mapped[str | None] = mapped_column() + user_id: Mapped[str | None] = mapped_column(ForeignKey("users.user_id", ondelete="CASCADE"), index=True) + created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=lambda: datetime.now(UTC)) + last_used_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True)) + expires_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True)) + is_active: Mapped[bool] = mapped_column(default=True) + # When true, requests authenticated with this key are logged with their computed + # cost but skip budget reservation/reconciliation: their spend is never written to + # User.spend and never gates enforcement. Default false keeps every existing key + # (and all keys minted before this column) on the normal enforced path. + exclude_from_budget: Mapped[bool] = mapped_column(default=False) + # Per-key override of the deployment-wide ``reject_user_mismatch`` setting. + # NULL = inherit (the default, and where every key predating this column + # stays), True = always reject a request naming a different ``user``, False = + # always accept it. The override only decides the 403: spend binds to this + # key's own user either way, so the client value stays a provider-side tag. + # False is for clients whose ``user`` is telemetry rather than an identity + # (Claude Code sends a per-session JSON blob); True lets a deployment that + # relaxed the check globally keep an individual key strict. + reject_user_mismatch: Mapped[bool | None] = mapped_column(default=None) + # Per-key override of the deployment-wide ``capture_agent_telemetry`` setting. + # NULL = inherit (the default), True = always store behavioral events from + # this key, False = always discard them. Usage capture/billing is unaffected + # either way; this only gates the content-free agent_telemetry row. + capture_agent_telemetry: Mapped[bool | None] = mapped_column(default=None) + # Per-key model allow-list. NULL = unrestricted (default; every key predating + # this column stays unrestricted), [] = deny all, a list = canonical + # instance:model entries (with instance:* / instance:prefix* wildcards). + allowed_models: Mapped[list[str] | None] = mapped_column(JSON) + + metadata_: Mapped[dict[str, Any]] = mapped_column("metadata", JSON, default=dict) + + user = relationship("User", back_populates="api_keys") + usage_logs = relationship("UsageLog", back_populates="api_key", passive_deletes=True) + + def to_dict(self) -> dict[str, Any]: + """Convert model to dictionary.""" + return { + "id": self.id, + "key_prefix": self.key_prefix, + "key_suffix": self.key_suffix, + "key_name": self.key_name, + "user_id": self.user_id, + "created_at": self.created_at.isoformat() if self.created_at else None, + "last_used_at": self.last_used_at.isoformat() if self.last_used_at else None, + "expires_at": self.expires_at.isoformat() if self.expires_at else None, + "is_active": self.is_active, + "exclude_from_budget": self.exclude_from_budget, + "reject_user_mismatch": self.reject_user_mismatch, + "capture_agent_telemetry": self.capture_agent_telemetry, + "allowed_models": self.allowed_models, + "metadata": self.metadata_, + } diff --git a/src/gateway/models/base.py b/src/gateway/models/base.py new file mode 100644 index 0000000000..d36f3e1ee0 --- /dev/null +++ b/src/gateway/models/base.py @@ -0,0 +1,122 @@ +"""The declarative base, and the column types and mixins the table modules share.""" + +import uuid +from datetime import UTC, datetime +from typing import Any + +from sqlalchemy import DateTime, func +from sqlalchemy.engine.interfaces import Dialect +from sqlalchemy.orm import DeclarativeBase +from sqlalchemy.types import TypeDecorator +from sqlmodel import Field, SQLModel + + +class Base(DeclarativeBase): + """Base class for SQLAlchemy models. + + Shares ``SQLModel.metadata`` so the reconciled control plane's SQLModel + tables (`gateway.models.tenancy`) and the gateway's own declarative tables + land in one collection. That is what lets Alembic keep a single + ``target_metadata``, and ``create_all``/``drop_all`` cover the whole schema, + without either style having to know the other exists. The two classes keep + separate declarative *registries*, so a same-named model on either side + (``User``, during the strangle) resolves unambiguously. + """ + + metadata = SQLModel.metadata + + +class UtcDateTime(TypeDecorator[datetime]): + """A timestamp that reads back UTC-aware on every engine. + + ``DateTime(timezone=True)`` alone is not enough, and the gap is the whole + reason this exists. PostgreSQL honors it and hands back an aware value; + SQLite has no timestamp type at all, so SQLAlchemy stores an ISO string and + the flag is a no-op, and a value written as ``datetime.now(UTC)`` reads back + with ``tzinfo=None``. A naive datetime then serializes with no offset, and a + browser parses an offset-less timestamp as **local** time, so every tenancy + timestamp in the dashboard would be wrong by the deployment's UTC offset on + the engine the OSS edition ships by default. + + Both directions are handled: an aware value is normalized to UTC before it + is stored, so a caller in another zone cannot write a wall-clock time that + means something else, and a naive value read back is stamped UTC, because + UTC is what everything here writes. + + The rendered DDL is exactly ``impl``'s, so this changes no migration and + ``compare_metadata`` stays clean. + """ + + impl = DateTime(timezone=True) + cache_ok = True + + def process_bind_param(self, value: datetime | None, dialect: Dialect) -> datetime | None: + if value is None: + return None + if value.utcoffset() is None: + # Refused rather than assumed. Reading a naive value back as UTC is + # safe, because UTC is what everything here writes; writing one is + # not, because the engines disagree about what it means. PostgreSQL + # interprets it in the *session* time zone, so the same value lands + # as a different instant depending on who connected, while SQLite + # stores the wall clock as written. Silently picking one is how a + # timestamp ends up hours off with nothing to show for it. + msg = "A tenancy timestamp must be timezone-aware; got a naive datetime" + raise ValueError(msg) + return value.astimezone(UTC) + + def process_result_value(self, value: datetime | None, dialect: Dialect) -> datetime | None: + if value is not None and value.tzinfo is None: + return value.replace(tzinfo=UTC) + return value + + +def _timestamp_field(*, default: Any = None, default_factory: Any = None, column_kwargs: dict[str, Any]) -> Any: + """Build a timezone-aware timestamp field. + + Two things are worked around here, once, instead of at five inheriting + tables. SQLModel's ``Field`` overloads type ``sa_type`` as a *class*, while + the type we want is an *instance* (the runtime accepts either and hands it + straight to ``Column``). And the type has to arrive as ``sa_type`` rather + than a ready-made ``sa_column``, because a ``Column`` instance declared on a + mixin cannot be attached to more than one table; ``sa_type`` plus kwargs + lets SQLModel build a fresh column per model. + """ + if default_factory is not None: + return Field( # type: ignore[call-overload] + default_factory=default_factory, + sa_type=UtcDateTime(), + sa_column_kwargs=column_kwargs, + ) + return Field( # type: ignore[call-overload] + default=default, + sa_type=UtcDateTime(), + sa_column_kwargs=column_kwargs, + ) + + +class PrimaryKeyMixin: + """A UUID primary key, rendered as CHAR(32) on SQLite and native on PostgreSQL.""" + + id: uuid.UUID = Field(default_factory=uuid.uuid4, primary_key=True) + + +class CreatedAtMixin: + """Creation timestamp, defaulted in Python and in the database.""" + + created_at: datetime = _timestamp_field( + default_factory=lambda: datetime.now(UTC), + column_kwargs={"server_default": func.now()}, + ) + + +class UpdatedAtMixin: + """Last-modification timestamp, stamped by the database on update. + + ``default=None`` and not merely a nullable annotation: without an explicit + default the field is *required* on the pydantic side, which a table class + hides (table models skip construction validation) and any schema inheriting + this mixin would not. + """ + + updated_at: datetime | None = _timestamp_field(default=None, column_kwargs={"onupdate": func.now()}) diff --git a/src/gateway/models/budgets.py b/src/gateway/models/budgets.py new file mode 100644 index 0000000000..f56c1bf65e --- /dev/null +++ b/src/gateway/models/budgets.py @@ -0,0 +1,431 @@ +"""ORM tables for budgets: limits, scoped ceilings, reservations, and workspace defaults.""" + +import uuid +from datetime import UTC, datetime +from decimal import Decimal +from typing import Any + +from sqlalchemy import BigInteger, CheckConstraint, DateTime, ForeignKey, Index, Uuid, false, text +from sqlalchemy.orm import Mapped, mapped_column, relationship + +from gateway.models.base import Base, UtcDateTime +from gateway.models.money import UsdCost + +# The largest token or request limit a route accepts, and the largest hold it +# will place. Well below the BIGINT ceiling those columns are, for the two +# reasons :data:`gateway.models.money.MAX_USD_LIMIT` keeps its own headroom. +# +# The wire cannot carry the type's maximum. JSON numbers are doubles, so +# ``9223372036854775807`` renders in the published schema as +# ``9.223372036854776e+18``, which is 9223372036854775808: a client sending +# exactly the maximum the spec advertises sends a value the column refuses. A +# quadrillion is exact as a double, so the schema says what it means. +# +# And the gate adds server-side. Every reserve evaluates +# ``current + reserved + held`` as BIGINT arithmetic, so a nonzero counter plus a +# hold near the type's ceiling overflows and answers with a 500 where a 403 was +# owed. Holds are clamped to this too, because the token estimate derives from a +# client-supplied output bound that nothing else limits. +# +# A quadrillion tokens is four orders of magnitude past any real allowance, and +# leaves the sum of three of them ~9000x inside the type. +MAX_COUNT_LIMIT = 1_000_000_000_000_000 + + +class Budget(Base): + """Budget model for spending limits.""" + + __tablename__ = "budgets" + __table_args__ = ( + # A period comes from one place or the other, never both, matching the + # rule ``scoped_budgets`` already enforced when it carried its own. Without + # it the pair encodes one concept twice and ``(86400, calendar_month)`` is + # storable and meaningless. + CheckConstraint( + "NOT (budget_duration_sec IS NOT NULL AND reset_alignment IS NOT NULL)", + name="ck_budgets_single_period_source", + ), + ) + + budget_id: Mapped[str] = mapped_column(primary_key=True, default=lambda: str(uuid.uuid4())) + name: Mapped[str | None] = mapped_column(default=None) + # Which tenant defined this budget, and therefore who may change it. + # + # NULL means the deployment's own: every budget predating `b7e1c4a9d2f5` + # reads NULL, and so does every one the otari-ai cutover migration mints, + # because that migration deliberately shares one budget per distinct + # (cap, period) shape across the ceilings it writes and a shape shared by two + # tenants' ceilings has no single owner. NULL is not "unowned and up for + # grabs": the organization-scoped surface never lists, offers or repoints one, + # so from a tenant's side it does not exist. + # + # Nullable and set only on the tenant-scoped create path, which is what keeps + # that migration working with no backfill: its preflight refuses on a + # *missing* column and its inserts name theirs explicitly, so a new nullable + # one is invisible to it. + organization_id: Mapped[uuid.UUID | None] = mapped_column( + Uuid(), + ForeignKey("organization.id", name="fk_budgets_organization_id", ondelete="CASCADE"), + default=None, + index=True, + ) + # Exact, like the counters it is compared against: the gate is + # ``spend + reserved <= max_budget``, and a cap stored as a binary float + # would decide a 403 against an amount an operator never typed + # (mozilla-ai/otari#691). + max_budget: Mapped[Decimal | None] = mapped_column(UsdCost()) + # The non-USD ceilings, independent of ``max_budget`` and of each other: a + # budget may cap dollars, tokens, requests, or any combination, and a NULL on + # an axis is unbounded there. Deliberately no "at least one limit" constraint, + # because a budget with every limit NULL is a named period that admits + # everything and predates these columns. + # + # BIGINT: a monthly token allowance for one organization outgrows a 32-bit + # counter, and the counters compared against these are the same width. + token_limit: Mapped[int | None] = mapped_column(BigInteger(), default=None) + request_limit: Mapped[int | None] = mapped_column(BigInteger(), default=None) + budget_duration_sec: Mapped[int | None] = mapped_column() + # Snap the window to a UTC calendar boundary instead of counting a fixed + # number of seconds, which is the only way to express a calendar month (2592000 + # seconds is a different, 1.5 percent more generous, product). It lives here + # rather than on the rows that enforce a budget because a limit and the period + # it is spent over are one product decision, and splitting them let a ceiling + # reset on a cadence the budget defining it had never heard of. + reset_alignment: Mapped[str | None] = mapped_column(default=None) + created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=lambda: datetime.now(UTC)) + updated_at: Mapped[datetime] = mapped_column( + DateTime(timezone=True), + default=lambda: datetime.now(UTC), + onupdate=lambda: datetime.now(UTC), + ) + + users = relationship("User", back_populates="budget") + reset_logs = relationship("BudgetResetLog", back_populates="budget") + + def to_dict(self) -> dict[str, Any]: + """Convert model to dictionary.""" + return { + "budget_id": self.budget_id, + "name": self.name, + "max_budget": self.max_budget, + "token_limit": self.token_limit, + "request_limit": self.request_limit, + "budget_duration_sec": self.budget_duration_sec, + "reset_alignment": self.reset_alignment, + "created_at": self.created_at.isoformat() if self.created_at else None, + "updated_at": self.updated_at.isoformat() if self.updated_at else None, + } + + +class BudgetResetLog(Base): + """Budget reset log model for tracking budget resets.""" + + __tablename__ = "budget_reset_logs" + + id: Mapped[int] = mapped_column(primary_key=True, autoincrement=True) + user_id: Mapped[str | None] = mapped_column(ForeignKey("users.user_id", ondelete="SET NULL"), index=True) + # Indexed: the reset-log drill-down filters on this column, and the table only + # grows, so an unindexed FK degrades that endpoint to a full scan over time. + budget_id: Mapped[str] = mapped_column(ForeignKey("budgets.budget_id"), index=True) + # The ledger's record of a counter that is now exact, so it is exact too: + # a float snapshot of an exact ``users.spend`` would no longer equal the + # spend it claims to have recorded. + previous_spend: Mapped[Decimal] = mapped_column(UsdCost()) + reset_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=lambda: datetime.now(UTC)) + next_reset_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True)) + + user = relationship("User", back_populates="reset_logs") + budget = relationship("Budget", back_populates="reset_logs") + + def to_dict(self) -> dict[str, Any]: + """Convert model to dictionary.""" + return { + "id": self.id, + "user_id": self.user_id, + "budget_id": self.budget_id, + "previous_spend": self.previous_spend, + "reset_at": self.reset_at.isoformat() if self.reset_at else None, + "next_reset_at": self.next_reset_at.isoformat() if self.next_reset_at else None, + } + + +class ScopedBudget(Base): + """A spending ceiling on one tenancy scope, optionally narrowed to one provider. + + Two axes. The identity axis is ``(scope_type, scope_id)``: who is capped, an + organization, a workspace, a member of either, or a single API key. The + resource axis is ``provider_key_id``: NULL caps spend across every provider, + a value narrows the cap to one provider instance. A request must pass every + row that applies to it, and each row is an independent ceiling with its own + counters and its own period window, unlike ``budgets``, where the window and + the counters live on the user. + + No limit is stored here. A limit is a property of the budget this names, + which is the only place in the schema that maps a cap to a figure, on any of + the three axes it can cap. + + ``scope_type`` is a plain string rather than a database enum so a new scope + needs no enum migration, and ``scope_id`` is a string so it holds both this + codebase's string ids (an API key's) and the platform's UUIDs. Nothing here + is a foreign key for the same reason: the rows a scope names live in four + different tables, and a provider instance may be configured in ``config.yml`` + and have no row at all. + + A row names a ``budgets`` row and holds the counters for spending it. The + limit and the period are read through the budget, never copied, so editing a + budget moves every ceiling that names it. That is deliberate: a budget is a + named thing an operator hands out, and the alternative was the same figure + typed once per place it applied. + + This table does not replace ``budgets``, and the two enforce differently. A + budget reached through ``users.budget_id`` is checked against + ``users.spend + users.reserved``, so N users sharing one each get the full + limit. A budget reached through a row here is checked against *this row's* + counters, so everyone the scope names draws on one allowance. Same budget, + two enforcement shapes, which is why both mechanisms exist. + """ + + __tablename__ = "scoped_budgets" + __table_args__ = ( + # PostgreSQL treats NULLs as distinct in a plain UNIQUE, so one index + # over the triple would enforce nothing on the aggregate rows (every one + # of them has a NULL key, so no two are ever "equal"). Two partial + # indexes instead: the narrowed rows are unique on the triple, and the + # aggregate rows are unique on the identity alone, which is what makes + # "one aggregate cap per scope" a real constraint. + Index( + "uq_scoped_budgets_scope_with_key", + "scope_type", + "scope_id", + "provider_key_id", + unique=True, + postgresql_where=text("provider_key_id IS NOT NULL"), + sqlite_where=text("provider_key_id IS NOT NULL"), + ), + Index( + "uq_scoped_budgets_scope_no_key", + "scope_type", + "scope_id", + unique=True, + postgresql_where=text("provider_key_id IS NULL"), + sqlite_where=text("provider_key_id IS NULL"), + ), + # The request path resolves rows by identity, so the lookup needs a + # non-partial index: neither unique index above covers a scan that spans + # narrowed and aggregate rows. + Index("ix_scoped_budgets_scope", "scope_type", "scope_id"), + ) + + id: Mapped[str] = mapped_column(primary_key=True, default=lambda: str(uuid.uuid4())) + scope_type: Mapped[str] = mapped_column() + scope_id: Mapped[str] = mapped_column() + provider_key_id: Mapped[str | None] = mapped_column(default=None) + name: Mapped[str | None] = mapped_column(default=None) + # The budget this ceiling enforces. NOT NULL: a ceiling with no budget caps + # nothing. The limit and the period are read through it rather than copied, so + # editing a budget moves every ceiling that names it, which is the point of a + # budget being a named thing rather than a number typed twice. + budget_id: Mapped[str] = mapped_column( + ForeignKey("budgets.budget_id", ondelete="RESTRICT"), nullable=False, index=True + ) + current_spend: Mapped[Decimal] = mapped_column(UsdCost(), default=Decimal(0), server_default="0") + # In-flight holds from reservations that have passed the gate but whose actual + # cost is not known yet. Headroom is ``max_budget - current_spend - + # reserved_spend``; a period roll zeroes ``current_spend`` only, so a hold + # taken before the roll is still released correctly after it. + reserved_spend: Mapped[Decimal] = mapped_column(UsdCost(), default=Decimal(0), server_default="0") + # One counter pair per non-USD axis the budget can cap, holding and settling + # exactly as the money pair above does. A period roll zeroes the ``current_*`` + # of all three axes and leaves every hold, so a hold taken before a roll is + # still released correctly after it. + current_tokens: Mapped[int] = mapped_column(BigInteger(), default=0, server_default="0") + reserved_tokens: Mapped[int] = mapped_column(BigInteger(), default=0, server_default="0") + current_requests: Mapped[int] = mapped_column(BigInteger(), default=0, server_default="0") + reserved_requests: Mapped[int] = mapped_column(BigInteger(), default=0, server_default="0") + period_start: Mapped[datetime | None] = mapped_column(DateTime(timezone=True), default=None) + period_end: Mapped[datetime | None] = mapped_column(DateTime(timezone=True), default=None) + created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=lambda: datetime.now(UTC)) + updated_at: Mapped[datetime] = mapped_column( + DateTime(timezone=True), + default=lambda: datetime.now(UTC), + onupdate=lambda: datetime.now(UTC), + ) + + +class BudgetReservation(Base): + """One in-flight budget hold, recorded as a row. + + ``users.reserved`` and ``scoped_budgets.reserved_spend`` stay the O(1) + counters the gate reads; this is the ledger behind them, and it exists for + the two things a counter cannot do (mozilla-ai/otari#742): + + * **Release becomes idempotent.** Without an identity, a second release for + the same request silently subtracts the hold twice. ``_release_reserved`` + clamps at zero, so that shows up not as an error but as an under-count of + live holds, which weakens the very overspend guarantee the reserve gate + exists to provide. The status transition here is what makes only the first + release do the work. + * **A leaked hold becomes reclaimable individually.** A failure between + reserve and settle used to leave an amount in the counter that could be + seen only in aggregate and released by nothing at all: the budget reset + zeroes ``spend`` and leaves ``reserved`` where it is. With a row it has an + owner, an age and a TTL. + + The row is written *after* the holds it records, never before. A hold with no + row is the pre-existing leak the sweep bounds; a row with no hold would have + the sweep release an amount nobody holds, under-counting the live ones. Of + the two inconsistent windows only one is safe, and this is it. + + Standalone mode only: hybrid mode reserves nothing locally, because the + platform holds against its own ledger. + """ + + __tablename__ = "budget_reservations" + __table_args__ = ( + # The global sweep's access path: active rows whose TTL has elapsed. + # Equality on ``status`` leads so the range scan on ``expires_at`` rides + # the same index. + Index("ix_budget_reservations_status_expires_at", "status", "expires_at"), + # The per-user reclaim's, which runs on every request that takes a hold. + # It has to lead on ``user_id``: given only the index above, the planner + # takes it and filters ``user_id``, so one user's reclaim pays for the + # whole deployment's backlog of expired rows. Leading on ``user_id`` also + # serves the FK cascade, so this replaces the plain index on that column + # rather than joining it. + Index("ix_budget_reservations_user_status_expires", "user_id", "status", "expires_at"), + ) + + id: Mapped[str] = mapped_column(primary_key=True, default=lambda: str(uuid.uuid4())) + # No ``index=True``: the composite in ``__table_args__`` leads on this column, + # so a plain index here would be a second, redundant one, and the migration + # deliberately does not create it. Declaring it anyway made a ``create_all`` + # schema and a migrated one disagree. + user_id: Mapped[str] = mapped_column(ForeignKey("users.user_id", ondelete="CASCADE"), nullable=False) + # What the per-user leg holds in ``users.reserved``. Zero when the request + # held only scoped ceilings (a user with no budget row still passes those). + estimate: Mapped[Decimal] = mapped_column(UsdCost(), default=Decimal(0), server_default="0") + # What the same leg holds on the other two axes. Recorded per axis because the + # sweep has to give back every axis a leaked hold took: a period roll zeroes + # ``current_*`` and deliberately leaves the holds, so a token hold nothing + # releases shrinks that ceiling for good. + token_estimate: Mapped[int] = mapped_column(BigInteger(), default=0, server_default="0") + request_estimate: Mapped[int] = mapped_column(BigInteger(), default=0, server_default="0") + # Whether the ``users.reserved`` write actually happened. Distinct from + # ``estimate > 0`` because a zero-cost request on an enforced budget still + # takes the hold, and the release has to match what the reserve did. + user_reserved: Mapped[bool] = mapped_column(default=False, server_default=false()) + # A plain string rather than a database enum, matching ``scoped_budgets.scope_type``: + # a new state should not need an enum migration. Values are the + # ``RESERVATION_*`` constants in gateway.services.budget_reservation_ledger. + status: Mapped[str] = mapped_column(default="active", server_default="active", nullable=False) + # After this instant a still-active row is treated as leaked and reclaimed. + expires_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), nullable=False) + created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=lambda: datetime.now(UTC)) + updated_at: Mapped[datetime] = mapped_column( + DateTime(timezone=True), + default=lambda: datetime.now(UTC), + onupdate=lambda: datetime.now(UTC), + ) + + +class BudgetReservationScope(Base): + """The hold one reservation placed on one scoped ceiling. + + ``scoped_budget_id`` is deliberately not a foreign key, following + ``ScopedBudget``'s own convention: a ceiling deleted while a request is in + flight leaves an orphan line that the release skips, rather than forcing the + delete to cascade into live holds. + + The amounts are stored per line, one per axis, even though today every + ceiling of a request holds the same figures. A ledger line that does not say + what it holds is not a ledger line, and reading the amounts from the parent + would silently become wrong the first time the two diverge. + """ + + __tablename__ = "budget_reservation_scopes" + __table_args__ = (Index("ix_budget_reservation_scopes_reservation_id", "reservation_id"),) + + id: Mapped[str] = mapped_column(primary_key=True, default=lambda: str(uuid.uuid4())) + reservation_id: Mapped[str] = mapped_column( + ForeignKey("budget_reservations.id", ondelete="CASCADE"), nullable=False + ) + scoped_budget_id: Mapped[str] = mapped_column(nullable=False) + amount: Mapped[Decimal] = mapped_column(UsdCost(), default=Decimal(0), server_default="0") + token_amount: Mapped[int] = mapped_column(BigInteger(), default=0, server_default="0") + request_amount: Mapped[int] = mapped_column(BigInteger(), default=0, server_default="0") + + +class WorkspaceBudgetDefault(Base): + """A workspace-level template for a per-member ``ScopedBudget``. + + ``scoped_budgets`` holds concrete ceilings; this table has no counters of + its own and enforces nothing directly. It is **materialized**: creating one + on a workspace that already has members, or a member joining a workspace + that already has one, stages a ``ScopedBudget(scope_type="workspace_member", + scope_id=)`` row for each (see + ``services/tenancy/workspace_budget_default_service.py``). A member with an + existing ceiling for the same ``provider_key_id`` is left alone; a + member-specific override always wins over the template. + + Same two-axis shape as ``ScopedBudget``: ``workspace_id`` is who the + template belongs to, ``provider_key_id`` optionally narrows it to one + provider instance (NULL applies to all of them). Unlike ``ScopedBudget``, + ``workspace_id`` is a real foreign key: a template has exactly one owner + and nothing else names it, so it is deleted with the workspace rather than + requiring the same explicit cleanup ``ScopedBudget`` needs (see + ``WorkspaceService._delete_scoped_budgets_for``). + """ + + __tablename__ = "workspace_budget_defaults" + __table_args__ = ( + # Same reasoning as ScopedBudget's two partial indexes: PostgreSQL and + # SQLite both treat NULLs as distinct in a plain UNIQUE, so a single + # index over the pair would enforce nothing on the aggregate (NULL-key) + # rows. + Index( + "uq_workspace_budget_defaults_with_key", + "workspace_id", + "provider_key_id", + unique=True, + postgresql_where=text("provider_key_id IS NOT NULL"), + sqlite_where=text("provider_key_id IS NOT NULL"), + ), + Index( + "uq_workspace_budget_defaults_no_key", + "workspace_id", + unique=True, + postgresql_where=text("provider_key_id IS NULL"), + sqlite_where=text("provider_key_id IS NULL"), + ), + ) + + id: Mapped[str] = mapped_column(primary_key=True, default=lambda: str(uuid.uuid4())) + workspace_id: Mapped[uuid.UUID] = mapped_column( + Uuid, ForeignKey("workspace.id", ondelete="CASCADE"), nullable=False, index=True + ) + provider_key_id: Mapped[str | None] = mapped_column(default=None) + # The budget this workspace hands to every member. NOT NULL: a default that + # names no budget is a template for nothing. ``RESTRICT`` because deleting a + # budget a workspace hands out should be refused and explained rather than + # silently withdraw the limit from every ceiling it materialized. + # + # The limit and the period live on the budget, not here, which is what lets + # the Budgets page say that a row is a workspace's default. ``provider_key_id`` + # stays on this side: which provider a workspace applies the budget to is a + # property of the assignment, and two workspaces may narrow one budget + # differently. + budget_id: Mapped[str] = mapped_column( + ForeignKey("budgets.budget_id", ondelete="RESTRICT"), nullable=False, index=True + ) + # ``UtcDateTime``, not ``DateTime(timezone=True)``: these two are serialized with + # ``.isoformat()`` (``WorkspaceMemberBudgetPolicyPublic.from_model``) for the + # dashboard, and on SQLite (this repo's default ``database_url``) + # a plain ``DateTime(timezone=True)`` round-trips naive, so the wire value + # would carry no offset and a browser would read it as local time. + # ``UtcDateTime.impl`` is ``DateTime(timezone=True)``, so the DDL is unchanged. + created_at: Mapped[datetime] = mapped_column(UtcDateTime(), default=lambda: datetime.now(UTC)) + updated_at: Mapped[datetime] = mapped_column( + UtcDateTime(), + default=lambda: datetime.now(UTC), + onupdate=lambda: datetime.now(UTC), + ) diff --git a/src/gateway/models/entities.py b/src/gateway/models/entities.py deleted file mode 100644 index e7b748e552..0000000000 --- a/src/gateway/models/entities.py +++ /dev/null @@ -1,2016 +0,0 @@ -import uuid -from datetime import UTC, datetime -from decimal import Decimal -from typing import Any - -from sqlalchemy import ( - JSON, - BigInteger, - CheckConstraint, - DateTime, - ForeignKey, - Index, - String, - Text, - UniqueConstraint, - Uuid, - false, - func, - text, - true, -) -from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column, relationship -from sqlmodel import SQLModel - -# The timezone-aware timestamp type the tenancy tables already use. Imported -# rather than redefined: it exists because the engines disagree about -# ``timezone=True``, and two copies of that reasoning would drift. -from gateway.models.money import UsdCost, UsdRate -from gateway.models.secret_fields import redact_secret_like_values -from gateway.models.tenancy import UtcDateTime - -# The vocabulary of ``ModelPricing.unit`` and ``OrganizationModelPricing.unit``. -# Every per-unit reader (``services/pricing_service`` helpers, the catalog) keys -# on these spellings, and the request schemas validate against them. -PRICING_UNITS: tuple[str, ...] = ("tokens", "requests", "images") - -# The vocabulary of ``origin`` on the same two tables. -PRICING_ORIGINS: tuple[str, ...] = ("config", "api", "migration") - - -class Base(DeclarativeBase): - """Base class for SQLAlchemy models. - - Shares ``SQLModel.metadata`` so the reconciled control plane's SQLModel - tables (`gateway.models.tenancy`) and the gateway's own declarative tables - land in one collection. That is what lets Alembic keep a single - ``target_metadata``, and ``create_all``/``drop_all`` cover the whole schema, - without either style having to know the other exists. The two classes keep - separate declarative *registries*, so a same-named model on either side - (``User``, during the strangle) resolves unambiguously. - """ - - metadata = SQLModel.metadata - - -def _epoch_seconds(value: datetime | None) -> int | None: - """Return a UTC epoch from a stored datetime. - - SQLite hands datetimes back naive; ``datetime.timestamp()`` would then read - them as local time and skew the epoch by the server's UTC offset. Treat a - naive value as the UTC it was stored as before converting. - """ - if value is None: - return None - if value.tzinfo is None: - value = value.replace(tzinfo=UTC) - return int(value.timestamp()) - - -class APIKey(Base): - """API Key model for authentication and authorization.""" - - __tablename__ = "api_keys" - - id: Mapped[str] = mapped_column(primary_key=True) - key_hash: Mapped[str] = mapped_column(unique=True, index=True) - # The workspace this row belongs to, and the canonical note for the three - # tables below that carry the same column. NOT NULL: a workspace is the unit - # the dashboard scopes by, so "no workspace" is never a real state, only an - # unmigrated one. Existing rows were backfilled onto the deployment's default - # workspace, which the same migration seeds when tenancy was never touched. - # RESTRICT rather than cascade: deleting a workspace must not silently take - # its keys, usage, aliases and policies with it. Which workspace a write - # lands in is resolved in `services/workspace_scope.py`. - workspace_id: Mapped[uuid.UUID] = mapped_column( - Uuid, ForeignKey("workspace.id", ondelete="RESTRICT"), nullable=False, index=True - ) - # Display-only leading characters of the plaintext key, kept so the dashboard can - # recognize a key after its one-time reveal. Nullable: keys minted before this - # column existed cannot be back-filled (the plaintext is unrecoverable). - key_prefix: Mapped[str | None] = mapped_column() - # Display-only trailing characters, stored so the dashboard can tell two keys - # apart when they share a prefix. Nullable for the same reason as ``key_prefix`` - # and for one more: every key minted before this column will show prefix-only - # forever, because the plaintext is unrecoverable. Named ``key_suffix`` rather - # than the ``last4`` its provider-credential counterpart uses, so that the pair - # on this table reads as a pair. - key_suffix: Mapped[str | None] = mapped_column() - key_name: Mapped[str | None] = mapped_column() - user_id: Mapped[str | None] = mapped_column(ForeignKey("users.user_id", ondelete="CASCADE"), index=True) - created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=lambda: datetime.now(UTC)) - last_used_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True)) - expires_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True)) - is_active: Mapped[bool] = mapped_column(default=True) - # When true, requests authenticated with this key are logged with their computed - # cost but skip budget reservation/reconciliation: their spend is never written to - # User.spend and never gates enforcement. Default false keeps every existing key - # (and all keys minted before this column) on the normal enforced path. - exclude_from_budget: Mapped[bool] = mapped_column(default=False) - # Per-key override of the deployment-wide ``reject_user_mismatch`` setting. - # NULL = inherit (the default, and where every key predating this column - # stays), True = always reject a request naming a different ``user``, False = - # always accept it. The override only decides the 403: spend binds to this - # key's own user either way, so the client value stays a provider-side tag. - # False is for clients whose ``user`` is telemetry rather than an identity - # (Claude Code sends a per-session JSON blob); True lets a deployment that - # relaxed the check globally keep an individual key strict. - reject_user_mismatch: Mapped[bool | None] = mapped_column(default=None) - # Per-key override of the deployment-wide ``capture_agent_telemetry`` setting. - # NULL = inherit (the default), True = always store behavioral events from - # this key, False = always discard them. Usage capture/billing is unaffected - # either way; this only gates the content-free agent_telemetry row. - capture_agent_telemetry: Mapped[bool | None] = mapped_column(default=None) - # Per-key model allow-list. NULL = unrestricted (default; every key predating - # this column stays unrestricted), [] = deny all, a list = canonical - # instance:model entries (with instance:* / instance:prefix* wildcards). - allowed_models: Mapped[list[str] | None] = mapped_column(JSON) - - metadata_: Mapped[dict[str, Any]] = mapped_column("metadata", JSON, default=dict) - - user = relationship("User", back_populates="api_keys") - usage_logs = relationship("UsageLog", back_populates="api_key", passive_deletes=True) - - def to_dict(self) -> dict[str, Any]: - """Convert model to dictionary.""" - return { - "id": self.id, - "key_prefix": self.key_prefix, - "key_suffix": self.key_suffix, - "key_name": self.key_name, - "user_id": self.user_id, - "created_at": self.created_at.isoformat() if self.created_at else None, - "last_used_at": self.last_used_at.isoformat() if self.last_used_at else None, - "expires_at": self.expires_at.isoformat() if self.expires_at else None, - "is_active": self.is_active, - "exclude_from_budget": self.exclude_from_budget, - "reject_user_mismatch": self.reject_user_mismatch, - "capture_agent_telemetry": self.capture_agent_telemetry, - "allowed_models": self.allowed_models, - "metadata": self.metadata_, - } - - -# The largest token or request limit a route accepts, and the largest hold it -# will place. Well below the BIGINT ceiling those columns are, for the two -# reasons :data:`gateway.models.money.MAX_USD_LIMIT` keeps its own headroom. -# -# The wire cannot carry the type's maximum. JSON numbers are doubles, so -# ``9223372036854775807`` renders in the published schema as -# ``9.223372036854776e+18``, which is 9223372036854775808: a client sending -# exactly the maximum the spec advertises sends a value the column refuses. A -# quadrillion is exact as a double, so the schema says what it means. -# -# And the gate adds server-side. Every reserve evaluates -# ``current + reserved + held`` as BIGINT arithmetic, so a nonzero counter plus a -# hold near the type's ceiling overflows and answers with a 500 where a 403 was -# owed. Holds are clamped to this too, because the token estimate derives from a -# client-supplied output bound that nothing else limits. -# -# A quadrillion tokens is four orders of magnitude past any real allowance, and -# leaves the sum of three of them ~9000x inside the type. -MAX_COUNT_LIMIT = 1_000_000_000_000_000 - - -class Budget(Base): - """Budget model for spending limits.""" - - __tablename__ = "budgets" - __table_args__ = ( - # A period comes from one place or the other, never both, matching the - # rule ``scoped_budgets`` already enforced when it carried its own. Without - # it the pair encodes one concept twice and ``(86400, calendar_month)`` is - # storable and meaningless. - CheckConstraint( - "NOT (budget_duration_sec IS NOT NULL AND reset_alignment IS NOT NULL)", - name="ck_budgets_single_period_source", - ), - ) - - budget_id: Mapped[str] = mapped_column(primary_key=True, default=lambda: str(uuid.uuid4())) - name: Mapped[str | None] = mapped_column(default=None) - # Which tenant defined this budget, and therefore who may change it. - # - # NULL means the deployment's own: every budget predating `b7e1c4a9d2f5` - # reads NULL, and so does every one the otari-ai cutover migration mints, - # because that migration deliberately shares one budget per distinct - # (cap, period) shape across the ceilings it writes and a shape shared by two - # tenants' ceilings has no single owner. NULL is not "unowned and up for - # grabs": the organization-scoped surface never lists, offers or repoints one, - # so from a tenant's side it does not exist. - # - # Nullable and set only on the tenant-scoped create path, which is what keeps - # that migration working with no backfill: its preflight refuses on a - # *missing* column and its inserts name theirs explicitly, so a new nullable - # one is invisible to it. - organization_id: Mapped[uuid.UUID | None] = mapped_column( - Uuid(), - ForeignKey("organization.id", name="fk_budgets_organization_id", ondelete="CASCADE"), - default=None, - index=True, - ) - # Exact, like the counters it is compared against: the gate is - # ``spend + reserved <= max_budget``, and a cap stored as a binary float - # would decide a 403 against an amount an operator never typed - # (mozilla-ai/otari#691). - max_budget: Mapped[Decimal | None] = mapped_column(UsdCost()) - # The non-USD ceilings, independent of ``max_budget`` and of each other: a - # budget may cap dollars, tokens, requests, or any combination, and a NULL on - # an axis is unbounded there. Deliberately no "at least one limit" constraint, - # because a budget with every limit NULL is a named period that admits - # everything and predates these columns. - # - # BIGINT: a monthly token allowance for one organization outgrows a 32-bit - # counter, and the counters compared against these are the same width. - token_limit: Mapped[int | None] = mapped_column(BigInteger(), default=None) - request_limit: Mapped[int | None] = mapped_column(BigInteger(), default=None) - budget_duration_sec: Mapped[int | None] = mapped_column() - # Snap the window to a UTC calendar boundary instead of counting a fixed - # number of seconds, which is the only way to express a calendar month (2592000 - # seconds is a different, 1.5 percent more generous, product). It lives here - # rather than on the rows that enforce a budget because a limit and the period - # it is spent over are one product decision, and splitting them let a ceiling - # reset on a cadence the budget defining it had never heard of. - reset_alignment: Mapped[str | None] = mapped_column(default=None) - created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=lambda: datetime.now(UTC)) - updated_at: Mapped[datetime] = mapped_column( - DateTime(timezone=True), - default=lambda: datetime.now(UTC), - onupdate=lambda: datetime.now(UTC), - ) - - users = relationship("User", back_populates="budget") - reset_logs = relationship("BudgetResetLog", back_populates="budget") - - def to_dict(self) -> dict[str, Any]: - """Convert model to dictionary.""" - return { - "budget_id": self.budget_id, - "name": self.name, - "max_budget": self.max_budget, - "token_limit": self.token_limit, - "request_limit": self.request_limit, - "budget_duration_sec": self.budget_duration_sec, - "reset_alignment": self.reset_alignment, - "created_at": self.created_at.isoformat() if self.created_at else None, - "updated_at": self.updated_at.isoformat() if self.updated_at else None, - } - - -class User(Base): - """User/Customer model for end-user tracking.""" - - __tablename__ = "users" - - user_id: Mapped[str] = mapped_column(primary_key=True) - alias: Mapped[str | None] = mapped_column() - # The spend ledger, exact to the micro-dollar like the ``usage_logs`` rows - # that sum into it (mozilla-ai/otari#691). As a float it drifted: four - # completions whose settled costs were each exact left this at - # 0.6619999999999999, and the drift accumulated across every reconcile until - # the budget reset. - spend: Mapped[Decimal] = mapped_column(UsdCost(), default=Decimal(0)) - # In-flight budget held by requests that have passed the budget gate but - # whose actual cost is not yet known. The effective committed amount is - # ``spend + reserved``; reservations are reconciled into ``spend`` (actual - # cost) on success or released on failure. See gateway.services.budget_service. - reserved: Mapped[Decimal] = mapped_column(UsdCost(), default=Decimal(0), server_default="0") - # The token and request counters, gated by the same budget's ``token_limit`` - # and ``request_limit`` the way the pair above is gated by ``max_budget``. - # Each axis names itself rather than extending the bare ``spend``/``reserved`` - # pair, which is USD and predates them. - current_tokens: Mapped[int] = mapped_column(BigInteger(), default=0, server_default="0") - reserved_tokens: Mapped[int] = mapped_column(BigInteger(), default=0, server_default="0") - current_requests: Mapped[int] = mapped_column(BigInteger(), default=0, server_default="0") - reserved_requests: Mapped[int] = mapped_column(BigInteger(), default=0, server_default="0") - # Indexed: the budgets list groups users by this column to build each budget's - # usage rollup, so an unindexed FK turns that page into a users table scan. - budget_id: Mapped[str | None] = mapped_column(ForeignKey("budgets.budget_id"), index=True) - # Default model access-list every one of this user's keys inherits when the - # key has no list of its own. null = unrestricted, [] = deny all, else - # canonical instance:model entries (see services/model_access.py). A key may - # narrow this default but never broaden it (validated on key write). - allowed_models: Mapped[list[str] | None] = mapped_column(JSON) - budget_started_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True)) - next_budget_reset_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True)) - blocked: Mapped[bool] = mapped_column(default=False) - deleted_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True), default=None, index=True) - created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=lambda: datetime.now(UTC)) - updated_at: Mapped[datetime] = mapped_column( - DateTime(timezone=True), - default=lambda: datetime.now(UTC), - onupdate=lambda: datetime.now(UTC), - ) - metadata_: Mapped[dict[str, Any]] = mapped_column("metadata", JSON, default=dict) - - budget = relationship("Budget", back_populates="users") - api_keys = relationship("APIKey", back_populates="user", passive_deletes=True) - usage_logs = relationship("UsageLog", back_populates="user", passive_deletes=True) - reset_logs = relationship("BudgetResetLog", back_populates="user", passive_deletes=True) - - def to_dict(self) -> dict[str, Any]: - """Convert model to dictionary.""" - return { - "user_id": self.user_id, - "alias": self.alias, - "spend": self.spend, - "reserved": self.reserved, - "budget_id": self.budget_id, - "allowed_models": self.allowed_models, - "budget_started_at": self.budget_started_at.isoformat() if self.budget_started_at else None, - "next_budget_reset_at": self.next_budget_reset_at.isoformat() if self.next_budget_reset_at else None, - "blocked": self.blocked, - "created_at": self.created_at.isoformat() if self.created_at else None, - "updated_at": self.updated_at.isoformat() if self.updated_at else None, - "metadata": self.metadata_, - } - - -class ModelAlias(Base): - """A display name that resolves to a real model selector. - - The runtime counterpart of the ``aliases:`` block in config.yml: same - meaning, but writable through the API. Pricing, budgets, and usage all key - on the resolved target, so nothing here is billed against ``name``. - - There are two scopes, and they are independent. ``workspace_id`` says which - tenant owns the alias: it resolves only for requests in that workspace, so - two workspaces can each point ``fast`` somewhere different. Within a - workspace, ``user_id`` narrows it further: ``NULL`` means every caller in - that workspace sees it, which is what every row predating the column is, and - a non-null ``user_id`` scopes it to that user, shadowing the workspace-wide - row of the same name for them alone. - - Uniqueness needs two constraints rather than one because SQLite and - PostgreSQL both treat NULLs as distinct in a unique index: the composite - constraint keeps one row per (workspace, name, user), and the partial index - keeps one workspace-wide row per (workspace, name), which the composite one - cannot, its ``user_id`` being NULL. The surrogate ``id`` exists only because - the natural key contains a nullable column, which a primary key cannot. - """ - - __tablename__ = "model_aliases" - __table_args__ = ( - # Workspace-scoped, so two workspaces can each hold a "fast" entry - # pointing somewhere different. Safe only because resolution is keyed by - # workspace too (``services/alias_service``); while that cache was keyed - # on name alone the second workspace's row silently shadowed the first at - # request time, which is why this constraint waited for it. - UniqueConstraint("workspace_id", "name", "user_id", name="uq_model_aliases_workspace_name_user"), - Index( - "uq_model_aliases_workspace_global_name", - "workspace_id", - "name", - unique=True, - sqlite_where=text("user_id IS NULL"), - postgresql_where=text("user_id IS NULL"), - ), - ) - - id: Mapped[str] = mapped_column(primary_key=True, default=lambda: str(uuid.uuid4())) - # No index of its own: both constraints above lead with `workspace_id` and - # carry `name` second, and a listing is always workspace-scoped. A third copy - # would be paid for on every write to serve reads that mostly do not happen, - # since resolution goes through the process-wide alias cache. - name: Mapped[str] = mapped_column() - target: Mapped[str] = mapped_column() - user_id: Mapped[str | None] = mapped_column(ForeignKey("users.user_id", ondelete="CASCADE"), index=True) - # The workspace this row belongs to; see `APIKey.workspace_id` for why. - workspace_id: Mapped[uuid.UUID] = mapped_column( - Uuid, ForeignKey("workspace.id", ondelete="RESTRICT"), nullable=False, index=True - ) - created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=lambda: datetime.now(UTC)) - updated_at: Mapped[datetime] = mapped_column( - DateTime(timezone=True), - default=lambda: datetime.now(UTC), - onupdate=lambda: datetime.now(UTC), - ) - - def to_dict(self) -> dict[str, Any]: - return { - "name": self.name, - "target": self.target, - "user_id": self.user_id, - "created_at": self.created_at.isoformat() if self.created_at else None, - "updated_at": self.updated_at.isoformat() if self.updated_at else None, - } - - -class RoutingPolicy(Base): - """A named routing policy, writable through the API. - - The runtime counterpart of the ``routing.policies`` block in config.yml. The - spec is stored as JSON rather than as columns because it is a nested, - versioned document (``select`` entries with conditions, ``on_failure``, - guardrails); flattening it into columns would mean a migration per - schema addition and would still need JSON for the conditions. It is validated - against :class:`gateway.models.routing.PolicySpec` on write and again on load, - so a row that predates a schema change surfaces as a startup warning rather - than as a request-time crash. - - Scoping mirrors :class:`ModelAlias` exactly, workspace included, and so does - the two-constraint uniqueness (SQLite and PostgreSQL both treat NULLs as - distinct in a unique index, so the composite constraint cannot keep one - *workspace-wide* row per name). A policy and an alias are the same concept at - different complexities, so it would be strange for their scoping rules to - differ. - """ - - __tablename__ = "routing_policies" - __table_args__ = ( - # Workspace-scoped for the same reason, and on the same precondition, as - # :class:`ModelAlias`: ``services/policy_store`` keys its cache by - # workspace, so two workspaces holding a "fast" policy each resolve their - # own rather than one shadowing the other. - UniqueConstraint("workspace_id", "name", "user_id", name="uq_routing_policies_workspace_name_user"), - Index( - "uq_routing_policies_workspace_global_name", - "workspace_id", - "name", - unique=True, - sqlite_where=text("user_id IS NULL"), - postgresql_where=text("user_id IS NULL"), - ), - ) - - id: Mapped[str] = mapped_column(primary_key=True, default=lambda: str(uuid.uuid4())) - name: Mapped[str] = mapped_column() - spec: Mapped[dict[str, Any]] = mapped_column(JSON) - user_id: Mapped[str | None] = mapped_column(ForeignKey("users.user_id", ondelete="CASCADE"), index=True) - # The workspace this row belongs to; see `APIKey.workspace_id` for why. - workspace_id: Mapped[uuid.UUID] = mapped_column( - Uuid, ForeignKey("workspace.id", ondelete="RESTRICT"), nullable=False, index=True - ) - created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=lambda: datetime.now(UTC)) - updated_at: Mapped[datetime] = mapped_column( - DateTime(timezone=True), - default=lambda: datetime.now(UTC), - onupdate=lambda: datetime.now(UTC), - ) - - def to_dict(self) -> dict[str, Any]: - return { - "name": self.name, - "spec": self.spec, - "user_id": self.user_id, - "created_at": self.created_at.isoformat() if self.created_at else None, - "updated_at": self.updated_at.isoformat() if self.updated_at else None, - } - - -class RuntimeSetting(Base): - """A persisted override for a runtime-toggleable config flag. - - A small key/value store for the handful of settings the dashboard can flip - at runtime (model discovery, default pricing). When a key is present it wins - over the config-file/env value and is applied on startup; when absent the - config value stands. The value is stored as a string ("true"/"false") so the - table can hold future non-boolean settings without a schema change. - """ - - __tablename__ = "runtime_settings" - - key: Mapped[str] = mapped_column(primary_key=True) - value: Mapped[str] = mapped_column() - updated_at: Mapped[datetime] = mapped_column( - DateTime(timezone=True), - default=lambda: datetime.now(UTC), - onupdate=lambda: datetime.now(UTC), - ) - - -class DashboardSession(Base): - """A server-side admin-dashboard sign-in session, held by one identity. - - Minted when an operator signs in to the dashboard with the master key: the - browser holds only an opaque token in an HttpOnly cookie and this table - stores the token's SHA-256 hash, so neither the master key nor a usable - session credential is ever persisted in JS-readable storage. Sessions - expire on a TTL and are revoked on sign-out and on master-key rotation. - - ``user_id`` is what lets a session resolve a caller rather than only prove - that the master key was presented once. It names a tenancy identity - (`models.tenancy.User`), whose ``active_organization_id`` is the - organization the session acts in, so a tenancy surface reads its scope off - the session. Master-key sign-in binds the session to the deployment's - bootstrap operator; a per-user sign-in flow binds it to whoever - authenticated. - - NOT NULL on purpose: a session that names nobody cannot answer "who is - calling", which is the whole point of the column, and the migration that - added it bound existing sessions to that same bootstrap operator. CASCADE - on the foreign key, so deleting an identity revokes its sessions rather - than leaving a live cookie pointing at a row that is gone. - """ - - __tablename__ = "dashboard_sessions" - - token_hash: Mapped[str] = mapped_column(primary_key=True) - user_id: Mapped[uuid.UUID] = mapped_column( - Uuid, ForeignKey("user.id", ondelete="CASCADE"), nullable=False, index=True - ) - created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=lambda: datetime.now(UTC)) - expires_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), index=True) - - -class PricingSnapshot(Base): - """An approved, source-tagged upstream pricing catalog.""" - - __tablename__ = "pricing_snapshots" - - source: Mapped[str] = mapped_column(primary_key=True) - snapshot: Mapped[str] = mapped_column(Text) - updated_at: Mapped[datetime] = mapped_column( - DateTime(timezone=True), - default=lambda: datetime.now(UTC), - onupdate=lambda: datetime.now(UTC), - ) - - -class PricingSnapshotHistory(Base): - """One accepted upstream pricing snapshot, kept after a later one replaces it. - - ``pricing_snapshots`` is the current state; this is the record. Written on - every accept, never updated. ``accepted_by`` says whether an operator - confirmed it or the scheduled refresh applied it on its own. - """ - - __tablename__ = "pricing_snapshot_history" - __table_args__ = (Index("ix_pricing_snapshot_history_source_accepted_at", "source", "accepted_at"),) - - id: Mapped[uuid.UUID] = mapped_column(Uuid, primary_key=True, default=uuid.uuid4) - source: Mapped[str] = mapped_column(String(64)) - accepted_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=lambda: datetime.now(UTC)) - accepted_by: Mapped[str] = mapped_column(String(32)) - model_count: Mapped[int] = mapped_column() - snapshot: Mapped[str] = mapped_column(Text) - - -class ProviderCredential(Base): - """A provider instance configured at runtime through the dashboard. - - The database counterpart of a ``providers:`` entry in config.yml: it is - merged over the config-file providers at runtime (see - ``provider_store_service``), with the stored row winning on an instance-name - collision. The API key is held encrypted (``secret_box``); ``last4`` is kept - in clear only so the UI can show which key is set without ever decrypting. - Standalone mode only, never used in the hybrid platform path. - """ - - __tablename__ = "provider_credentials" - - instance: Mapped[str] = mapped_column(primary_key=True) - provider_type: Mapped[str | None] = mapped_column() - api_base: Mapped[str | None] = mapped_column() - encrypted_api_key: Mapped[str | None] = mapped_column() - last4: Mapped[str | None] = mapped_column() - client_args: Mapped[dict[str, Any]] = mapped_column("client_args", JSON, default=dict) - created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=lambda: datetime.now(UTC)) - updated_at: Mapped[datetime] = mapped_column( - DateTime(timezone=True), - default=lambda: datetime.now(UTC), - onupdate=lambda: datetime.now(UTC), - ) - - def to_public_dict(self) -> dict[str, Any]: - """Serialize for the API. Never includes the secret, only ``last4``. - - ``client_args`` is masked by key name the same way - ``OrgProviderKey.to_public`` masks its own: a standalone Bedrock instance - keeps its ``aws_secret_access_key`` there, so the field this table holds - in clear is as much a credential as ``encrypted_api_key`` is, and it must - not round-trip over the API either. - """ - return { - "instance": self.instance, - "provider_type": self.provider_type, - "api_base": self.api_base, - "last4": self.last4, - "client_args": redact_secret_like_values(self.client_args) or {}, - "created_at": self.created_at.isoformat() if self.created_at else None, - "updated_at": self.updated_at.isoformat() if self.updated_at else None, - } - - -class SearchToolCredential(Base): - """A ``POST /v1/search`` tool configured at runtime through the dashboard. - - The database counterpart of a ``search_tools:`` entry in config.yml: it is - merged over the config-file tools at runtime (see - ``search_tool_store_service``), with the stored row winning on a name - collision, exactly as ``ProviderCredential`` does for providers. The API key - is held encrypted (``secret_box``) and is optional, because a ``searxng`` - backend is normally keyless; ``last4`` is kept in clear only so the UI can - show which key is set without ever decrypting. Standalone mode only. - """ - - __tablename__ = "search_tool_credentials" - - name: Mapped[str] = mapped_column(primary_key=True) - provider: Mapped[str] = mapped_column() - api_base: Mapped[str | None] = mapped_column() - encrypted_api_key: Mapped[str | None] = mapped_column() - last4: Mapped[str | None] = mapped_column() - # Named for its unit; the config-file key it stands in for is plain ``timeout``, - # and ``to_public_dict`` / the overlay entry both use that name. - timeout_seconds: Mapped[float | None] = mapped_column() - options: Mapped[dict[str, Any]] = mapped_column("options", JSON, default=dict) - created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=lambda: datetime.now(UTC)) - updated_at: Mapped[datetime] = mapped_column( - DateTime(timezone=True), - default=lambda: datetime.now(UTC), - onupdate=lambda: datetime.now(UTC), - ) - - def to_public_dict(self) -> dict[str, Any]: - """Serialize for the API. Never includes the secret, only ``last4``. - - ``options`` is masked by key name for the reason ``ProviderCredential`` - gives above: it is free-form backend configuration, so a second - credential an operator put there is not echoed back either. - """ - return { - "name": self.name, - "provider": self.provider, - "api_base": self.api_base, - "last4": self.last4, - "timeout": self.timeout_seconds, - "options": redact_secret_like_values(self.options) or {}, - "created_at": self.created_at.isoformat() if self.created_at else None, - "updated_at": self.updated_at.isoformat() if self.updated_at else None, - } - - -class ModelPricing(Base): - """Model pricing configuration.""" - - __tablename__ = "model_pricing" - - model_key: Mapped[str] = mapped_column(primary_key=True) - effective_at: Mapped[datetime] = mapped_column( - DateTime(timezone=True), - primary_key=True, - default=lambda: datetime.now(UTC), - ) - input_price_per_million: Mapped[Decimal] = mapped_column(UsdRate()) - output_price_per_million: Mapped[Decimal] = mapped_column(UsdRate()) - # Nullable: providers without prompt caching (or models without a - # discounted cache rate) leave these unset. When set, the cost - # calculation prices cache_read_tokens / cache_write_tokens at these - # per-million-token rates, following the provider inclusion convention - # (see log_usage in _pipeline.py). - cache_read_price_per_million: Mapped[Decimal | None] = mapped_column(UsdRate(), nullable=True) - cache_write_price_per_million: Mapped[Decimal | None] = mapped_column(UsdRate(), nullable=True) - cache_write_1h_price_per_million: Mapped[Decimal | None] = mapped_column(UsdRate(), nullable=True) - # Ordered threshold rules. Each rule applies its supplied rates to the - # entire request once ``total_input_tokens`` reaches ``min_input_tokens``. - pricing_tiers: Mapped[list[dict[str, Any]]] = mapped_column(JSON, default=list) - # What ``input_price_per_million`` is a rate per: ``tokens`` for a model, - # ``requests`` for a gateway-run tool or a moderation call, ``images`` for - # image generation. The rate columns are shared by all three and a reader - # cannot tell which from the number, so the row says (``PRICING_UNITS``). - unit: Mapped[str] = mapped_column(String(16), default="tokens", server_default="tokens") - # Which path wrote the row: ``config`` (the file's ``pricing:`` block), - # ``api`` (``POST /v1/pricing``), or ``migration``. NULL on a row written - # before origins were recorded, which is a real answer and not a default. - origin: Mapped[str | None] = mapped_column(String(16), nullable=True) - created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=lambda: datetime.now(UTC)) - updated_at: Mapped[datetime] = mapped_column( - DateTime(timezone=True), - default=lambda: datetime.now(UTC), - onupdate=lambda: datetime.now(UTC), - ) - - def to_dict(self) -> dict[str, Any]: - """Convert model to dictionary.""" - return { - "model_key": self.model_key, - "effective_at": self.effective_at.isoformat() if self.effective_at else None, - "input_price_per_million": self.input_price_per_million, - "output_price_per_million": self.output_price_per_million, - "cache_read_price_per_million": self.cache_read_price_per_million, - "cache_write_price_per_million": self.cache_write_price_per_million, - "cache_write_1h_price_per_million": self.cache_write_1h_price_per_million, - "pricing_tiers": self.pricing_tiers, - "unit": self.unit, - "origin": self.origin, - "created_at": self.created_at.isoformat() if self.created_at else None, - "updated_at": self.updated_at.isoformat() if self.updated_at else None, - } - - -class UsageLog(Base): - """Usage log model for tracking API requests.""" - - __tablename__ = "usage_logs" - __table_args__ = ( - Index("ix_usage_logs_user_id_timestamp", "user_id", "timestamp"), - # Supports the activity-log viewer's primary "show errors, newest-first" - # query. status is low-cardinality; model is high-cardinality and left - # unindexed on purpose. - Index("ix_usage_logs_status_timestamp", "status", "timestamp"), - # Supports the setup guide's two questions about one workspace: has any - # request in it ever succeeded (oldest first), and what did the last one - # do (newest first). Both filter a workspace, a source and a status and - # then order by time, which the workspace-only and status-first indexes - # above can each answer only halfway: on a deployment with real traffic - # the guide would otherwise scan the workspace's rows on every dashboard - # load, and where usage is imported as well most of those rows are the - # wrong source anyway. Equality columns first, the ordering column last. - Index( - "ix_usage_logs_workspace_source_status_timestamp", - "workspace_id", - "source", - "status", - "timestamp", - ), - # Idempotency for imported usage: re-submitting the same (source, - # source_event_id) must not create a second row. Gateway-originated rows - # keep source_event_id NULL, and SQL treats NULLs as distinct on both - # SQLite and Postgres, so many (gateway, NULL) rows coexist freely. - UniqueConstraint("source", "source_event_id", name="uq_usage_logs_source_event"), - ) - - id: Mapped[str] = mapped_column(primary_key=True, default=lambda: str(uuid.uuid4())) - # The workspace this row belongs to; see `APIKey.workspace_id` for why. - workspace_id: Mapped[uuid.UUID] = mapped_column( - Uuid, ForeignKey("workspace.id", ondelete="RESTRICT"), nullable=False, index=True - ) - api_key_id: Mapped[str | None] = mapped_column(ForeignKey("api_keys.id", ondelete="SET NULL"), index=True) - user_id: Mapped[str | None] = mapped_column(ForeignKey("users.user_id", ondelete="SET NULL"), index=True) - timestamp: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=lambda: datetime.now(UTC), index=True) - - model: Mapped[str] = mapped_column() - provider: Mapped[str | None] = mapped_column() - endpoint: Mapped[str] = mapped_column() - - # Provenance. "gateway" for requests Otari served itself; a source slug (e.g. - # "claude_code") for usage imported through POST /v1/usage/external-events. A row - # backfilled from hosted history keeps its origin's slug behind a legacy prefix - # ("otari-ai:gateway", "otari-ai:claude_code"), so asking whether this deployment - # served a row means asking about the slug behind that prefix: core/usage_source. - # source_event_id is the upstream event id used for idempotent import (NULL for - # gateway rows); source_label carries optional session/project attribution. - source: Mapped[str] = mapped_column(default="gateway", index=True) - source_event_id: Mapped[str | None] = mapped_column() - source_label: Mapped[str | None] = mapped_column() - # Whether this row's cost participates in budget enforcement. True for normal - # gateway rows; false for imported usage and for rows from keys flagged - # exclude_from_budget. False rows are recorded (and appear in cost analytics) - # but their cost is never written to User.spend. - counts_toward_budget: Mapped[bool] = mapped_column(default=True) - - prompt_tokens: Mapped[int | None] = mapped_column() - completion_tokens: Mapped[int | None] = mapped_column() - total_tokens: Mapped[int | None] = mapped_column() - cache_read_tokens: Mapped[int | None] = mapped_column() - cache_write_tokens: Mapped[int | None] = mapped_column() - cache_write_1h_tokens: Mapped[int | None] = mapped_column() - # Which cached-token convention the counts above were reported under: True - # when the cache buckets are already inside ``prompt_tokens`` (OpenAI shape), - # False when they are additive to it (Anthropic / Claude Code shape). Written - # by settlement from ``GatewayUsage.cache_tokens_in_prompt`` and by the - # external-usage ingest from the value the submitter sent, so a row can be - # repriced under the convention it was recorded with rather than one inferred - # from the numbers, which cannot tell the two apart. - # - # Nullable, and deliberately not defaulted: "not recorded" and "inclusive" are - # different answers. Rows written before this column existed are NULL, and - # repricing falls back to recovering the convention from ``billing_meters`` - # for exactly those (see ``usage_admin_service._row_cache_tokens_included``). - # A default would make every historical row claim a convention nothing - # checked, and mis-price the half that were the other one. - cache_tokens_in_prompt: Mapped[bool | None] = mapped_column() - billing_meters: Mapped[dict[str, Any] | None] = mapped_column(JSON) - pricing_breakdown: Mapped[list[dict[str, Any]] | None] = mapped_column(JSON) - # The settled amount, and the accounting truth for this row - # (mozilla-ai/otari-ai#1751). Exact to the micro-dollar; see - # ``models/money.py`` for what that costs on each engine. - cost: Mapped[Decimal | None] = mapped_column(UsdCost()) - - # Why ``cost`` is the amount it is, which the row cannot re-derive on its own: - # ``pricing_source`` names the price list that settled it ("organization", - # "managed", "genai_prices"), ``pricing_reference`` identifies the entry in it - # (a pricing row's id, or a ``provider:model`` key), ``pricing_effective_at`` - # is when that rate took effect, and ``pricing_version`` pins the revision of - # the list. ``calculated_at`` is when the amount was priced, which is not - # ``timestamp`` (when the request ran): usage settled or repriced later moves - # the two apart. - # - # All nullable with no backfill. The gateway's own settlement does not record - # provenance, so these are written by the hosted-usage backfill - # (mozilla-ai/otari-ai#1798) from the platform's ``gateway_usage_settlement`` - # row, and null reads correctly as "not recorded". The lengths mirror that - # table's columns rather than this file's usual unbounded strings, so a value - # copied across always fits. - # - # ``pricing_source`` speaks the platform's settlement vocabulary, the values - # ``_platform.SettledCost.pricing_source`` already carries on the hybrid wire - # (echoed to callers as ``usage.pricing_source``). It is not the same field as - # the one on a listed model in ``api/routes/models.py`` ("configured", - # "default", "dynamic", "none"), which says where a price list entry came from - # in this deployment rather than what settled one row's amount. - pricing_source: Mapped[str | None] = mapped_column(String(32)) - pricing_reference: Mapped[str | None] = mapped_column(String(511)) - pricing_effective_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True)) - pricing_version: Mapped[str | None] = mapped_column(String(255)) - calculated_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True)) - - # "success", "error", or "absorbed". ``absorbed`` is a failed attempt that a - # routing policy recovered from by trying the next candidate: the request - # itself succeeded (or failed on a later attempt), so counting it as an error - # would make a working fallback chain look like an outage. Every error metric - # in the product counts ``status == "error"`` exactly, and ``request_count`` - # excludes absorbed rows, because a request that took two attempts is still one - # request. - status: Mapped[str] = mapped_column() - error_message: Mapped[str | None] = mapped_column() - - # Routing attribution. All nullable: a request that named a plain model was not - # routed through a policy, and null reads correctly as exactly that. - # - # `policy_name` is the name the caller sent. `selection_reason` says why this - # candidate was chosen ("default", "condition:", "on_failure", - # "router:"). `attempt_position` and `attempt_count` locate the row in - # the plan, so "served on attempt 2 of 3" is a query rather than a log grep. - # `request_group_id` ties a request's rows together, which is what makes the - # absorbed attempts findable from the row that served. - policy_name: Mapped[str | None] = mapped_column(index=True) - selection_reason: Mapped[str | None] = mapped_column() - attempt_position: Mapped[int | None] = mapped_column() - attempt_count: Mapped[int | None] = mapped_column() - request_group_id: Mapped[str | None] = mapped_column(index=True) - - # HTTP status that classifies a failure, so failures can be grouped with a - # GROUP BY instead of substring-matching provider-specific error prose. It is - # the status the provider returned when it sent one (an upstream 401 stays - # visible as a credential fault even though the caller sees the generic 502 - # that keeps gateway config out of the response), otherwise the gateway's own - # rejection or classification code (402 missing pricing, 422 tool-loop cap, - # 504 timeout, 502 unreachable). Nullable: historical rows predate the column, - # a successful request has no failure to classify, and some failures carry no - # HTTP status at all (e.g. a stream that ended without usage data). - status_code: Mapped[int | None] = mapped_column() - - # Total server-side wall-clock for the request, in milliseconds. Nullable: - # historical rows predate the column, and some write paths (batch jobs, - # provider-never-reached rejections) have no meaningful request duration. - latency_ms: Mapped[int | None] = mapped_column() - - # Milliseconds from request start to the first streamed chunk. Nullable: - # non-streaming requests have no first chunk, historical rows predate the - # column, and a stream that failed before yielding anything never reached one. - # - # ``started_at`` is taken in the handler preamble, so on a routing plan the - # serving row's value also carries every earlier attempt's setup time. - # Nothing in the column says so; a percentile keyed by the serving model - # attributes failover time to the model that actually served. - # - # Hybrid (platform-fallback) streams never write this column at all: every - # settlement callback in build_streaming_response returns before reaching - # log_usage on that path, and run_streaming_with_fallback passes db=None. - ttft_ms: Mapped[int | None] = mapped_column() - - api_key = relationship("APIKey", back_populates="usage_logs") - user = relationship("User", back_populates="usage_logs") - - def to_dict(self) -> dict[str, Any]: - """Convert model to dictionary.""" - return { - "id": self.id, - "api_key_id": self.api_key_id, - "user_id": self.user_id, - "timestamp": self.timestamp.isoformat() if self.timestamp else None, - "model": self.model, - "endpoint": self.endpoint, - "source": self.source, - "source_label": self.source_label, - "counts_toward_budget": self.counts_toward_budget, - "prompt_tokens": self.prompt_tokens, - "completion_tokens": self.completion_tokens, - "total_tokens": self.total_tokens, - "cache_read_tokens": self.cache_read_tokens, - "cache_write_tokens": self.cache_write_tokens, - "cache_write_1h_tokens": self.cache_write_1h_tokens, - "cache_tokens_in_prompt": self.cache_tokens_in_prompt, - "billing_meters": self.billing_meters, - "pricing_breakdown": self.pricing_breakdown, - "cost": self.cost, - "status": self.status, - "error_message": self.error_message, - "status_code": self.status_code, - "latency_ms": self.latency_ms, - "policy_name": self.policy_name, - "selection_reason": self.selection_reason, - "attempt_position": self.attempt_position, - "attempt_count": self.attempt_count, - "request_group_id": self.request_group_id, - } - - -class AgentTelemetry(Base): - """Content-free outcome metrics and behavioral events from coding agents.""" - - __tablename__ = "agent_telemetry" - __table_args__ = ( - UniqueConstraint("source", "dedup_key", name="uq_agent_telemetry_source_dedup"), - Index("ix_agent_telemetry_user_id_timestamp", "user_id", "timestamp"), - # Read-time cumulative-to-delta derivation orders one series' points by time. - Index("ix_agent_telemetry_series_timestamp", "series_key", "timestamp"), - ) - - id: Mapped[str] = mapped_column(primary_key=True, default=lambda: str(uuid.uuid4())) - api_key_id: Mapped[str | None] = mapped_column(ForeignKey("api_keys.id", ondelete="SET NULL"), index=True) - user_id: Mapped[str | None] = mapped_column(ForeignKey("users.user_id", ondelete="SET NULL"), index=True) - timestamp: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=lambda: datetime.now(UTC), index=True) - name: Mapped[str] = mapped_column() - tool_name: Mapped[str | None] = mapped_column() - decision: Mapped[str | None] = mapped_column() - success: Mapped[bool | None] = mapped_column() - duration_ms: Mapped[int | None] = mapped_column() - status_code: Mapped[int | None] = mapped_column() - prompt_length: Mapped[int | None] = mapped_column() - source: Mapped[str] = mapped_column(index=True) - session_label: Mapped[str | None] = mapped_column() - dedup_key: Mapped[str] = mapped_column() - created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=lambda: datetime.now(UTC)) - - # Outcome-metric columns. Populated only on a metric row (``kind="metric"``), - # NULL on a behavioral one, which is the inverse of the allow-list columns - # above. ``value`` is stored exactly as OTLP reported it (a running total or - # an increment, per ``temporality``); the read endpoints do the delta - # arithmetic, so nothing is normalized at ingest. ``series_key`` is the pure - # OTLP series identity (name plus attributes), which is what makes a - # dimensioned metric two series rather than one. - kind: Mapped[str | None] = mapped_column() - value: Mapped[float | None] = mapped_column() - temporality: Mapped[str | None] = mapped_column() - series_start: Mapped[datetime | None] = mapped_column(DateTime(timezone=True)) - series_key: Mapped[str | None] = mapped_column() - - -class FileObject(Base): - """Uploaded file metadata for the OpenAI-compatible /v1/files API. - - The raw bytes live in a pluggable blob backend (see - gateway.services.file_store); this row holds metadata plus the backend - ``storage_ref`` used to fetch them. Files are scoped to ``user_id`` for - tenant isolation and soft-deleted via ``deleted_at``. ``workspace_id`` is a - second, independent axis: it says which workspace the upload was made in, so - a key confined to one workspace never reaches another's files even when the - same user holds keys in both. - """ - - __tablename__ = "file_objects" - - id: Mapped[str] = mapped_column(primary_key=True, default=lambda: f"file-{uuid.uuid4().hex}") - # Always set to the authenticated user; non-null enforces the user-scoping - # contract at the schema level. CASCADE removes a user's files on delete. - user_id: Mapped[str] = mapped_column(ForeignKey("users.user_id", ondelete="CASCADE"), index=True) - # The workspace this row belongs to; see `APIKey.workspace_id` for why it is - # NOT NULL and RESTRICT rather than nullable and cascading. Existing rows were - # backfilled onto the deployment's default workspace, which is also where a - # master-key upload lands. - workspace_id: Mapped[uuid.UUID] = mapped_column( - Uuid, ForeignKey("workspace.id", ondelete="RESTRICT"), nullable=False, index=True - ) - filename: Mapped[str] = mapped_column() - mime_type: Mapped[str] = mapped_column() - bytes: Mapped[int] = mapped_column() - purpose: Mapped[str] = mapped_column(default="user_data") - storage_ref: Mapped[str] = mapped_column() - created_at: Mapped[datetime] = mapped_column( - DateTime(timezone=True), default=lambda: datetime.now(UTC), index=True - ) - expires_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True)) - deleted_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True), default=None, index=True) - - metadata_: Mapped[dict[str, Any]] = mapped_column("metadata", JSON, default=dict) - - def to_dict(self) -> dict[str, Any]: - """Convert to the OpenAI file object shape.""" - return { - "id": self.id, - "object": "file", - "bytes": self.bytes, - "created_at": _epoch_seconds(self.created_at), - "expires_at": _epoch_seconds(self.expires_at), - "filename": self.filename, - "purpose": self.purpose, - } - - -class BatchRecord(Base): - """Ownership and accounting record for an asynchronous batch job. - - Written at creation time so results accounting can be made idempotent (bill - and log once, on the first completed retrieval), the batch cost can be folded - into ``users.spend``, and ownership can be enforced without depending on the - provider round-tripping the ``otari_user_id`` metadata marker. Batches created - before this table existed carry no record and fall back to the - metadata-anchored ownership path in ``api/routes/batches.py``. ``workspace_id`` - additionally anchors which workspace's organization-scoped provider key - (otari#643) lifecycle calls should resolve credentials from. - """ - - __tablename__ = "batches" - - # Provider-assigned batch id (globally unique per provider), used as the - # lookup key on retrieve/cancel/results. - id: Mapped[str] = mapped_column(primary_key=True) - # Instance/provider name the batch was created against (echoed to clients). - provider: Mapped[str] = mapped_column() - # Billed owner, stamped from the authenticated principal at creation. Non-null: - # this record is the strict ownership anchor, so it must always name an owner. - # CASCADE: deleting the user drops the ownership record (the user's keys are - # gone too, and usage_logs remain the billing history). - user_id: Mapped[str] = mapped_column( - ForeignKey("users.user_id", ondelete="CASCADE"), nullable=False, index=True - ) - # SET NULL: a key may be revoked while its batch is still in flight. - api_key_id: Mapped[str | None] = mapped_column(ForeignKey("api_keys.id", ondelete="SET NULL"), index=True) - # The workspace this batch was CREATED in (otari#643 follow-up), so - # lifecycle calls (retrieve/cancel/results) can resolve organization-scoped - # credentials from the batch's own origin rather than the retriever's - # current workspace: a master-key or legitimately cross-workspace retrieval - # would otherwise use the wrong organization's key, or find none, exactly - # the failure `api_key_id` going NULL on key revocation already risks for - # ownership. Nullable and SET NULL, not RESTRICT: batches created before - # this column existed carry NULL here and fall back to the caller's own - # workspace in `api/routes/batches.py`, and a workspace deleted out from - # under an in-flight batch must not block that delete. - workspace_id: Mapped[uuid.UUID | None] = mapped_column( - Uuid, ForeignKey("workspace.id", ondelete="SET NULL"), index=True - ) - model: Mapped[str] = mapped_column() - created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=lambda: datetime.now(UTC)) - # NULL until the first completed results retrieval accounts the batch; the - # atomic NULL -> now transition is the idempotency gate for billing/logging. - results_accounted_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True)) - - -class RoutingMemory(Base): - """One record per scored example: a prompt embedding plus the quality each - candidate model earned on it. - - The kNN router (:mod:`gateway.services.routing.knn`) retrieves the nearest - neighbors of an incoming request's task embedding within one user's records - and votes on the cheapest candidate that is still good enough. One record is - one example (one prompt), so the vote is over distinct prompts; ``qualities`` - maps each model to its ``[0, 1]`` score for this prompt, keyed on canonical - ``instance:model`` so a candidate's spelling never decides whether it matches - (the router canonicalizes what it reads, so older rows keyed on another - spelling still match). Records are written by the preference-collection flow, - never by live traffic (passive learning is a fast-follow). - - Vectors are stored as a JSON list of floats for SQLite/PostgreSQL - portability and scanned linearly in Python. That holds into the low thousands - of records per user (the ``router_max_records_per_user`` cap); larger pools - need an indexed vector store. - ``embedding_model`` tags each row so changing the embedding model invalidates - stale vectors instead of mixing incomparable spaces. - - Scoped by ``user_id``, which is the identity the request is routed and billed - under, so one user's examples never steer another's traffic. CASCADE: the - records are derived training data, worthless once the user is gone. - ``workspace_id`` narrows that further: the router reads one (user, workspace) - partition, so a user who holds keys in two workspaces does not have one - workspace's labels steering the other's traffic. - """ - - __tablename__ = "routing_memory" - __table_args__ = ( - # Every read filters on the workspace as well as the user, so the - # workspace leads: the same three shapes, one partition narrower. - Index("ix_routing_memory_workspace_user_model", "workspace_id", "user_id", "embedding_model"), - Index("ix_routing_memory_workspace_user_created", "workspace_id", "user_id", "created_at"), - # A task-scoped read filters on all four; without this it walks every - # record the user has for the embedding model before partitioning. - Index( - "ix_routing_memory_workspace_user_model_task", - "workspace_id", - "user_id", - "embedding_model", - "task_id", - ), - ) - - id: Mapped[str] = mapped_column(primary_key=True, default=lambda: str(uuid.uuid4())) - user_id: Mapped[str] = mapped_column( - ForeignKey("users.user_id", ondelete="CASCADE"), nullable=False, index=True - ) - # The workspace this row belongs to; see `APIKey.workspace_id` for why. - workspace_id: Mapped[uuid.UUID] = mapped_column( - Uuid, ForeignKey("workspace.id", ondelete="RESTRICT"), nullable=False, index=True - ) - embedding_model: Mapped[str] = mapped_column() - embedding: Mapped[list[float]] = mapped_column(JSON) - qualities: Mapped[dict[str, float]] = mapped_column(JSON) - task_id: Mapped[str | None] = mapped_column(default=None, index=True) - label_source: Mapped[str] = mapped_column(default="human") - created_at: Mapped[datetime] = mapped_column( - DateTime(timezone=True), default=lambda: datetime.now(UTC), index=True - ) - - def to_dict(self) -> dict[str, Any]: - """Convert model to dictionary. - - The embedding itself is deliberately left out: it is thousands of floats - that no management surface renders, and the prompt it came from is on the - :class:`RouterPreference` audit row. - """ - return { - "id": self.id, - "user_id": self.user_id, - "workspace_id": str(self.workspace_id), - "embedding_model": self.embedding_model, - "qualities": self.qualities, - "task_id": self.task_id, - "label_source": self.label_source, - "created_at": self.created_at.isoformat() if self.created_at else None, - } - - -class RouterPreference(Base): - """An audit record of one preference-collection scoring. - - Each ``/v1/routing/preferences/rank`` submission writes one row here for - provenance plus one :class:`RoutingMemory` row. The routing-memory row keeps - only the embedding, so this is where the prompt text and the raw per-model - scores live: enough to recompute the memory if the scoring changes, and to - tell a human label from a judge's. - - ``workspace_id`` matches the :class:`RoutingMemory` row written beside it, so - the audit trail partitions exactly the way the training data does. - """ - - __tablename__ = "router_preferences" - __table_args__ = ( - Index("ix_router_preferences_workspace_user_created", "workspace_id", "user_id", "created_at"), - ) - - id: Mapped[str] = mapped_column(primary_key=True, default=lambda: str(uuid.uuid4())) - user_id: Mapped[str] = mapped_column( - ForeignKey("users.user_id", ondelete="CASCADE"), nullable=False, index=True - ) - # The workspace this row belongs to; see `APIKey.workspace_id` for why. - workspace_id: Mapped[uuid.UUID] = mapped_column( - Uuid, ForeignKey("workspace.id", ondelete="RESTRICT"), nullable=False, index=True - ) - prompt: Mapped[str] = mapped_column() - task_id: Mapped[str | None] = mapped_column(default=None) - scores: Mapped[dict[str, float]] = mapped_column(JSON) - label_source: Mapped[str] = mapped_column(default="human") - created_at: Mapped[datetime] = mapped_column( - DateTime(timezone=True), default=lambda: datetime.now(UTC), index=True - ) - - def to_dict(self) -> dict[str, Any]: - """Convert model to dictionary.""" - return { - "id": self.id, - "user_id": self.user_id, - "workspace_id": str(self.workspace_id), - "prompt": self.prompt, - "task_id": self.task_id, - "scores": self.scores, - "label_source": self.label_source, - "created_at": self.created_at.isoformat() if self.created_at else None, - } - - -class BudgetResetLog(Base): - """Budget reset log model for tracking budget resets.""" - - __tablename__ = "budget_reset_logs" - - id: Mapped[int] = mapped_column(primary_key=True, autoincrement=True) - user_id: Mapped[str | None] = mapped_column(ForeignKey("users.user_id", ondelete="SET NULL"), index=True) - # Indexed: the reset-log drill-down filters on this column, and the table only - # grows, so an unindexed FK degrades that endpoint to a full scan over time. - budget_id: Mapped[str] = mapped_column(ForeignKey("budgets.budget_id"), index=True) - # The ledger's record of a counter that is now exact, so it is exact too: - # a float snapshot of an exact ``users.spend`` would no longer equal the - # spend it claims to have recorded. - previous_spend: Mapped[Decimal] = mapped_column(UsdCost()) - reset_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=lambda: datetime.now(UTC)) - next_reset_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True)) - - user = relationship("User", back_populates="reset_logs") - budget = relationship("Budget", back_populates="reset_logs") - - def to_dict(self) -> dict[str, Any]: - """Convert model to dictionary.""" - return { - "id": self.id, - "user_id": self.user_id, - "budget_id": self.budget_id, - "previous_spend": self.previous_spend, - "reset_at": self.reset_at.isoformat() if self.reset_at else None, - "next_reset_at": self.next_reset_at.isoformat() if self.next_reset_at else None, - } - - -class ScopedBudget(Base): - """A spending ceiling on one tenancy scope, optionally narrowed to one provider. - - Two axes. The identity axis is ``(scope_type, scope_id)``: who is capped, an - organization, a workspace, a member of either, or a single API key. The - resource axis is ``provider_key_id``: NULL caps spend across every provider, - a value narrows the cap to one provider instance. A request must pass every - row that applies to it, and each row is an independent ceiling with its own - counters and its own period window, unlike ``budgets``, where the window and - the counters live on the user. - - No limit is stored here. A limit is a property of the budget this names, - which is the only place in the schema that maps a cap to a figure, on any of - the three axes it can cap. - - ``scope_type`` is a plain string rather than a database enum so a new scope - needs no enum migration, and ``scope_id`` is a string so it holds both this - codebase's string ids (an API key's) and the platform's UUIDs. Nothing here - is a foreign key for the same reason: the rows a scope names live in four - different tables, and a provider instance may be configured in ``config.yml`` - and have no row at all. - - A row names a ``budgets`` row and holds the counters for spending it. The - limit and the period are read through the budget, never copied, so editing a - budget moves every ceiling that names it. That is deliberate: a budget is a - named thing an operator hands out, and the alternative was the same figure - typed once per place it applied. - - This table does not replace ``budgets``, and the two enforce differently. A - budget reached through ``users.budget_id`` is checked against - ``users.spend + users.reserved``, so N users sharing one each get the full - limit. A budget reached through a row here is checked against *this row's* - counters, so everyone the scope names draws on one allowance. Same budget, - two enforcement shapes, which is why both mechanisms exist. - """ - - __tablename__ = "scoped_budgets" - __table_args__ = ( - # PostgreSQL treats NULLs as distinct in a plain UNIQUE, so one index - # over the triple would enforce nothing on the aggregate rows (every one - # of them has a NULL key, so no two are ever "equal"). Two partial - # indexes instead: the narrowed rows are unique on the triple, and the - # aggregate rows are unique on the identity alone, which is what makes - # "one aggregate cap per scope" a real constraint. - Index( - "uq_scoped_budgets_scope_with_key", - "scope_type", - "scope_id", - "provider_key_id", - unique=True, - postgresql_where=text("provider_key_id IS NOT NULL"), - sqlite_where=text("provider_key_id IS NOT NULL"), - ), - Index( - "uq_scoped_budgets_scope_no_key", - "scope_type", - "scope_id", - unique=True, - postgresql_where=text("provider_key_id IS NULL"), - sqlite_where=text("provider_key_id IS NULL"), - ), - # The request path resolves rows by identity, so the lookup needs a - # non-partial index: neither unique index above covers a scan that spans - # narrowed and aggregate rows. - Index("ix_scoped_budgets_scope", "scope_type", "scope_id"), - ) - - id: Mapped[str] = mapped_column(primary_key=True, default=lambda: str(uuid.uuid4())) - scope_type: Mapped[str] = mapped_column() - scope_id: Mapped[str] = mapped_column() - provider_key_id: Mapped[str | None] = mapped_column(default=None) - name: Mapped[str | None] = mapped_column(default=None) - # The budget this ceiling enforces. NOT NULL: a ceiling with no budget caps - # nothing. The limit and the period are read through it rather than copied, so - # editing a budget moves every ceiling that names it, which is the point of a - # budget being a named thing rather than a number typed twice. - budget_id: Mapped[str] = mapped_column( - ForeignKey("budgets.budget_id", ondelete="RESTRICT"), nullable=False, index=True - ) - current_spend: Mapped[Decimal] = mapped_column(UsdCost(), default=Decimal(0), server_default="0") - # In-flight holds from reservations that have passed the gate but whose actual - # cost is not known yet. Headroom is ``max_budget - current_spend - - # reserved_spend``; a period roll zeroes ``current_spend`` only, so a hold - # taken before the roll is still released correctly after it. - reserved_spend: Mapped[Decimal] = mapped_column(UsdCost(), default=Decimal(0), server_default="0") - # One counter pair per non-USD axis the budget can cap, holding and settling - # exactly as the money pair above does. A period roll zeroes the ``current_*`` - # of all three axes and leaves every hold, so a hold taken before a roll is - # still released correctly after it. - current_tokens: Mapped[int] = mapped_column(BigInteger(), default=0, server_default="0") - reserved_tokens: Mapped[int] = mapped_column(BigInteger(), default=0, server_default="0") - current_requests: Mapped[int] = mapped_column(BigInteger(), default=0, server_default="0") - reserved_requests: Mapped[int] = mapped_column(BigInteger(), default=0, server_default="0") - period_start: Mapped[datetime | None] = mapped_column(DateTime(timezone=True), default=None) - period_end: Mapped[datetime | None] = mapped_column(DateTime(timezone=True), default=None) - created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=lambda: datetime.now(UTC)) - updated_at: Mapped[datetime] = mapped_column( - DateTime(timezone=True), - default=lambda: datetime.now(UTC), - onupdate=lambda: datetime.now(UTC), - ) - - -class BudgetReservation(Base): - """One in-flight budget hold, recorded as a row. - - ``users.reserved`` and ``scoped_budgets.reserved_spend`` stay the O(1) - counters the gate reads; this is the ledger behind them, and it exists for - the two things a counter cannot do (mozilla-ai/otari#742): - - * **Release becomes idempotent.** Without an identity, a second release for - the same request silently subtracts the hold twice. ``_release_reserved`` - clamps at zero, so that shows up not as an error but as an under-count of - live holds, which weakens the very overspend guarantee the reserve gate - exists to provide. The status transition here is what makes only the first - release do the work. - * **A leaked hold becomes reclaimable individually.** A failure between - reserve and settle used to leave an amount in the counter that could be - seen only in aggregate and released by nothing at all: the budget reset - zeroes ``spend`` and leaves ``reserved`` where it is. With a row it has an - owner, an age and a TTL. - - The row is written *after* the holds it records, never before. A hold with no - row is the pre-existing leak the sweep bounds; a row with no hold would have - the sweep release an amount nobody holds, under-counting the live ones. Of - the two inconsistent windows only one is safe, and this is it. - - Standalone mode only: hybrid mode reserves nothing locally, because the - platform holds against its own ledger. - """ - - __tablename__ = "budget_reservations" - __table_args__ = ( - # The global sweep's access path: active rows whose TTL has elapsed. - # Equality on ``status`` leads so the range scan on ``expires_at`` rides - # the same index. - Index("ix_budget_reservations_status_expires_at", "status", "expires_at"), - # The per-user reclaim's, which runs on every request that takes a hold. - # It has to lead on ``user_id``: given only the index above, the planner - # takes it and filters ``user_id``, so one user's reclaim pays for the - # whole deployment's backlog of expired rows. Leading on ``user_id`` also - # serves the FK cascade, so this replaces the plain index on that column - # rather than joining it. - Index("ix_budget_reservations_user_status_expires", "user_id", "status", "expires_at"), - ) - - id: Mapped[str] = mapped_column(primary_key=True, default=lambda: str(uuid.uuid4())) - # No ``index=True``: the composite in ``__table_args__`` leads on this column, - # so a plain index here would be a second, redundant one, and the migration - # deliberately does not create it. Declaring it anyway made a ``create_all`` - # schema and a migrated one disagree. - user_id: Mapped[str] = mapped_column(ForeignKey("users.user_id", ondelete="CASCADE"), nullable=False) - # What the per-user leg holds in ``users.reserved``. Zero when the request - # held only scoped ceilings (a user with no budget row still passes those). - estimate: Mapped[Decimal] = mapped_column(UsdCost(), default=Decimal(0), server_default="0") - # What the same leg holds on the other two axes. Recorded per axis because the - # sweep has to give back every axis a leaked hold took: a period roll zeroes - # ``current_*`` and deliberately leaves the holds, so a token hold nothing - # releases shrinks that ceiling for good. - token_estimate: Mapped[int] = mapped_column(BigInteger(), default=0, server_default="0") - request_estimate: Mapped[int] = mapped_column(BigInteger(), default=0, server_default="0") - # Whether the ``users.reserved`` write actually happened. Distinct from - # ``estimate > 0`` because a zero-cost request on an enforced budget still - # takes the hold, and the release has to match what the reserve did. - user_reserved: Mapped[bool] = mapped_column(default=False, server_default=false()) - # A plain string rather than a database enum, matching ``scoped_budgets.scope_type``: - # a new state should not need an enum migration. Values are the - # ``RESERVATION_*`` constants in gateway.services.budget_reservation_ledger. - status: Mapped[str] = mapped_column(default="active", server_default="active", nullable=False) - # After this instant a still-active row is treated as leaked and reclaimed. - expires_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), nullable=False) - created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=lambda: datetime.now(UTC)) - updated_at: Mapped[datetime] = mapped_column( - DateTime(timezone=True), - default=lambda: datetime.now(UTC), - onupdate=lambda: datetime.now(UTC), - ) - - -class BudgetReservationScope(Base): - """The hold one reservation placed on one scoped ceiling. - - ``scoped_budget_id`` is deliberately not a foreign key, following - ``ScopedBudget``'s own convention: a ceiling deleted while a request is in - flight leaves an orphan line that the release skips, rather than forcing the - delete to cascade into live holds. - - The amounts are stored per line, one per axis, even though today every - ceiling of a request holds the same figures. A ledger line that does not say - what it holds is not a ledger line, and reading the amounts from the parent - would silently become wrong the first time the two diverge. - """ - - __tablename__ = "budget_reservation_scopes" - __table_args__ = (Index("ix_budget_reservation_scopes_reservation_id", "reservation_id"),) - - id: Mapped[str] = mapped_column(primary_key=True, default=lambda: str(uuid.uuid4())) - reservation_id: Mapped[str] = mapped_column( - ForeignKey("budget_reservations.id", ondelete="CASCADE"), nullable=False - ) - scoped_budget_id: Mapped[str] = mapped_column(nullable=False) - amount: Mapped[Decimal] = mapped_column(UsdCost(), default=Decimal(0), server_default="0") - token_amount: Mapped[int] = mapped_column(BigInteger(), default=0, server_default="0") - request_amount: Mapped[int] = mapped_column(BigInteger(), default=0, server_default="0") - - -class OrganizationModelPricing(Base): - """One organization's rate for a model, sitting above the deployment price list. - - ``model_pricing`` carries no tenancy column: it is one price list for the - whole deployment. This table is the layer above it, so an organization can - price the models it uses at its own negotiated rates while every other - organization, and every model it has not overridden, keeps resolving exactly - as before. Resolution order is override, then deployment row, then the - genai-prices dataset (`services.pricing_service.find_model_pricing`). - - **Keyed on ``model_key``, not a split provider and model.** The platform's - equivalent table (`otari-ai` ``organization_model_pricing``) carries - ``provider`` and ``model`` as separate columns. Here the whole pricing chain - keys on one ``provider:model`` string, and that string is not always a - provider and a model: a pricing key names a provider *instance* - (``home_lab:llama-3``, over ``provider_type: openai``) and sometimes no model - at all (``otari:web_search``). Splitting it would make an override - unmatchable for exactly the keys an operator is most likely to have priced by - hand, so the override keys the same way the row it overrides does. - - **An interval, where ``model_pricing`` carries a version series.** A price in - ``model_pricing`` is ``(model_key, effective_at)`` and a later row silently - shadows an earlier one, which is the right shape for a catalog an operator - re-imports. An override is a commitment for a period, so it carries both ends - and overlapping periods for one model are refused rather than shadowed - (`services.organization_pricing_service`). ``effective_to`` NULL means open - ended. - - **The overlap rule is enforced in the service, not by the database.** The - natural constraint is a PostgreSQL ``EXCLUDE`` over a ``tstzrange``, and the - platform has one. SQLite has no exclusion constraint and no range type, and - it is what the OSS edition ships by default, so a database-side rule would - hold on one engine and be a comment on the other. The unique index below is - what both engines can enforce: it stops the exact-duplicate start, which is - the collision two concurrent writers actually produce, while a partial - overlap between two simultaneous inserts remains a narrow race the service's - check can lose. Single-writer configuration traffic, and a wrong rate is - visible and correctable rather than silent. - - Rates are the exact ``UsdRate`` type ``ModelPricing`` uses, and the two - tables carry it together (#661, one migration over both). An override - resolves *into* a transient ``ModelPricing``, so one implementation of the - cost math prices both, which it could not if they disagreed about the type - of money. - """ - - __tablename__ = "organization_model_pricing" - __table_args__ = ( - # One index, doing both jobs, because they want the same columns in the - # same order. As a constraint it refuses two rows for one key that begin - # at the same instant, which is the part of the overlap rule either - # engine can hold (see the class docstring for why the rest is in the - # service). As an index it serves the resolution lookup: the two equality - # columns lead, so the request path gets a prefix scan, and - # ``effective_from`` trails so picking the newest applicable period is - # index-ordered rather than a sort. - Index( - "uq_organization_model_pricing_period_start", - "organization_id", - "model_key", - "effective_from", - unique=True, - ), - # An inverted period would resolve for no instant at all, so it is a - # storage error rather than a pricing decision. Equal ends are refused - # too: a zero-width period is the same silent nothing. - CheckConstraint( - "effective_to IS NULL OR effective_to > effective_from", - name="ck_organization_model_pricing_period_ordered", - ), - # Negative money prices a request as a credit. The service rejects it - # with a message; these are the backstop for a writer that is not the - # service, and they are spelled out per column because a single check - # over all five would not say which rate was wrong. - CheckConstraint( - "input_price_per_million >= 0", - name="ck_organization_model_pricing_input_non_negative", - ), - CheckConstraint( - "output_price_per_million >= 0", - name="ck_organization_model_pricing_output_non_negative", - ), - CheckConstraint( - "cache_read_price_per_million IS NULL OR cache_read_price_per_million >= 0", - name="ck_organization_model_pricing_cache_read_non_negative", - ), - CheckConstraint( - "cache_write_price_per_million IS NULL OR cache_write_price_per_million >= 0", - name="ck_organization_model_pricing_cache_write_non_negative", - ), - CheckConstraint( - "cache_write_1h_price_per_million IS NULL OR cache_write_1h_price_per_million >= 0", - name="ck_organization_model_pricing_cache_write_1h_non_negative", - ), - ) - - id: Mapped[uuid.UUID] = mapped_column(Uuid, primary_key=True, default=uuid.uuid4) - # CASCADE, where the request-plane tables above use RESTRICT. Those are kept - # because a workspace's spend history must survive it; an override is - # configuration, and an organization's rates mean nothing once the - # organization is gone. The usage rows priced under it keep their own settled - # cost, so deleting this loses no accounting. - # No ``index=True``: the composite lookup index below leads on this column, - # so a plain one on it would be a second index serving queries the first - # already answers, paid for on every write. - organization_id: Mapped[uuid.UUID] = mapped_column( - Uuid, ForeignKey("organization.id", ondelete="CASCADE"), nullable=False - ) - model_key: Mapped[str] = mapped_column() - input_price_per_million: Mapped[Decimal] = mapped_column(UsdRate()) - output_price_per_million: Mapped[Decimal] = mapped_column(UsdRate()) - # Nullable for the same reason ``ModelPricing``'s are: a provider without - # prompt caching, or a model with no discounted cache rate, leaves them unset - # and the cost calculation falls back the way it already does. - cache_read_price_per_million: Mapped[Decimal | None] = mapped_column(UsdRate(), nullable=True) - cache_write_price_per_million: Mapped[Decimal | None] = mapped_column(UsdRate(), nullable=True) - cache_write_1h_price_per_million: Mapped[Decimal | None] = mapped_column(UsdRate(), nullable=True) - # Same shape and same ``min_input_tokens`` key as ``ModelPricing``, so the - # transient row an override resolves into needs no tier translation. - pricing_tiers: Mapped[list[dict[str, Any]]] = mapped_column(JSON, default=list) - # The same two columns ``ModelPricing`` carries, for the same reasons: an - # override is read as a ``ModelPricing`` and has to say what it is per. - unit: Mapped[str] = mapped_column(String(16), default="tokens", server_default="tokens") - origin: Mapped[str | None] = mapped_column(String(16), nullable=True) - # ``UtcDateTime``, not ``DateTime(timezone=True)``, and this is the one place - # in this file where that distinction is load-bearing. The flag is a no-op on - # SQLite, which is what ``core/config.py`` defaults ``database_url`` to, so a - # plain column reads back naive there and this table's timestamps are the - # ones that go out over the wire: ``OrganizationModelPricingPublic`` would - # serialize them with no offset, a browser parses an offset-less date-time as - # *local*, and the Edit dialog would then round-trip the period shifted by - # the reader's UTC offset on every save. ``UtcDateTime.impl`` is - # ``DateTime(timezone=True)``, so the DDL and the migration are unchanged; it - # normalizes on the way in and stamps UTC on the way out. - # - # ``ModelPricing`` above keeps the plain column because nothing renders its - # ``effective_at`` into an editable control; the transient row an override - # resolves into is stamped in ``_override_as_model_pricing`` for the cost - # path, which is a different fix for a different reader. - effective_from: Mapped[datetime] = mapped_column( - UtcDateTime(), - default=lambda: datetime.now(UTC), - ) - # NULL means open ended, which is the common case: an organization sets a - # rate and it applies until something replaces it. - effective_to: Mapped[datetime | None] = mapped_column(UtcDateTime(), default=None) - created_at: Mapped[datetime] = mapped_column(UtcDateTime(), default=lambda: datetime.now(UTC)) - updated_at: Mapped[datetime] = mapped_column( - UtcDateTime(), - default=lambda: datetime.now(UTC), - onupdate=lambda: datetime.now(UTC), - ) - - -class WorkspaceBudgetDefault(Base): - """A workspace-level template for a per-member ``ScopedBudget``. - - ``scoped_budgets`` holds concrete ceilings; this table has no counters of - its own and enforces nothing directly. It is **materialized**: creating one - on a workspace that already has members, or a member joining a workspace - that already has one, stages a ``ScopedBudget(scope_type="workspace_member", - scope_id=)`` row for each (see - ``services/tenancy/workspace_budget_default_service.py``). A member with an - existing ceiling for the same ``provider_key_id`` is left alone; a - member-specific override always wins over the template. - - Same two-axis shape as ``ScopedBudget``: ``workspace_id`` is who the - template belongs to, ``provider_key_id`` optionally narrows it to one - provider instance (NULL applies to all of them). Unlike ``ScopedBudget``, - ``workspace_id`` is a real foreign key: a template has exactly one owner - and nothing else names it, so it is deleted with the workspace rather than - requiring the same explicit cleanup ``ScopedBudget`` needs (see - ``WorkspaceService._delete_scoped_budgets_for``). - """ - - __tablename__ = "workspace_budget_defaults" - __table_args__ = ( - # Same reasoning as ScopedBudget's two partial indexes: PostgreSQL and - # SQLite both treat NULLs as distinct in a plain UNIQUE, so a single - # index over the pair would enforce nothing on the aggregate (NULL-key) - # rows. - Index( - "uq_workspace_budget_defaults_with_key", - "workspace_id", - "provider_key_id", - unique=True, - postgresql_where=text("provider_key_id IS NOT NULL"), - sqlite_where=text("provider_key_id IS NOT NULL"), - ), - Index( - "uq_workspace_budget_defaults_no_key", - "workspace_id", - unique=True, - postgresql_where=text("provider_key_id IS NULL"), - sqlite_where=text("provider_key_id IS NULL"), - ), - ) - - id: Mapped[str] = mapped_column(primary_key=True, default=lambda: str(uuid.uuid4())) - workspace_id: Mapped[uuid.UUID] = mapped_column( - Uuid, ForeignKey("workspace.id", ondelete="CASCADE"), nullable=False, index=True - ) - provider_key_id: Mapped[str | None] = mapped_column(default=None) - # The budget this workspace hands to every member. NOT NULL: a default that - # names no budget is a template for nothing. ``RESTRICT`` because deleting a - # budget a workspace hands out should be refused and explained rather than - # silently withdraw the limit from every ceiling it materialized. - # - # The limit and the period live on the budget, not here, which is what lets - # the Budgets page say that a row is a workspace's default. ``provider_key_id`` - # stays on this side: which provider a workspace applies the budget to is a - # property of the assignment, and two workspaces may narrow one budget - # differently. - budget_id: Mapped[str] = mapped_column( - ForeignKey("budgets.budget_id", ondelete="RESTRICT"), nullable=False, index=True - ) - # ``UtcDateTime``, not ``DateTime(timezone=True)``: these two are serialized with - # ``.isoformat()`` (``WorkspaceMemberBudgetPolicyPublic.from_model``) for the - # dashboard, and on SQLite (this repo's default ``database_url``) - # a plain ``DateTime(timezone=True)`` round-trips naive, so the wire value - # would carry no offset and a browser would read it as local time. - # ``UtcDateTime.impl`` is ``DateTime(timezone=True)``, so the DDL is unchanged. - created_at: Mapped[datetime] = mapped_column(UtcDateTime(), default=lambda: datetime.now(UTC)) - updated_at: Mapped[datetime] = mapped_column( - UtcDateTime(), - default=lambda: datetime.now(UTC), - onupdate=lambda: datetime.now(UTC), - ) - - -class WorkspaceActivationState(Base): - """What the dashboard's first-request setup guide remembers about a workspace. - - The guide walks a workspace from "no traffic" to its first successful - request (`services/tenancy/workspace_activation_service.py`). Only what - cannot be observed elsewhere is stored here: whether someone dismissed it, - when it last handed out a key, and which key that was. Whether the workspace - has *activated* is deliberately not a column, because ``usage_logs`` already - records it: the first successful gateway request in the workspace is the - evidence, so there is no second copy of it to backfill or to disagree with - the Activity page. - - Ported from the platform's ``workspace_activation_state`` / - ``workspace_activation_experience_state`` pair - (`otari-ai` `backend/app/models/workspace_activation.py`), which does carry - the attempt telemetry as columns, because its usage pipeline is asynchronous - and crosses services. Here the usage row is written by this process into this - database, so the derivation is exact. - - One row per workspace, not per workspace and viewer: the guide is about a - workspace's first request, so dismissing it says "this workspace is set up, - stop offering the guide" for everyone who can manage it. - """ - - __tablename__ = "workspace_activation_state" - - workspace_id: Mapped[uuid.UUID] = mapped_column( - Uuid, ForeignKey("workspace.id", ondelete="CASCADE"), primary_key=True - ) - # When the guide first and last minted an API key for this workspace. The - # first is what an operator reads as "when was this offered"; the last is - # what makes a rotation visible next to the key it rotated. - first_presented_at: Mapped[datetime | None] = mapped_column(UtcDateTime(), default=None) - last_presented_at: Mapped[datetime | None] = mapped_column(UtcDateTime(), default=None) - # Set by Skip, and permanent: the guide is a first-run offer, so a workspace - # that turned it down is not asked again on the next page load. - dismissed_at: Mapped[datetime | None] = mapped_column(UtcDateTime(), default=None) - # The key the guide issued, rotated in place on each presentation so a - # workspace collects one "Setup guide" key rather than one per page load. - # ``SET NULL`` because deleting that key from the Keys page is a legitimate - # thing to do, and it must not take this row (or the dismissal on it) with it. - api_key_id: Mapped[str | None] = mapped_column( - ForeignKey("api_keys.id", ondelete="SET NULL"), default=None, index=True - ) - # ``UtcDateTime`` rather than ``DateTime(timezone=True)`` for the same reason - # ``WorkspaceBudgetDefault`` above uses it: on SQLite, which this edition - # ships by default, the plain type round-trips naive and a browser would read - # the value as local time. - created_at: Mapped[datetime] = mapped_column(UtcDateTime(), default=lambda: datetime.now(UTC)) - updated_at: Mapped[datetime] = mapped_column( - UtcDateTime(), - default=lambda: datetime.now(UTC), - onupdate=lambda: datetime.now(UTC), - ) - - -class WorkspaceMcpServer(Base): - """One MCP server a workspace has configured, referenced by id from a request. - - Ported from otari-ai's ``mcp_server`` table (otari#658). A request names - stored servers with ``mcp_server_ids``; hybrid mode resolves those ids - through the platform and standalone mode resolves them here, against the - workspace the request's key belongs to. There is no deployment-wide MCP - server list for these rows to narrow, which is why MCP is the stated - exception to the "a workspace row never grants" rule in - ``src/gateway/AGENTS.md``. - - ``encrypted_token`` holds the server's bearer token, Fernet-encrypted with - ``OTARI_SECRET_KEY`` (``services/secret_box.py``), the same treatment - ``ProviderCredential.encrypted_api_key`` gets. Nothing serializes it: the - public shape carries ``has_token`` and no prefix or suffix of the value, - because unlike a provider key's ``last4`` there is no operator workflow - here that needs to tell two tokens apart at a glance. - - ``enabled`` is a workspace-level off switch that keeps the row and its - token: a disabled server is skipped at resolve rather than refusing the - request, so a caller whose stored id list outlives one server's - decommissioning still gets the rest. - - CASCADE, not the ``RESTRICT`` the request-plane tables above use: this is a - workspace-owned configuration row, like ``workspace_budget_defaults``, with - no meaning once its workspace is gone. - """ - - __tablename__ = "workspace_mcp_servers" - __table_args__ = ( - # Duplicate names within one workspace are rejected at the database, not - # only in the service layer, so two concurrent creates cannot both land - # (otari#658's third Definition-of-Done item). The name is what an - # operator recognizes a server by and what the tool loop labels its - # tools with, so collapsing two onto one name would silently hide a - # server. - UniqueConstraint("workspace_id", "name", name="uq_workspace_mcp_servers_workspace_name"), - ) - - id: Mapped[uuid.UUID] = mapped_column(Uuid, primary_key=True, default=uuid.uuid4) - workspace_id: Mapped[uuid.UUID] = mapped_column( - Uuid, ForeignKey("workspace.id", ondelete="CASCADE"), nullable=False, index=True - ) - name: Mapped[str] = mapped_column(nullable=False) - url: Mapped[str] = mapped_column(nullable=False) - encrypted_token: Mapped[str | None] = mapped_column(Text, default=None) - purpose_hint: Mapped[str | None] = mapped_column(Text, default=None) - allowed_tools: Mapped[list[str] | None] = mapped_column(JSON, default=None) - enabled: Mapped[bool] = mapped_column(default=True, nullable=False) - # ``UtcDateTime`` for the same reason ``WorkspaceBudgetDefault``'s are: these - # go over the wire and a naive SQLite round-trip would drop the offset. - created_at: Mapped[datetime] = mapped_column(UtcDateTime(), default=lambda: datetime.now(UTC)) - updated_at: Mapped[datetime] = mapped_column( - UtcDateTime(), - default=lambda: datetime.now(UTC), - onupdate=lambda: datetime.now(UTC), - ) - - -class WorkspaceCodeExecutionPolicy(Base): - """A workspace's policy over the deployment-wide code-execution sandbox. - - The sandbox itself stays deployment-wide (``sandbox_url`` and its - credential are operator concerns and never move here, see - ``src/gateway/AGENTS.md``); this row says who on that deployment may ask - for it and within which limits. Resolved at admission by - ``prepare_gateway_tools`` and applied to the tool loop, the standalone - counterpart of the hybrid path's ``/gateway/code-execution/resolve``. - - A row may only *narrow*: ``enabled=False`` refuses the tool for this - workspace, and the two limits are floored against the values a request - would otherwise get. No row means no narrowing, which is what keeps a - deployment that configures nothing behaving as it did (#655/#678). - - ``workspace_id`` is the primary key: a workspace has one policy or none, - so there is nothing else to identify a row by. It is a real foreign key - with ``CASCADE``, like ``workspace_budget_defaults``: nothing else names - the row, so it rides the workspace's own delete. - - ``image`` and ``tools`` reach the same two decisions the hosted - ``CodeExecutionConfig`` carries (#740). Neither breaks the rule above: - ``image`` may only name something the deployment's operator has already - curated into ``sandbox_allowed_session_images``, so a workspace picks from an - operator's shelf rather than pointing the gateway at an image of its own, - and ``tools`` may only remove tool kinds from what the sandbox backend - already serves. - """ - - __tablename__ = "workspace_code_execution_policies" - __table_args__ = ( - # Both limits are ceilings that get floored into an effective value, so - # zero or negative is a storage error rather than a stricter policy: it - # would floor the loop to nothing runnable while reading as configured. - # The request schemas refuse it first; these are the backstop for a - # writer that is not the service. - CheckConstraint( - "max_iterations IS NULL OR max_iterations > 0", - name="ck_workspace_code_execution_policies_max_iterations_positive", - ), - CheckConstraint( - "exec_timeout_s IS NULL OR exec_timeout_s > 0", - name="ck_workspace_code_execution_policies_exec_timeout_positive", - ), - ) - - workspace_id: Mapped[uuid.UUID] = mapped_column( - Uuid, ForeignKey("workspace.id", ondelete="CASCADE"), primary_key=True - ) - enabled: Mapped[bool] = mapped_column(default=True, nullable=False) - # NULL means "no workspace default": the request's own hint, then the - # deployment's, then the backend's built-in, exactly as today. - default_purpose_hint: Mapped[str | None] = mapped_column(Text, default=None) - # Both NULL-able ceilings, applied with ``min`` against what the request - # would otherwise get, so a value above the deployment ceiling narrows - # nothing rather than raising it. - max_iterations: Mapped[int | None] = mapped_column(default=None) - exec_timeout_s: Mapped[int | None] = mapped_column(default=None) - # NULL means "no workspace image": whatever the deployment names in - # ``sandbox_session_image``, and failing that whatever the sandbox backend runs by - # default, which is what every request got before this column existed. - # ``String(255)`` rather than ``Text`` to match the hosted column's own - # bound; an image reference that long is already pathological. - image: Mapped[str | None] = mapped_column(String(255), default=None) - # NULL means "no workspace tool allow-list": the backend offers what it - # offers. A stored list is an intersection, never a union, so it can only - # take tool kinds away. JSON rather than a child table for the same reason - # ``WorkspaceWebSearchConfig`` stores its domain lists that way: short, read - # whole, and nothing queries into it. - tools: Mapped[list[str] | None] = mapped_column(JSON, default=None) - # ``UtcDateTime`` for the same reason ``WorkspaceBudgetDefault`` uses it: - # these are serialized with ``.isoformat()`` for the dashboard, and a plain - # ``DateTime(timezone=True)`` round-trips naive on SQLite. - created_at: Mapped[datetime] = mapped_column(UtcDateTime(), default=lambda: datetime.now(UTC)) - updated_at: Mapped[datetime] = mapped_column( - UtcDateTime(), - default=lambda: datetime.now(UTC), - onupdate=lambda: datetime.now(UTC), - ) - - -class WorkspaceWebSearchConfig(Base): - """A workspace's configuration over the deployment-wide web-search backend. - - The backend itself stays deployment-wide (``web_search_url`` and the - credential the adapter in front of it holds are operator concerns and never - move here, see ``src/gateway/AGENTS.md``); this row says which workspaces - may reach it and how their searches are constrained. Resolved at admission - by ``prepare_gateway_tools``, the standalone counterpart of the hybrid - path's ``/gateway/web-search/resolve``. - - A row may only *narrow*: ``enabled=False`` refuses ``otari_web_search`` for - the workspace, ``max_results`` is floored against what the request asked - for, ``blocked_domains`` is added to the request's own block-list, and - ``allowed_domains`` intersects the request's. No row means no narrowing, - which is what keeps a deployment that configures nothing behaving as it did - (#655/#678). - - ``workspace_id`` is the primary key, and a real foreign key with - ``CASCADE``, for the same reasons as :class:`WorkspaceCodeExecutionPolicy` - next door: one row per workspace, and nothing else names it. - - There is deliberately no ``provider`` column, which the hosted config - carries: on this deployment the operator picks the backend by pointing - ``web_search_url`` somewhere, so a provider named here would either be inert - or would ask the gateway to reach an endpoint the operator did not choose, - which is the one thing the narrowing rule forbids. - """ - - __tablename__ = "workspace_web_search_configs" - __table_args__ = ( - # ``max_results`` is floored into an effective value, so zero or less is - # a storage error rather than a stricter policy: it would ask for a - # search that can return nothing while reading as configured. The - # request schema refuses it first; this is the backstop for a writer - # that is not the service. - CheckConstraint( - "max_results IS NULL OR max_results > 0", - name="ck_workspace_web_search_configs_max_results_positive", - ), - ) - - workspace_id: Mapped[uuid.UUID] = mapped_column( - Uuid, ForeignKey("workspace.id", ondelete="CASCADE"), primary_key=True - ) - # ``server_default`` mirrors the migration so autogenerate sees no drift, and - # so a row written by anything other than this mapping still gets a value. - enabled: Mapped[bool] = mapped_column(default=True, nullable=False, server_default=true()) - # NULL means "no workspace ceiling": the request's own value, then the - # deployment's, then the backend's built-in, exactly as today. - max_results: Mapped[int | None] = mapped_column(default=None) - # NULL means "no workspace default": the request's own hint, then the - # deployment's, then the backend's built-in. - purpose_hint: Mapped[str | None] = mapped_column(Text, default=None) - # Two domain lists and an opaque provider bag, stored as JSON for the same - # reason the hosted table does: they are short, they are read whole, and - # nothing queries into them. ``JSON`` rather than ``JSONB`` to match every - # other JSON column here, which has to work on SQLite too. - allowed_domains: Mapped[list[str] | None] = mapped_column(JSON, default=None) - blocked_domains: Mapped[list[str] | None] = mapped_column(JSON, default=None) - # Provider-specific knobs (Tavily's ``search_depth``, say). Opaque here and - # forwarded to the backend, which is what lets a new provider need no - # migration; the adapter in front of it whitelists what it understands. - provider_options: Mapped[dict[str, Any] | None] = mapped_column(JSON, default=None) - # ``UtcDateTime`` for the same reason ``WorkspaceCodeExecutionPolicy`` uses - # it: these are serialized with ``.isoformat()`` for the dashboard, and a - # plain ``DateTime(timezone=True)`` round-trips naive on SQLite. The Python - # default is what every write here uses; ``server_default`` is the backstop - # for a writer that is not this mapping, matching ``workspace`` itself. - created_at: Mapped[datetime] = mapped_column( - UtcDateTime(), default=lambda: datetime.now(UTC), server_default=func.now() - ) - updated_at: Mapped[datetime] = mapped_column( - UtcDateTime(), - default=lambda: datetime.now(UTC), - onupdate=lambda: datetime.now(UTC), - server_default=func.now(), - ) - - -class OrganizationGuardrail(Base): - """A guardrail an organization runs over the requests of its workspaces. - - The plane *above* the deployment-wide guardrail settings, not a replacement - for them: ``guardrails_url`` stays in ``runtime_settings`` and a deployment - that configures no organization guardrails behaves exactly as it did - (otari#654). A row here is a check the organization mandates; it is merged - into the effective guardrail list at admission by ``prepare_gateway_tools`` - the same way a routing policy's mandate already is, so an organization can - only ever add a check or tighten one a caller asked for. - - That is what keeps this inside the rule ``src/gateway/AGENTS.md`` records - from #655/#678: a mandated guardrail can only make *fewer* requests succeed, - never more, whichever endpoint it names. Which is also why the entry may - carry its own ``url`` and credential where a workspace code-execution policy - may not: the sandbox is a capability a workspace would be acquiring, and a - guardrail is a restriction the organization is accepting. A caller can - already point a request-body guardrail at a URL of their own - (``models/guardrails.GuardrailConfig.url``, SSRF-checked on the request - path), so storing one here grants nothing that was not already reachable. - - ``profile`` is unique per organization rather than a nickname being unique, - which is where this parts company with the hosted - ``organization_guardrail_key`` (unique on ``(organization_id, nickname)``, - so one profile may be configured twice). The effective guardrail set on this - request path is keyed by profile, because ``merge_guardrail_layers`` has - always merged that way; two rows of one profile could therefore never both - run, and one would silently win. - """ - - __tablename__ = "organization_guardrails" - __table_args__ = (UniqueConstraint("organization_id", "profile", name="uq_organization_guardrails_org_profile"),) - - id: Mapped[uuid.UUID] = mapped_column(Uuid, primary_key=True, default=uuid.uuid4) - organization_id: Mapped[uuid.UUID] = mapped_column( - Uuid, ForeignKey("organization.id", ondelete="CASCADE"), nullable=False, index=True - ) - profile: Mapped[str] = mapped_column(nullable=False) - # NULL means "use the deployment's guardrails_url", which is the ordinary - # case: an organization that runs its own any-guardrail deployment names it - # here, and then the credential below is what authenticates to it. - url: Mapped[str | None] = mapped_column(default=None) - encrypted_credential: Mapped[str | None] = mapped_column(Text, default=None) - mode: Mapped[str] = mapped_column(default="monitor", nullable=False) - on_unavailable: Mapped[str] = mapped_column(default="block", nullable=False) - validate_kwargs: Mapped[dict[str, Any] | None] = mapped_column(JSON, default=None) - # The organization's own kill switch. A disabled entry runs nowhere, - # whatever its scope says, so an organization can stop a guardrail without - # losing the credential and the workspace list it took to set up. - enabled: Mapped[bool] = mapped_column(default=True, nullable=False) - # The inheritance rule otari#654 asks for, and the hosted plane's - # ``is_org_default`` under a name that says what it does: true means every - # workspace of the organization runs this, including one created tomorrow, - # and the scope rows below are not consulted. False means it runs only in - # the workspaces named there, and a new workspace inherits nothing. - applies_to_all_workspaces: Mapped[bool] = mapped_column(default=False, nullable=False) - # ``UtcDateTime`` for the reason its neighbors use it: these are serialized - # with ``.isoformat()`` for the dashboard, and a plain ``DateTime(timezone=True)`` - # round-trips naive on SQLite. - created_at: Mapped[datetime] = mapped_column(UtcDateTime(), default=lambda: datetime.now(UTC)) - updated_at: Mapped[datetime] = mapped_column( - UtcDateTime(), - default=lambda: datetime.now(UTC), - onupdate=lambda: datetime.now(UTC), - ) - - -class OrganizationGuardrailWorkspace(Base): - """One workspace an organization guardrail is scoped to. - - Membership only: a row means "this guardrail runs in this workspace", and - its absence means it does not. The hosted plane instead carries a - ``disabled`` flag on the equivalent row and admits three states, two of - which resolve to off; there is nothing here for a third state to record, - because the scope is the organization's to set and a workspace has no veto - over it (a veto would widen what succeeds, which #655/#678 does not allow). - - Ignored entirely when the guardrail's ``applies_to_all_workspaces`` is set, - so rows left behind by flipping that on are inert rather than contradictory. - - Both sides cascade: the pairing has no meaning once either end is gone. - """ - - __tablename__ = "organization_guardrail_workspaces" - - organization_guardrail_id: Mapped[uuid.UUID] = mapped_column( - Uuid, ForeignKey("organization_guardrails.id", ondelete="CASCADE"), primary_key=True - ) - workspace_id: Mapped[uuid.UUID] = mapped_column( - Uuid, ForeignKey("workspace.id", ondelete="CASCADE"), primary_key=True, index=True - ) - created_at: Mapped[datetime] = mapped_column(UtcDateTime(), default=lambda: datetime.now(UTC)) diff --git a/src/gateway/models/guardrails.py b/src/gateway/models/guardrails.py index 56f9da03ff..bc5f166b3e 100644 --- a/src/gateway/models/guardrails.py +++ b/src/gateway/models/guardrails.py @@ -11,13 +11,21 @@ operator-controlled guardrails service (``otari-anyguardrails-container``, which exposes ``POST /validate``), and strips the field before forwarding the request upstream. Omit the field entirely → no guardrail runs. + +Also holds the organization guardrail tables. """ from __future__ import annotations +import uuid +from datetime import UTC, datetime from typing import Any, Literal from pydantic import BaseModel, Field +from sqlalchemy import JSON, ForeignKey, Text, UniqueConstraint, Uuid +from sqlalchemy.orm import Mapped, mapped_column + +from gateway.models.base import Base, UtcDateTime GuardrailDirection = Literal["input", "output"] @@ -71,3 +79,96 @@ class GuardrailConfig(BaseModel): validate_kwargs: dict[str, Any] = Field(default_factory=dict) """Extra kwargs forwarded to the guardrails service ``/validate`` call, merged on top of the profile's own ``validate_kwargs`` server-side.""" + + +class OrganizationGuardrail(Base): + """A guardrail an organization runs over the requests of its workspaces. + + The plane *above* the deployment-wide guardrail settings, not a replacement + for them: ``guardrails_url`` stays in ``runtime_settings`` and a deployment + that configures no organization guardrails behaves exactly as it did + (otari#654). A row here is a check the organization mandates; it is merged + into the effective guardrail list at admission by ``prepare_gateway_tools`` + the same way a routing policy's mandate already is, so an organization can + only ever add a check or tighten one a caller asked for. + + That is what keeps this inside the rule ``src/gateway/AGENTS.md`` records + from #655/#678: a mandated guardrail can only make *fewer* requests succeed, + never more, whichever endpoint it names. Which is also why the entry may + carry its own ``url`` and credential where a workspace code-execution policy + may not: the sandbox is a capability a workspace would be acquiring, and a + guardrail is a restriction the organization is accepting. A caller can + already point a request-body guardrail at a URL of their own + (``models/guardrails.GuardrailConfig.url``, SSRF-checked on the request + path), so storing one here grants nothing that was not already reachable. + + ``profile`` is unique per organization rather than a nickname being unique, + which is where this parts company with the hosted + ``organization_guardrail_key`` (unique on ``(organization_id, nickname)``, + so one profile may be configured twice). The effective guardrail set on this + request path is keyed by profile, because ``merge_guardrail_layers`` has + always merged that way; two rows of one profile could therefore never both + run, and one would silently win. + """ + + __tablename__ = "organization_guardrails" + __table_args__ = (UniqueConstraint("organization_id", "profile", name="uq_organization_guardrails_org_profile"),) + + id: Mapped[uuid.UUID] = mapped_column(Uuid, primary_key=True, default=uuid.uuid4) + organization_id: Mapped[uuid.UUID] = mapped_column( + Uuid, ForeignKey("organization.id", ondelete="CASCADE"), nullable=False, index=True + ) + profile: Mapped[str] = mapped_column(nullable=False) + # NULL means "use the deployment's guardrails_url", which is the ordinary + # case: an organization that runs its own any-guardrail deployment names it + # here, and then the credential below is what authenticates to it. + url: Mapped[str | None] = mapped_column(default=None) + encrypted_credential: Mapped[str | None] = mapped_column(Text, default=None) + mode: Mapped[str] = mapped_column(default="monitor", nullable=False) + on_unavailable: Mapped[str] = mapped_column(default="block", nullable=False) + validate_kwargs: Mapped[dict[str, Any] | None] = mapped_column(JSON, default=None) + # The organization's own kill switch. A disabled entry runs nowhere, + # whatever its scope says, so an organization can stop a guardrail without + # losing the credential and the workspace list it took to set up. + enabled: Mapped[bool] = mapped_column(default=True, nullable=False) + # The inheritance rule otari#654 asks for, and the hosted plane's + # ``is_org_default`` under a name that says what it does: true means every + # workspace of the organization runs this, including one created tomorrow, + # and the scope rows below are not consulted. False means it runs only in + # the workspaces named there, and a new workspace inherits nothing. + applies_to_all_workspaces: Mapped[bool] = mapped_column(default=False, nullable=False) + # Gotcha: a plain DateTime(timezone=True) reads back naive on SQLite. The dashboard + # then shows it as local time. + created_at: Mapped[datetime] = mapped_column(UtcDateTime(), default=lambda: datetime.now(UTC)) + updated_at: Mapped[datetime] = mapped_column( + UtcDateTime(), + default=lambda: datetime.now(UTC), + onupdate=lambda: datetime.now(UTC), + ) + + +class OrganizationGuardrailWorkspace(Base): + """One workspace an organization guardrail is scoped to. + + Membership only: a row means "this guardrail runs in this workspace", and + its absence means it does not. The hosted plane instead carries a + ``disabled`` flag on the equivalent row and admits three states, two of + which resolve to off; there is nothing here for a third state to record, + because the scope is the organization's to set and a workspace has no veto + over it (a veto would widen what succeeds, which #655/#678 does not allow). + + Ignored entirely when the guardrail's ``applies_to_all_workspaces`` is set, + so rows left behind by flipping that on are inert rather than contradictory. + + Both sides cascade: the pairing has no meaning once either end is gone. + """ + + __tablename__ = "organization_guardrail_workspaces" + + organization_guardrail_id: Mapped[uuid.UUID] = mapped_column( + Uuid, ForeignKey("organization_guardrails.id", ondelete="CASCADE"), primary_key=True + ) + workspace_id: Mapped[uuid.UUID] = mapped_column( + Uuid, ForeignKey("workspace.id", ondelete="CASCADE"), primary_key=True, index=True + ) + created_at: Mapped[datetime] = mapped_column(UtcDateTime(), default=lambda: datetime.now(UTC)) diff --git a/src/gateway/models/inference.py b/src/gateway/models/inference.py new file mode 100644 index 0000000000..b3c41afb71 --- /dev/null +++ b/src/gateway/models/inference.py @@ -0,0 +1,58 @@ +"""ORM table for asynchronous batch jobs.""" + +import uuid +from datetime import UTC, datetime + +from sqlalchemy import DateTime, ForeignKey, Uuid +from sqlalchemy.orm import Mapped, mapped_column + +from gateway.models.base import Base + + +class BatchRecord(Base): + """Ownership and accounting record for an asynchronous batch job. + + Written at creation time so results accounting can be made idempotent (bill + and log once, on the first completed retrieval), the batch cost can be folded + into ``users.spend``, and ownership can be enforced without depending on the + provider round-tripping the ``otari_user_id`` metadata marker. Batches created + before this table existed carry no record and fall back to the + metadata-anchored ownership path in ``api/routes/batches.py``. ``workspace_id`` + additionally anchors which workspace's organization-scoped provider key + (otari#643) lifecycle calls should resolve credentials from. + """ + + __tablename__ = "batches" + + # Provider-assigned batch id (globally unique per provider), used as the + # lookup key on retrieve/cancel/results. + id: Mapped[str] = mapped_column(primary_key=True) + # Instance/provider name the batch was created against (echoed to clients). + provider: Mapped[str] = mapped_column() + # Billed owner, stamped from the authenticated principal at creation. Non-null: + # this record is the strict ownership anchor, so it must always name an owner. + # CASCADE: deleting the user drops the ownership record (the user's keys are + # gone too, and usage_logs remain the billing history). + user_id: Mapped[str] = mapped_column( + ForeignKey("users.user_id", ondelete="CASCADE"), nullable=False, index=True + ) + # SET NULL: a key may be revoked while its batch is still in flight. + api_key_id: Mapped[str | None] = mapped_column(ForeignKey("api_keys.id", ondelete="SET NULL"), index=True) + # The workspace this batch was CREATED in (otari#643 follow-up), so + # lifecycle calls (retrieve/cancel/results) can resolve organization-scoped + # credentials from the batch's own origin rather than the retriever's + # current workspace: a master-key or legitimately cross-workspace retrieval + # would otherwise use the wrong organization's key, or find none, exactly + # the failure `api_key_id` going NULL on key revocation already risks for + # ownership. Nullable and SET NULL, not RESTRICT: batches created before + # this column existed carry NULL here and fall back to the caller's own + # workspace in `api/routes/batches.py`, and a workspace deleted out from + # under an in-flight batch must not block that delete. + workspace_id: Mapped[uuid.UUID | None] = mapped_column( + Uuid, ForeignKey("workspace.id", ondelete="SET NULL"), index=True + ) + model: Mapped[str] = mapped_column() + created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=lambda: datetime.now(UTC)) + # NULL until the first completed results retrieval accounts the batch; the + # atomic NULL -> now transition is the idempotency gate for billing/logging. + results_accounted_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True)) diff --git a/src/gateway/models/platform.py b/src/gateway/models/platform.py new file mode 100644 index 0000000000..1f71bddc34 --- /dev/null +++ b/src/gateway/models/platform.py @@ -0,0 +1,29 @@ +"""ORM table for deployment settings the dashboard can change at runtime.""" + +from datetime import UTC, datetime + +from sqlalchemy import DateTime +from sqlalchemy.orm import Mapped, mapped_column + +from gateway.models.base import Base + + +class RuntimeSetting(Base): + """A persisted override for a runtime-toggleable config flag. + + A small key/value store for the handful of settings the dashboard can flip + at runtime (model discovery, default pricing). When a key is present it wins + over the config-file/env value and is applied on startup; when absent the + config value stands. The value is stored as a string ("true"/"false") so the + table can hold future non-boolean settings without a schema change. + """ + + __tablename__ = "runtime_settings" + + key: Mapped[str] = mapped_column(primary_key=True) + value: Mapped[str] = mapped_column() + updated_at: Mapped[datetime] = mapped_column( + DateTime(timezone=True), + default=lambda: datetime.now(UTC), + onupdate=lambda: datetime.now(UTC), + ) diff --git a/src/gateway/models/playground.py b/src/gateway/models/playground.py index c4256ae4bc..f5fd84c27c 100644 --- a/src/gateway/models/playground.py +++ b/src/gateway/models/playground.py @@ -45,7 +45,7 @@ are credentials, and none of these tables holds one. Style follows ``models/tenancy.py`` and ``models/provider_keys.py``: SQLModel -rather than `entities.py`'s declarative style, because the ``Public`` schemas +rather than the declarative ``Base`` style, because the ``Public`` schemas below are the endpoint contracts the generated dashboard client is built from, and no ``relationship()`` is declared (lazy loading raises ``MissingGreenlet`` on an ``AsyncSession``), so the routes join explicitly. @@ -58,7 +58,7 @@ from sqlalchemy import Column, Index, Text, UniqueConstraint from sqlmodel import Field, SQLModel -from gateway.models.tenancy import CreatedAtMixin, PrimaryKeyMixin, UpdatedAtMixin +from gateway.models.base import CreatedAtMixin, PrimaryKeyMixin, UpdatedAtMixin # What one save may carry. Each ceiling is enforced at the request schema, so an # oversized save is a 422 naming the field rather than a database error, and the diff --git a/src/gateway/models/pricing.py b/src/gateway/models/pricing.py new file mode 100644 index 0000000000..a3b2ba7299 --- /dev/null +++ b/src/gateway/models/pricing.py @@ -0,0 +1,262 @@ +"""ORM tables for pricing: the deployment price list, organization overrides, and upstream snapshots.""" + +import uuid +from datetime import UTC, datetime +from decimal import Decimal +from typing import Any + +from sqlalchemy import JSON, CheckConstraint, DateTime, ForeignKey, Index, String, Text, Uuid +from sqlalchemy.orm import Mapped, mapped_column + +from gateway.models.base import Base, UtcDateTime +from gateway.models.money import UsdRate + +# The vocabulary of ``ModelPricing.unit`` and ``OrganizationModelPricing.unit``. +# Every per-unit reader (``services/pricing_service`` helpers, the catalog) keys +# on these spellings, and the request schemas validate against them. +PRICING_UNITS: tuple[str, ...] = ("tokens", "requests", "images") + +# The vocabulary of ``origin`` on the same two tables. +PRICING_ORIGINS: tuple[str, ...] = ("config", "api", "migration") + + +class PricingSnapshot(Base): + """An approved, source-tagged upstream pricing catalog.""" + + __tablename__ = "pricing_snapshots" + + source: Mapped[str] = mapped_column(primary_key=True) + snapshot: Mapped[str] = mapped_column(Text) + updated_at: Mapped[datetime] = mapped_column( + DateTime(timezone=True), + default=lambda: datetime.now(UTC), + onupdate=lambda: datetime.now(UTC), + ) + + +class PricingSnapshotHistory(Base): + """One accepted upstream pricing snapshot, kept after a later one replaces it. + + ``pricing_snapshots`` is the current state; this is the record. Written on + every accept, never updated. ``accepted_by`` says whether an operator + confirmed it or the scheduled refresh applied it on its own. + """ + + __tablename__ = "pricing_snapshot_history" + __table_args__ = (Index("ix_pricing_snapshot_history_source_accepted_at", "source", "accepted_at"),) + + id: Mapped[uuid.UUID] = mapped_column(Uuid, primary_key=True, default=uuid.uuid4) + source: Mapped[str] = mapped_column(String(64)) + accepted_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=lambda: datetime.now(UTC)) + accepted_by: Mapped[str] = mapped_column(String(32)) + model_count: Mapped[int] = mapped_column() + snapshot: Mapped[str] = mapped_column(Text) + + +class ModelPricing(Base): + """Model pricing configuration.""" + + __tablename__ = "model_pricing" + + model_key: Mapped[str] = mapped_column(primary_key=True) + effective_at: Mapped[datetime] = mapped_column( + DateTime(timezone=True), + primary_key=True, + default=lambda: datetime.now(UTC), + ) + input_price_per_million: Mapped[Decimal] = mapped_column(UsdRate()) + output_price_per_million: Mapped[Decimal] = mapped_column(UsdRate()) + # Nullable: providers without prompt caching (or models without a + # discounted cache rate) leave these unset. When set, the cost + # calculation prices cache_read_tokens / cache_write_tokens at these + # per-million-token rates, following the provider inclusion convention + # (see log_usage in _pipeline.py). + cache_read_price_per_million: Mapped[Decimal | None] = mapped_column(UsdRate(), nullable=True) + cache_write_price_per_million: Mapped[Decimal | None] = mapped_column(UsdRate(), nullable=True) + cache_write_1h_price_per_million: Mapped[Decimal | None] = mapped_column(UsdRate(), nullable=True) + # Ordered threshold rules. Each rule applies its supplied rates to the + # entire request once ``total_input_tokens`` reaches ``min_input_tokens``. + pricing_tiers: Mapped[list[dict[str, Any]]] = mapped_column(JSON, default=list) + # What ``input_price_per_million`` is a rate per: ``tokens`` for a model, + # ``requests`` for a gateway-run tool or a moderation call, ``images`` for + # image generation. The rate columns are shared by all three and a reader + # cannot tell which from the number, so the row says (``PRICING_UNITS``). + unit: Mapped[str] = mapped_column(String(16), default="tokens", server_default="tokens") + # Which path wrote the row: ``config`` (the file's ``pricing:`` block), + # ``api`` (``POST /v1/pricing``), or ``migration``. NULL on a row written + # before origins were recorded, which is a real answer and not a default. + origin: Mapped[str | None] = mapped_column(String(16), nullable=True) + created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=lambda: datetime.now(UTC)) + updated_at: Mapped[datetime] = mapped_column( + DateTime(timezone=True), + default=lambda: datetime.now(UTC), + onupdate=lambda: datetime.now(UTC), + ) + + def to_dict(self) -> dict[str, Any]: + """Convert model to dictionary.""" + return { + "model_key": self.model_key, + "effective_at": self.effective_at.isoformat() if self.effective_at else None, + "input_price_per_million": self.input_price_per_million, + "output_price_per_million": self.output_price_per_million, + "cache_read_price_per_million": self.cache_read_price_per_million, + "cache_write_price_per_million": self.cache_write_price_per_million, + "cache_write_1h_price_per_million": self.cache_write_1h_price_per_million, + "pricing_tiers": self.pricing_tiers, + "unit": self.unit, + "origin": self.origin, + "created_at": self.created_at.isoformat() if self.created_at else None, + "updated_at": self.updated_at.isoformat() if self.updated_at else None, + } + + +class OrganizationModelPricing(Base): + """One organization's rate for a model, sitting above the deployment price list. + + ``model_pricing`` carries no tenancy column: it is one price list for the + whole deployment. This table is the layer above it, so an organization can + price the models it uses at its own negotiated rates while every other + organization, and every model it has not overridden, keeps resolving exactly + as before. Resolution order is override, then deployment row, then the + genai-prices dataset (`services.pricing_service.find_model_pricing`). + + **Keyed on ``model_key``, not a split provider and model.** The platform's + equivalent table (`otari-ai` ``organization_model_pricing``) carries + ``provider`` and ``model`` as separate columns. Here the whole pricing chain + keys on one ``provider:model`` string, and that string is not always a + provider and a model: a pricing key names a provider *instance* + (``home_lab:llama-3``, over ``provider_type: openai``) and sometimes no model + at all (``otari:web_search``). Splitting it would make an override + unmatchable for exactly the keys an operator is most likely to have priced by + hand, so the override keys the same way the row it overrides does. + + **An interval, where ``model_pricing`` carries a version series.** A price in + ``model_pricing`` is ``(model_key, effective_at)`` and a later row silently + shadows an earlier one, which is the right shape for a catalog an operator + re-imports. An override is a commitment for a period, so it carries both ends + and overlapping periods for one model are refused rather than shadowed + (`services.organization_pricing_service`). ``effective_to`` NULL means open + ended. + + **The overlap rule is enforced in the service, not by the database.** The + natural constraint is a PostgreSQL ``EXCLUDE`` over a ``tstzrange``, and the + platform has one. SQLite has no exclusion constraint and no range type, and + it is what the OSS edition ships by default, so a database-side rule would + hold on one engine and be a comment on the other. The unique index below is + what both engines can enforce: it stops the exact-duplicate start, which is + the collision two concurrent writers actually produce, while a partial + overlap between two simultaneous inserts remains a narrow race the service's + check can lose. Single-writer configuration traffic, and a wrong rate is + visible and correctable rather than silent. + + Rates are the exact ``UsdRate`` type ``ModelPricing`` uses, and the two + tables carry it together (#661, one migration over both). An override + resolves *into* a transient ``ModelPricing``, so one implementation of the + cost math prices both, which it could not if they disagreed about the type + of money. + """ + + __tablename__ = "organization_model_pricing" + __table_args__ = ( + # One index, doing both jobs, because they want the same columns in the + # same order. As a constraint it refuses two rows for one key that begin + # at the same instant, which is the part of the overlap rule either + # engine can hold (see the class docstring for why the rest is in the + # service). As an index it serves the resolution lookup: the two equality + # columns lead, so the request path gets a prefix scan, and + # ``effective_from`` trails so picking the newest applicable period is + # index-ordered rather than a sort. + Index( + "uq_organization_model_pricing_period_start", + "organization_id", + "model_key", + "effective_from", + unique=True, + ), + # An inverted period would resolve for no instant at all, so it is a + # storage error rather than a pricing decision. Equal ends are refused + # too: a zero-width period is the same silent nothing. + CheckConstraint( + "effective_to IS NULL OR effective_to > effective_from", + name="ck_organization_model_pricing_period_ordered", + ), + # Negative money prices a request as a credit. The service rejects it + # with a message; these are the backstop for a writer that is not the + # service, and they are spelled out per column because a single check + # over all five would not say which rate was wrong. + CheckConstraint( + "input_price_per_million >= 0", + name="ck_organization_model_pricing_input_non_negative", + ), + CheckConstraint( + "output_price_per_million >= 0", + name="ck_organization_model_pricing_output_non_negative", + ), + CheckConstraint( + "cache_read_price_per_million IS NULL OR cache_read_price_per_million >= 0", + name="ck_organization_model_pricing_cache_read_non_negative", + ), + CheckConstraint( + "cache_write_price_per_million IS NULL OR cache_write_price_per_million >= 0", + name="ck_organization_model_pricing_cache_write_non_negative", + ), + CheckConstraint( + "cache_write_1h_price_per_million IS NULL OR cache_write_1h_price_per_million >= 0", + name="ck_organization_model_pricing_cache_write_1h_non_negative", + ), + ) + + id: Mapped[uuid.UUID] = mapped_column(Uuid, primary_key=True, default=uuid.uuid4) + # CASCADE: an override means nothing without its organization, and usage rows + # keep their own settled cost. + # No ``index=True``: the composite lookup index below leads on this column, + # so a plain one on it would be a second index serving queries the first + # already answers, paid for on every write. + organization_id: Mapped[uuid.UUID] = mapped_column( + Uuid, ForeignKey("organization.id", ondelete="CASCADE"), nullable=False + ) + model_key: Mapped[str] = mapped_column() + input_price_per_million: Mapped[Decimal] = mapped_column(UsdRate()) + output_price_per_million: Mapped[Decimal] = mapped_column(UsdRate()) + # Nullable for the same reason ``ModelPricing``'s are: a provider without + # prompt caching, or a model with no discounted cache rate, leaves them unset + # and the cost calculation falls back the way it already does. + cache_read_price_per_million: Mapped[Decimal | None] = mapped_column(UsdRate(), nullable=True) + cache_write_price_per_million: Mapped[Decimal | None] = mapped_column(UsdRate(), nullable=True) + cache_write_1h_price_per_million: Mapped[Decimal | None] = mapped_column(UsdRate(), nullable=True) + # Same shape and same ``min_input_tokens`` key as ``ModelPricing``, so the + # transient row an override resolves into needs no tier translation. + pricing_tiers: Mapped[list[dict[str, Any]]] = mapped_column(JSON, default=list) + # The same two columns ``ModelPricing`` carries, for the same reasons: an + # override is read as a ``ModelPricing`` and has to say what it is per. + unit: Mapped[str] = mapped_column(String(16), default="tokens", server_default="tokens") + origin: Mapped[str | None] = mapped_column(String(16), nullable=True) + # ``UtcDateTime``, not ``DateTime(timezone=True)``, and this is the one place + # in this file where that distinction is load-bearing. The flag is a no-op on + # SQLite, which is what ``core/config.py`` defaults ``database_url`` to, so a + # plain column reads back naive there and this table's timestamps are the + # ones that go out over the wire: ``OrganizationModelPricingPublic`` would + # serialize them with no offset, a browser parses an offset-less date-time as + # *local*, and the Edit dialog would then round-trip the period shifted by + # the reader's UTC offset on every save. ``UtcDateTime.impl`` is + # ``DateTime(timezone=True)``, so the DDL and the migration are unchanged; it + # normalizes on the way in and stamps UTC on the way out. + # + # ``ModelPricing`` above keeps the plain column because nothing renders its + # ``effective_at`` into an editable control; the transient row an override + # resolves into is stamped in ``_override_as_model_pricing`` for the cost + # path, which is a different fix for a different reader. + effective_from: Mapped[datetime] = mapped_column( + UtcDateTime(), + default=lambda: datetime.now(UTC), + ) + # NULL means open ended, which is the common case: an organization sets a + # rate and it applies until something replaces it. + effective_to: Mapped[datetime | None] = mapped_column(UtcDateTime(), default=None) + created_at: Mapped[datetime] = mapped_column(UtcDateTime(), default=lambda: datetime.now(UTC)) + updated_at: Mapped[datetime] = mapped_column( + UtcDateTime(), + default=lambda: datetime.now(UTC), + onupdate=lambda: datetime.now(UTC), + ) diff --git a/src/gateway/models/provider_keys.py b/src/gateway/models/provider_keys.py index c77b046df7..c9f9b2f712 100644 --- a/src/gateway/models/provider_keys.py +++ b/src/gateway/models/provider_keys.py @@ -14,7 +14,7 @@ workspace that has an org-scoped key. See mozilla-ai/otari#643. Three tables, named to avoid a collision that already exists in this -codebase: ``ScopedBudget.provider_key_id`` (`models/entities.py`) already +codebase: ``ScopedBudget.provider_key_id`` (`models/budgets.py`) already means "an instance-name string, no FK". These tables use ``org_provider_key`` throughout so no column here is ever ambiguously named ``provider_key_id``. @@ -35,8 +35,8 @@ ``(workspace, key)`` pair means every model is allowed; one or more rows narrows it to exactly those. -Style follows ``models/tenancy.py``: SQLModel (not `entities.py`'s declarative -style) because these are tenancy-scoped tables sharing its mixins and +Style follows ``models/tenancy.py``: SQLModel (not the declarative ``Base`` +style) because these are tenancy-scoped tables sharing the same mixins and ``UtcDateTime`` timestamp handling, and no ``relationship()`` is declared (lazy loading raises ``MissingGreenlet`` on an ``AsyncSession``); repositories join explicitly. @@ -57,8 +57,8 @@ from sqlalchemy import JSON, Column, ForeignKeyConstraint, Index, UniqueConstraint, text from sqlmodel import Field, SQLModel +from gateway.models.base import CreatedAtMixin, PrimaryKeyMixin, UpdatedAtMixin, _timestamp_field from gateway.models.secret_fields import redact_secret_like_values -from gateway.models.tenancy import CreatedAtMixin, PrimaryKeyMixin, UpdatedAtMixin, _timestamp_field # ``client_args`` is arbitrary JSON, and this gateway's own Bedrock support is # the reason a credential-shaped entry in it cannot simply be rejected outright: @@ -82,7 +82,7 @@ class OrgProviderKeyCreateRequest(SQLModel): The plaintext key is never stored as sent: the service encrypts it (`services/secret_box.py`) and keeps only the ciphertext and ``last4``, - the same convention `entities.ProviderCredential` already uses. + the same convention `providers.ProviderCredential` already uses. """ provider: str = Field(max_length=255) diff --git a/src/gateway/models/providers.py b/src/gateway/models/providers.py new file mode 100644 index 0000000000..a99f96ece0 --- /dev/null +++ b/src/gateway/models/providers.py @@ -0,0 +1,127 @@ +"""ORM tables for providers: provider instances configured at runtime, and model aliases.""" + +import uuid +from datetime import UTC, datetime +from typing import Any + +from sqlalchemy import JSON, DateTime, ForeignKey, Index, UniqueConstraint, Uuid, text +from sqlalchemy.orm import Mapped, mapped_column + +from gateway.models.base import Base +from gateway.models.secret_fields import redact_secret_like_values + + +class ModelAlias(Base): + """A display name that resolves to a real model selector. + + The runtime counterpart of the ``aliases:`` block in config.yml: same + meaning, but writable through the API. Pricing, budgets, and usage all key + on the resolved target, so nothing here is billed against ``name``. + + There are two scopes, and they are independent. ``workspace_id`` says which + tenant owns the alias: it resolves only for requests in that workspace, so + two workspaces can each point ``fast`` somewhere different. Within a + workspace, ``user_id`` narrows it further: ``NULL`` means every caller in + that workspace sees it, which is what every row predating the column is, and + a non-null ``user_id`` scopes it to that user, shadowing the workspace-wide + row of the same name for them alone. + + Uniqueness needs two constraints rather than one because SQLite and + PostgreSQL both treat NULLs as distinct in a unique index: the composite + constraint keeps one row per (workspace, name, user), and the partial index + keeps one workspace-wide row per (workspace, name), which the composite one + cannot, its ``user_id`` being NULL. The surrogate ``id`` exists only because + the natural key contains a nullable column, which a primary key cannot. + """ + + __tablename__ = "model_aliases" + __table_args__ = ( + # Workspace-scoped, so two workspaces can each hold a "fast" entry + # pointing somewhere different. Safe only because resolution is keyed by + # workspace too (``services/alias_service``); while that cache was keyed + # on name alone the second workspace's row silently shadowed the first at + # request time, which is why this constraint waited for it. + UniqueConstraint("workspace_id", "name", "user_id", name="uq_model_aliases_workspace_name_user"), + Index( + "uq_model_aliases_workspace_global_name", + "workspace_id", + "name", + unique=True, + sqlite_where=text("user_id IS NULL"), + postgresql_where=text("user_id IS NULL"), + ), + ) + + id: Mapped[str] = mapped_column(primary_key=True, default=lambda: str(uuid.uuid4())) + # No index of its own: both constraints above lead with `workspace_id` and + # carry `name` second, and a listing is always workspace-scoped. A third copy + # would be paid for on every write to serve reads that mostly do not happen, + # since resolution goes through the process-wide alias cache. + name: Mapped[str] = mapped_column() + target: Mapped[str] = mapped_column() + user_id: Mapped[str | None] = mapped_column(ForeignKey("users.user_id", ondelete="CASCADE"), index=True) + # The workspace this row belongs to; see `APIKey.workspace_id` for why. + workspace_id: Mapped[uuid.UUID] = mapped_column( + Uuid, ForeignKey("workspace.id", ondelete="RESTRICT"), nullable=False, index=True + ) + created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=lambda: datetime.now(UTC)) + updated_at: Mapped[datetime] = mapped_column( + DateTime(timezone=True), + default=lambda: datetime.now(UTC), + onupdate=lambda: datetime.now(UTC), + ) + + def to_dict(self) -> dict[str, Any]: + return { + "name": self.name, + "target": self.target, + "user_id": self.user_id, + "created_at": self.created_at.isoformat() if self.created_at else None, + "updated_at": self.updated_at.isoformat() if self.updated_at else None, + } + + +class ProviderCredential(Base): + """A provider instance configured at runtime through the dashboard. + + The database counterpart of a ``providers:`` entry in config.yml: it is + merged over the config-file providers at runtime (see + ``provider_store_service``), with the stored row winning on an instance-name + collision. The API key is held encrypted (``secret_box``); ``last4`` is kept + in clear only so the UI can show which key is set without ever decrypting. + Standalone mode only, never used in the hybrid platform path. + """ + + __tablename__ = "provider_credentials" + + instance: Mapped[str] = mapped_column(primary_key=True) + provider_type: Mapped[str | None] = mapped_column() + api_base: Mapped[str | None] = mapped_column() + encrypted_api_key: Mapped[str | None] = mapped_column() + last4: Mapped[str | None] = mapped_column() + client_args: Mapped[dict[str, Any]] = mapped_column("client_args", JSON, default=dict) + created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=lambda: datetime.now(UTC)) + updated_at: Mapped[datetime] = mapped_column( + DateTime(timezone=True), + default=lambda: datetime.now(UTC), + onupdate=lambda: datetime.now(UTC), + ) + + def to_public_dict(self) -> dict[str, Any]: + """Serialize for the API. Never includes the secret, only ``last4``. + + ``client_args`` is masked by key name the same way + ``OrgProviderKey.to_public`` masks its own: a standalone Bedrock instance + keeps its ``aws_secret_access_key`` there, so the field this table holds + in clear is as much a credential as ``encrypted_api_key`` is, and it must + not round-trip over the API either. + """ + return { + "instance": self.instance, + "provider_type": self.provider_type, + "api_base": self.api_base, + "last4": self.last4, + "client_args": redact_secret_like_values(self.client_args) or {}, + "created_at": self.created_at.isoformat() if self.created_at else None, + "updated_at": self.updated_at.isoformat() if self.updated_at else None, + } diff --git a/src/gateway/models/routing.py b/src/gateway/models/routing.py index 7010e6531c..e4373ba814 100644 --- a/src/gateway/models/routing.py +++ b/src/gateway/models/routing.py @@ -22,21 +22,32 @@ ``extra="ignore"`` (a pydantic-settings default that also swallows stray env vars), so a typo'd key inside a policy would otherwise vanish and the policy would quietly not do what it says. + +Also holds the routing tables. """ from __future__ import annotations import math +import uuid +from datetime import UTC, datetime from typing import Any, Literal from pydantic import BaseModel, ConfigDict, Field, model_validator +from sqlalchemy import JSON, DateTime, ForeignKey, Index, UniqueConstraint, Uuid, text +from sqlalchemy.orm import Mapped, mapped_column + +from gateway.models.base import Base __all__ = [ "MAX_CANDIDATES", "WEIGHTED_BACKEND", "PolicyGuardrail", "PolicySpec", + "RouterPreference", "RoutingConfig", + "RoutingMemory", + "RoutingPolicy", "SelectEntry", "Threshold", "WhenClause", @@ -491,3 +502,195 @@ class RoutingConfig(BaseModel): ), ) policies: dict[str, PolicySpec] = Field(default_factory=dict) + + +class RoutingPolicy(Base): + """A named routing policy, writable through the API. + + The runtime counterpart of the ``routing.policies`` block in config.yml. The + spec is stored as JSON rather than as columns because it is a nested, + versioned document (``select`` entries with conditions, ``on_failure``, + guardrails); flattening it into columns would mean a migration per + schema addition and would still need JSON for the conditions. It is validated + against :class:`gateway.models.routing.PolicySpec` on write and again on load, + so a row that predates a schema change surfaces as a startup warning rather + than as a request-time crash. + + Scoping mirrors :class:`gateway.models.providers.ModelAlias` exactly, + workspace included, and so does the two-constraint uniqueness (SQLite and + PostgreSQL both treat NULLs as distinct in a unique index, so the composite + constraint cannot keep one *workspace-wide* row per name). A policy and an alias are the same concept at + different complexities, so it would be strange for their scoping rules to + differ. + """ + + __tablename__ = "routing_policies" + __table_args__ = ( + # Workspace-scoped for the same reason, and on the same precondition, as + # :class:`gateway.models.providers.ModelAlias`: ``services/policy_store`` + # keys its cache by workspace, so two workspaces holding a "fast" policy + # each resolve their own rather than one shadowing the other. + UniqueConstraint("workspace_id", "name", "user_id", name="uq_routing_policies_workspace_name_user"), + Index( + "uq_routing_policies_workspace_global_name", + "workspace_id", + "name", + unique=True, + sqlite_where=text("user_id IS NULL"), + postgresql_where=text("user_id IS NULL"), + ), + ) + + id: Mapped[str] = mapped_column(primary_key=True, default=lambda: str(uuid.uuid4())) + name: Mapped[str] = mapped_column() + spec: Mapped[dict[str, Any]] = mapped_column(JSON) + user_id: Mapped[str | None] = mapped_column(ForeignKey("users.user_id", ondelete="CASCADE"), index=True) + # The workspace this row belongs to; see `APIKey.workspace_id` for why. + workspace_id: Mapped[uuid.UUID] = mapped_column( + Uuid, ForeignKey("workspace.id", ondelete="RESTRICT"), nullable=False, index=True + ) + created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=lambda: datetime.now(UTC)) + updated_at: Mapped[datetime] = mapped_column( + DateTime(timezone=True), + default=lambda: datetime.now(UTC), + onupdate=lambda: datetime.now(UTC), + ) + + def to_dict(self) -> dict[str, Any]: + return { + "name": self.name, + "spec": self.spec, + "user_id": self.user_id, + "created_at": self.created_at.isoformat() if self.created_at else None, + "updated_at": self.updated_at.isoformat() if self.updated_at else None, + } + + +class RoutingMemory(Base): + """One record per scored example: a prompt embedding plus the quality each + candidate model earned on it. + + The kNN router (:mod:`gateway.services.routing.knn`) retrieves the nearest + neighbors of an incoming request's task embedding within one user's records + and votes on the cheapest candidate that is still good enough. One record is + one example (one prompt), so the vote is over distinct prompts; ``qualities`` + maps each model to its ``[0, 1]`` score for this prompt, keyed on canonical + ``instance:model`` so a candidate's spelling never decides whether it matches + (the router canonicalizes what it reads, so older rows keyed on another + spelling still match). Records are written by the preference-collection flow, + never by live traffic (passive learning is a fast-follow). + + Vectors are stored as a JSON list of floats for SQLite/PostgreSQL + portability and scanned linearly in Python. That holds into the low thousands + of records per user (the ``router_max_records_per_user`` cap); larger pools + need an indexed vector store. + ``embedding_model`` tags each row so changing the embedding model invalidates + stale vectors instead of mixing incomparable spaces. + + Scoped by ``user_id``, which is the identity the request is routed and billed + under, so one user's examples never steer another's traffic. CASCADE: the + records are derived training data, worthless once the user is gone. + ``workspace_id`` narrows that further: the router reads one (user, workspace) + partition, so a user who holds keys in two workspaces does not have one + workspace's labels steering the other's traffic. + """ + + __tablename__ = "routing_memory" + __table_args__ = ( + # Every read filters on the workspace as well as the user, so the + # workspace leads: the same three shapes, one partition narrower. + Index("ix_routing_memory_workspace_user_model", "workspace_id", "user_id", "embedding_model"), + Index("ix_routing_memory_workspace_user_created", "workspace_id", "user_id", "created_at"), + # A task-scoped read filters on all four; without this it walks every + # record the user has for the embedding model before partitioning. + Index( + "ix_routing_memory_workspace_user_model_task", + "workspace_id", + "user_id", + "embedding_model", + "task_id", + ), + ) + + id: Mapped[str] = mapped_column(primary_key=True, default=lambda: str(uuid.uuid4())) + user_id: Mapped[str] = mapped_column( + ForeignKey("users.user_id", ondelete="CASCADE"), nullable=False, index=True + ) + # The workspace this row belongs to; see `APIKey.workspace_id` for why. + workspace_id: Mapped[uuid.UUID] = mapped_column( + Uuid, ForeignKey("workspace.id", ondelete="RESTRICT"), nullable=False, index=True + ) + embedding_model: Mapped[str] = mapped_column() + embedding: Mapped[list[float]] = mapped_column(JSON) + qualities: Mapped[dict[str, float]] = mapped_column(JSON) + task_id: Mapped[str | None] = mapped_column(default=None, index=True) + label_source: Mapped[str] = mapped_column(default="human") + created_at: Mapped[datetime] = mapped_column( + DateTime(timezone=True), default=lambda: datetime.now(UTC), index=True + ) + + def to_dict(self) -> dict[str, Any]: + """Convert model to dictionary. + + The embedding itself is deliberately left out: it is thousands of floats + that no management surface renders, and the prompt it came from is on the + :class:`RouterPreference` audit row. + """ + return { + "id": self.id, + "user_id": self.user_id, + "workspace_id": str(self.workspace_id), + "embedding_model": self.embedding_model, + "qualities": self.qualities, + "task_id": self.task_id, + "label_source": self.label_source, + "created_at": self.created_at.isoformat() if self.created_at else None, + } + + +class RouterPreference(Base): + """An audit record of one preference-collection scoring. + + Each ``/v1/routing/preferences/rank`` submission writes one row here for + provenance plus one :class:`RoutingMemory` row. The routing-memory row keeps + only the embedding, so this is where the prompt text and the raw per-model + scores live: enough to recompute the memory if the scoring changes, and to + tell a human label from a judge's. + + ``workspace_id`` matches the :class:`RoutingMemory` row written beside it, so + the audit trail partitions exactly the way the training data does. + """ + + __tablename__ = "router_preferences" + __table_args__ = ( + Index("ix_router_preferences_workspace_user_created", "workspace_id", "user_id", "created_at"), + ) + + id: Mapped[str] = mapped_column(primary_key=True, default=lambda: str(uuid.uuid4())) + user_id: Mapped[str] = mapped_column( + ForeignKey("users.user_id", ondelete="CASCADE"), nullable=False, index=True + ) + # The workspace this row belongs to; see `APIKey.workspace_id` for why. + workspace_id: Mapped[uuid.UUID] = mapped_column( + Uuid, ForeignKey("workspace.id", ondelete="RESTRICT"), nullable=False, index=True + ) + prompt: Mapped[str] = mapped_column() + task_id: Mapped[str | None] = mapped_column(default=None) + scores: Mapped[dict[str, float]] = mapped_column(JSON) + label_source: Mapped[str] = mapped_column(default="human") + created_at: Mapped[datetime] = mapped_column( + DateTime(timezone=True), default=lambda: datetime.now(UTC), index=True + ) + + def to_dict(self) -> dict[str, Any]: + """Convert model to dictionary.""" + return { + "id": self.id, + "user_id": self.user_id, + "workspace_id": str(self.workspace_id), + "prompt": self.prompt, + "task_id": self.task_id, + "scores": self.scores, + "label_source": self.label_source, + "created_at": self.created_at.isoformat() if self.created_at else None, + } diff --git a/src/gateway/models/tenancy.py b/src/gateway/models/tenancy.py index bd7bf3402a..974fe05c56 100644 --- a/src/gateway/models/tenancy.py +++ b/src/gateway/models/tenancy.py @@ -5,17 +5,10 @@ graph: an ``organization`` owns ``workspace`` rows, and ``user`` rows join both through ``organization_member`` and ``workspace_member``. -**Why SQLModel here and plain SQLAlchemy in `entities.py`.** SQLModel is -SQLAlchemy underneath, so these classes bind to the gateway's ``AsyncSession`` -unchanged while keeping the ``Create``/``Update``/``Public`` schema layer the -routes and the generated dashboard client are built on. Converting them to -`entities.py`'s declarative style would have rewritten every endpoint contract -in the slice for no behavioral gain. `entities.py` stays as it is: the two -styles coexist deliberately, and new *gateway* tables still belong there. - -**One MetaData, two styles.** `entities.py`'s ``Base`` shares -``SQLModel.metadata`` (see the note there), so Alembic, ``create_all`` and -``drop_all`` see one schema no matter which style declared a table. +**Two styles.** Tables whose ``Create``/``Update``/``Public`` schemas are endpoint +contracts use SQLModel. ``DashboardSession`` and ``WorkspaceActivationState`` +have no such contract, so they use the declarative ``Base``. Both share +``SQLModel.metadata``. Three deliberate departures from the platform's models, applied on arrival: @@ -25,7 +18,7 @@ ``datetime.now(UTC)`` into it: the offset is silently dropped on the way in, and the value reads back as local-looking UTC. That is a latent bug, not a style difference, so it is fixed here rather than carried. ``timezone=True`` - alone does not fix it, which is why ``UtcDateTime`` below exists: PostgreSQL + alone does not fix it, which is why ``UtcDateTime`` exists: PostgreSQL honors the flag and SQLite ignores it, and SQLite is what the OSS edition ships by default, so on that engine the departure would have been a comment rather than a behavior. ``tests/unit/test_tenancy_timestamps.py`` is what @@ -60,7 +53,7 @@ import uuid from datetime import UTC, datetime, timedelta -from typing import Any, Literal +from typing import Literal from pydantic import field_validator from sqlalchemy import ( @@ -75,10 +68,11 @@ func, text, ) -from sqlalchemy.engine.interfaces import Dialect -from sqlalchemy.types import TypeDecorator +from sqlalchemy.orm import Mapped, mapped_column from sqlmodel import Field, SQLModel +from gateway.models.base import Base, CreatedAtMixin, PrimaryKeyMixin, UpdatedAtMixin, UtcDateTime, _timestamp_field + ORGANIZATION_MEMBER_ROLES = {"owner", "admin", "member", "viewer"} ORGANIZATION_MEMBER_STATUSES = {"active", "invited", "suspended"} WORKSPACE_MEMBER_ROLES = {"owner", "admin", "member", "viewer"} @@ -121,102 +115,6 @@ def _validate_membership(value: str, *, allowed: set[str], kind: str) -> str: return value -class UtcDateTime(TypeDecorator[datetime]): - """A timestamp that reads back UTC-aware on every engine. - - ``DateTime(timezone=True)`` alone is not enough, and the gap is the whole - reason this exists. PostgreSQL honors it and hands back an aware value; - SQLite has no timestamp type at all, so SQLAlchemy stores an ISO string and - the flag is a no-op, and a value written as ``datetime.now(UTC)`` reads back - with ``tzinfo=None``. A naive datetime then serializes with no offset, and a - browser parses an offset-less timestamp as **local** time, so every tenancy - timestamp in the dashboard would be wrong by the deployment's UTC offset on - the engine the OSS edition ships by default. - - Both directions are handled: an aware value is normalized to UTC before it - is stored, so a caller in another zone cannot write a wall-clock time that - means something else, and a naive value read back is stamped UTC, because - UTC is what everything here writes. - - The rendered DDL is exactly ``impl``'s, so this changes no migration and - ``compare_metadata`` stays clean. - """ - - impl = DateTime(timezone=True) - cache_ok = True - - def process_bind_param(self, value: datetime | None, dialect: Dialect) -> datetime | None: - if value is None: - return None - if value.utcoffset() is None: - # Refused rather than assumed. Reading a naive value back as UTC is - # safe, because UTC is what everything here writes; writing one is - # not, because the engines disagree about what it means. PostgreSQL - # interprets it in the *session* time zone, so the same value lands - # as a different instant depending on who connected, while SQLite - # stores the wall clock as written. Silently picking one is how a - # timestamp ends up hours off with nothing to show for it. - msg = "A tenancy timestamp must be timezone-aware; got a naive datetime" - raise ValueError(msg) - return value.astimezone(UTC) - - def process_result_value(self, value: datetime | None, dialect: Dialect) -> datetime | None: - if value is not None and value.tzinfo is None: - return value.replace(tzinfo=UTC) - return value - - -def _timestamp_field(*, default: Any = None, default_factory: Any = None, column_kwargs: dict[str, Any]) -> Any: - """Build a timezone-aware timestamp field. - - Two things are worked around here, once, instead of at five inheriting - tables. SQLModel's ``Field`` overloads type ``sa_type`` as a *class*, while - the type we want is an *instance* (the runtime accepts either and hands it - straight to ``Column``). And the type has to arrive as ``sa_type`` rather - than a ready-made ``sa_column``, because a ``Column`` instance declared on a - mixin cannot be attached to more than one table; ``sa_type`` plus kwargs - lets SQLModel build a fresh column per model. - """ - if default_factory is not None: - return Field( # type: ignore[call-overload] - default_factory=default_factory, - sa_type=UtcDateTime(), - sa_column_kwargs=column_kwargs, - ) - return Field( # type: ignore[call-overload] - default=default, - sa_type=UtcDateTime(), - sa_column_kwargs=column_kwargs, - ) - - -class PrimaryKeyMixin: - """A UUID primary key, rendered as CHAR(32) on SQLite and native on PostgreSQL.""" - - id: uuid.UUID = Field(default_factory=uuid.uuid4, primary_key=True) - - -class CreatedAtMixin: - """Creation timestamp, defaulted in Python and in the database.""" - - created_at: datetime = _timestamp_field( - default_factory=lambda: datetime.now(UTC), - column_kwargs={"server_default": func.now()}, - ) - - -class UpdatedAtMixin: - """Last-modification timestamp, stamped by the database on update. - - ``default=None`` and not merely a nullable annotation: without an explicit - default the field is *required* on the pydantic side, which a table class - hides (table models skip construction validation) and any schema inheriting - this mixin would not. - """ - - updated_at: datetime | None = _timestamp_field(default=None, column_kwargs={"onupdate": func.now()}) - - # ============================================================================= # Identity # ============================================================================= @@ -252,7 +150,7 @@ class UserCreate(UserBase): class User(UserBase, PrimaryKeyMixin, CreatedAtMixin, UpdatedAtMixin, table=True): """An identity in the reconciled control plane. - Not to be confused with `entities.User`, the gateway's own string-keyed + Not to be confused with `users.User`, the gateway's own string-keyed per-request spend identity, which is what keys, budgets, and usage attach to. Both exist, and how they converge is no longer settled: otari-ai#1719 made otari's schema the survivor, which retired the pre-flip plan of re-parenting @@ -1509,6 +1407,94 @@ class OAuthPendingState(SQLModel, table=True): expires_at: datetime = Field(sa_type=UtcDateTime(), index=True) # type: ignore[call-overload] +class DashboardSession(Base): + """A server-side admin-dashboard sign-in session, held by one identity. + + Minted when an operator signs in to the dashboard with the master key: the + browser holds only an opaque token in an HttpOnly cookie and this table + stores the token's SHA-256 hash, so neither the master key nor a usable + session credential is ever persisted in JS-readable storage. Sessions + expire on a TTL and are revoked on sign-out and on master-key rotation. + + ``user_id`` is what lets a session resolve a caller rather than only prove + that the master key was presented once. It names a tenancy identity + (`models.tenancy.User`), whose ``active_organization_id`` is the + organization the session acts in, so a tenancy surface reads its scope off + the session. Master-key sign-in binds the session to the deployment's + bootstrap operator; a per-user sign-in flow binds it to whoever + authenticated. + + NOT NULL on purpose: a session that names nobody cannot answer "who is + calling", which is the whole point of the column, and the migration that + added it bound existing sessions to that same bootstrap operator. CASCADE + on the foreign key, so deleting an identity revokes its sessions rather + than leaving a live cookie pointing at a row that is gone. + """ + + __tablename__ = "dashboard_sessions" + + token_hash: Mapped[str] = mapped_column(primary_key=True) + user_id: Mapped[uuid.UUID] = mapped_column( + Uuid, ForeignKey("user.id", ondelete="CASCADE"), nullable=False, index=True + ) + created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=lambda: datetime.now(UTC)) + expires_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), index=True) + + +class WorkspaceActivationState(Base): + """What the dashboard's first-request setup guide remembers about a workspace. + + The guide walks a workspace from "no traffic" to its first successful + request (`services/tenancy/workspace_activation_service.py`). Only what + cannot be observed elsewhere is stored here: whether someone dismissed it, + when it last handed out a key, and which key that was. Whether the workspace + has *activated* is deliberately not a column, because ``usage_logs`` already + records it: the first successful gateway request in the workspace is the + evidence, so there is no second copy of it to backfill or to disagree with + the Activity page. + + Ported from the platform's ``workspace_activation_state`` / + ``workspace_activation_experience_state`` pair + (`otari-ai` `backend/app/models/workspace_activation.py`), which does carry + the attempt telemetry as columns, because its usage pipeline is asynchronous + and crosses services. Here the usage row is written by this process into this + database, so the derivation is exact. + + One row per workspace, not per workspace and viewer: the guide is about a + workspace's first request, so dismissing it says "this workspace is set up, + stop offering the guide" for everyone who can manage it. + """ + + __tablename__ = "workspace_activation_state" + + workspace_id: Mapped[uuid.UUID] = mapped_column( + Uuid, ForeignKey("workspace.id", ondelete="CASCADE"), primary_key=True + ) + # When the guide first and last minted an API key for this workspace. The + # first is what an operator reads as "when was this offered"; the last is + # what makes a rotation visible next to the key it rotated. + first_presented_at: Mapped[datetime | None] = mapped_column(UtcDateTime(), default=None) + last_presented_at: Mapped[datetime | None] = mapped_column(UtcDateTime(), default=None) + # Set by Skip, and permanent: the guide is a first-run offer, so a workspace + # that turned it down is not asked again on the next page load. + dismissed_at: Mapped[datetime | None] = mapped_column(UtcDateTime(), default=None) + # The key the guide issued, rotated in place on each presentation so a + # workspace collects one "Setup guide" key rather than one per page load. + # ``SET NULL`` because deleting that key from the Keys page is a legitimate + # thing to do, and it must not take this row (or the dismissal on it) with it. + api_key_id: Mapped[str | None] = mapped_column( + ForeignKey("api_keys.id", ondelete="SET NULL"), default=None, index=True + ) + # Gotcha: a plain DateTime(timezone=True) reads back naive on SQLite. The dashboard + # then shows it as local time. + created_at: Mapped[datetime] = mapped_column(UtcDateTime(), default=lambda: datetime.now(UTC)) + updated_at: Mapped[datetime] = mapped_column( + UtcDateTime(), + default=lambda: datetime.now(UTC), + onupdate=lambda: datetime.now(UTC), + ) + + __all__ = [ "DeploymentAdminAccessPublic", "DeploymentUserOrganizationPublic", @@ -1542,6 +1528,7 @@ class OAuthPendingState(SQLModel, table=True): "CallerOrganizationMembershipPublic", "CallerOrganizationMembershipsPublic", "CallerWorkspaceMembershipPublic", + "DashboardSession", "Invitation", "InvitationCreate", "InvitationPreviewPublic", @@ -1586,6 +1573,7 @@ class OAuthPendingState(SQLModel, table=True): "WebAuthnCredentialsPublic", "Workspace", "WorkspaceActivationClassification", + "WorkspaceActivationState", "WorkspaceAssignmentRequest", "WorkspaceCreate", "WorkspaceMember", diff --git a/src/gateway/models/tools.py b/src/gateway/models/tools.py new file mode 100644 index 0000000000..48bd2390b6 --- /dev/null +++ b/src/gateway/models/tools.py @@ -0,0 +1,342 @@ +"""ORM tables for gateway-run tools: search credentials, uploaded files, and workspace tool policies.""" + +import uuid +from datetime import UTC, datetime +from typing import Any + +from sqlalchemy import JSON, CheckConstraint, DateTime, ForeignKey, String, Text, UniqueConstraint, Uuid, func, true +from sqlalchemy.orm import Mapped, mapped_column + +from gateway.models.base import Base, UtcDateTime +from gateway.models.secret_fields import redact_secret_like_values + + +def _epoch_seconds(value: datetime | None) -> int | None: + """Return a UTC epoch from a stored datetime. + + SQLite hands datetimes back naive; ``datetime.timestamp()`` would then read + them as local time and skew the epoch by the server's UTC offset. Treat a + naive value as the UTC it was stored as before converting. + """ + if value is None: + return None + if value.tzinfo is None: + value = value.replace(tzinfo=UTC) + return int(value.timestamp()) + + +class SearchToolCredential(Base): + """A ``POST /v1/search`` tool configured at runtime through the dashboard. + + The database counterpart of a ``search_tools:`` entry in config.yml: it is + merged over the config-file tools at runtime (see + ``search_tool_store_service``), with the stored row winning on a name + collision, exactly as ``ProviderCredential`` does for providers. The API key + is held encrypted (``secret_box``) and is optional, because a ``searxng`` + backend is normally keyless; ``last4`` is kept in clear only so the UI can + show which key is set without ever decrypting. Standalone mode only. + """ + + __tablename__ = "search_tool_credentials" + + name: Mapped[str] = mapped_column(primary_key=True) + provider: Mapped[str] = mapped_column() + api_base: Mapped[str | None] = mapped_column() + encrypted_api_key: Mapped[str | None] = mapped_column() + last4: Mapped[str | None] = mapped_column() + # Named for its unit; the config-file key it stands in for is plain ``timeout``, + # and ``to_public_dict`` / the overlay entry both use that name. + timeout_seconds: Mapped[float | None] = mapped_column() + options: Mapped[dict[str, Any]] = mapped_column("options", JSON, default=dict) + created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=lambda: datetime.now(UTC)) + updated_at: Mapped[datetime] = mapped_column( + DateTime(timezone=True), + default=lambda: datetime.now(UTC), + onupdate=lambda: datetime.now(UTC), + ) + + def to_public_dict(self) -> dict[str, Any]: + """Serialize for the API. Never includes the secret, only ``last4``. + + ``options`` is free-form, so it is masked by key name in case it holds a + credential. + """ + return { + "name": self.name, + "provider": self.provider, + "api_base": self.api_base, + "last4": self.last4, + "timeout": self.timeout_seconds, + "options": redact_secret_like_values(self.options) or {}, + "created_at": self.created_at.isoformat() if self.created_at else None, + "updated_at": self.updated_at.isoformat() if self.updated_at else None, + } + + +class FileObject(Base): + """Uploaded file metadata for the OpenAI-compatible /v1/files API. + + The raw bytes live in a pluggable blob backend (see + gateway.services.file_store); this row holds metadata plus the backend + ``storage_ref`` used to fetch them. Files are scoped to ``user_id`` for + tenant isolation and soft-deleted via ``deleted_at``. ``workspace_id`` is a + second, independent axis: it says which workspace the upload was made in, so + a key confined to one workspace never reaches another's files even when the + same user holds keys in both. + """ + + __tablename__ = "file_objects" + + id: Mapped[str] = mapped_column(primary_key=True, default=lambda: f"file-{uuid.uuid4().hex}") + # Always set to the authenticated user; non-null enforces the user-scoping + # contract at the schema level. CASCADE removes a user's files on delete. + user_id: Mapped[str] = mapped_column(ForeignKey("users.user_id", ondelete="CASCADE"), index=True) + # The workspace this row belongs to; see `APIKey.workspace_id` for why it is + # NOT NULL and RESTRICT rather than nullable and cascading. Existing rows were + # backfilled onto the deployment's default workspace, which is also where a + # master-key upload lands. + workspace_id: Mapped[uuid.UUID] = mapped_column( + Uuid, ForeignKey("workspace.id", ondelete="RESTRICT"), nullable=False, index=True + ) + filename: Mapped[str] = mapped_column() + mime_type: Mapped[str] = mapped_column() + bytes: Mapped[int] = mapped_column() + purpose: Mapped[str] = mapped_column(default="user_data") + storage_ref: Mapped[str] = mapped_column() + created_at: Mapped[datetime] = mapped_column( + DateTime(timezone=True), default=lambda: datetime.now(UTC), index=True + ) + expires_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True)) + deleted_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True), default=None, index=True) + + metadata_: Mapped[dict[str, Any]] = mapped_column("metadata", JSON, default=dict) + + def to_dict(self) -> dict[str, Any]: + """Convert to the OpenAI file object shape.""" + return { + "id": self.id, + "object": "file", + "bytes": self.bytes, + "created_at": _epoch_seconds(self.created_at), + "expires_at": _epoch_seconds(self.expires_at), + "filename": self.filename, + "purpose": self.purpose, + } + + +class WorkspaceMcpServer(Base): + """One MCP server a workspace has configured, referenced by id from a request. + + Ported from otari-ai's ``mcp_server`` table (otari#658). A request names + stored servers with ``mcp_server_ids``; hybrid mode resolves those ids + through the platform and standalone mode resolves them here, against the + workspace the request's key belongs to. There is no deployment-wide MCP + server list for these rows to narrow, which is why MCP is the stated + exception to the "a workspace row never grants" rule in + ``src/gateway/AGENTS.md``. + + ``encrypted_token`` holds the server's bearer token, Fernet-encrypted with + ``OTARI_SECRET_KEY`` (``services/secret_box.py``), the same treatment + ``ProviderCredential.encrypted_api_key`` gets. Nothing serializes it: the + public shape carries ``has_token`` and no prefix or suffix of the value, + because unlike a provider key's ``last4`` there is no operator workflow + here that needs to tell two tokens apart at a glance. + + ``enabled`` is a workspace-level off switch that keeps the row and its + token: a disabled server is skipped at resolve rather than refusing the + request, so a caller whose stored id list outlives one server's + decommissioning still gets the rest. + + CASCADE: a workspace-owned configuration row means nothing without its + workspace. + """ + + __tablename__ = "workspace_mcp_servers" + __table_args__ = ( + # Duplicate names within one workspace are rejected at the database, not + # only in the service layer, so two concurrent creates cannot both land + # (otari#658's third Definition-of-Done item). The name is what an + # operator recognizes a server by and what the tool loop labels its + # tools with, so collapsing two onto one name would silently hide a + # server. + UniqueConstraint("workspace_id", "name", name="uq_workspace_mcp_servers_workspace_name"), + ) + + id: Mapped[uuid.UUID] = mapped_column(Uuid, primary_key=True, default=uuid.uuid4) + workspace_id: Mapped[uuid.UUID] = mapped_column( + Uuid, ForeignKey("workspace.id", ondelete="CASCADE"), nullable=False, index=True + ) + name: Mapped[str] = mapped_column(nullable=False) + url: Mapped[str] = mapped_column(nullable=False) + encrypted_token: Mapped[str | None] = mapped_column(Text, default=None) + purpose_hint: Mapped[str | None] = mapped_column(Text, default=None) + allowed_tools: Mapped[list[str] | None] = mapped_column(JSON, default=None) + enabled: Mapped[bool] = mapped_column(default=True, nullable=False) + # ``UtcDateTime`` for the same reason ``WorkspaceBudgetDefault``'s are: these + # go over the wire and a naive SQLite round-trip would drop the offset. + created_at: Mapped[datetime] = mapped_column(UtcDateTime(), default=lambda: datetime.now(UTC)) + updated_at: Mapped[datetime] = mapped_column( + UtcDateTime(), + default=lambda: datetime.now(UTC), + onupdate=lambda: datetime.now(UTC), + ) + + +class WorkspaceCodeExecutionPolicy(Base): + """A workspace's policy over the deployment-wide code-execution sandbox. + + The sandbox itself stays deployment-wide (``sandbox_url`` and its + credential are operator concerns and never move here, see + ``src/gateway/AGENTS.md``); this row says who on that deployment may ask + for it and within which limits. Resolved at admission by + ``prepare_gateway_tools`` and applied to the tool loop, the standalone + counterpart of the hybrid path's ``/gateway/code-execution/resolve``. + + A row may only *narrow*: ``enabled=False`` refuses the tool for this + workspace, and the two limits are floored against the values a request + would otherwise get. No row means no narrowing, which is what keeps a + deployment that configures nothing behaving as it did (#655/#678). + + ``workspace_id`` is the primary key: a workspace has one policy or none, + so there is nothing else to identify a row by. It is a real foreign key + with ``CASCADE``, like ``workspace_budget_defaults``: nothing else names + the row, so it rides the workspace's own delete. + + ``image`` and ``tools`` reach the same two decisions the hosted + ``CodeExecutionConfig`` carries (#740). Neither breaks the rule above: + ``image`` may only name something the deployment's operator has already + curated into ``sandbox_allowed_session_images``, so a workspace picks from an + operator's shelf rather than pointing the gateway at an image of its own, + and ``tools`` may only remove tool kinds from what the sandbox backend + already serves. + """ + + __tablename__ = "workspace_code_execution_policies" + __table_args__ = ( + # Both limits are ceilings that get floored into an effective value, so + # zero or negative is a storage error rather than a stricter policy: it + # would floor the loop to nothing runnable while reading as configured. + # The request schemas refuse it first; these are the backstop for a + # writer that is not the service. + CheckConstraint( + "max_iterations IS NULL OR max_iterations > 0", + name="ck_workspace_code_execution_policies_max_iterations_positive", + ), + CheckConstraint( + "exec_timeout_s IS NULL OR exec_timeout_s > 0", + name="ck_workspace_code_execution_policies_exec_timeout_positive", + ), + ) + + workspace_id: Mapped[uuid.UUID] = mapped_column( + Uuid, ForeignKey("workspace.id", ondelete="CASCADE"), primary_key=True + ) + enabled: Mapped[bool] = mapped_column(default=True, nullable=False) + # NULL means "no workspace default": the request's own hint, then the + # deployment's, then the backend's built-in, exactly as today. + default_purpose_hint: Mapped[str | None] = mapped_column(Text, default=None) + # Both NULL-able ceilings, applied with ``min`` against what the request + # would otherwise get, so a value above the deployment ceiling narrows + # nothing rather than raising it. + max_iterations: Mapped[int | None] = mapped_column(default=None) + exec_timeout_s: Mapped[int | None] = mapped_column(default=None) + # NULL means "no workspace image": whatever the deployment names in + # ``sandbox_session_image``, and failing that whatever the sandbox backend runs by + # default, which is what every request got before this column existed. + # ``String(255)`` rather than ``Text`` to match the hosted column's own + # bound; an image reference that long is already pathological. + image: Mapped[str | None] = mapped_column(String(255), default=None) + # NULL means "no workspace tool allow-list": the backend offers what it + # offers. A stored list is an intersection, never a union, so it can only + # take tool kinds away. JSON rather than a child table for the same reason + # ``WorkspaceWebSearchConfig`` stores its domain lists that way: short, read + # whole, and nothing queries into it. + tools: Mapped[list[str] | None] = mapped_column(JSON, default=None) + # ``UtcDateTime`` for the same reason ``WorkspaceBudgetDefault`` uses it: + # these are serialized with ``.isoformat()`` for the dashboard, and a plain + # ``DateTime(timezone=True)`` round-trips naive on SQLite. + created_at: Mapped[datetime] = mapped_column(UtcDateTime(), default=lambda: datetime.now(UTC)) + updated_at: Mapped[datetime] = mapped_column( + UtcDateTime(), + default=lambda: datetime.now(UTC), + onupdate=lambda: datetime.now(UTC), + ) + + +class WorkspaceWebSearchConfig(Base): + """A workspace's configuration over the deployment-wide web-search backend. + + The backend itself stays deployment-wide (``web_search_url`` and the + credential the adapter in front of it holds are operator concerns and never + move here, see ``src/gateway/AGENTS.md``); this row says which workspaces + may reach it and how their searches are constrained. Resolved at admission + by ``prepare_gateway_tools``, the standalone counterpart of the hybrid + path's ``/gateway/web-search/resolve``. + + A row may only *narrow*: ``enabled=False`` refuses ``otari_web_search`` for + the workspace, ``max_results`` is floored against what the request asked + for, ``blocked_domains`` is added to the request's own block-list, and + ``allowed_domains`` intersects the request's. No row means no narrowing, + which is what keeps a deployment that configures nothing behaving as it did + (#655/#678). + + ``workspace_id`` is the primary key, and a real foreign key with + ``CASCADE``, for the same reasons as :class:`WorkspaceCodeExecutionPolicy` + next door: one row per workspace, and nothing else names it. + + There is deliberately no ``provider`` column, which the hosted config + carries: on this deployment the operator picks the backend by pointing + ``web_search_url`` somewhere, so a provider named here would either be inert + or would ask the gateway to reach an endpoint the operator did not choose, + which is the one thing the narrowing rule forbids. + """ + + __tablename__ = "workspace_web_search_configs" + __table_args__ = ( + # ``max_results`` is floored into an effective value, so zero or less is + # a storage error rather than a stricter policy: it would ask for a + # search that can return nothing while reading as configured. The + # request schema refuses it first; this is the backstop for a writer + # that is not the service. + CheckConstraint( + "max_results IS NULL OR max_results > 0", + name="ck_workspace_web_search_configs_max_results_positive", + ), + ) + + workspace_id: Mapped[uuid.UUID] = mapped_column( + Uuid, ForeignKey("workspace.id", ondelete="CASCADE"), primary_key=True + ) + # ``server_default`` mirrors the migration so autogenerate sees no drift, and + # so a row written by anything other than this mapping still gets a value. + enabled: Mapped[bool] = mapped_column(default=True, nullable=False, server_default=true()) + # NULL means "no workspace ceiling": the request's own value, then the + # deployment's, then the backend's built-in, exactly as today. + max_results: Mapped[int | None] = mapped_column(default=None) + # NULL means "no workspace default": the request's own hint, then the + # deployment's, then the backend's built-in. + purpose_hint: Mapped[str | None] = mapped_column(Text, default=None) + # Two domain lists and an opaque provider bag, stored as JSON for the same + # reason the hosted table does: they are short, they are read whole, and + # nothing queries into them. ``JSON`` rather than ``JSONB`` to match every + # other JSON column here, which has to work on SQLite too. + allowed_domains: Mapped[list[str] | None] = mapped_column(JSON, default=None) + blocked_domains: Mapped[list[str] | None] = mapped_column(JSON, default=None) + # Provider-specific knobs (Tavily's ``search_depth``, say). Opaque here and + # forwarded to the backend, which is what lets a new provider need no + # migration; the adapter in front of it whitelists what it understands. + provider_options: Mapped[dict[str, Any] | None] = mapped_column(JSON, default=None) + # ``UtcDateTime`` for the same reason ``WorkspaceCodeExecutionPolicy`` uses + # it: these are serialized with ``.isoformat()`` for the dashboard, and a + # plain ``DateTime(timezone=True)`` round-trips naive on SQLite. The Python + # default is what every write here uses; ``server_default`` is the backstop + # for a writer that is not this mapping, matching ``workspace`` itself. + created_at: Mapped[datetime] = mapped_column( + UtcDateTime(), default=lambda: datetime.now(UTC), server_default=func.now() + ) + updated_at: Mapped[datetime] = mapped_column( + UtcDateTime(), + default=lambda: datetime.now(UTC), + onupdate=lambda: datetime.now(UTC), + server_default=func.now(), + ) diff --git a/src/gateway/models/usage.py b/src/gateway/models/usage.py new file mode 100644 index 0000000000..62197394db --- /dev/null +++ b/src/gateway/models/usage.py @@ -0,0 +1,262 @@ +"""ORM tables for usage rows and coding-agent telemetry.""" + +import uuid +from datetime import UTC, datetime +from decimal import Decimal +from typing import Any + +from sqlalchemy import JSON, DateTime, ForeignKey, Index, String, UniqueConstraint, Uuid +from sqlalchemy.orm import Mapped, mapped_column, relationship + +from gateway.models.base import Base +from gateway.models.money import UsdCost + + +class UsageLog(Base): + """Usage log model for tracking API requests.""" + + __tablename__ = "usage_logs" + __table_args__ = ( + Index("ix_usage_logs_user_id_timestamp", "user_id", "timestamp"), + # Supports the activity-log viewer's primary "show errors, newest-first" + # query. status is low-cardinality; model is high-cardinality and left + # unindexed on purpose. + Index("ix_usage_logs_status_timestamp", "status", "timestamp"), + # Supports the setup guide's two questions about one workspace: has any + # request in it ever succeeded (oldest first), and what did the last one + # do (newest first). Both filter a workspace, a source and a status and + # then order by time, which the workspace-only and status-first indexes + # above can each answer only halfway: on a deployment with real traffic + # the guide would otherwise scan the workspace's rows on every dashboard + # load, and where usage is imported as well most of those rows are the + # wrong source anyway. Equality columns first, the ordering column last. + Index( + "ix_usage_logs_workspace_source_status_timestamp", + "workspace_id", + "source", + "status", + "timestamp", + ), + # Idempotency for imported usage: re-submitting the same (source, + # source_event_id) must not create a second row. Gateway-originated rows + # keep source_event_id NULL, and SQL treats NULLs as distinct on both + # SQLite and Postgres, so many (gateway, NULL) rows coexist freely. + UniqueConstraint("source", "source_event_id", name="uq_usage_logs_source_event"), + ) + + id: Mapped[str] = mapped_column(primary_key=True, default=lambda: str(uuid.uuid4())) + # The workspace this row belongs to; see `APIKey.workspace_id` for why. + workspace_id: Mapped[uuid.UUID] = mapped_column( + Uuid, ForeignKey("workspace.id", ondelete="RESTRICT"), nullable=False, index=True + ) + api_key_id: Mapped[str | None] = mapped_column(ForeignKey("api_keys.id", ondelete="SET NULL"), index=True) + user_id: Mapped[str | None] = mapped_column(ForeignKey("users.user_id", ondelete="SET NULL"), index=True) + timestamp: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=lambda: datetime.now(UTC), index=True) + + model: Mapped[str] = mapped_column() + provider: Mapped[str | None] = mapped_column() + endpoint: Mapped[str] = mapped_column() + + # Provenance. "gateway" for requests Otari served itself; a source slug (e.g. + # "claude_code") for usage imported through POST /v1/usage/external-events. A row + # backfilled from hosted history keeps its origin's slug behind a legacy prefix + # ("otari-ai:gateway", "otari-ai:claude_code"), so asking whether this deployment + # served a row means asking about the slug behind that prefix: core/usage_source. + # source_event_id is the upstream event id used for idempotent import (NULL for + # gateway rows); source_label carries optional session/project attribution. + source: Mapped[str] = mapped_column(default="gateway", index=True) + source_event_id: Mapped[str | None] = mapped_column() + source_label: Mapped[str | None] = mapped_column() + # Whether this row's cost participates in budget enforcement. True for normal + # gateway rows; false for imported usage and for rows from keys flagged + # exclude_from_budget. False rows are recorded (and appear in cost analytics) + # but their cost is never written to User.spend. + counts_toward_budget: Mapped[bool] = mapped_column(default=True) + + prompt_tokens: Mapped[int | None] = mapped_column() + completion_tokens: Mapped[int | None] = mapped_column() + total_tokens: Mapped[int | None] = mapped_column() + cache_read_tokens: Mapped[int | None] = mapped_column() + cache_write_tokens: Mapped[int | None] = mapped_column() + cache_write_1h_tokens: Mapped[int | None] = mapped_column() + # Which cached-token convention the counts above were reported under: True + # when the cache buckets are already inside ``prompt_tokens`` (OpenAI shape), + # False when they are additive to it (Anthropic / Claude Code shape). Written + # by settlement from ``GatewayUsage.cache_tokens_in_prompt`` and by the + # external-usage ingest from the value the submitter sent, so a row can be + # repriced under the convention it was recorded with rather than one inferred + # from the numbers, which cannot tell the two apart. + # + # Nullable, and deliberately not defaulted: "not recorded" and "inclusive" are + # different answers. Rows written before this column existed are NULL, and + # repricing falls back to recovering the convention from ``billing_meters`` + # for exactly those (see ``usage_admin_service._row_cache_tokens_included``). + # A default would make every historical row claim a convention nothing + # checked, and mis-price the half that were the other one. + cache_tokens_in_prompt: Mapped[bool | None] = mapped_column() + billing_meters: Mapped[dict[str, Any] | None] = mapped_column(JSON) + pricing_breakdown: Mapped[list[dict[str, Any]] | None] = mapped_column(JSON) + # The settled amount, and the accounting truth for this row + # (mozilla-ai/otari-ai#1751). Exact to the micro-dollar; see + # ``models/money.py`` for what that costs on each engine. + cost: Mapped[Decimal | None] = mapped_column(UsdCost()) + + # Why ``cost`` is the amount it is, which the row cannot re-derive on its own: + # ``pricing_source`` names the price list that settled it ("organization", + # "managed", "genai_prices"), ``pricing_reference`` identifies the entry in it + # (a pricing row's id, or a ``provider:model`` key), ``pricing_effective_at`` + # is when that rate took effect, and ``pricing_version`` pins the revision of + # the list. ``calculated_at`` is when the amount was priced, which is not + # ``timestamp`` (when the request ran): usage settled or repriced later moves + # the two apart. + # + # All nullable with no backfill. The gateway's own settlement does not record + # provenance, so these are written by the hosted-usage backfill + # (mozilla-ai/otari-ai#1798) from the platform's ``gateway_usage_settlement`` + # row, and null reads correctly as "not recorded". The lengths mirror that + # table's columns rather than this file's usual unbounded strings, so a value + # copied across always fits. + # + # ``pricing_source`` speaks the platform's settlement vocabulary, the values + # ``_platform.SettledCost.pricing_source`` already carries on the hybrid wire + # (echoed to callers as ``usage.pricing_source``). It is not the same field as + # the one on a listed model in ``api/routes/models.py`` ("configured", + # "default", "dynamic", "none"), which says where a price list entry came from + # in this deployment rather than what settled one row's amount. + pricing_source: Mapped[str | None] = mapped_column(String(32)) + pricing_reference: Mapped[str | None] = mapped_column(String(511)) + pricing_effective_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True)) + pricing_version: Mapped[str | None] = mapped_column(String(255)) + calculated_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True)) + + # "success", "error", or "absorbed". ``absorbed`` is a failed attempt that a + # routing policy recovered from by trying the next candidate: the request + # itself succeeded (or failed on a later attempt), so counting it as an error + # would make a working fallback chain look like an outage. Every error metric + # in the product counts ``status == "error"`` exactly, and ``request_count`` + # excludes absorbed rows, because a request that took two attempts is still one + # request. + status: Mapped[str] = mapped_column() + error_message: Mapped[str | None] = mapped_column() + + # Routing attribution. All nullable: a request that named a plain model was not + # routed through a policy, and null reads correctly as exactly that. + # + # `policy_name` is the name the caller sent. `selection_reason` says why this + # candidate was chosen ("default", "condition:", "on_failure", + # "router:"). `attempt_position` and `attempt_count` locate the row in + # the plan, so "served on attempt 2 of 3" is a query rather than a log grep. + # `request_group_id` ties a request's rows together, which is what makes the + # absorbed attempts findable from the row that served. + policy_name: Mapped[str | None] = mapped_column(index=True) + selection_reason: Mapped[str | None] = mapped_column() + attempt_position: Mapped[int | None] = mapped_column() + attempt_count: Mapped[int | None] = mapped_column() + request_group_id: Mapped[str | None] = mapped_column(index=True) + + # HTTP status that classifies a failure, so failures can be grouped with a + # GROUP BY instead of substring-matching provider-specific error prose. It is + # the status the provider returned when it sent one (an upstream 401 stays + # visible as a credential fault even though the caller sees the generic 502 + # that keeps gateway config out of the response), otherwise the gateway's own + # rejection or classification code (402 missing pricing, 422 tool-loop cap, + # 504 timeout, 502 unreachable). Nullable: historical rows predate the column, + # a successful request has no failure to classify, and some failures carry no + # HTTP status at all (e.g. a stream that ended without usage data). + status_code: Mapped[int | None] = mapped_column() + + # Total server-side wall-clock for the request, in milliseconds. Nullable: + # historical rows predate the column, and some write paths (batch jobs, + # provider-never-reached rejections) have no meaningful request duration. + latency_ms: Mapped[int | None] = mapped_column() + + # Milliseconds from request start to the first streamed chunk. Nullable: + # non-streaming requests have no first chunk, historical rows predate the + # column, and a stream that failed before yielding anything never reached one. + # + # ``started_at`` is taken in the handler preamble, so on a routing plan the + # serving row's value also carries every earlier attempt's setup time. + # Nothing in the column says so; a percentile keyed by the serving model + # attributes failover time to the model that actually served. + # + # Hybrid (platform-fallback) streams never write this column at all: every + # settlement callback in build_streaming_response returns before reaching + # log_usage on that path, and run_streaming_with_fallback passes db=None. + ttft_ms: Mapped[int | None] = mapped_column() + + api_key = relationship("APIKey", back_populates="usage_logs") + user = relationship("User", back_populates="usage_logs") + + def to_dict(self) -> dict[str, Any]: + """Convert model to dictionary.""" + return { + "id": self.id, + "api_key_id": self.api_key_id, + "user_id": self.user_id, + "timestamp": self.timestamp.isoformat() if self.timestamp else None, + "model": self.model, + "endpoint": self.endpoint, + "source": self.source, + "source_label": self.source_label, + "counts_toward_budget": self.counts_toward_budget, + "prompt_tokens": self.prompt_tokens, + "completion_tokens": self.completion_tokens, + "total_tokens": self.total_tokens, + "cache_read_tokens": self.cache_read_tokens, + "cache_write_tokens": self.cache_write_tokens, + "cache_write_1h_tokens": self.cache_write_1h_tokens, + "cache_tokens_in_prompt": self.cache_tokens_in_prompt, + "billing_meters": self.billing_meters, + "pricing_breakdown": self.pricing_breakdown, + "cost": self.cost, + "status": self.status, + "error_message": self.error_message, + "status_code": self.status_code, + "latency_ms": self.latency_ms, + "policy_name": self.policy_name, + "selection_reason": self.selection_reason, + "attempt_position": self.attempt_position, + "attempt_count": self.attempt_count, + "request_group_id": self.request_group_id, + } + + +class AgentTelemetry(Base): + """Content-free outcome metrics and behavioral events from coding agents.""" + + __tablename__ = "agent_telemetry" + __table_args__ = ( + UniqueConstraint("source", "dedup_key", name="uq_agent_telemetry_source_dedup"), + Index("ix_agent_telemetry_user_id_timestamp", "user_id", "timestamp"), + # Read-time cumulative-to-delta derivation orders one series' points by time. + Index("ix_agent_telemetry_series_timestamp", "series_key", "timestamp"), + ) + + id: Mapped[str] = mapped_column(primary_key=True, default=lambda: str(uuid.uuid4())) + api_key_id: Mapped[str | None] = mapped_column(ForeignKey("api_keys.id", ondelete="SET NULL"), index=True) + user_id: Mapped[str | None] = mapped_column(ForeignKey("users.user_id", ondelete="SET NULL"), index=True) + timestamp: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=lambda: datetime.now(UTC), index=True) + name: Mapped[str] = mapped_column() + tool_name: Mapped[str | None] = mapped_column() + decision: Mapped[str | None] = mapped_column() + success: Mapped[bool | None] = mapped_column() + duration_ms: Mapped[int | None] = mapped_column() + status_code: Mapped[int | None] = mapped_column() + prompt_length: Mapped[int | None] = mapped_column() + source: Mapped[str] = mapped_column(index=True) + session_label: Mapped[str | None] = mapped_column() + dedup_key: Mapped[str] = mapped_column() + created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=lambda: datetime.now(UTC)) + + # Outcome-metric columns. Populated only on a metric row (``kind="metric"``), + # NULL on a behavioral one, which is the inverse of the allow-list columns + # above. ``value`` is stored exactly as OTLP reported it (a running total or + # an increment, per ``temporality``); the read endpoints do the delta + # arithmetic, so nothing is normalized at ingest. ``series_key`` is the pure + # OTLP series identity (name plus attributes), which is what makes a + # dimensioned metric two series rather than one. + kind: Mapped[str | None] = mapped_column() + value: Mapped[float | None] = mapped_column() + temporality: Mapped[str | None] = mapped_column() + series_start: Mapped[datetime | None] = mapped_column(DateTime(timezone=True)) + series_key: Mapped[str | None] = mapped_column() diff --git a/src/gateway/models/users.py b/src/gateway/models/users.py new file mode 100644 index 0000000000..4523f730d2 --- /dev/null +++ b/src/gateway/models/users.py @@ -0,0 +1,84 @@ +"""ORM table for the gateway's billing identity. + +API keys, budgets, and usage rows attach to this ``User``. Gotcha: another model +class named ``User`` is the dashboard sign-in identity. +""" + +from datetime import UTC, datetime +from decimal import Decimal +from typing import Any + +from sqlalchemy import JSON, BigInteger, DateTime, ForeignKey +from sqlalchemy.orm import Mapped, mapped_column, relationship + +from gateway.models.base import Base +from gateway.models.money import UsdCost + + +class User(Base): + """User/Customer model for end-user tracking.""" + + __tablename__ = "users" + + user_id: Mapped[str] = mapped_column(primary_key=True) + alias: Mapped[str | None] = mapped_column() + # The spend ledger, exact to the micro-dollar like the ``usage_logs`` rows + # that sum into it (mozilla-ai/otari#691). As a float it drifted: four + # completions whose settled costs were each exact left this at + # 0.6619999999999999, and the drift accumulated across every reconcile until + # the budget reset. + spend: Mapped[Decimal] = mapped_column(UsdCost(), default=Decimal(0)) + # In-flight budget held by requests that have passed the budget gate but + # whose actual cost is not yet known. The effective committed amount is + # ``spend + reserved``; reservations are reconciled into ``spend`` (actual + # cost) on success or released on failure. See gateway.services.budget_service. + reserved: Mapped[Decimal] = mapped_column(UsdCost(), default=Decimal(0), server_default="0") + # The token and request counters, gated by the same budget's ``token_limit`` + # and ``request_limit`` the way the pair above is gated by ``max_budget``. + # Each axis names itself rather than extending the bare ``spend``/``reserved`` + # pair, which is USD and predates them. + current_tokens: Mapped[int] = mapped_column(BigInteger(), default=0, server_default="0") + reserved_tokens: Mapped[int] = mapped_column(BigInteger(), default=0, server_default="0") + current_requests: Mapped[int] = mapped_column(BigInteger(), default=0, server_default="0") + reserved_requests: Mapped[int] = mapped_column(BigInteger(), default=0, server_default="0") + # Indexed: the budgets list groups users by this column to build each budget's + # usage rollup, so an unindexed FK turns that page into a users table scan. + budget_id: Mapped[str | None] = mapped_column(ForeignKey("budgets.budget_id"), index=True) + # Default model access-list every one of this user's keys inherits when the + # key has no list of its own. null = unrestricted, [] = deny all, else + # canonical instance:model entries (see services/model_access.py). A key may + # narrow this default but never broaden it (validated on key write). + allowed_models: Mapped[list[str] | None] = mapped_column(JSON) + budget_started_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True)) + next_budget_reset_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True)) + blocked: Mapped[bool] = mapped_column(default=False) + deleted_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True), default=None, index=True) + created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=lambda: datetime.now(UTC)) + updated_at: Mapped[datetime] = mapped_column( + DateTime(timezone=True), + default=lambda: datetime.now(UTC), + onupdate=lambda: datetime.now(UTC), + ) + metadata_: Mapped[dict[str, Any]] = mapped_column("metadata", JSON, default=dict) + + budget = relationship("Budget", back_populates="users") + api_keys = relationship("APIKey", back_populates="user", passive_deletes=True) + usage_logs = relationship("UsageLog", back_populates="user", passive_deletes=True) + reset_logs = relationship("BudgetResetLog", back_populates="user", passive_deletes=True) + + def to_dict(self) -> dict[str, Any]: + """Convert model to dictionary.""" + return { + "user_id": self.user_id, + "alias": self.alias, + "spend": self.spend, + "reserved": self.reserved, + "budget_id": self.budget_id, + "allowed_models": self.allowed_models, + "budget_started_at": self.budget_started_at.isoformat() if self.budget_started_at else None, + "next_budget_reset_at": self.next_budget_reset_at.isoformat() if self.next_budget_reset_at else None, + "blocked": self.blocked, + "created_at": self.created_at.isoformat() if self.created_at else None, + "updated_at": self.updated_at.isoformat() if self.updated_at else None, + "metadata": self.metadata_, + } diff --git a/src/gateway/repositories/users_repository.py b/src/gateway/repositories/users_repository.py index 242ebb1ad6..46dd9b6eda 100644 --- a/src/gateway/repositories/users_repository.py +++ b/src/gateway/repositories/users_repository.py @@ -7,8 +7,10 @@ from sqlalchemy.sql.elements import ColumnElement from sqlmodel import col -from gateway.models.entities import APIKey, UsageLog, User +from gateway.models.api_keys import APIKey from gateway.models.tenancy import OrganizationMember, Workspace +from gateway.models.usage import UsageLog +from gateway.models.users import User # The owner a key falls back to when it is created without a user_id (the API's # convenience path, and the first-run bootstrap key). One shared, visible, diff --git a/src/gateway/services/agent_telemetry_service.py b/src/gateway/services/agent_telemetry_service.py index f302426e0a..84269d4142 100644 --- a/src/gateway/services/agent_telemetry_service.py +++ b/src/gateway/services/agent_telemetry_service.py @@ -13,7 +13,7 @@ from sqlalchemy.ext.asyncio import AsyncSession -from gateway.models.entities import APIKey +from gateway.models.api_keys import APIKey from gateway.ports.telemetry_storage_port import IngestResult, TelemetryRecord, TelemetryStoragePort from gateway.repositories.users_repository import get_active_user diff --git a/src/gateway/services/alias_service.py b/src/gateway/services/alias_service.py index b6465e5edf..626c52b8b7 100644 --- a/src/gateway/services/alias_service.py +++ b/src/gateway/services/alias_service.py @@ -49,7 +49,7 @@ from gateway.core.config import GatewayConfig from gateway.core.database import create_session from gateway.log_config import logger -from gateway.models.entities import ModelAlias +from gateway.models.providers import ModelAlias from gateway.services.workspace_scope import lookup_default_workspace_id # How long a worker may serve a stale alias map before refreshing. A new alias diff --git a/src/gateway/services/batch_service.py b/src/gateway/services/batch_service.py index bdb244c690..c29d52f0e8 100644 --- a/src/gateway/services/batch_service.py +++ b/src/gateway/services/batch_service.py @@ -15,7 +15,7 @@ from sqlalchemy.exc import SQLAlchemyError from gateway.log_config import logger -from gateway.models.entities import BatchRecord +from gateway.models.inference import BatchRecord if TYPE_CHECKING: import uuid diff --git a/src/gateway/services/bootstrap_service.py b/src/gateway/services/bootstrap_service.py index d21bb6b608..42d3c198c3 100644 --- a/src/gateway/services/bootstrap_service.py +++ b/src/gateway/services/bootstrap_service.py @@ -7,7 +7,7 @@ from gateway.auth import generate_api_key, hash_key, key_prefix, key_suffix from gateway.core.config import GatewayConfig from gateway.log_config import log_secret -from gateway.models.entities import APIKey +from gateway.models.api_keys import APIKey from gateway.repositories.users_repository import get_or_create_default_user from gateway.services.workspace_scope import default_workspace_id diff --git a/src/gateway/services/budget_reservation_ledger.py b/src/gateway/services/budget_reservation_ledger.py index fd241ed728..25c1418f5a 100644 --- a/src/gateway/services/budget_reservation_ledger.py +++ b/src/gateway/services/budget_reservation_ledger.py @@ -57,7 +57,8 @@ from gateway.core.database import create_session from gateway.log_config import logger -from gateway.models.entities import BudgetReservation, BudgetReservationScope, User +from gateway.models.budgets import BudgetReservation, BudgetReservationScope +from gateway.models.users import User from gateway.services.scoped_budget_service import release as release_scoped if TYPE_CHECKING: diff --git a/src/gateway/services/budget_retiming.py b/src/gateway/services/budget_retiming.py index 50aa5c2610..09f4e7eb4c 100644 --- a/src/gateway/services/budget_retiming.py +++ b/src/gateway/services/budget_retiming.py @@ -32,7 +32,7 @@ from sqlalchemy import update from sqlalchemy.ext.asyncio import AsyncSession -from gateway.models.entities import ScopedBudget +from gateway.models.budgets import ScopedBudget from gateway.services.budget_periods import period_window __all__ = ["cadence_of", "retime_ceilings_for_budget"] diff --git a/src/gateway/services/budget_service.py b/src/gateway/services/budget_service.py index d3deb70ef2..d9bc4289e7 100644 --- a/src/gateway/services/budget_service.py +++ b/src/gateway/services/budget_service.py @@ -17,8 +17,10 @@ from gateway.core.metered_pricing import estimate_metered_cost from gateway.log_config import logger from gateway.metrics import REGISTRY, Counter -from gateway.models.entities import MAX_COUNT_LIMIT, Budget, BudgetResetLog, ModelPricing, User +from gateway.models.budgets import MAX_COUNT_LIMIT, Budget, BudgetResetLog from gateway.models.money import to_usd +from gateway.models.pricing import ModelPricing +from gateway.models.users import User from gateway.repositories.users_repository import get_active_user from gateway.services import budget_reservation_ledger as ledger from gateway.services.budget_periods import budget_window diff --git a/src/gateway/services/dashboard_session_service.py b/src/gateway/services/dashboard_session_service.py index f12309e5c3..7c20e617b5 100644 --- a/src/gateway/services/dashboard_session_service.py +++ b/src/gateway/services/dashboard_session_service.py @@ -41,8 +41,8 @@ from gateway.core.config import GatewayConfig from gateway.core.database import create_session from gateway.log_config import logger -from gateway.models.entities import DashboardSession, RuntimeSetting -from gateway.models.tenancy import User +from gateway.models.platform import RuntimeSetting +from gateway.models.tenancy import DashboardSession, User from gateway.services.master_key_service import hash_master_key SESSION_COOKIE_NAME = "otari_dashboard_session" diff --git a/src/gateway/services/external_usage_service.py b/src/gateway/services/external_usage_service.py index 12c62b63cb..4c64a3b127 100644 --- a/src/gateway/services/external_usage_service.py +++ b/src/gateway/services/external_usage_service.py @@ -30,7 +30,10 @@ from gateway.core.config import API_ROOT from gateway.core.metered_pricing import BillableUsage, ChargeLine, billable_usage, price_billable_usage from gateway.log_config import logger -from gateway.models.entities import APIKey, ModelPricing, UsageLog, User +from gateway.models.api_keys import APIKey +from gateway.models.pricing import ModelPricing +from gateway.models.usage import UsageLog +from gateway.models.users import User from gateway.services.pricing_service import ( OverridePeriod, default_model_pricing, diff --git a/src/gateway/services/file_service.py b/src/gateway/services/file_service.py index 49abcdafc8..cbbd44bf7d 100644 --- a/src/gateway/services/file_service.py +++ b/src/gateway/services/file_service.py @@ -14,7 +14,7 @@ from sqlalchemy import select from sqlalchemy.ext.asyncio import AsyncSession -from gateway.models.entities import FileObject +from gateway.models.tools import FileObject from gateway.services.file_store import FileStore diff --git a/src/gateway/services/file_store.py b/src/gateway/services/file_store.py index b8af3bdfd8..765fe4b6ea 100644 --- a/src/gateway/services/file_store.py +++ b/src/gateway/services/file_store.py @@ -1,7 +1,7 @@ """Pluggable blob storage for uploaded file bytes. The ``/v1/files`` API stores file *metadata* in the database (see -``gateway.models.entities.FileObject``) and the raw *bytes* here, keyed by an +``gateway.models.tools.FileObject``) and the raw *bytes* here, keyed by an opaque ``storage_ref``. Keeping bytes out of the relational store lets large uploads live on a filesystem / object store while the DB stays lean. diff --git a/src/gateway/services/log_writer.py b/src/gateway/services/log_writer.py index baf0db1b42..728359cf66 100644 --- a/src/gateway/services/log_writer.py +++ b/src/gateway/services/log_writer.py @@ -9,7 +9,7 @@ from gateway.core.database import DATABASE_ERRORS, create_log_session from gateway.log_config import logger from gateway.metrics import REGISTRY, Counter, Gauge, Histogram -from gateway.models.entities import UsageLog +from gateway.models.usage import UsageLog QUEUE_DEPTH = Gauge( "gateway_usage_log_queue_depth", diff --git a/src/gateway/services/maintenance_mode_service.py b/src/gateway/services/maintenance_mode_service.py index 6c38118077..2d4c8f6cbb 100644 --- a/src/gateway/services/maintenance_mode_service.py +++ b/src/gateway/services/maintenance_mode_service.py @@ -41,7 +41,7 @@ from sqlalchemy import select from sqlalchemy.ext.asyncio import AsyncSession -from gateway.models.entities import RuntimeSetting +from gateway.models.platform import RuntimeSetting # The ``runtime_settings`` row this flag lives in. Deliberately absent from # ``runtime_settings_service._SPECS``: see the module docstring. diff --git a/src/gateway/services/master_key_service.py b/src/gateway/services/master_key_service.py index 74021d9940..8d76454f15 100644 --- a/src/gateway/services/master_key_service.py +++ b/src/gateway/services/master_key_service.py @@ -24,7 +24,7 @@ from gateway.core.config import GatewayConfig from gateway.log_config import log_secret, logger -from gateway.models.entities import RuntimeSetting +from gateway.models.platform import RuntimeSetting # Stored in runtime_settings; ignored by runtime_settings_service (not a SETTABLE_KEY). MASTER_KEY_HASH_KEY = "master_key_hash" diff --git a/src/gateway/services/merged_catalog_service.py b/src/gateway/services/merged_catalog_service.py index 57a5f5f2a7..68d7335663 100644 --- a/src/gateway/services/merged_catalog_service.py +++ b/src/gateway/services/merged_catalog_service.py @@ -22,8 +22,9 @@ from gateway.core.config import GatewayConfig from gateway.log_config import logger -from gateway.models.entities import APIKey, ModelPricing +from gateway.models.api_keys import APIKey from gateway.models.money import as_float +from gateway.models.pricing import ModelPricing from gateway.models.pricing_schemas import PricingTier from gateway.models.routing import PolicySpec from gateway.models.tenancy import User as TenancyUser diff --git a/src/gateway/services/model_access.py b/src/gateway/services/model_access.py index e6c45bd60b..2b14d8e2ee 100644 --- a/src/gateway/services/model_access.py +++ b/src/gateway/services/model_access.py @@ -20,7 +20,8 @@ from sqlalchemy.ext.asyncio import AsyncSession from gateway.core.config import GatewayConfig -from gateway.models.entities import APIKey, User +from gateway.models.api_keys import APIKey +from gateway.models.users import User from gateway.services.alias_service import all_alias_names from gateway.services.provider_kwargs import split_selector diff --git a/src/gateway/services/organization_pricing_service.py b/src/gateway/services/organization_pricing_service.py index ea665d67e2..2efb7c2b12 100644 --- a/src/gateway/services/organization_pricing_service.py +++ b/src/gateway/services/organization_pricing_service.py @@ -15,7 +15,7 @@ which the OSS edition ships by default, has neither exclusion constraints nor range types. So the rule is checked here and the schema holds the part both engines can (a unique index on the period start); see - `models.entities.OrganizationModelPricing` for the race that leaves. + `models.pricing.OrganizationModelPricing` for the race that leaves. - **Only a management role may write.** Rates decide what every member of the organization is billed, so this is the same owner-or-admin gate the rest of the organization surface uses, delegated to ``OrganizationService`` rather than @@ -44,8 +44,8 @@ from sqlalchemy.ext.asyncio import AsyncSession from gateway.core.config import GatewayConfig -from gateway.models.entities import OrganizationModelPricing from gateway.models.money import to_usd, to_usd_or_none +from gateway.models.pricing import OrganizationModelPricing from gateway.models.tenancy import User as TenancyUser from gateway.ports.model_provider_port import HostedAccessDeniedError, ModelProviderPort from gateway.repositories.tenancy import ( diff --git a/src/gateway/services/playground_service.py b/src/gateway/services/playground_service.py index a88816314e..1c953bceef 100644 --- a/src/gateway/services/playground_service.py +++ b/src/gateway/services/playground_service.py @@ -26,12 +26,6 @@ from gateway.core.config import GatewayConfig from gateway.core.env import otari_env -from gateway.models.entities import ( - User, - WorkspaceCodeExecutionPolicy, - WorkspaceMcpServer, - WorkspaceWebSearchConfig, -) from gateway.models.playground import ( MAX_FAVORITE_MODELS, MAX_SAVED_COMPARISONS, @@ -50,6 +44,8 @@ PlaygroundMessagePublic, ) from gateway.models.tenancy import User as TenancyUser +from gateway.models.tools import WorkspaceCodeExecutionPolicy, WorkspaceMcpServer, WorkspaceWebSearchConfig +from gateway.models.users import User from gateway.repositories.users_repository import get_or_create_attribution_user from gateway.services.tenancy import OrganizationService from gateway.services.tenancy.authorization import resolve_workspace_in_organization diff --git a/src/gateway/services/policy_store.py b/src/gateway/services/policy_store.py index fa987cf0a5..26f22ea2fd 100644 --- a/src/gateway/services/policy_store.py +++ b/src/gateway/services/policy_store.py @@ -45,8 +45,7 @@ from gateway.core.config import GatewayConfig from gateway.core.database import create_session from gateway.log_config import logger -from gateway.models.entities import RoutingPolicy -from gateway.models.routing import PolicySpec +from gateway.models.routing import PolicySpec, RoutingPolicy from gateway.services.workspace_scope import lookup_default_workspace_id __all__ = [ diff --git a/src/gateway/services/pricing_init_service.py b/src/gateway/services/pricing_init_service.py index f8ecc9d304..fcdc3d92be 100644 --- a/src/gateway/services/pricing_init_service.py +++ b/src/gateway/services/pricing_init_service.py @@ -7,7 +7,7 @@ from gateway.core.config import API_ROOT, GatewayConfig from gateway.core.env import otari_env from gateway.log_config import logger -from gateway.models.entities import ModelPricing +from gateway.models.pricing import ModelPricing from gateway.services.pricing_service import ( GATEWAY_TOOL_PRICING_PROVIDER, find_model_pricing, diff --git a/src/gateway/services/pricing_refresh_service.py b/src/gateway/services/pricing_refresh_service.py index e919503382..62ce3fbd98 100644 --- a/src/gateway/services/pricing_refresh_service.py +++ b/src/gateway/services/pricing_refresh_service.py @@ -17,7 +17,7 @@ from gateway.core.config import GatewayConfig from gateway.core.database import create_session from gateway.log_config import logger -from gateway.models.entities import PricingSnapshot, PricingSnapshotHistory +from gateway.models.pricing import PricingSnapshot, PricingSnapshotHistory from gateway.services.pricing_service import normalize_effective_at, reset_price_cache _PREVIEW_CHANGE_LIMIT = 100 diff --git a/src/gateway/services/pricing_service.py b/src/gateway/services/pricing_service.py index 94353c67fb..0ee5fc3acf 100644 --- a/src/gateway/services/pricing_service.py +++ b/src/gateway/services/pricing_service.py @@ -14,7 +14,7 @@ from gateway.core.config import API_ROOT from gateway.core.metered_pricing import meter_cost, quantize_cost, to_decimal from gateway.log_config import logger -from gateway.models.entities import ModelPricing, OrganizationModelPricing +from gateway.models.pricing import ModelPricing, OrganizationModelPricing # A zero-token usage is enough to resolve a model's per-million rates from # genai-prices without depending on real token counts. diff --git a/src/gateway/services/provider_store_service.py b/src/gateway/services/provider_store_service.py index 70d7ffa8b7..ab9031886b 100644 --- a/src/gateway/services/provider_store_service.py +++ b/src/gateway/services/provider_store_service.py @@ -29,7 +29,7 @@ from gateway.core.config import GatewayConfig from gateway.core.database import create_session from gateway.log_config import logger -from gateway.models.entities import ProviderCredential +from gateway.models.providers import ProviderCredential from gateway.models.secret_fields import restore_redacted_values from gateway.services.secret_box import ( SecretBoxUnavailableError, diff --git a/src/gateway/services/routing/knn.py b/src/gateway/services/routing/knn.py index 390282baf9..fd830b6856 100644 --- a/src/gateway/services/routing/knn.py +++ b/src/gateway/services/routing/knn.py @@ -16,7 +16,7 @@ The store is a linear cosine scan over the records the requesting user has in the requesting workspace, held in the gateway DB -(:class:`gateway.models.entities.RoutingMemory`). That holds into the low +(:class:`gateway.models.routing.RoutingMemory`). That holds into the low thousands of records per partition (the ``router_max_records_per_user`` cap, which bounds one user's records in one workspace, since that is what a decision loads); larger pools need an indexed vector store. Records carry an @@ -47,7 +47,7 @@ from gateway.core.database import create_session from gateway.log_config import logger -from gateway.models.entities import RoutingMemory +from gateway.models.routing import RoutingMemory from gateway.services.pricing_service import find_model_pricing from gateway.services.provider_kwargs import resolve_provider_selector from gateway.services.routing.backends import RoutingContext, RoutingDecision diff --git a/src/gateway/services/runtime_settings_service.py b/src/gateway/services/runtime_settings_service.py index 811a239fd4..cd1d8f992f 100644 --- a/src/gateway/services/runtime_settings_service.py +++ b/src/gateway/services/runtime_settings_service.py @@ -45,7 +45,7 @@ GatewayConfig, ) from gateway.log_config import logger -from gateway.models.entities import RuntimeSetting +from gateway.models.platform import RuntimeSetting from gateway.services.pricing_service import configure_default_pricing MODEL_DISCOVERY = "model_discovery" diff --git a/src/gateway/services/scoped_budget_service.py b/src/gateway/services/scoped_budget_service.py index c70ece507b..867378dabb 100644 --- a/src/gateway/services/scoped_budget_service.py +++ b/src/gateway/services/scoped_budget_service.py @@ -42,7 +42,8 @@ from sqlmodel import col from gateway.log_config import logger -from gateway.models.entities import APIKey, Budget, ScopedBudget +from gateway.models.api_keys import APIKey +from gateway.models.budgets import Budget, ScopedBudget from gateway.models.tenancy import OrganizationMember, Workspace, WorkspaceMember from gateway.services.budget_periods import ( ALIGN_DAY, diff --git a/src/gateway/services/search_tool_store_service.py b/src/gateway/services/search_tool_store_service.py index 038d03825b..8956a9e599 100644 --- a/src/gateway/services/search_tool_store_service.py +++ b/src/gateway/services/search_tool_store_service.py @@ -30,8 +30,8 @@ from gateway.core.config import GatewayConfig from gateway.core.database import create_session from gateway.log_config import logger -from gateway.models.entities import SearchToolCredential from gateway.models.secret_fields import restore_redacted_values +from gateway.models.tools import SearchToolCredential from gateway.services.secret_box import ( SecretBoxUnavailableError, SecretDecryptionError, diff --git a/src/gateway/services/tenancy/organization_budget_service.py b/src/gateway/services/tenancy/organization_budget_service.py index 9524060446..78b5c71a8b 100644 --- a/src/gateway/services/tenancy/organization_budget_service.py +++ b/src/gateway/services/tenancy/organization_budget_service.py @@ -63,10 +63,11 @@ from sqlalchemy.sql.elements import ColumnElement from sqlmodel import col -from gateway.models.entities import MAX_COUNT_LIMIT, APIKey, Budget, ScopedBudget, WorkspaceBudgetDefault -from gateway.models.entities import User as GatewayUser +from gateway.models.api_keys import APIKey +from gateway.models.budgets import MAX_COUNT_LIMIT, Budget, ScopedBudget, WorkspaceBudgetDefault from gateway.models.money import MAX_USD_LIMIT, as_float, to_usd_or_none from gateway.models.tenancy import Organization, OrganizationMember, User, Workspace, WorkspaceMember +from gateway.models.users import User as GatewayUser from gateway.services.budget_periods import ResetAlignment, period_window from gateway.services.budget_retiming import cadence_of, retime_ceilings_for_budget from gateway.services.tenancy.errors import ( diff --git a/src/gateway/services/tenancy/organization_guardrail_service.py b/src/gateway/services/tenancy/organization_guardrail_service.py index d4c4ad697b..bea86418ae 100644 --- a/src/gateway/services/tenancy/organization_guardrail_service.py +++ b/src/gateway/services/tenancy/organization_guardrail_service.py @@ -59,8 +59,7 @@ from sqlalchemy.exc import IntegrityError from sqlalchemy.ext.asyncio import AsyncSession -from gateway.models.entities import OrganizationGuardrail, OrganizationGuardrailWorkspace -from gateway.models.guardrails import GuardrailConfig +from gateway.models.guardrails import GuardrailConfig, OrganizationGuardrail, OrganizationGuardrailWorkspace from gateway.models.secret_fields import redact_secret_like_values, restore_redacted_values from gateway.models.tenancy import User from gateway.repositories.tenancy import WorkspaceRepository @@ -116,7 +115,7 @@ class OrganizationGuardrailCreate(BaseModel): ``credential`` is never stored as sent: it is encrypted with ``OTARI_SECRET_KEY`` and only the ciphertext is kept, the same convention - `entities.WorkspaceMcpServer` and `entities.ProviderCredential` use. It is + `tools.WorkspaceMcpServer` and `providers.ProviderCredential` use. It is sent to the endpoint as ``Authorization: Bearer`` when the guardrail runs, so it authenticates this gateway to the guardrails service the entry names. A guardrail *vendor's* own key is not this: the guardrails service builds diff --git a/src/gateway/services/tenancy/provisioning_service.py b/src/gateway/services/tenancy/provisioning_service.py index 39ad52f581..c4c0af40c9 100644 --- a/src/gateway/services/tenancy/provisioning_service.py +++ b/src/gateway/services/tenancy/provisioning_service.py @@ -31,7 +31,7 @@ from sqlmodel import col from gateway.log_config import logger -from gateway.models.entities import RuntimeSetting +from gateway.models.platform import RuntimeSetting from gateway.models.tenancy import Organization, User from gateway.repositories.tenancy import ( OrganizationMemberRepository, diff --git a/src/gateway/services/tenancy/workspace_activation_service.py b/src/gateway/services/tenancy/workspace_activation_service.py index 8a675b4f70..16a80961f2 100644 --- a/src/gateway/services/tenancy/workspace_activation_service.py +++ b/src/gateway/services/tenancy/workspace_activation_service.py @@ -49,9 +49,10 @@ from gateway.auth.models import generate_api_key, hash_key, key_prefix, key_suffix from gateway.core.config import GatewayConfig from gateway.core.usage_source import integration_traffic, served_here -from gateway.models.entities import APIKey, UsageLog, WorkspaceActivationState +from gateway.models.api_keys import APIKey from gateway.models.money import as_float -from gateway.models.tenancy import User, Workspace +from gateway.models.tenancy import User, Workspace, WorkspaceActivationState +from gateway.models.usage import UsageLog from gateway.repositories.users_repository import get_or_create_attribution_user from gateway.services.tenancy import authorization from gateway.services.tenancy.errors import ( diff --git a/src/gateway/services/tenancy/workspace_budget_default_service.py b/src/gateway/services/tenancy/workspace_budget_default_service.py index e2f765fa04..bdcc16103f 100644 --- a/src/gateway/services/tenancy/workspace_budget_default_service.py +++ b/src/gateway/services/tenancy/workspace_budget_default_service.py @@ -32,7 +32,7 @@ from sqlalchemy.ext.asyncio import AsyncSession from sqlmodel import col -from gateway.models.entities import Budget, ScopedBudget, WorkspaceBudgetDefault +from gateway.models.budgets import Budget, ScopedBudget, WorkspaceBudgetDefault from gateway.models.money import as_float from gateway.models.tenancy import User, Workspace, WorkspaceMember from gateway.repositories.tenancy import WorkspaceMemberRepository, WorkspaceRepository diff --git a/src/gateway/services/tenancy/workspace_code_execution_policy_service.py b/src/gateway/services/tenancy/workspace_code_execution_policy_service.py index cefa6abfee..ef8789b78a 100644 --- a/src/gateway/services/tenancy/workspace_code_execution_policy_service.py +++ b/src/gateway/services/tenancy/workspace_code_execution_policy_service.py @@ -52,8 +52,8 @@ from sqlalchemy.exc import IntegrityError, SQLAlchemyError from sqlalchemy.ext.asyncio import AsyncSession -from gateway.models.entities import WorkspaceCodeExecutionPolicy from gateway.models.tenancy import User, Workspace +from gateway.models.tools import WorkspaceCodeExecutionPolicy from gateway.services.mcp_loop import MAX_TOOL_ITERATIONS_CAP from gateway.services.sandbox_backend import ( CODE_EXECUTION_TOOL_NAME, diff --git a/src/gateway/services/tenancy/workspace_mcp_server_service.py b/src/gateway/services/tenancy/workspace_mcp_server_service.py index d97ba18c90..f167beef57 100644 --- a/src/gateway/services/tenancy/workspace_mcp_server_service.py +++ b/src/gateway/services/tenancy/workspace_mcp_server_service.py @@ -49,9 +49,9 @@ from sqlalchemy.exc import IntegrityError from sqlalchemy.ext.asyncio import AsyncSession -from gateway.models.entities import WorkspaceMcpServer from gateway.models.mcp import McpServerConfig, ResolvedMcpServer from gateway.models.tenancy import User +from gateway.models.tools import WorkspaceMcpServer from gateway.repositories.tenancy import WorkspaceRepository from gateway.services.secret_box import ( SecretBoxUnavailableError, @@ -93,7 +93,7 @@ class WorkspaceMcpServerCreate(BaseModel): ``authorization_token`` is never stored as sent: it is encrypted with ``OTARI_SECRET_KEY`` and only the ciphertext is kept, the same convention - `entities.ProviderCredential` and `OrgProviderKey` already use. + `providers.ProviderCredential` and `OrgProviderKey` already use. """ name: str = Field(min_length=1, max_length=128, description="Label for the server, unique within the workspace") diff --git a/src/gateway/services/tenancy/workspace_service.py b/src/gateway/services/tenancy/workspace_service.py index 2672d209e9..99f1c454e1 100644 --- a/src/gateway/services/tenancy/workspace_service.py +++ b/src/gateway/services/tenancy/workspace_service.py @@ -22,7 +22,7 @@ from sqlalchemy.ext.asyncio import AsyncSession from sqlmodel import col -from gateway.models.entities import ScopedBudget +from gateway.models.budgets import ScopedBudget from gateway.models.tenancy import ( MANAGEMENT_ROLES, WORKSPACE_MEMBER_ROLES, diff --git a/src/gateway/services/tenancy/workspace_web_search_service.py b/src/gateway/services/tenancy/workspace_web_search_service.py index 260362bdb0..0bb3e60e6f 100644 --- a/src/gateway/services/tenancy/workspace_web_search_service.py +++ b/src/gateway/services/tenancy/workspace_web_search_service.py @@ -57,8 +57,8 @@ from sqlalchemy.exc import IntegrityError, SQLAlchemyError from sqlalchemy.ext.asyncio import AsyncSession -from gateway.models.entities import WorkspaceWebSearchConfig from gateway.models.tenancy import User, Workspace +from gateway.models.tools import WorkspaceWebSearchConfig from gateway.services.tenancy import authorization from gateway.services.tenancy.errors import WorkspaceWebSearchDomainsExcludedError from gateway.services.tenancy.organization_service import OrganizationService diff --git a/src/gateway/services/tool_settings_service.py b/src/gateway/services/tool_settings_service.py index 90c7fe8219..23cde095e5 100644 --- a/src/gateway/services/tool_settings_service.py +++ b/src/gateway/services/tool_settings_service.py @@ -33,7 +33,7 @@ from gateway.core.config import GatewayConfig from gateway.core.env import otari_env from gateway.log_config import logger -from gateway.models.entities import RuntimeSetting +from gateway.models.platform import RuntimeSetting from gateway.services.runtime_settings_service import SettingValue WEB_SEARCH_URL = "web_search_url" diff --git a/src/gateway/services/usage_admin_service.py b/src/gateway/services/usage_admin_service.py index c3f5ffaa66..ef04026b5f 100644 --- a/src/gateway/services/usage_admin_service.py +++ b/src/gateway/services/usage_admin_service.py @@ -32,7 +32,8 @@ from gateway.core.sql import MAX_FILTER_VALUES, match_any, utc_bound from gateway.core.usage_source import not_served_here from gateway.log_config import logger -from gateway.models.entities import ModelPricing, UsageLog +from gateway.models.pricing import ModelPricing +from gateway.models.usage import UsageLog from gateway.services.tool_usage import TOOL_METER_NAMESPACE # Cap on an explicit id list. Page selections drive the id path and the largest diff --git a/src/gateway/services/workspace_scope.py b/src/gateway/services/workspace_scope.py index 8e03676889..0d684a42ad 100644 --- a/src/gateway/services/workspace_scope.py +++ b/src/gateway/services/workspace_scope.py @@ -39,7 +39,7 @@ from sqlalchemy.ext.asyncio import AsyncSession from sqlmodel import col -from gateway.models.entities import APIKey +from gateway.models.api_keys import APIKey from gateway.models.tenancy import Organization, Workspace from gateway.services.tenancy.provisioning_service import ( DEFAULT_ORGANIZATION_NAME, diff --git a/tests/integration/test_agent_telemetry_admin.py b/tests/integration/test_agent_telemetry_admin.py index fc93cbf5d7..268f80f1f6 100644 --- a/tests/integration/test_agent_telemetry_admin.py +++ b/tests/integration/test_agent_telemetry_admin.py @@ -9,7 +9,9 @@ from conftest import seed_workspace_id from gateway.core.config import API_ROOT -from gateway.models.entities import AgentTelemetry, APIKey, User +from gateway.models.api_keys import APIKey +from gateway.models.usage import AgentTelemetry +from gateway.models.users import User DELETE_PATH = f"{API_ROOT}/agent-telemetry" diff --git a/tests/integration/test_agent_telemetry_read.py b/tests/integration/test_agent_telemetry_read.py index cf341fe1cc..7b2e371a28 100644 --- a/tests/integration/test_agent_telemetry_read.py +++ b/tests/integration/test_agent_telemetry_read.py @@ -18,7 +18,9 @@ from gateway.core.config import API_ROOT, GatewayConfig from gateway.core.database import reset_db from gateway.main import create_app -from gateway.models.entities import AgentTelemetry, APIKey, UsageLog, User +from gateway.models.api_keys import APIKey +from gateway.models.usage import AgentTelemetry, UsageLog +from gateway.models.users import User SUMMARY_PATH = f"{API_ROOT}/agent-telemetry/summary" COUNT_PATH = f"{API_ROOT}/agent-telemetry/count" diff --git a/tests/integration/test_atomic_spend_update.py b/tests/integration/test_atomic_spend_update.py index 2a943abd73..8961c47cef 100644 --- a/tests/integration/test_atomic_spend_update.py +++ b/tests/integration/test_atomic_spend_update.py @@ -11,7 +11,7 @@ import pytest from sqlalchemy.ext.asyncio import AsyncSession -from gateway.models.entities import User +from gateway.models.users import User from gateway.services.budget_service import ReservationHandle, reconcile_reservation diff --git a/tests/integration/test_budget_dashboard.py b/tests/integration/test_budget_dashboard.py index 70f26c0732..213b532c11 100644 --- a/tests/integration/test_budget_dashboard.py +++ b/tests/integration/test_budget_dashboard.py @@ -9,8 +9,9 @@ from sqlalchemy.orm import Session from gateway.core.config import API_ROOT -from gateway.models.entities import BudgetResetLog, ScopedBudget, User, WorkspaceBudgetDefault +from gateway.models.budgets import BudgetResetLog, ScopedBudget, WorkspaceBudgetDefault from gateway.models.tenancy import Organization, Workspace +from gateway.models.users import User def _make_budget(client: TestClient, headers: dict[str, str], max_budget: float | None = 100.0) -> str: diff --git a/tests/integration/test_budget_race_condition.py b/tests/integration/test_budget_race_condition.py index fbaa9407fd..52d3088592 100644 --- a/tests/integration/test_budget_race_condition.py +++ b/tests/integration/test_budget_race_condition.py @@ -13,7 +13,9 @@ import pytest from fastapi import HTTPException -from gateway.models.entities import MAX_COUNT_LIMIT, Budget, ModelPricing, User +from gateway.models.budgets import MAX_COUNT_LIMIT, Budget +from gateway.models.pricing import ModelPricing +from gateway.models.users import User from gateway.repositories.users_repository import get_active_user from gateway.services.budget_service import ( estimate_cost, diff --git a/tests/integration/test_budget_reservation_ledger.py b/tests/integration/test_budget_reservation_ledger.py index 47612e2ec9..e89f3cb734 100644 --- a/tests/integration/test_budget_reservation_ledger.py +++ b/tests/integration/test_budget_reservation_ledger.py @@ -23,7 +23,8 @@ from sqlalchemy.exc import SQLAlchemyError from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine -from gateway.models.entities import Budget, BudgetReservation, BudgetReservationScope, ScopedBudget, User +from gateway.models.budgets import Budget, BudgetReservation, BudgetReservationScope, ScopedBudget +from gateway.models.users import User from gateway.services import budget_reservation_ledger as ledger from gateway.services.budget_service import ( increase_reservation, diff --git a/tests/integration/test_catalog_endpoint.py b/tests/integration/test_catalog_endpoint.py index 045387fd18..f474d2b284 100644 --- a/tests/integration/test_catalog_endpoint.py +++ b/tests/integration/test_catalog_endpoint.py @@ -17,8 +17,8 @@ from sqlalchemy.orm import Session from gateway.core.config import API_ROOT, GatewayConfig -from gateway.models.entities import DashboardSession, OrganizationModelPricing -from gateway.models.tenancy import Organization, OrganizationMember, User +from gateway.models.pricing import OrganizationModelPricing +from gateway.models.tenancy import DashboardSession, Organization, OrganizationMember, User from gateway.services import model_catalog_service as mcs from gateway.services.dashboard_session_service import SESSION_COOKIE_NAME, hash_session_token @@ -475,8 +475,8 @@ def test_a_signed_in_caller_sees_their_own_usage_of_an_offering( """The listed rate is what a token costs; this is what the tokens cost.""" from sqlmodel import select - from gateway.models.entities import UsageLog from gateway.models.tenancy import Workspace + from gateway.models.usage import UsageLog # The master key acts in the default workspace, which boot provisioned. assert priced.get(f"{API_ROOT}/organizations/me", headers=master_header).status_code == status.HTTP_200_OK diff --git a/tests/integration/test_deployment_operator_gate.py b/tests/integration/test_deployment_operator_gate.py index 840ab4672b..8605a5e3d8 100644 --- a/tests/integration/test_deployment_operator_gate.py +++ b/tests/integration/test_deployment_operator_gate.py @@ -33,8 +33,7 @@ from sqlmodel import col from gateway.core.config import API_ROOT, GatewayConfig -from gateway.models.entities import DashboardSession -from gateway.models.tenancy import Organization, OrganizationMember, User +from gateway.models.tenancy import DashboardSession, Organization, OrganizationMember, User from gateway.services.dashboard_session_service import SESSION_COOKIE_NAME, hash_session_token # One probe per deployment-wide router family, each the cheapest request that diff --git a/tests/integration/test_deployment_user_administration.py b/tests/integration/test_deployment_user_administration.py index 5bdd0655dc..06b384b24b 100644 --- a/tests/integration/test_deployment_user_administration.py +++ b/tests/integration/test_deployment_user_administration.py @@ -22,8 +22,9 @@ from sqlmodel import col from gateway.core.config import API_ROOT, GatewayConfig -from gateway.models.entities import DashboardSession, RuntimeSetting +from gateway.models.platform import RuntimeSetting from gateway.models.tenancy import ( + DashboardSession, DeploymentUserUpdateRequest, Organization, OrganizationMember, diff --git a/tests/integration/test_exact_budget_ledger.py b/tests/integration/test_exact_budget_ledger.py index 01476f41b9..44fa483e66 100644 --- a/tests/integration/test_exact_budget_ledger.py +++ b/tests/integration/test_exact_budget_ledger.py @@ -24,7 +24,9 @@ from sqlalchemy.orm import Session from gateway.core.config import API_KEY_HEADER, API_ROOT -from gateway.models.entities import Budget, ScopedBudget, UsageLog, User +from gateway.models.budgets import Budget, ScopedBudget +from gateway.models.usage import UsageLog +from gateway.models.users import User from gateway.services.budget_service import ReservationHandle, reconcile_reservation, reserve_budget from gateway.services.scoped_budget_service import ApplicableBudget, reserve, settle diff --git a/tests/integration/test_exact_money_columns.py b/tests/integration/test_exact_money_columns.py index 9fd2d71cb5..5327b99b8f 100644 --- a/tests/integration/test_exact_money_columns.py +++ b/tests/integration/test_exact_money_columns.py @@ -29,8 +29,9 @@ from conftest import seed_workspace_id from gateway.core.metered_pricing import calculate_metered_cost from gateway.core.usage import GatewayUsage -from gateway.models.entities import ModelPricing, OrganizationModelPricing, UsageLog +from gateway.models.pricing import ModelPricing, OrganizationModelPricing from gateway.models.tenancy import Organization +from gateway.models.usage import UsageLog from gateway.services.pricing_service import default_model_pricing # Rates an operator's price list actually holds, taken from the catalog the diff --git a/tests/integration/test_exclude_from_budget.py b/tests/integration/test_exclude_from_budget.py index 553ec1be57..6a78b14d56 100644 --- a/tests/integration/test_exclude_from_budget.py +++ b/tests/integration/test_exclude_from_budget.py @@ -14,7 +14,8 @@ from sqlalchemy.orm import Session from gateway.core.config import API_KEY_HEADER, API_ROOT -from gateway.models.entities import UsageLog, User +from gateway.models.usage import UsageLog +from gateway.models.users import User from .conftest import MODEL_NAME diff --git a/tests/integration/test_external_usage_events.py b/tests/integration/test_external_usage_events.py index b5011e7646..6e264dd767 100644 --- a/tests/integration/test_external_usage_events.py +++ b/tests/integration/test_external_usage_events.py @@ -15,8 +15,11 @@ from sqlalchemy.orm import Session from gateway.core.config import API_ROOT -from gateway.models.entities import OrganizationModelPricing, RuntimeSetting, UsageLog, User +from gateway.models.platform import RuntimeSetting +from gateway.models.pricing import OrganizationModelPricing from gateway.models.tenancy import Organization, OrganizationMember, Workspace +from gateway.models.usage import UsageLog +from gateway.models.users import User from gateway.services.tenancy.provisioning_service import BOOTSTRAP_IDENTITY_KEY _SRC = "claude_code" diff --git a/tests/integration/test_files_endpoint.py b/tests/integration/test_files_endpoint.py index 061d0ab8ed..56dc240262 100644 --- a/tests/integration/test_files_endpoint.py +++ b/tests/integration/test_files_endpoint.py @@ -24,7 +24,7 @@ from sqlalchemy.orm import Session from gateway.core.config import API_ROOT -from gateway.models.entities import FileObject +from gateway.models.tools import FileObject from gateway.services.file_extractors import ExtractionResult from gateway.services.file_store import LocalDirFileStore diff --git a/tests/integration/test_growth_signal_lifecycle.py b/tests/integration/test_growth_signal_lifecycle.py index a1f84dafbc..5992c0595c 100644 --- a/tests/integration/test_growth_signal_lifecycle.py +++ b/tests/integration/test_growth_signal_lifecycle.py @@ -28,8 +28,7 @@ from sqlmodel import col from gateway.core.config import API_ROOT, GatewayConfig -from gateway.models.entities import DashboardSession -from gateway.models.tenancy import Organization, OrganizationMember, User, Workspace, WorkspaceMember +from gateway.models.tenancy import DashboardSession, Organization, OrganizationMember, User, Workspace, WorkspaceMember from gateway.ports.growth_signal_port import GrowthActivationEvent, GrowthSignalPort from gateway.services.dashboard_session_service import SESSION_COOKIE_NAME, hash_session_token diff --git a/tests/integration/test_hosted_credential_request_path.py b/tests/integration/test_hosted_credential_request_path.py index 0c482eaae1..c723161aba 100644 --- a/tests/integration/test_hosted_credential_request_path.py +++ b/tests/integration/test_hosted_credential_request_path.py @@ -21,7 +21,7 @@ from gateway.core.config import API_KEY_HEADER, API_ROOT, GatewayConfig from gateway.log_config import logger as gateway_logger -from gateway.models.entities import User +from gateway.models.users import User from .conftest import build_test_client diff --git a/tests/integration/test_log_usage_commit_scope.py b/tests/integration/test_log_usage_commit_scope.py index 48f77e53be..cafa3fa370 100644 --- a/tests/integration/test_log_usage_commit_scope.py +++ b/tests/integration/test_log_usage_commit_scope.py @@ -11,7 +11,8 @@ from gateway.api.routes.chat import log_usage from gateway.core.usage import GatewayUsage -from gateway.models.entities import ModelPricing, UsageLog +from gateway.models.pricing import ModelPricing +from gateway.models.usage import UsageLog @dataclass diff --git a/tests/integration/test_mcp_stored_server_endpoints.py b/tests/integration/test_mcp_stored_server_endpoints.py index 02e57b6b5f..af0c1a5dec 100644 --- a/tests/integration/test_mcp_stored_server_endpoints.py +++ b/tests/integration/test_mcp_stored_server_endpoints.py @@ -30,9 +30,11 @@ from gateway.core.config import API_ROOT from gateway.core.database import release_session from gateway.inflight import InFlightRegistry -from gateway.models.entities import APIKey, User, WorkspaceMcpServer +from gateway.models.api_keys import APIKey from gateway.models.mcp import ResolvedMcpServer from gateway.models.tenancy import Organization, Workspace +from gateway.models.tools import WorkspaceMcpServer +from gateway.models.users import User from gateway.services import mcp_stateless from gateway.services.secret_box import encrypt_secret, generate_secret_key diff --git a/tests/integration/test_messages_streaming_usage.py b/tests/integration/test_messages_streaming_usage.py index d673a3c283..6473ff9797 100644 --- a/tests/integration/test_messages_streaming_usage.py +++ b/tests/integration/test_messages_streaming_usage.py @@ -35,7 +35,7 @@ from sqlalchemy.orm import Session from gateway.core.config import API_ROOT -from gateway.models.entities import UsageLog +from gateway.models.usage import UsageLog from .conftest import MODEL_NAME diff --git a/tests/integration/test_migrated_key_auth.py b/tests/integration/test_migrated_key_auth.py index e1245e6204..1de3be18d0 100644 --- a/tests/integration/test_migrated_key_auth.py +++ b/tests/integration/test_migrated_key_auth.py @@ -15,7 +15,7 @@ from sqlalchemy.orm import Session from gateway.core.config import API_KEY_HEADER, API_ROOT -from gateway.models.entities import APIKey +from gateway.models.api_keys import APIKey # The shape otari-ai mints: it fails both the ``gw-``/``gw_`` prefix check and the # ``gw[-_][A-Za-z0-9_-]+`` charset check the old validator applied. diff --git a/tests/integration/test_organization_budgets.py b/tests/integration/test_organization_budgets.py index 1345f886d3..97e797ee2c 100644 --- a/tests/integration/test_organization_budgets.py +++ b/tests/integration/test_organization_budgets.py @@ -29,17 +29,10 @@ from sqlmodel import col from gateway.core.config import API_ROOT -from gateway.models.entities import ( - APIKey, - Budget, - BudgetResetLog, - ScopedBudget, - WorkspaceBudgetDefault, -) -from gateway.models.entities import ( - User as ApiUser, -) +from gateway.models.api_keys import APIKey +from gateway.models.budgets import Budget, BudgetResetLog, ScopedBudget, WorkspaceBudgetDefault from gateway.models.tenancy import Organization, OrganizationMember, User, Workspace, WorkspaceMember +from gateway.models.users import User as ApiUser from gateway.repositories.tenancy import ( OrganizationMemberRepository, OrganizationRepository, diff --git a/tests/integration/test_organization_guardrails.py b/tests/integration/test_organization_guardrails.py index 8eaf00a8cb..4e7300ad36 100644 --- a/tests/integration/test_organization_guardrails.py +++ b/tests/integration/test_organization_guardrails.py @@ -19,7 +19,7 @@ from sqlalchemy import select from sqlalchemy.ext.asyncio import AsyncSession -from gateway.models.entities import OrganizationGuardrail, OrganizationGuardrailWorkspace +from gateway.models.guardrails import OrganizationGuardrail, OrganizationGuardrailWorkspace from gateway.models.tenancy import Organization, User, Workspace from gateway.repositories.tenancy import ( OrganizationMemberRepository, diff --git a/tests/integration/test_organization_member_keys.py b/tests/integration/test_organization_member_keys.py index d9b5804968..921fc2f8f3 100644 --- a/tests/integration/test_organization_member_keys.py +++ b/tests/integration/test_organization_member_keys.py @@ -28,9 +28,9 @@ from gateway.auth.models import generate_api_key, hash_key, key_prefix from gateway.core.config import API_ROOT -from gateway.models.entities import APIKey, DashboardSession -from gateway.models.entities import User as BillingUser -from gateway.models.tenancy import Organization, OrganizationMember, User, Workspace, WorkspaceMember +from gateway.models.api_keys import APIKey +from gateway.models.tenancy import DashboardSession, Organization, OrganizationMember, User, Workspace, WorkspaceMember +from gateway.models.users import User as BillingUser from gateway.services.dashboard_session_service import SESSION_COOKIE_NAME, hash_session_token _PREFIX = f"{API_ROOT}/organizations/me/keys" diff --git a/tests/integration/test_organization_pricing_routes.py b/tests/integration/test_organization_pricing_routes.py index d01418f179..09a41f872c 100644 --- a/tests/integration/test_organization_pricing_routes.py +++ b/tests/integration/test_organization_pricing_routes.py @@ -23,8 +23,9 @@ from sqlmodel import col from gateway.core.config import API_ROOT, GatewayConfig -from gateway.models.entities import APIKey, DashboardSession, ModelPricing, OrganizationModelPricing -from gateway.models.tenancy import Organization, OrganizationMember, User, Workspace +from gateway.models.api_keys import APIKey +from gateway.models.pricing import ModelPricing, OrganizationModelPricing +from gateway.models.tenancy import DashboardSession, Organization, OrganizationMember, User, Workspace from gateway.ports.model_provider_port import HostedAccessDeniedError, HostedCredential, ModelProviderPort from gateway.repositories.tenancy import ( OrganizationMemberRepository, diff --git a/tests/integration/test_organization_routing_policies_scope.py b/tests/integration/test_organization_routing_policies_scope.py index fb3f267864..acbebd404e 100644 --- a/tests/integration/test_organization_routing_policies_scope.py +++ b/tests/integration/test_organization_routing_policies_scope.py @@ -25,8 +25,8 @@ from sqlalchemy.orm import Session from gateway.core.config import API_ROOT -from gateway.models.entities import DashboardSession, RoutingPolicy -from gateway.models.tenancy import Organization, OrganizationMember, User, Workspace, WorkspaceMember +from gateway.models.routing import RoutingPolicy +from gateway.models.tenancy import DashboardSession, Organization, OrganizationMember, User, Workspace, WorkspaceMember from gateway.services.dashboard_session_service import SESSION_COOKIE_NAME, hash_session_token _SCOPED_PATH = f"{API_ROOT}/organizations/me/routing-policies" diff --git a/tests/integration/test_organization_routing_write_scope.py b/tests/integration/test_organization_routing_write_scope.py index 9ff9c9758b..bd87cf9e24 100644 --- a/tests/integration/test_organization_routing_write_scope.py +++ b/tests/integration/test_organization_routing_write_scope.py @@ -31,9 +31,8 @@ from sqlalchemy.orm import Session from gateway.core.config import API_ROOT -from gateway.models.entities import DashboardSession from gateway.models.provider_keys import OrgProviderKey -from gateway.models.tenancy import Organization, OrganizationMember, User, Workspace, WorkspaceMember +from gateway.models.tenancy import DashboardSession, Organization, OrganizationMember, User, Workspace, WorkspaceMember from gateway.services.dashboard_session_service import SESSION_COOKIE_NAME, hash_session_token from gateway.services.secret_box import encrypt_secret, generate_secret_key diff --git a/tests/integration/test_organization_usage_scope.py b/tests/integration/test_organization_usage_scope.py index b40fd3b0cb..e524f510ef 100644 --- a/tests/integration/test_organization_usage_scope.py +++ b/tests/integration/test_organization_usage_scope.py @@ -37,8 +37,8 @@ from sqlmodel import col from gateway.core.config import API_ROOT -from gateway.models.entities import DashboardSession, UsageLog -from gateway.models.tenancy import Organization, OrganizationMember, User, Workspace, WorkspaceMember +from gateway.models.tenancy import DashboardSession, Organization, OrganizationMember, User, Workspace, WorkspaceMember +from gateway.models.usage import UsageLog from gateway.services.dashboard_session_service import SESSION_COOKIE_NAME, hash_session_token # Every read this router serves. Parametrized rather than asserted once, because diff --git a/tests/integration/test_otlp.py b/tests/integration/test_otlp.py index d4de253eed..2fa3bbcbd8 100644 --- a/tests/integration/test_otlp.py +++ b/tests/integration/test_otlp.py @@ -18,7 +18,8 @@ from sqlalchemy.orm import Session from gateway.core.config import API_ROOT -from gateway.models.entities import UsageLog, User +from gateway.models.usage import UsageLog +from gateway.models.users import User def _usd(tokens: int, rate_per_million: str) -> Decimal: diff --git a/tests/integration/test_otlp_logs_behavioral.py b/tests/integration/test_otlp_logs_behavioral.py index a1d9b53027..b31b70f2de 100644 --- a/tests/integration/test_otlp_logs_behavioral.py +++ b/tests/integration/test_otlp_logs_behavioral.py @@ -7,7 +7,8 @@ from gateway.api.routes.otlp import _MAX_EVENTS_PER_EXPORT from gateway.core.config import API_ROOT -from gateway.models.entities import AgentTelemetry, UsageLog, User +from gateway.models.usage import AgentTelemetry, UsageLog +from gateway.models.users import User from .otlp_helpers import log_record, logs_export diff --git a/tests/integration/test_otlp_metrics.py b/tests/integration/test_otlp_metrics.py index 48c2ad4c2a..01597bdb3e 100644 --- a/tests/integration/test_otlp_metrics.py +++ b/tests/integration/test_otlp_metrics.py @@ -18,7 +18,8 @@ from gateway.api.routes.otlp import _MAX_METRIC_DATA_POINTS from gateway.core.config import API_ROOT -from gateway.models.entities import AgentTelemetry, UsageLog, User +from gateway.models.usage import AgentTelemetry, UsageLog +from gateway.models.users import User from .otlp_helpers import gauge_metric, metrics_export, metrics_export_protobuf, number_point, sum_metric diff --git a/tests/integration/test_playground.py b/tests/integration/test_playground.py index 05d731ca96..81565f4d84 100644 --- a/tests/integration/test_playground.py +++ b/tests/integration/test_playground.py @@ -39,9 +39,10 @@ from gateway.core.config import API_ROOT from gateway.core.usage_source import PLAYGROUND_USAGE_ENDPOINT, SERVED_HERE_SLUG -from gateway.models.entities import DashboardSession, UsageLog, WorkspaceWebSearchConfig -from gateway.models.entities import User as BillingUser -from gateway.models.tenancy import Organization, OrganizationMember, User, Workspace, WorkspaceMember +from gateway.models.tenancy import DashboardSession, Organization, OrganizationMember, User, Workspace, WorkspaceMember +from gateway.models.tools import WorkspaceWebSearchConfig +from gateway.models.usage import UsageLog +from gateway.models.users import User as BillingUser from gateway.services.dashboard_session_service import SESSION_COOKIE_NAME, hash_session_token from .conftest import MODEL_NAME diff --git a/tests/integration/test_pricing_config.py b/tests/integration/test_pricing_config.py index f6fcfc49fc..be6f65e0de 100644 --- a/tests/integration/test_pricing_config.py +++ b/tests/integration/test_pricing_config.py @@ -14,7 +14,7 @@ from gateway.core.config import API_ROOT, GatewayConfig, PricingConfig from gateway.db import ModelPricing, get_db from gateway.main import create_app -from gateway.models.entities import UsageLog +from gateway.models.usage import UsageLog from .conftest import build_async_session_override diff --git a/tests/integration/test_pricing_provenance_columns.py b/tests/integration/test_pricing_provenance_columns.py index e71dc34738..0c843dd2b0 100644 --- a/tests/integration/test_pricing_provenance_columns.py +++ b/tests/integration/test_pricing_provenance_columns.py @@ -26,7 +26,7 @@ from sqlalchemy.orm import Session from conftest import seed_workspace_id -from gateway.models.entities import UsageLog +from gateway.models.usage import UsageLog _RAN_AT = datetime(2026, 8, 1, 9, 30, tzinfo=UTC) # Later than _RAN_AT on purpose: an amount can be settled or repriced well after diff --git a/tests/integration/test_pricing_refresh_endpoint.py b/tests/integration/test_pricing_refresh_endpoint.py index b43b19513b..6afcf1d20a 100644 --- a/tests/integration/test_pricing_refresh_endpoint.py +++ b/tests/integration/test_pricing_refresh_endpoint.py @@ -101,7 +101,7 @@ def test_a_pending_update_is_previewed_without_fetching_and_accepting_it_is_reme db_session_factory: Callable[[], Session], ) -> None: """What the scheduled refresh leaves behind under the review policy.""" - from gateway.models.entities import PricingSnapshot + from gateway.models.pricing import PricingSnapshot from gateway.services.pricing_refresh_service import GENAI_PRICES_PENDING_SOURCE, reset_price_refresh_state session = db_session_factory() @@ -138,7 +138,7 @@ def test_the_history_keeps_only_the_newest_snapshots( monkeypatch: pytest.MonkeyPatch, ) -> None: """Each accept is the whole dataset, so the history is a window.""" - from gateway.models.entities import PricingSnapshot + from gateway.models.pricing import PricingSnapshot from gateway.services import pricing_refresh_service as refresh monkeypatch.setattr(refresh, "PRICING_SNAPSHOT_HISTORY_KEEP", 2) diff --git a/tests/integration/test_pricing_startup_warning.py b/tests/integration/test_pricing_startup_warning.py index 62ee3334be..9de9c27402 100644 --- a/tests/integration/test_pricing_startup_warning.py +++ b/tests/integration/test_pricing_startup_warning.py @@ -8,7 +8,7 @@ from gateway.core.config import GatewayConfig from gateway.log_config import logger as gateway_logger -from gateway.models.entities import ModelPricing +from gateway.models.pricing import ModelPricing from gateway.services.pricing_init_service import ( warn_if_require_pricing_without_pricing, warn_if_search_tools_lack_flat_pricing, diff --git a/tests/integration/test_provider_credentials_api.py b/tests/integration/test_provider_credentials_api.py index ff0519fcfa..78deb53663 100644 --- a/tests/integration/test_provider_credentials_api.py +++ b/tests/integration/test_provider_credentials_api.py @@ -13,7 +13,7 @@ from gateway.api.routes import providers as providers_route from gateway.core.config import API_ROOT -from gateway.models.entities import ProviderCredential +from gateway.models.providers import ProviderCredential from gateway.services.model_discovery_service import ProviderDiscovery from gateway.services.provider_store_service import reset_provider_cache from gateway.services.secret_box import decrypt_secret, generate_secret_key diff --git a/tests/integration/test_schema_metadata_parity.py b/tests/integration/test_schema_metadata_parity.py index 108569148e..01ba1d0acd 100644 --- a/tests/integration/test_schema_metadata_parity.py +++ b/tests/integration/test_schema_metadata_parity.py @@ -27,7 +27,7 @@ from sqlalchemy.engine import make_url import gateway.models # noqa: F401 # populates the metadata with every model module -from gateway.models.entities import Base +from gateway.models.base import Base if TYPE_CHECKING: from collections.abc import Generator diff --git a/tests/integration/test_scoped_budgets.py b/tests/integration/test_scoped_budgets.py index 221c032159..b841e6054e 100644 --- a/tests/integration/test_scoped_budgets.py +++ b/tests/integration/test_scoped_budgets.py @@ -19,9 +19,11 @@ from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine from gateway.core.config import API_ROOT -from gateway.models.entities import APIKey, Budget, ScopedBudget, User +from gateway.models.api_keys import APIKey +from gateway.models.budgets import Budget, ScopedBudget from gateway.models.tenancy import Organization, OrganizationMember, Workspace, WorkspaceMember from gateway.models.tenancy import User as TenancyUser +from gateway.models.users import User from gateway.services.budget_service import ( ReservationHandle, increase_reservation, diff --git a/tests/integration/test_session_model_catalog_scope.py b/tests/integration/test_session_model_catalog_scope.py index d9ae06ac50..3226771c43 100644 --- a/tests/integration/test_session_model_catalog_scope.py +++ b/tests/integration/test_session_model_catalog_scope.py @@ -29,9 +29,8 @@ from sqlalchemy.orm import Session from gateway.core.config import API_ROOT -from gateway.models.entities import DashboardSession from gateway.models.provider_keys import OrgProviderKey, WorkspaceProviderModelRestriction -from gateway.models.tenancy import Organization, OrganizationMember, User, Workspace, WorkspaceMember +from gateway.models.tenancy import DashboardSession, Organization, OrganizationMember, User, Workspace, WorkspaceMember from gateway.services.dashboard_session_service import SESSION_COOKIE_NAME, hash_session_token from gateway.services.secret_box import encrypt_secret, generate_secret_key diff --git a/tests/integration/test_streaming_precommit_refund.py b/tests/integration/test_streaming_precommit_refund.py index d91856cbb0..1b84350f38 100644 --- a/tests/integration/test_streaming_precommit_refund.py +++ b/tests/integration/test_streaming_precommit_refund.py @@ -21,7 +21,8 @@ from sqlalchemy.orm import Session from gateway.core.config import API_ROOT -from gateway.models.entities import UsageLog, User +from gateway.models.usage import UsageLog +from gateway.models.users import User from .conftest import MODEL_NAME diff --git a/tests/integration/test_streaming_ttft.py b/tests/integration/test_streaming_ttft.py index 445796a7ad..0ceae75aad 100644 --- a/tests/integration/test_streaming_ttft.py +++ b/tests/integration/test_streaming_ttft.py @@ -25,7 +25,7 @@ from sqlalchemy.orm import Session from gateway.core.config import API_ROOT -from gateway.models.entities import UsageLog +from gateway.models.usage import UsageLog from .conftest import MODEL_NAME diff --git a/tests/integration/test_tenancy_api.py b/tests/integration/test_tenancy_api.py index 5aeba882a8..c299f69e91 100644 --- a/tests/integration/test_tenancy_api.py +++ b/tests/integration/test_tenancy_api.py @@ -18,7 +18,7 @@ from sqlmodel import col from gateway.core.config import API_ROOT -from gateway.models.entities import RuntimeSetting +from gateway.models.platform import RuntimeSetting from gateway.models.tenancy import ( MAX_WORKSPACE_ASSIGNMENTS, Organization, diff --git a/tests/integration/test_tenancy_attribution.py b/tests/integration/test_tenancy_attribution.py index 5e8df55793..3b16cea13a 100644 --- a/tests/integration/test_tenancy_attribution.py +++ b/tests/integration/test_tenancy_attribution.py @@ -17,7 +17,7 @@ from sqlalchemy.orm import Session from gateway.core.config import API_ROOT -from gateway.models.entities import User as GatewayUser +from gateway.models.users import User as GatewayUser def _add_member(client: TestClient, headers: dict[str, str], email: str) -> dict[str, Any]: diff --git a/tests/integration/test_tenancy_races.py b/tests/integration/test_tenancy_races.py index 42242cd661..097c68cc8a 100644 --- a/tests/integration/test_tenancy_races.py +++ b/tests/integration/test_tenancy_races.py @@ -21,13 +21,14 @@ from gateway.auth.models import hash_key from gateway.core.config import GatewayConfig -from gateway.models.entities import APIKey, WorkspaceActivationState +from gateway.models.api_keys import APIKey from gateway.models.tenancy import ( ActiveOrganizationMemberCreateRequest, ActiveOrganizationMemberUpdateRequest, # noqa: E402 InviteOrganizationMemberRequest, Organization, User, + WorkspaceActivationState, WorkspaceAssignmentRequest, WorkspaceCreate, ) diff --git a/tests/integration/test_timezone_consistency.py b/tests/integration/test_timezone_consistency.py index 90142ff54b..40dd3eafd3 100644 --- a/tests/integration/test_timezone_consistency.py +++ b/tests/integration/test_timezone_consistency.py @@ -4,7 +4,7 @@ from sqlalchemy.ext.asyncio import AsyncSession from gateway.api.routes.chat import log_usage -from gateway.models.entities import UsageLog +from gateway.models.usage import UsageLog @pytest.mark.asyncio diff --git a/tests/integration/test_tool_billing_settlement.py b/tests/integration/test_tool_billing_settlement.py index 01aaa9a1d4..4616de0bff 100644 --- a/tests/integration/test_tool_billing_settlement.py +++ b/tests/integration/test_tool_billing_settlement.py @@ -27,7 +27,8 @@ from sqlalchemy.orm import Session from gateway.core.config import API_KEY_HEADER, API_ROOT -from gateway.models.entities import UsageLog, User +from gateway.models.usage import UsageLog +from gateway.models.users import User from gateway.services.tool_usage import TOOL_METER_NAMESPACE from .conftest import MODEL_NAME diff --git a/tests/integration/test_tool_settings_tenant_read.py b/tests/integration/test_tool_settings_tenant_read.py index 95f9875211..453b0b4680 100644 --- a/tests/integration/test_tool_settings_tenant_read.py +++ b/tests/integration/test_tool_settings_tenant_read.py @@ -24,8 +24,7 @@ from sqlalchemy.orm import Session from gateway.core.config import API_ROOT -from gateway.models.entities import DashboardSession -from gateway.models.tenancy import Organization, OrganizationMember, User +from gateway.models.tenancy import DashboardSession, Organization, OrganizationMember, User from gateway.services.dashboard_session_service import SESSION_COOKIE_NAME, hash_session_token _PATH = f"{API_ROOT}/tool-settings" diff --git a/tests/integration/test_transaction_safety.py b/tests/integration/test_transaction_safety.py index 75f89f8227..4e3e2321fe 100644 --- a/tests/integration/test_transaction_safety.py +++ b/tests/integration/test_transaction_safety.py @@ -9,7 +9,8 @@ from sqlalchemy.ext.asyncio import AsyncSession from gateway.core.config import API_KEY_HEADER, API_ROOT -from gateway.models.entities import Budget, User +from gateway.models.budgets import Budget +from gateway.models.users import User from gateway.services.budget_service import _cas_reset_user_budget, _is_model_free diff --git a/tests/integration/test_usage_admin.py b/tests/integration/test_usage_admin.py index db6c8f80ae..2ee6f45d4f 100644 --- a/tests/integration/test_usage_admin.py +++ b/tests/integration/test_usage_admin.py @@ -16,7 +16,8 @@ from gateway.core.config import API_ROOT from gateway.core.sql import MAX_FILTER_VALUES from gateway.core.usage_source import SERVED_HERE_SLUG, SERVED_HERE_SOURCES -from gateway.models.entities import UsageLog, User +from gateway.models.usage import UsageLog +from gateway.models.users import User DELETE_PATH = f"{API_ROOT}/usage" SET_PRICE_PATH = f"{API_ROOT}/usage/set-price" diff --git a/tests/integration/test_usage_endpoint.py b/tests/integration/test_usage_endpoint.py index 6a701fce16..7812b27b7c 100644 --- a/tests/integration/test_usage_endpoint.py +++ b/tests/integration/test_usage_endpoint.py @@ -10,7 +10,9 @@ from conftest import seed_workspace_id from gateway.core.config import API_ROOT -from gateway.models.entities import APIKey, UsageLog, User +from gateway.models.api_keys import APIKey +from gateway.models.usage import UsageLog +from gateway.models.users import User USAGE_PATH = f"{API_ROOT}/usage" diff --git a/tests/integration/test_usage_summary.py b/tests/integration/test_usage_summary.py index e73c7e3806..183a92e120 100644 --- a/tests/integration/test_usage_summary.py +++ b/tests/integration/test_usage_summary.py @@ -19,7 +19,9 @@ from conftest import seed_workspace_id from gateway.core.config import API_ROOT from gateway.core.sql import MAX_FILTER_VALUES -from gateway.models.entities import APIKey, UsageLog, User +from gateway.models.api_keys import APIKey +from gateway.models.usage import UsageLog +from gateway.models.users import User SUMMARY_PATH = f"{API_ROOT}/usage/summary" SERIES_PATH = f"{API_ROOT}/usage/series" diff --git a/tests/integration/test_usage_tracking.py b/tests/integration/test_usage_tracking.py index a80e985337..c3e1911dd4 100644 --- a/tests/integration/test_usage_tracking.py +++ b/tests/integration/test_usage_tracking.py @@ -11,7 +11,8 @@ from sqlalchemy.orm import Session from gateway.core.config import API_ROOT -from gateway.models.entities import UsageLog, User +from gateway.models.usage import UsageLog +from gateway.models.users import User from .conftest import MODEL_NAME diff --git a/tests/integration/test_user_delete_preserve_logs.py b/tests/integration/test_user_delete_preserve_logs.py index 7790328047..e10fad7982 100644 --- a/tests/integration/test_user_delete_preserve_logs.py +++ b/tests/integration/test_user_delete_preserve_logs.py @@ -11,7 +11,10 @@ from gateway.adapters.telemetry_storage_adapter import DatabaseTelemetryStorageAdapter from gateway.core.config import API_KEY_HEADER, API_ROOT -from gateway.models.entities import APIKey, BudgetResetLog, UsageLog, User +from gateway.models.api_keys import APIKey +from gateway.models.budgets import BudgetResetLog +from gateway.models.usage import UsageLog +from gateway.models.users import User from .conftest import MODEL_NAME diff --git a/tests/integration/test_users_organization_scope.py b/tests/integration/test_users_organization_scope.py index 1975bb7959..cfcf5044d0 100644 --- a/tests/integration/test_users_organization_scope.py +++ b/tests/integration/test_users_organization_scope.py @@ -31,9 +31,11 @@ from sqlalchemy.orm import Session from gateway.core.config import API_ROOT -from gateway.models.entities import APIKey, DashboardSession, UsageLog, User -from gateway.models.tenancy import Organization, OrganizationMember, Workspace +from gateway.models.api_keys import APIKey +from gateway.models.tenancy import DashboardSession, Organization, OrganizationMember, Workspace from gateway.models.tenancy import User as TenancyUser +from gateway.models.usage import UsageLog +from gateway.models.users import User from gateway.services.dashboard_session_service import SESSION_COOKIE_NAME, hash_session_token # The request-plane ids this suite reasons about. Each names the join that puts diff --git a/tests/integration/test_web_search_config_enforcement.py b/tests/integration/test_web_search_config_enforcement.py index 1cbf630d82..74d9a93e12 100644 --- a/tests/integration/test_web_search_config_enforcement.py +++ b/tests/integration/test_web_search_config_enforcement.py @@ -29,7 +29,7 @@ from sqlalchemy.orm import Session from gateway.core.config import API_KEY_HEADER, API_ROOT -from gateway.models.entities import WorkspaceWebSearchConfig +from gateway.models.tools import WorkspaceWebSearchConfig _SEARCH_URL = "http://127.0.0.1:9998/search" _REQUEST = { diff --git a/tests/integration/test_workspace_activation.py b/tests/integration/test_workspace_activation.py index 6cfd1c2cb7..db029b814a 100644 --- a/tests/integration/test_workspace_activation.py +++ b/tests/integration/test_workspace_activation.py @@ -18,8 +18,9 @@ from gateway.auth.models import hash_key from gateway.core.config import GatewayConfig -from gateway.models.entities import APIKey, UsageLog, WorkspaceActivationState -from gateway.models.tenancy import Organization, User, Workspace +from gateway.models.api_keys import APIKey +from gateway.models.tenancy import Organization, User, Workspace, WorkspaceActivationState +from gateway.models.usage import UsageLog from gateway.repositories.tenancy import ( OrganizationMemberRepository, OrganizationRepository, diff --git a/tests/integration/test_workspace_code_execution_policy.py b/tests/integration/test_workspace_code_execution_policy.py index 2ff6c619ac..816e78de8b 100644 --- a/tests/integration/test_workspace_code_execution_policy.py +++ b/tests/integration/test_workspace_code_execution_policy.py @@ -16,8 +16,8 @@ from pydantic import ValidationError from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine -from gateway.models.entities import WorkspaceCodeExecutionPolicy from gateway.models.tenancy import Organization, User, Workspace +from gateway.models.tools import WorkspaceCodeExecutionPolicy from gateway.repositories.tenancy import ( OrganizationMemberRepository, OrganizationRepository, diff --git a/tests/integration/test_workspace_mcp_servers.py b/tests/integration/test_workspace_mcp_servers.py index 6f3c417f14..905155f562 100644 --- a/tests/integration/test_workspace_mcp_servers.py +++ b/tests/integration/test_workspace_mcp_servers.py @@ -28,9 +28,9 @@ from gateway.api.routes._pipeline import RequestContext, prepare_gateway_tools from gateway.api.routes.chat import ChatCompletionRequest from gateway.core.config import GatewayConfig -from gateway.models.entities import WorkspaceMcpServer from gateway.models.mcp import MAX_MCP_SERVER_IDS, McpServerConfig from gateway.models.tenancy import Organization, User, Workspace +from gateway.models.tools import WorkspaceMcpServer from gateway.repositories.tenancy import ( OrganizationMemberRepository, OrganizationRepository, diff --git a/tests/integration/test_workspace_member_budget_policies.py b/tests/integration/test_workspace_member_budget_policies.py index 3f2f096d47..f55fb191d5 100644 --- a/tests/integration/test_workspace_member_budget_policies.py +++ b/tests/integration/test_workspace_member_budget_policies.py @@ -18,7 +18,7 @@ from sqlalchemy import select from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine -from gateway.models.entities import Budget, ScopedBudget, WorkspaceBudgetDefault +from gateway.models.budgets import Budget, ScopedBudget, WorkspaceBudgetDefault from gateway.models.money import as_float from gateway.models.tenancy import ( ActiveOrganizationMemberCreateRequest, diff --git a/tests/integration/test_workspace_scope.py b/tests/integration/test_workspace_scope.py index 381f853872..0c0fc772e6 100644 --- a/tests/integration/test_workspace_scope.py +++ b/tests/integration/test_workspace_scope.py @@ -14,7 +14,8 @@ from sqlalchemy.orm import Session from gateway.core.config import API_ROOT -from gateway.models.entities import APIKey, UsageLog +from gateway.models.api_keys import APIKey +from gateway.models.usage import UsageLog def _default_workspace(client: TestClient, headers: dict[str, str]) -> str: diff --git a/tests/integration/test_workspace_web_search.py b/tests/integration/test_workspace_web_search.py index cf228a74c5..49722e6b2e 100644 --- a/tests/integration/test_workspace_web_search.py +++ b/tests/integration/test_workspace_web_search.py @@ -15,8 +15,8 @@ import pytest_asyncio from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine -from gateway.models.entities import WorkspaceWebSearchConfig from gateway.models.tenancy import Organization, User, Workspace +from gateway.models.tools import WorkspaceWebSearchConfig from gateway.repositories.tenancy import ( OrganizationMemberRepository, OrganizationRepository, diff --git a/tests/unit/test_agent_telemetry_ingest.py b/tests/unit/test_agent_telemetry_ingest.py index 84c64b525d..fdf64b6f4f 100644 --- a/tests/unit/test_agent_telemetry_ingest.py +++ b/tests/unit/test_agent_telemetry_ingest.py @@ -13,7 +13,7 @@ import pytest from sqlalchemy.ext.asyncio import AsyncSession -from gateway.models.entities import APIKey +from gateway.models.api_keys import APIKey from gateway.ports.telemetry_storage_port import IngestResult, TelemetryRecord, TelemetryStoragePort from gateway.services.agent_telemetry_service import ingest diff --git a/tests/unit/test_budget_cost_estimate.py b/tests/unit/test_budget_cost_estimate.py index 7febd793d8..077d4c2578 100644 --- a/tests/unit/test_budget_cost_estimate.py +++ b/tests/unit/test_budget_cost_estimate.py @@ -9,7 +9,8 @@ from decimal import Decimal from typing import Any -from gateway.models.entities import MAX_COUNT_LIMIT, Budget, ModelPricing +from gateway.models.budgets import MAX_COUNT_LIMIT, Budget +from gateway.models.pricing import ModelPricing from gateway.services.budget_service import _blocked_axis, estimate_cost, estimate_tokens diff --git a/tests/unit/test_check_architecture.py b/tests/unit/test_check_architecture.py index ee5cb7d39a..63e5b5f6aa 100644 --- a/tests/unit/test_check_architecture.py +++ b/tests/unit/test_check_architecture.py @@ -31,7 +31,7 @@ def _write(src_root: Path, relative_path: str, content: str) -> Path: def test_service_importing_models_is_clean(tmp_path: Path) -> None: - file_path = _write(tmp_path, "gateway/services/thing.py", "from gateway.models.entities import User\n") + file_path = _write(tmp_path, "gateway/services/thing.py", "from gateway.models.users import User\n") assert check.check_file(file_path, tmp_path) == [] diff --git a/tests/unit/test_compute_cost.py b/tests/unit/test_compute_cost.py index a363069e81..a424574541 100644 --- a/tests/unit/test_compute_cost.py +++ b/tests/unit/test_compute_cost.py @@ -32,7 +32,7 @@ from gateway.api.routes._pipeline import _compute_cost from gateway.core.config import PricingConfig, PricingTierConfig from gateway.core.usage import GatewayUsage -from gateway.models.entities import ModelPricing +from gateway.models.pricing import ModelPricing def _usd(tokens: int, rate_per_million: str) -> Decimal: diff --git a/tests/unit/test_content_normalizer.py b/tests/unit/test_content_normalizer.py index ce947fb309..2addd326fd 100644 --- a/tests/unit/test_content_normalizer.py +++ b/tests/unit/test_content_normalizer.py @@ -13,7 +13,7 @@ from any_llm.types.completion import CompletionUsage from gateway.core.config import GatewayConfig -from gateway.models.entities import FileObject +from gateway.models.tools import FileObject from gateway.services import content_normalizer as cn from gateway.services.content_normalizer import normalize_messages from gateway.services.file_extractors import ExtractionResult diff --git a/tests/unit/test_dashboard_session.py b/tests/unit/test_dashboard_session.py index 27a176b110..cf9c3ab622 100644 --- a/tests/unit/test_dashboard_session.py +++ b/tests/unit/test_dashboard_session.py @@ -25,7 +25,7 @@ from gateway.api.routes import auth_session as auth_session_route from gateway.core.config import API_ROOT, GatewayConfig from gateway.main import create_app -from gateway.models.entities import DashboardSession +from gateway.models.tenancy import DashboardSession from gateway.services import dashboard_session_service, master_key_service from gateway.services.dashboard_session_service import SESSION_COOKIE_NAME from gateway.services.tenancy.provisioning_service import BOOTSTRAP_IDENTITY_KEY diff --git a/tests/unit/test_exact_money_schema_chain.py b/tests/unit/test_exact_money_schema_chain.py index 29015bb45e..a11ac5a62d 100644 --- a/tests/unit/test_exact_money_schema_chain.py +++ b/tests/unit/test_exact_money_schema_chain.py @@ -29,7 +29,8 @@ import gateway.models # noqa: F401 (registers every table on the shared metadata) from gateway.core.metered_pricing import COST_QUANTUM, calculate_token_cost -from gateway.models.entities import ModelPricing, UsageLog +from gateway.models.pricing import ModelPricing +from gateway.models.usage import UsageLog from gateway.services.pricing_service import default_model_pricing _ALEMBIC_DIR = Path(__file__).resolve().parents[2] / "alembic" diff --git a/tests/unit/test_gateway_rejection_logging_best_effort.py b/tests/unit/test_gateway_rejection_logging_best_effort.py index a394b0b24c..b500884e1c 100644 --- a/tests/unit/test_gateway_rejection_logging_best_effort.py +++ b/tests/unit/test_gateway_rejection_logging_best_effort.py @@ -15,7 +15,7 @@ import pytest from gateway.api.routes._pipeline import log_gateway_rejection -from gateway.models.entities import UsageLog +from gateway.models.usage import UsageLog class _BoomWriter: diff --git a/tests/unit/test_key_fingerprint_schemas.py b/tests/unit/test_key_fingerprint_schemas.py index b70715d67d..463692789f 100644 --- a/tests/unit/test_key_fingerprint_schemas.py +++ b/tests/unit/test_key_fingerprint_schemas.py @@ -11,7 +11,7 @@ from datetime import UTC, datetime from gateway.api.routes.keys import CreateKeyResponse, KeyInfo -from gateway.models.entities import APIKey +from gateway.models.api_keys import APIKey from gateway.services.tenancy.workspace_activation_service import ActivationApiKeyPublic diff --git a/tests/unit/test_knn_router.py b/tests/unit/test_knn_router.py index 1bdf62c2f3..6c9fc9850a 100644 --- a/tests/unit/test_knn_router.py +++ b/tests/unit/test_knn_router.py @@ -28,7 +28,7 @@ from gateway.api.routes._helpers import conversation_opening_text, first_user_text, latest_user_text from gateway.core.config import GatewayConfig -from gateway.models.entities import RoutingMemory +from gateway.models.routing import RoutingMemory from gateway.services import alias_service from gateway.services.routing import knn from gateway.services.routing.backends import RoutingContext diff --git a/tests/unit/test_log_writer.py b/tests/unit/test_log_writer.py index 1f00d9b358..f96dc13ee1 100644 --- a/tests/unit/test_log_writer.py +++ b/tests/unit/test_log_writer.py @@ -5,7 +5,7 @@ import pytest from sqlalchemy.exc import SQLAlchemyError -from gateway.models.entities import UsageLog +from gateway.models.usage import UsageLog from gateway.services.log_writer import SingleLogWriter diff --git a/tests/unit/test_master_key_service.py b/tests/unit/test_master_key_service.py index 7f2879dfb0..c9eba157e6 100644 --- a/tests/unit/test_master_key_service.py +++ b/tests/unit/test_master_key_service.py @@ -14,7 +14,7 @@ from gateway.api import deps from gateway.core.config import API_ROOT, GatewayConfig from gateway.main import create_app -from gateway.models.entities import RuntimeSetting +from gateway.models.platform import RuntimeSetting from gateway.services import master_key_service from gateway.services.master_key_service import ( MASTER_KEY_HASH_KEY, diff --git a/tests/unit/test_models_default_pricing.py b/tests/unit/test_models_default_pricing.py index ff31c608dc..43801b5978 100644 --- a/tests/unit/test_models_default_pricing.py +++ b/tests/unit/test_models_default_pricing.py @@ -6,7 +6,7 @@ import pytest from gateway.core.config import GatewayConfig -from gateway.models.entities import ModelPricing +from gateway.models.pricing import ModelPricing from gateway.services.merged_catalog_service import ModelObject, ModelPricingInfo, alias_model, apply_default_pricing from gateway.services.pricing_service import configure_default_pricing, configure_provider_types diff --git a/tests/unit/test_organization_pricing_resolution.py b/tests/unit/test_organization_pricing_resolution.py index 42d26266a7..c924b5a7a0 100644 --- a/tests/unit/test_organization_pricing_resolution.py +++ b/tests/unit/test_organization_pricing_resolution.py @@ -24,7 +24,7 @@ import gateway.models # noqa: F401 (registers every table on the shared metadata) from gateway.core.config import GatewayConfig -from gateway.models.entities import ModelPricing, OrganizationModelPricing +from gateway.models.pricing import ModelPricing, OrganizationModelPricing from gateway.models.tenancy import Organization from gateway.services.external_usage_service import _load_pricing_index, _resolve_pricing from gateway.services.organization_pricing_service import ( diff --git a/tests/unit/test_password_reset.py b/tests/unit/test_password_reset.py index 147087d70c..db5692fd41 100644 --- a/tests/unit/test_password_reset.py +++ b/tests/unit/test_password_reset.py @@ -20,7 +20,7 @@ from gateway.core.config import API_ROOT, GatewayConfig from gateway.log_config import logger as gateway_logger from gateway.main import create_app -from gateway.models.entities import DashboardSession +from gateway.models.tenancy import DashboardSession MASTER_KEY = "sk-test-master" PASSWORD = "a-real-password" # pragma: allowlist secret diff --git a/tests/unit/test_password_sign_in.py b/tests/unit/test_password_sign_in.py index daa0dc1701..845b52a8d8 100644 --- a/tests/unit/test_password_sign_in.py +++ b/tests/unit/test_password_sign_in.py @@ -26,7 +26,7 @@ from gateway.core.config import API_ROOT, GatewayConfig from gateway.log_config import logger as gateway_logger from gateway.main import create_app -from gateway.models.entities import DashboardSession +from gateway.models.tenancy import DashboardSession from gateway.services.dashboard_session_service import SESSION_COOKIE_NAME from gateway.services.password_service import MAX_PASSWORD_BYTES, MIN_PASSWORD_LENGTH, hash_password from gateway.services.tenancy.provisioning_service import BOOTSTRAP_IDENTITY_KEY diff --git a/tests/unit/test_pricing_provenance_schema_chain.py b/tests/unit/test_pricing_provenance_schema_chain.py index 30880ce6a8..d2dbac02e7 100644 --- a/tests/unit/test_pricing_provenance_schema_chain.py +++ b/tests/unit/test_pricing_provenance_schema_chain.py @@ -28,7 +28,7 @@ from sqlmodel import SQLModel import gateway.models # noqa: F401 (registers every table on the shared metadata) -from gateway.models.entities import UsageLog +from gateway.models.usage import UsageLog _ALEMBIC_DIR = Path(__file__).resolve().parents[2] / "alembic" _PROVENANCE_REVISION = "a9c4e2b6d8f1" diff --git a/tests/unit/test_pricing_refresh_service.py b/tests/unit/test_pricing_refresh_service.py index 1a962ad923..f9b0ae1e32 100644 --- a/tests/unit/test_pricing_refresh_service.py +++ b/tests/unit/test_pricing_refresh_service.py @@ -11,7 +11,7 @@ from sqlalchemy.ext.asyncio import AsyncSession import gateway.services.pricing_refresh_service as pricing_refresh_service -from gateway.models.entities import PricingSnapshot, PricingSnapshotHistory +from gateway.models.pricing import PricingSnapshot, PricingSnapshotHistory _PERSISTED_SNAPSHOT = ( '[{"id":"test","name":"Test","api_pattern":"","models":[' diff --git a/tests/unit/test_pricing_snapshot_history_schema_chain.py b/tests/unit/test_pricing_snapshot_history_schema_chain.py index 0d8a125ba4..bb250f7447 100644 --- a/tests/unit/test_pricing_snapshot_history_schema_chain.py +++ b/tests/unit/test_pricing_snapshot_history_schema_chain.py @@ -13,7 +13,7 @@ from sqlmodel import SQLModel import gateway.models # noqa: F401 (registers every table on the shared metadata) -from gateway.models.entities import PricingSnapshotHistory +from gateway.models.pricing import PricingSnapshotHistory _ALEMBIC_DIR = Path(__file__).resolve().parents[2] / "alembic" _REVISION = "b2d4f6a8c0e2" diff --git a/tests/unit/test_pricing_unit_conventions.py b/tests/unit/test_pricing_unit_conventions.py index 2770b9fa21..a71d21cebf 100644 --- a/tests/unit/test_pricing_unit_conventions.py +++ b/tests/unit/test_pricing_unit_conventions.py @@ -11,7 +11,7 @@ import pytest -from gateway.models.entities import ModelPricing +from gateway.models.pricing import ModelPricing from gateway.services.pricing_service import flat_request_cost, input_token_cost, per_image_cost diff --git a/tests/unit/test_pricing_unit_origin_schema_chain.py b/tests/unit/test_pricing_unit_origin_schema_chain.py index b21af59443..1fcf0076e8 100644 --- a/tests/unit/test_pricing_unit_origin_schema_chain.py +++ b/tests/unit/test_pricing_unit_origin_schema_chain.py @@ -22,7 +22,7 @@ from sqlmodel import SQLModel import gateway.models # noqa: F401 (registers every table on the shared metadata) -from gateway.models.entities import ModelPricing +from gateway.models.pricing import ModelPricing _ALEMBIC_DIR = Path(__file__).resolve().parents[2] / "alembic" _REVISION = "c7e9a1b3d5f7" diff --git a/tests/unit/test_provider_store_service.py b/tests/unit/test_provider_store_service.py index 55df081a56..962d1cbd82 100644 --- a/tests/unit/test_provider_store_service.py +++ b/tests/unit/test_provider_store_service.py @@ -7,7 +7,7 @@ import pytest from gateway.core.config import GatewayConfig -from gateway.models.entities import ProviderCredential +from gateway.models.providers import ProviderCredential from gateway.services import provider_store_service as store from gateway.services.provider_store_service import apply_to_config, reset_provider_cache from gateway.services.secret_box import ( diff --git a/tests/unit/test_runtime_settings_service.py b/tests/unit/test_runtime_settings_service.py index ba3e9e05bd..086452a193 100644 --- a/tests/unit/test_runtime_settings_service.py +++ b/tests/unit/test_runtime_settings_service.py @@ -7,7 +7,8 @@ from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine from gateway.core.config import GatewayConfig -from gateway.models.entities import Base, RuntimeSetting +from gateway.models.base import Base +from gateway.models.platform import RuntimeSetting from gateway.services.pricing_service import configure_default_pricing, default_pricing_enabled from gateway.services.runtime_settings_service import ( DEFAULT_PRICING, diff --git a/tests/unit/test_search_tool_store_service.py b/tests/unit/test_search_tool_store_service.py index 6c92a42abf..fbde5fb231 100644 --- a/tests/unit/test_search_tool_store_service.py +++ b/tests/unit/test_search_tool_store_service.py @@ -7,7 +7,7 @@ import pytest from gateway.core.config import GatewayConfig -from gateway.models.entities import SearchToolCredential +from gateway.models.tools import SearchToolCredential from gateway.services import search_tool_store_service as store from gateway.services.search_backend import resolve_search_tool from gateway.services.search_tool_store_service import ( diff --git a/tests/unit/test_secret_fields.py b/tests/unit/test_secret_fields.py index 3230663dd7..2dbee74c2f 100644 --- a/tests/unit/test_secret_fields.py +++ b/tests/unit/test_secret_fields.py @@ -19,12 +19,14 @@ import uuid from datetime import UTC, datetime -from gateway.models.entities import OrganizationGuardrail, ProviderCredential, SearchToolCredential +from gateway.models.guardrails import OrganizationGuardrail +from gateway.models.providers import ProviderCredential from gateway.models.secret_fields import ( REDACTED_VALUE, redact_secret_like_values, restore_redacted_values, ) +from gateway.models.tools import SearchToolCredential from gateway.services.tenancy.organization_guardrail_service import OrganizationGuardrailPublic diff --git a/tests/unit/test_tenancy_schema_chain.py b/tests/unit/test_tenancy_schema_chain.py index a472279b44..6aefd87eea 100644 --- a/tests/unit/test_tenancy_schema_chain.py +++ b/tests/unit/test_tenancy_schema_chain.py @@ -39,7 +39,7 @@ from sqlmodel import SQLModel import gateway.models # noqa: F401 (registers every table on the shared metadata) -from gateway.models.tenancy import UtcDateTime +from gateway.models.base import UtcDateTime _ALEMBIC_DIR = Path(__file__).resolve().parents[2] / "alembic" _TENANCY_REVISION = "c4b6d8e0f2a3" @@ -192,7 +192,7 @@ def test_upgrade_downgrade_upgrade_round_trips(sqlite_at_head: tuple[Config, Eng def test_naming_one_model_module_registers_them_all() -> None: """``Base.metadata`` is whole however few model modules the caller imported. - ``alembic/env.py`` names only ``gateway.models.entities`` and relies on the + ``alembic/env.py`` names only ``gateway.models.base`` and relies on the package ``__init__`` to pull in the rest. If that import chain breaks, the metadata silently loses the tenancy tables and autogenerate proposes ``DROP TABLE`` for them, which is data-loss-class and invisible until @@ -200,7 +200,7 @@ def test_naming_one_model_module_registers_them_all() -> None: a test runs in this one every model module is already imported. """ source = ( - "from gateway.models.entities import Base;" + "from gateway.models.base import Base;" "import json,sys;" "sys.stdout.write(json.dumps(sorted(Base.metadata.tables)))" ) diff --git a/tests/unit/test_usage_cache_tokens.py b/tests/unit/test_usage_cache_tokens.py index 4df32d6e4b..c9f2eb9ee4 100644 --- a/tests/unit/test_usage_cache_tokens.py +++ b/tests/unit/test_usage_cache_tokens.py @@ -16,7 +16,7 @@ from gateway.api.routes.messages import _messages_stream_usage, _MessagesAdapter, _requested_cache_write_ttl from gateway.api.routes.responses import _usage_to_completion_usage from gateway.core.usage import GatewayUsage, cache_read_tokens_of, cache_write_1h_tokens_of, cache_write_tokens_of -from gateway.models.entities import UsageLog +from gateway.models.usage import UsageLog from gateway.services.usage_admin_service import _row_cache_tokens_included diff --git a/tests/unit/test_usage_log_metering_pool.py b/tests/unit/test_usage_log_metering_pool.py index 1f984ef4e8..cc3033f62d 100644 --- a/tests/unit/test_usage_log_metering_pool.py +++ b/tests/unit/test_usage_log_metering_pool.py @@ -16,7 +16,7 @@ from gateway.core.config import GatewayConfig from gateway.core.database import init_db, reset_db -from gateway.models.entities import UsageLog +from gateway.models.usage import UsageLog from gateway.services import log_writer as log_writer_module from gateway.services.log_writer import BatchLogWriter, SingleLogWriter diff --git a/web/src/client/schema.ts b/web/src/client/schema.ts index 18d3af746b..834043ed6d 100644 --- a/web/src/client/schema.ts +++ b/web/src/client/schema.ts @@ -8654,7 +8654,7 @@ export interface components { * * The plaintext key is never stored as sent: the service encrypts it * (`services/secret_box.py`) and keeps only the ciphertext and ``last4``, - * the same convention `entities.ProviderCredential` already uses. + * the same convention `providers.ProviderCredential` already uses. */ OrgProviderKeyCreateRequest: { /** Api Base */ @@ -8957,7 +8957,7 @@ export interface components { * * ``credential`` is never stored as sent: it is encrypted with * ``OTARI_SECRET_KEY`` and only the ciphertext is kept, the same convention - * `entities.WorkspaceMcpServer` and `entities.ProviderCredential` use. It is + * `tools.WorkspaceMcpServer` and `providers.ProviderCredential` use. It is * sent to the endpoint as ``Authorization: Bearer`` when the guardrail runs, * so it authenticates this gateway to the guardrails service the entry names. * A guardrail *vendor's* own key is not this: the guardrails service builds @@ -12434,7 +12434,7 @@ export interface components { * * ``authorization_token`` is never stored as sent: it is encrypted with * ``OTARI_SECRET_KEY`` and only the ciphertext is kept, the same convention - * `entities.ProviderCredential` and `OrgProviderKey` already use. + * `providers.ProviderCredential` and `OrgProviderKey` already use. */ WorkspaceMcpServerCreate: { /**