diff --git a/backend/api/ai_settings.py b/backend/api/ai_settings.py index 10a6293..37915ec 100644 --- a/backend/api/ai_settings.py +++ b/backend/api/ai_settings.py @@ -14,6 +14,8 @@ class AIProviderSettingsRequest(BaseModel): clearApiKey: bool = False adminApiKey: str | None = None clearAdminApiKey: bool = False + providerProjectId: str = "" + providerApiKeyId: str = "" keyPolicy: str | None = None assistMode: str = "auto" defaultModel: str = "" diff --git a/backend/services/ai_cost_reconciliation.py b/backend/services/ai_cost_reconciliation.py index ed9c459..251927e 100644 --- a/backend/services/ai_cost_reconciliation.py +++ b/backend/services/ai_cost_reconciliation.py @@ -18,6 +18,11 @@ class UsageCostEvent: provider: str model_name: str estimated_cost_usd: Decimal + input_tokens: int = 0 + cached_input_tokens: int = 0 + output_tokens: int = 0 + provider_project_id: str = "" + provider_api_key_id: str = "" @dataclass(frozen=True) diff --git a/backend/services/openai_assist_accounting.py b/backend/services/openai_assist_accounting.py index 3740ba4..14f60f8 100644 --- a/backend/services/openai_assist_accounting.py +++ b/backend/services/openai_assist_accounting.py @@ -32,6 +32,8 @@ def _assist_event_base( "decision_reason": decision_reason, "paid_by": settings["paidBy"], "provider_key_owner_user_id": settings.get("providerKeyOwnerUserId"), + "provider_project_id": settings.get("providerProjectId") or "", + "provider_api_key_id": settings.get("providerApiKeyId") or "", } def _estimate_openai_cost(self, settings: dict[str, Any], usage: dict[str, int]) -> float: diff --git a/db/ai_provider_store.py b/db/ai_provider_store.py index ebf5412..83af54b 100644 --- a/db/ai_provider_store.py +++ b/db/ai_provider_store.py @@ -326,6 +326,8 @@ def record_ai_assist_event( estimated_cost: float = 0.0, paid_by: str = "unknown", provider_key_owner_user_id: int | None = None, + provider_project_id: str = "", + provider_api_key_id: str = "", success: bool = True, error_message: str | None = None, decision_reason: str = "", @@ -351,6 +353,8 @@ def record_ai_assist_event( float(estimated_cost or 0), paid_by, provider_key_owner_user_id, + str(provider_project_id or "").strip(), + str(provider_api_key_id or "").strip(), bool(success), str(error_message or "")[:1000] if error_message else None, str(decision_reason or "")[:1000], @@ -359,9 +363,21 @@ def record_ai_assist_event( ).fetchone() return int(row["id"]) if row else None - def list_openai_events_for_reconciliation(self, *, start_time: datetime, end_time: datetime) -> list[UsageCostEvent]: + def list_openai_events_for_reconciliation( + self, + *, + start_time: datetime, + end_time: datetime, + project_ids: list[str] | None = None, + api_key_ids: list[str] | None = None, + ) -> list[UsageCostEvent]: + project_filter = self._filter_or_none(project_ids) + api_key_filter = self._filter_or_none(api_key_ids) with self.database.connection() as conn: - rows = conn.execute(load_query("ai_assist_events_for_reconciliation.sql"), (start_time, end_time)).fetchall() + rows = conn.execute( + load_query("ai_assist_events_for_reconciliation.sql"), + (start_time, end_time, project_filter, project_filter, api_key_filter, api_key_filter), + ).fetchall() return [ UsageCostEvent( event_id=int(row["id"]), @@ -369,10 +385,79 @@ def list_openai_events_for_reconciliation(self, *, start_time: datetime, end_tim provider=row["provider"], model_name=row["model_name"] or "", estimated_cost_usd=Decimal(str(row["estimated_cost"] or "0")), + input_tokens=int(row.get("input_tokens") or 0), + cached_input_tokens=int(row.get("cached_input_tokens") or 0), + output_tokens=int(row.get("output_tokens") or 0), + provider_project_id=row.get("provider_project_id") or "", + provider_api_key_id=row.get("provider_api_key_id") or "", ) for row in rows ] + def reprice_openai_events_for_reconciliation( + self, + *, + start_time: datetime, + end_time: datetime, + project_ids: list[str] | None = None, + api_key_ids: list[str] | None = None, + ) -> int: + project_filter = self._filter_or_none(project_ids) + api_key_filter = self._filter_or_none(api_key_ids) + with self.database.connection() as conn: + rows = conn.execute( + load_query("ai_assist_events_for_reconciliation.sql"), + (start_time, end_time, project_filter, project_filter, api_key_filter, api_key_filter), + ).fetchall() + update_query = load_query("ai_assist_event_estimated_cost_update.sql") + updated = 0 + for row in rows: + paid_by = row.get("paid_by") or "system" + user_id = row.get("provider_key_owner_user_id") if paid_by == "user" else row.get("user_id") + estimate = self.estimate_cost( + provider="openai", + model_name=row.get("model_name") or "", + input_tokens=int(row.get("input_tokens") or 0), + cached_input_tokens=int(row.get("cached_input_tokens") or 0), + output_tokens=int(row.get("output_tokens") or 0), + billing_scope=paid_by, + entity_id=row.get("entity_id") if paid_by == "entity" else None, + user_id=user_id if paid_by == "user" else None, + ) + conn.execute(update_query, (Decimal(str(estimate)), row["id"])) + updated += 1 + return updated + + def clear_openai_reconciliation_window( + self, + *, + start_time: datetime, + end_time: datetime, + project_ids: list[str] | None = None, + api_key_ids: list[str] | None = None, + ) -> None: + project_filter = self._filter_or_none(project_ids) + api_key_filter = self._filter_or_none(api_key_ids) + with self.database.connection() as conn: + conn.execute( + load_query("ai_assist_event_reconciliation_clear.sql"), + (start_time, end_time, project_filter, project_filter, api_key_filter, api_key_filter), + ) + + def openai_reconciliation_filters(self) -> dict[str, list[str]]: + project_ids: list[str] = [] + api_key_ids: list[str] = [] + with self.database.connection() as conn: + rows = conn.execute(load_query("openai_reconciliation_configured_filters.sql")).fetchall() + for row in rows: + project_id = str(row.get("provider_project_id") or "").strip() + api_key_id = str(row.get("provider_api_key_id") or "").strip() + if project_id and project_id not in project_ids: + project_ids.append(project_id) + if api_key_id and api_key_id not in api_key_ids: + api_key_ids.append(api_key_id) + return {"projectIds": project_ids, "apiKeyIds": api_key_ids} + def create_cost_reconciliation_run( self, *, @@ -564,6 +649,8 @@ def save_system_settings(self, payload: dict[str, Any], updated_by: int | None, key_value, self.encryption_secret(), preview, + self._billing_id(payload.get("providerProjectId")), + self._billing_id(payload.get("providerApiKeyId")), ids["assist_mode_id"], str(payload.get("defaultModel") or ""), updated_by, @@ -609,6 +696,8 @@ def save_entity_settings(self, entity_id: int, payload: dict[str, Any], updated_ key_value, self.encryption_secret(), preview, + self._billing_id(payload.get("providerProjectId")), + self._billing_id(payload.get("providerApiKeyId")), ids["key_policy_id"], ids["assist_mode_id"], str(payload.get("defaultModel") or ""), @@ -646,6 +735,8 @@ def save_user_settings(self, user_id: int, payload: dict[str, Any], provider: st key_value, self.encryption_secret(), preview, + self._billing_id(payload.get("providerProjectId")), + self._billing_id(payload.get("providerApiKeyId")), ids["key_policy_id"], ids["assist_mode_id"], str(payload.get("defaultModel") or ""), @@ -723,6 +814,8 @@ def _secret_row( "adminApiKey": row.get("admin_api_key") or "", "keyPreview": row.get("key_preview") or "", "adminKeyPreview": row.get("admin_key_preview") or "", + "providerProjectId": row.get("provider_project_id") or "", + "providerApiKeyId": row.get("provider_api_key_id") or "", "keyPolicy": row.get("key_policy") or ("system" if scope == "system" else scope), "assistMode": row.get("assist_mode") or "auto", "defaultModel": row.get("default_model") or "", @@ -748,6 +841,8 @@ def _resolved_provider( "modelName": row.get("defaultModel") or default_model, "paidBy": paid_by, "providerKeyOwnerUserId": owner_user_id if paid_by == "user" else None, + "providerProjectId": row.get("providerProjectId") or "", + "providerApiKeyId": row.get("providerApiKeyId") or "", "pricingScope": paid_by, "pricingEntityId": entity_id if paid_by == "entity" else None, "pricingUserId": user_id if paid_by == "user" else None, @@ -776,6 +871,8 @@ def _settings_row(self, row: dict[str, Any] | None, *, scope: str, provider: str "hasAdminApiKey": False, "keyPreview": "", "adminKeyPreview": "", + "providerProjectId": "", + "providerApiKeyId": "", "keyPolicy": "system" if scope == "system" else ("entity" if scope == "entity" else "user_when_available"), "assistMode": "auto", "defaultModel": "", @@ -795,6 +892,8 @@ def _settings_row(self, row: dict[str, Any] | None, *, scope: str, provider: str "hasAdminApiKey": bool(row.get("admin_key_preview")), "keyPreview": row["key_preview"] or "", "adminKeyPreview": row.get("admin_key_preview") or "", + "providerProjectId": row.get("provider_project_id") or "", + "providerApiKeyId": row.get("provider_api_key_id") or "", "keyPolicy": row.get("key_policy") or ("system" if scope == "system" else "entity"), "assistMode": row["assist_mode"], "defaultModel": row["default_model"] or "", @@ -840,6 +939,8 @@ def _usage_event_row(self, row: dict[str, Any]) -> dict[str, Any]: "paidBy": row.get("paid_by") or "unknown", "providerKeyOwnerUserId": row.get("provider_key_owner_user_id"), "providerKeyOwnerUsername": row.get("provider_key_owner_username"), + "providerProjectId": row.get("provider_project_id") or "", + "providerApiKeyId": row.get("provider_api_key_id") or "", "success": bool(row.get("success")), "errorMessage": row.get("error_message"), "decisionReason": row.get("decision_reason") or "", @@ -902,3 +1003,16 @@ def _context_label(context_type: Any, context_id: Any) -> str: if not context_id: return str(context_type) return f"{context_type}:{str(context_id)[:8]}" + + @staticmethod + def _billing_id(value: Any) -> str: + return str(value or "").strip()[:120] + + @staticmethod + def _filter_or_none(values: list[str] | None) -> list[str] | None: + clean: list[str] = [] + for value in values or []: + item = str(value or "").strip() + if item and item not in clean: + clean.append(item) + return clean or None diff --git a/db/migrations/071_openai_reconciliation_billing_ids.sql b/db/migrations/071_openai_reconciliation_billing_ids.sql new file mode 100644 index 0000000..93bc547 --- /dev/null +++ b/db/migrations/071_openai_reconciliation_billing_ids.sql @@ -0,0 +1,112 @@ +BEGIN; + +ALTER TABLE system_ai_provider_settings + ADD COLUMN IF NOT EXISTS provider_project_id text NOT NULL DEFAULT '', + ADD COLUMN IF NOT EXISTS provider_api_key_id text NOT NULL DEFAULT ''; + +ALTER TABLE entity_ai_provider_settings + ADD COLUMN IF NOT EXISTS provider_project_id text NOT NULL DEFAULT '', + ADD COLUMN IF NOT EXISTS provider_api_key_id text NOT NULL DEFAULT ''; + +ALTER TABLE user_ai_provider_settings + ADD COLUMN IF NOT EXISTS provider_project_id text NOT NULL DEFAULT '', + ADD COLUMN IF NOT EXISTS provider_api_key_id text NOT NULL DEFAULT ''; + +ALTER TABLE ai_assist_events + ADD COLUMN IF NOT EXISTS provider_project_id text NOT NULL DEFAULT '', + ADD COLUMN IF NOT EXISTS provider_api_key_id text NOT NULL DEFAULT ''; + +CREATE INDEX IF NOT EXISTS ai_assist_events_provider_billing_ids_idx + ON ai_assist_events (provider_type_id, provider_project_id, provider_api_key_id, created_at DESC); + +UPDATE ai_model_pricing pricing + SET input_per_million = 5.00, + cached_input_per_million = 0.50, + output_per_million = 30.00, + updated_at = now() + FROM ai_provider_types provider + WHERE pricing.provider_type_id = provider.id + AND provider.code = 'openai' + AND pricing.model_name = 'gpt-5-chat-latest'; + +UPDATE ai_assist_events ev + SET provider_project_id = coalesce(nullif(s.provider_project_id, ''), ev.provider_project_id), + provider_api_key_id = coalesce(nullif(s.provider_api_key_id, ''), ev.provider_api_key_id) + FROM system_ai_provider_settings s + JOIN ai_provider_types provider ON provider.id = s.provider_type_id + WHERE ev.provider_type_id = provider.id + AND provider.code = 'openai' + AND ev.paid_by = 'system'; + +UPDATE ai_assist_events ev + SET provider_project_id = coalesce(nullif(s.provider_project_id, ''), ev.provider_project_id), + provider_api_key_id = coalesce(nullif(s.provider_api_key_id, ''), ev.provider_api_key_id) + FROM entity_ai_provider_settings s + JOIN ai_provider_types provider ON provider.id = s.provider_type_id + WHERE ev.provider_type_id = provider.id + AND provider.code = 'openai' + AND ev.paid_by = 'entity' + AND ev.entity_id = s.entity_id; + +UPDATE ai_assist_events ev + SET provider_project_id = coalesce(nullif(s.provider_project_id, ''), ev.provider_project_id), + provider_api_key_id = coalesce(nullif(s.provider_api_key_id, ''), ev.provider_api_key_id) + FROM user_ai_provider_settings s + JOIN ai_provider_types provider ON provider.id = s.provider_type_id + WHERE ev.provider_type_id = provider.id + AND provider.code = 'openai' + AND ev.paid_by = 'user' + AND ev.provider_key_owner_user_id = s.user_id; + +DROP VIEW IF EXISTS ai_assist_events_expanded; + +CREATE VIEW ai_assist_events_expanded AS +SELECT ev.id, + ev.created_at, + ev.entity_id, + en.name AS entity_name, + ev.user_id, + u.username, + coalesce(nullif(u.display_name, ''), nullif(u.nickname, ''), u.username::text) AS user_label, + ev.provider_key_owner_user_id, + key_owner.username AS provider_key_owner_username, + coalesce(nullif(key_owner.display_name, ''), nullif(key_owner.nickname, ''), key_owner.username::text) AS provider_key_owner_label, + ev.provider_project_id, + ev.provider_api_key_id, + ev.provider_type_id, + provider.code AS provider_code, + provider.display_name AS provider_name, + ev.task_type_id, + task.code AS task_code, + task.display_name AS task_name, + ev.model_name, + ev.context_type, + ev.context_id, + ev.round_number, + ev.round_count, + ev.input_tokens, + ev.cached_input_tokens, + ev.output_tokens, + ev.estimated_cost, + ev.final_cost, + ev.cost_status, + ev.cost_discrepancy, + ev.reconciliation_run_id, + ev.allocation_method, + ev.paid_by, + ev.success, + ev.error_message, + ev.decision_reason, + ev.latency_ms +FROM ai_assist_events ev +LEFT JOIN entities en ON en.id = ev.entity_id +LEFT JOIN users u ON u.id = ev.user_id +LEFT JOIN users key_owner ON key_owner.id = ev.provider_key_owner_user_id +LEFT JOIN ai_provider_types provider ON provider.id = ev.provider_type_id +LEFT JOIN ai_task_types task ON task.id = ev.task_type_id; + +INSERT INTO schema_migrations (version, name) +VALUES (71, 'openai_reconciliation_billing_ids') +ON CONFLICT (version) DO NOTHING; + +COMMIT; diff --git a/db/queries/ai_assist_event_estimated_cost_update.sql b/db/queries/ai_assist_event_estimated_cost_update.sql new file mode 100644 index 0000000..97cdb75 --- /dev/null +++ b/db/queries/ai_assist_event_estimated_cost_update.sql @@ -0,0 +1,4 @@ +UPDATE ai_assist_events + SET estimated_cost = %s + WHERE id = %s +RETURNING id; diff --git a/db/queries/ai_assist_event_insert.sql b/db/queries/ai_assist_event_insert.sql index 73070b3..3be2e1e 100644 --- a/db/queries/ai_assist_event_insert.sql +++ b/db/queries/ai_assist_event_insert.sql @@ -14,6 +14,8 @@ INSERT INTO ai_assist_events ( estimated_cost, paid_by, provider_key_owner_user_id, + provider_project_id, + provider_api_key_id, success, error_message, decision_reason, @@ -38,6 +40,8 @@ VALUES ( %s, %s, %s, + %s, + %s, %s ) RETURNING id; diff --git a/db/queries/ai_assist_event_reconciliation_clear.sql b/db/queries/ai_assist_event_reconciliation_clear.sql new file mode 100644 index 0000000..4b08ce6 --- /dev/null +++ b/db/queries/ai_assist_event_reconciliation_clear.sql @@ -0,0 +1,14 @@ +UPDATE ai_assist_events ev + SET final_cost = NULL, + cost_status = 'estimated', + cost_discrepancy = 0, + reconciliation_run_id = NULL, + allocation_method = '' + FROM ai_provider_types provider + WHERE ev.provider_type_id = provider.id + AND provider.code = 'openai' + AND ev.success = true + AND ev.created_at >= %s + AND ev.created_at < %s + AND (%s::text[] IS NULL OR ev.provider_project_id = '' OR ev.provider_project_id = ANY(%s::text[])) + AND (%s::text[] IS NULL OR ev.provider_api_key_id = '' OR ev.provider_api_key_id = ANY(%s::text[])); diff --git a/db/queries/ai_assist_events_for_reconciliation.sql b/db/queries/ai_assist_events_for_reconciliation.sql index 7ae25b7..493d06b 100644 --- a/db/queries/ai_assist_events_for_reconciliation.sql +++ b/db/queries/ai_assist_events_for_reconciliation.sql @@ -1,17 +1,25 @@ SELECT ev.id, ev.created_at, + ev.entity_id, + ev.user_id, provider.code AS provider, ev.model_name, + ev.provider_project_id, + ev.provider_api_key_id, ev.input_tokens, ev.cached_input_tokens, ev.output_tokens, ev.estimated_cost, ev.final_cost, - ev.cost_status + ev.cost_status, + ev.paid_by, + ev.provider_key_owner_user_id FROM ai_assist_events ev JOIN ai_provider_types provider ON provider.id = ev.provider_type_id WHERE provider.code = 'openai' AND ev.success = true AND ev.created_at >= %s AND ev.created_at < %s + AND (%s::text[] IS NULL OR ev.provider_project_id = '' OR ev.provider_project_id = ANY(%s::text[])) + AND (%s::text[] IS NULL OR ev.provider_api_key_id = '' OR ev.provider_api_key_id = ANY(%s::text[])) ORDER BY ev.created_at ASC, ev.id ASC; diff --git a/db/queries/ai_assist_usage_events.sql b/db/queries/ai_assist_usage_events.sql index 443f45c..8fb1c4e 100644 --- a/db/queries/ai_assist_usage_events.sql +++ b/db/queries/ai_assist_usage_events.sql @@ -24,6 +24,8 @@ SELECT ev.id, ev.paid_by, ev.provider_key_owner_user_id, key_owner.username AS provider_key_owner_username, + ev.provider_project_id, + ev.provider_api_key_id, ev.success, ev.error_message, ev.decision_reason, diff --git a/db/queries/ai_provider_backup_rows.sql b/db/queries/ai_provider_backup_rows.sql index f1c6dff..5dcb236 100644 --- a/db/queries/ai_provider_backup_rows.sql +++ b/db/queries/ai_provider_backup_rows.sql @@ -6,6 +6,8 @@ SELECT 'system' AS scope, s.encrypted_admin_api_key, s.key_preview, s.admin_key_preview, + s.provider_project_id, + s.provider_api_key_id, 'system' AS key_policy, am.code AS assist_mode, s.default_model, @@ -26,6 +28,8 @@ SELECT 'entity' AS scope, NULL::text AS encrypted_admin_api_key, s.key_preview, NULL::text AS admin_key_preview, + s.provider_project_id, + s.provider_api_key_id, kp.code AS key_policy, am.code AS assist_mode, s.default_model, @@ -47,6 +51,8 @@ SELECT 'user' AS scope, NULL::text AS encrypted_admin_api_key, s.key_preview, NULL::text AS admin_key_preview, + s.provider_project_id, + s.provider_api_key_id, kp.code AS key_policy, am.code AS assist_mode, s.default_model, diff --git a/db/queries/ai_provider_entity_encrypted_restore.sql b/db/queries/ai_provider_entity_encrypted_restore.sql index e339832..91844cc 100644 --- a/db/queries/ai_provider_entity_encrypted_restore.sql +++ b/db/queries/ai_provider_entity_encrypted_restore.sql @@ -4,6 +4,8 @@ INSERT INTO entity_ai_provider_settings ( enabled, encrypted_api_key, key_preview, + provider_project_id, + provider_api_key_id, key_policy_id, assist_mode_id, default_model, @@ -13,7 +15,7 @@ INSERT INTO entity_ai_provider_settings ( updated_by, updated_at ) -SELECT %s, p.id, %s, %s, %s, kp.id, am.id, %s, %s, %s, %s, %s, now() +SELECT %s, p.id, %s, %s, %s, %s, %s, kp.id, am.id, %s, %s, %s, %s, %s, now() FROM ai_provider_types p JOIN ai_key_policies kp ON kp.code = %s JOIN ai_assist_modes am ON am.code = %s @@ -22,6 +24,8 @@ ON CONFLICT (entity_id, provider_type_id) DO UPDATE SET enabled = EXCLUDED.enabled, encrypted_api_key = EXCLUDED.encrypted_api_key, key_preview = EXCLUDED.key_preview, + provider_project_id = EXCLUDED.provider_project_id, + provider_api_key_id = EXCLUDED.provider_api_key_id, key_policy_id = EXCLUDED.key_policy_id, assist_mode_id = EXCLUDED.assist_mode_id, default_model = EXCLUDED.default_model, diff --git a/db/queries/ai_provider_system_encrypted_restore.sql b/db/queries/ai_provider_system_encrypted_restore.sql index a2dae49..a7dba25 100644 --- a/db/queries/ai_provider_system_encrypted_restore.sql +++ b/db/queries/ai_provider_system_encrypted_restore.sql @@ -5,12 +5,14 @@ INSERT INTO system_ai_provider_settings ( encrypted_admin_api_key, key_preview, admin_key_preview, + provider_project_id, + provider_api_key_id, assist_mode_id, default_model, updated_by, updated_at ) -SELECT p.id, %s, %s, %s, %s, %s, am.id, %s, %s, now() +SELECT p.id, %s, %s, %s, %s, %s, %s, %s, am.id, %s, %s, now() FROM ai_provider_types p JOIN ai_assist_modes am ON am.code = %s WHERE p.code = %s @@ -20,6 +22,8 @@ ON CONFLICT (provider_type_id) DO UPDATE SET encrypted_admin_api_key = EXCLUDED.encrypted_admin_api_key, key_preview = EXCLUDED.key_preview, admin_key_preview = EXCLUDED.admin_key_preview, + provider_project_id = EXCLUDED.provider_project_id, + provider_api_key_id = EXCLUDED.provider_api_key_id, assist_mode_id = EXCLUDED.assist_mode_id, default_model = EXCLUDED.default_model, updated_by = EXCLUDED.updated_by, diff --git a/db/queries/ai_provider_user_encrypted_restore.sql b/db/queries/ai_provider_user_encrypted_restore.sql index 1347443..16a5358 100644 --- a/db/queries/ai_provider_user_encrypted_restore.sql +++ b/db/queries/ai_provider_user_encrypted_restore.sql @@ -4,6 +4,8 @@ INSERT INTO user_ai_provider_settings ( enabled, encrypted_api_key, key_preview, + provider_project_id, + provider_api_key_id, key_policy_id, assist_mode_id, default_model, @@ -12,7 +14,7 @@ INSERT INTO user_ai_provider_settings ( stop_percent, updated_at ) -SELECT %s, p.id, %s, %s, %s, kp.id, am.id, %s, %s, %s, %s, now() +SELECT %s, p.id, %s, %s, %s, %s, %s, kp.id, am.id, %s, %s, %s, %s, now() FROM ai_provider_types p JOIN ai_key_policies kp ON kp.code = %s JOIN ai_assist_modes am ON am.code = %s @@ -21,6 +23,8 @@ ON CONFLICT (user_id, provider_type_id) DO UPDATE SET enabled = EXCLUDED.enabled, encrypted_api_key = EXCLUDED.encrypted_api_key, key_preview = EXCLUDED.key_preview, + provider_project_id = EXCLUDED.provider_project_id, + provider_api_key_id = EXCLUDED.provider_api_key_id, key_policy_id = EXCLUDED.key_policy_id, assist_mode_id = EXCLUDED.assist_mode_id, default_model = EXCLUDED.default_model, diff --git a/db/queries/entity_ai_provider_get.sql b/db/queries/entity_ai_provider_get.sql index 49226bc..052ebba 100644 --- a/db/queries/entity_ai_provider_get.sql +++ b/db/queries/entity_ai_provider_get.sql @@ -2,6 +2,8 @@ SELECT p.code AS provider_code, s.entity_id, s.enabled, s.key_preview, + s.provider_project_id, + s.provider_api_key_id, kp.code AS key_policy, am.code AS assist_mode, s.default_model, diff --git a/db/queries/entity_ai_provider_secret_get.sql b/db/queries/entity_ai_provider_secret_get.sql index e5ef686..9ff3f5e 100644 --- a/db/queries/entity_ai_provider_secret_get.sql +++ b/db/queries/entity_ai_provider_secret_get.sql @@ -5,6 +5,8 @@ SELECT p.code AS provider_code, ELSE pgp_sym_decrypt(decode(s.encrypted_api_key, 'base64'), %s::text) END AS api_key, s.key_preview, + s.provider_project_id, + s.provider_api_key_id, kp.code AS key_policy, am.code AS assist_mode, s.default_model, diff --git a/db/queries/entity_ai_provider_upsert.sql b/db/queries/entity_ai_provider_upsert.sql index 88dff7d..14560b3 100644 --- a/db/queries/entity_ai_provider_upsert.sql +++ b/db/queries/entity_ai_provider_upsert.sql @@ -4,6 +4,8 @@ INSERT INTO entity_ai_provider_settings ( enabled, encrypted_api_key, key_preview, + provider_project_id, + provider_api_key_id, key_policy_id, assist_mode_id, default_model, @@ -30,6 +32,8 @@ VALUES ( %s, %s, %s, + %s, + %s, now() ) ON CONFLICT (entity_id, provider_type_id) DO UPDATE SET @@ -42,6 +46,8 @@ ON CONFLICT (entity_id, provider_type_id) DO UPDATE SET WHEN %s::boolean THEN EXCLUDED.key_preview ELSE entity_ai_provider_settings.key_preview END, + provider_project_id = EXCLUDED.provider_project_id, + provider_api_key_id = EXCLUDED.provider_api_key_id, key_policy_id = EXCLUDED.key_policy_id, assist_mode_id = EXCLUDED.assist_mode_id, default_model = EXCLUDED.default_model, diff --git a/db/queries/openai_reconciliation_configured_filters.sql b/db/queries/openai_reconciliation_configured_filters.sql new file mode 100644 index 0000000..15d09f3 --- /dev/null +++ b/db/queries/openai_reconciliation_configured_filters.sql @@ -0,0 +1,14 @@ +SELECT provider_project_id, provider_api_key_id + FROM system_ai_provider_settings s + JOIN ai_provider_types p ON p.id = s.provider_type_id + WHERE p.code = 'openai' +UNION ALL +SELECT provider_project_id, provider_api_key_id + FROM entity_ai_provider_settings s + JOIN ai_provider_types p ON p.id = s.provider_type_id + WHERE p.code = 'openai' +UNION ALL +SELECT provider_project_id, provider_api_key_id + FROM user_ai_provider_settings s + JOIN ai_provider_types p ON p.id = s.provider_type_id + WHERE p.code = 'openai'; diff --git a/db/queries/system_ai_provider_get.sql b/db/queries/system_ai_provider_get.sql index a08dcd9..d053d4f 100644 --- a/db/queries/system_ai_provider_get.sql +++ b/db/queries/system_ai_provider_get.sql @@ -2,6 +2,8 @@ SELECT p.code AS provider_code, s.enabled, s.key_preview, s.admin_key_preview, + s.provider_project_id, + s.provider_api_key_id, am.code AS assist_mode, s.default_model, s.updated_at diff --git a/db/queries/system_ai_provider_secret_get.sql b/db/queries/system_ai_provider_secret_get.sql index 1bc2b46..3f87351 100644 --- a/db/queries/system_ai_provider_secret_get.sql +++ b/db/queries/system_ai_provider_secret_get.sql @@ -10,6 +10,8 @@ SELECT p.code AS provider_code, END AS admin_api_key, s.key_preview, s.admin_key_preview, + s.provider_project_id, + s.provider_api_key_id, am.code AS assist_mode, s.default_model FROM system_ai_provider_settings s diff --git a/db/queries/system_ai_provider_upsert.sql b/db/queries/system_ai_provider_upsert.sql index a55af53..0cd2101 100644 --- a/db/queries/system_ai_provider_upsert.sql +++ b/db/queries/system_ai_provider_upsert.sql @@ -3,6 +3,8 @@ INSERT INTO system_ai_provider_settings ( enabled, encrypted_api_key, key_preview, + provider_project_id, + provider_api_key_id, assist_mode_id, default_model, updated_by, @@ -20,6 +22,8 @@ VALUES ( %s, %s, %s, + %s, + %s, now() ) ON CONFLICT (provider_type_id) DO UPDATE SET @@ -32,6 +36,8 @@ ON CONFLICT (provider_type_id) DO UPDATE SET WHEN %s::boolean THEN EXCLUDED.key_preview ELSE system_ai_provider_settings.key_preview END, + provider_project_id = EXCLUDED.provider_project_id, + provider_api_key_id = EXCLUDED.provider_api_key_id, assist_mode_id = EXCLUDED.assist_mode_id, default_model = EXCLUDED.default_model, updated_by = EXCLUDED.updated_by, diff --git a/db/queries/user_ai_provider_get.sql b/db/queries/user_ai_provider_get.sql index 821572e..22bb5f3 100644 --- a/db/queries/user_ai_provider_get.sql +++ b/db/queries/user_ai_provider_get.sql @@ -2,6 +2,8 @@ SELECT p.code AS provider_code, s.user_id, s.enabled, s.key_preview, + s.provider_project_id, + s.provider_api_key_id, kp.code AS key_policy, am.code AS assist_mode, s.default_model, diff --git a/db/queries/user_ai_provider_secret_get.sql b/db/queries/user_ai_provider_secret_get.sql index 654fe5d..daa634b 100644 --- a/db/queries/user_ai_provider_secret_get.sql +++ b/db/queries/user_ai_provider_secret_get.sql @@ -5,6 +5,8 @@ SELECT p.code AS provider_code, ELSE pgp_sym_decrypt(decode(s.encrypted_api_key, 'base64'), %s::text) END AS api_key, s.key_preview, + s.provider_project_id, + s.provider_api_key_id, kp.code AS key_policy, am.code AS assist_mode, s.default_model, diff --git a/db/queries/user_ai_provider_upsert.sql b/db/queries/user_ai_provider_upsert.sql index be4bcb3..9b0a5f4 100644 --- a/db/queries/user_ai_provider_upsert.sql +++ b/db/queries/user_ai_provider_upsert.sql @@ -4,6 +4,8 @@ INSERT INTO user_ai_provider_settings ( enabled, encrypted_api_key, key_preview, + provider_project_id, + provider_api_key_id, key_policy_id, assist_mode_id, default_model, @@ -28,6 +30,8 @@ VALUES ( %s, %s, %s, + %s, + %s, now() ) ON CONFLICT (user_id, provider_type_id) DO UPDATE SET @@ -40,6 +44,8 @@ ON CONFLICT (user_id, provider_type_id) DO UPDATE SET WHEN %s::boolean THEN EXCLUDED.key_preview ELSE user_ai_provider_settings.key_preview END, + provider_project_id = EXCLUDED.provider_project_id, + provider_api_key_id = EXCLUDED.provider_api_key_id, key_policy_id = EXCLUDED.key_policy_id, assist_mode_id = EXCLUDED.assist_mode_id, default_model = EXCLUDED.default_model, diff --git a/frontend/src/components/AIProviderSettingsPanel.tsx b/frontend/src/components/AIProviderSettingsPanel.tsx index ca75fea..d948498 100644 --- a/frontend/src/components/AIProviderSettingsPanel.tsx +++ b/frontend/src/components/AIProviderSettingsPanel.tsx @@ -4,6 +4,7 @@ import { SectionHeader } from "./SectionHeader"; import { AIProviderPricingOverrides } from "./AIProviderPricingOverrides"; import { useAIProviderSettingsForm } from "../hooks/useAIProviderSettingsForm"; import { AIProviderBudgetFields } from "./aiProvider/AIProviderBudgetFields"; +import { AIProviderBillingFields } from "./aiProvider/AIProviderBillingFields"; import { AIProviderKeyFields } from "./aiProvider/AIProviderKeyFields"; import { AIProviderMainFields } from "./aiProvider/AIProviderMainFields"; import { AIProviderModelRefresh } from "./aiProvider/AIProviderModelRefresh"; @@ -90,6 +91,7 @@ export function AIProviderSettingsPanel({ onClearApiKeyChange={setClearAdminApiKey} /> ) : null} + {loadModels ? ( void; +}) { + return ( +
+ + +
+ ); +} diff --git a/frontend/src/hooks/useAIProviderSettingsForm.ts b/frontend/src/hooks/useAIProviderSettingsForm.ts index 0b6de44..d486105 100644 --- a/frontend/src/hooks/useAIProviderSettingsForm.ts +++ b/frontend/src/hooks/useAIProviderSettingsForm.ts @@ -104,6 +104,8 @@ export function useAIProviderSettingsForm({ clearApiKey, adminApiKey: adminApiKey.trim() || undefined, clearAdminApiKey, + providerProjectId: settings.providerProjectId.trim(), + providerApiKeyId: settings.providerApiKeyId.trim(), keyPolicy: settings.keyPolicy, assistMode: settings.assistMode, defaultModel: settings.defaultModel, diff --git a/frontend/src/libs/aiProvider/defaultSettings.ts b/frontend/src/libs/aiProvider/defaultSettings.ts index d89d58d..41f721c 100644 --- a/frontend/src/libs/aiProvider/defaultSettings.ts +++ b/frontend/src/libs/aiProvider/defaultSettings.ts @@ -8,6 +8,8 @@ export const defaultAIProviderSettings: AIProviderSettings = { hasAdminApiKey: false, keyPreview: "", adminKeyPreview: "", + providerProjectId: "", + providerApiKeyId: "", keyPolicy: "user_when_available", assistMode: "auto", defaultModel: "", diff --git a/frontend/src/types/ai.ts b/frontend/src/types/ai.ts index 4f2224c..7a938b1 100644 --- a/frontend/src/types/ai.ts +++ b/frontend/src/types/ai.ts @@ -36,6 +36,8 @@ export interface AIProviderSettings { hasAdminApiKey: boolean; keyPreview: string; adminKeyPreview: string; + providerProjectId: string; + providerApiKeyId: string; keyPolicy: string; assistMode: string; defaultModel: string; @@ -52,6 +54,8 @@ export interface AIProviderSettingsPayload { clearApiKey?: boolean; adminApiKey?: string; clearAdminApiKey?: boolean; + providerProjectId?: string; + providerApiKeyId?: string; keyPolicy?: string; assistMode: string; defaultModel: string; @@ -96,6 +100,8 @@ export interface AIUsageEvent { paidBy: string; providerKeyOwnerUserId?: number | null; providerKeyOwnerUsername?: string | null; + providerProjectId?: string; + providerApiKeyId?: string; decisionReason?: string; latencyMs?: number; success: boolean; diff --git a/tests/test_ai_cost_reconciliation.py b/tests/test_ai_cost_reconciliation.py index 169fbd1..464bc77 100644 --- a/tests/test_ai_cost_reconciliation.py +++ b/tests/test_ai_cost_reconciliation.py @@ -132,6 +132,8 @@ def test_usage_event_row_exposes_reconciliation_fields(): "paid_by": "system", "provider_key_owner_user_id": None, "provider_key_owner_username": None, + "provider_project_id": "proj_123", + "provider_api_key_id": "key_456", "success": True, "error_message": None, "decision_reason": "test", @@ -146,6 +148,8 @@ def test_usage_event_row_exposes_reconciliation_fields(): assert event["costStatus"] == "adjusted" assert event["costDiscrepancy"] == 0.0025 assert event["reconciliationRunId"] == "11111111-1111-1111-1111-111111111111" + assert event["providerProjectId"] == "proj_123" + assert event["providerApiKeyId"] == "key_456" def test_usage_summary_row_exposes_estimated_actual_and_verified_costs(): diff --git a/tests/test_ai_provider_store.py b/tests/test_ai_provider_store.py index 708df94..1fe6b4e 100644 --- a/tests/test_ai_provider_store.py +++ b/tests/test_ai_provider_store.py @@ -91,6 +91,8 @@ def test_ai_provider_settings_row_exposes_admin_key_preview(): "enabled": True, "key_preview": "sk-live...1234", "admin_key_preview": "sk-admin...5678", + "provider_project_id": "proj_123", + "provider_api_key_id": "key_456", "key_policy": "system", "assist_mode": "auto", "default_model": "gpt-5-chat-latest", @@ -108,9 +110,36 @@ def test_ai_provider_settings_row_exposes_admin_key_preview(): assert settings["hasAdminApiKey"] is True assert settings["adminKeyPreview"] == "sk-admin...5678" + assert settings["providerProjectId"] == "proj_123" + assert settings["providerApiKeyId"] == "key_456" assert "adminApiKey" not in settings +def test_resolved_provider_carries_billing_ids(): + settings = AIProviderStore(database=None, config_path="config/config.yaml")._resolved_provider( + row={ + "provider": "openai", + "apiKey": "sk-live", + "assistMode": "auto", + "defaultModel": "gpt-5-chat-latest", + "providerProjectId": "proj_123", + "providerApiKeyId": "key_456", + "monthlyBudget": 0, + "warnPercent": 80, + "stopPercent": 100, + "scope": "system", + }, + paid_by="system", + owner_user_id=None, + default_model="fallback", + entity_id=None, + user_id=None, + ) + + assert settings["providerProjectId"] == "proj_123" + assert settings["providerApiKeyId"] == "key_456" + + def test_ai_provider_admin_api_key_for_provider_uses_system_secret_row(monkeypatch): store = AIProviderStore(database=None, config_path="config/config.yaml") monkeypatch.setattr( diff --git a/tools/ai_key_ops.py b/tools/ai_key_ops.py index 61fd444..f5d8dc1 100644 --- a/tools/ai_key_ops.py +++ b/tools/ai_key_ops.py @@ -64,6 +64,8 @@ def restore_keys(database_url: str, input_path: Path) -> int: row.get("encrypted_admin_api_key") or "", row.get("key_preview") or "", row.get("admin_key_preview") or "", + row.get("provider_project_id") or "", + row.get("provider_api_key_id") or "", row.get("default_model") or "", row.get("updated_by"), row.get("assist_mode") or "auto", @@ -78,6 +80,8 @@ def restore_keys(database_url: str, input_path: Path) -> int: bool(row.get("enabled")), row.get("encrypted_api_key") or "", row.get("key_preview") or "", + row.get("provider_project_id") or "", + row.get("provider_api_key_id") or "", row.get("default_model") or "", float(row.get("monthly_budget") or 0), int(row.get("warn_percent") or 80), @@ -96,6 +100,8 @@ def restore_keys(database_url: str, input_path: Path) -> int: bool(row.get("enabled")), row.get("encrypted_api_key") or "", row.get("key_preview") or "", + row.get("provider_project_id") or "", + row.get("provider_api_key_id") or "", row.get("default_model") or "", float(row.get("monthly_budget") or 0), int(row.get("warn_percent") or 80), diff --git a/tools/reconcile_openai_costs.py b/tools/reconcile_openai_costs.py index ddbf481..8e9bd90 100755 --- a/tools/reconcile_openai_costs.py +++ b/tools/reconcile_openai_costs.py @@ -62,6 +62,8 @@ def main() -> int: parser.add_argument("--group-by", action="append", choices=["project_id", "line_item", "api_key_id"], default=[]) parser.add_argument("--project-id", action="append", default=[]) parser.add_argument("--api-key-id", action="append", default=[]) + parser.add_argument("--skip-configured-filters", action="store_true", help="Do not default to stored provider project/API key IDs.") + parser.add_argument("--skip-reprice", action="store_true", help="Do not recalculate stored event estimates before reconciliation.") parser.add_argument("--dry-run", action="store_true", help="Fetch and calculate but do not persist final costs.") args = parser.parse_args() @@ -83,6 +85,12 @@ def main() -> int: return 2 client = OpenAIOrganizationUsageClient(admin_api_key=admin_key, base_url=args.base_url) allocation_method = "estimated_cost_proportional" + configured_filters = {"projectIds": [], "apiKeyIds": []} if args.skip_configured_filters else store.openai_reconciliation_filters() + project_ids = args.project_id or configured_filters["projectIds"] + api_key_ids = args.api_key_id or configured_filters["apiKeyIds"] + group_by = list(args.group_by) + if api_key_ids and "api_key_id" not in group_by: + group_by.append("api_key_id") run_id = "" if not args.dry_run: @@ -95,16 +103,35 @@ def main() -> int: allocation_method=allocation_method, ) try: + repriced_count = 0 + if not args.skip_reprice and not args.dry_run: + store.clear_openai_reconciliation_window( + start_time=args.start, + end_time=args.end, + project_ids=project_ids, + api_key_ids=api_key_ids, + ) + repriced_count = store.reprice_openai_events_for_reconciliation( + start_time=args.start, + end_time=args.end, + project_ids=project_ids, + api_key_ids=api_key_ids, + ) payload = client.fetch_costs( start_time=unix_seconds(args.start), end_time=unix_seconds(args.end), bucket_width=args.bucket_width, - group_by=args.group_by, - project_ids=args.project_id, - api_key_ids=args.api_key_id, + group_by=group_by, + project_ids=project_ids, + api_key_ids=api_key_ids, ) verified_cost = verified_cost_from_openai_payload(payload) - events = store.list_openai_events_for_reconciliation(start_time=args.start, end_time=args.end) + events = store.list_openai_events_for_reconciliation( + start_time=args.start, + end_time=args.end, + project_ids=project_ids, + api_key_ids=api_key_ids, + ) effective_run_id = run_id or "dry-run" updates = reconcile_cost_bucket(events, verified_cost_usd=verified_cost, reconciliation_run_id=effective_run_id) estimated_cost = sum((update.estimated_cost_usd for update in updates), Decimal("0")) @@ -116,11 +143,23 @@ def main() -> int: verified_cost=verified_cost, estimated_cost=estimated_cost, event_count=len(updates), - raw_provider_payload=payload, + raw_provider_payload={ + "filters": { + "projectIds": project_ids, + "apiKeyIds": api_key_ids, + "groupBy": group_by, + "repricedEvents": repriced_count, + }, + "payload": payload, + }, ) print(f"OpenAI reconciliation {'dry run' if args.dry_run else 'complete'}") print(f"Window: {args.start.isoformat()} to {args.end.isoformat()}") + if project_ids or api_key_ids: + print(f"Filters: projects={len(project_ids)} api_keys={len(api_key_ids)} group_by={','.join(group_by) or 'none'}") + if not args.skip_reprice and not args.dry_run: + print(f"Repriced events: {repriced_count}") print(f"Verified cost: ${verified_cost:.8f}") print(f"Estimated event cost: ${estimated_cost:.8f}") print(f"Events reconciled: {len(updates)}")