From 01807b570476387e8b2900f6c0f9a83bd7bd7e11 Mon Sep 17 00:00:00 2001 From: KayanoLiam Date: Sat, 24 Oct 2026 13:29:39 +0800 Subject: [PATCH] fix(ai): Preserve provider-specific model metadata --- src/sentry/relay/config/ai_model_costs.py | 4 +- src/sentry/tasks/ai_agent_monitoring.py | 65 +++++++++++----- .../sentry/tasks/test_ai_agent_monitoring.py | 76 +++++++++++++++++++ 3 files changed, 123 insertions(+), 22 deletions(-) diff --git a/src/sentry/relay/config/ai_model_costs.py b/src/sentry/relay/config/ai_model_costs.py index 849c0e442793..0567746e25e3 100644 --- a/src/sentry/relay/config/ai_model_costs.py +++ b/src/sentry/relay/config/ai_model_costs.py @@ -1,5 +1,5 @@ import logging -from typing import Required, TypedDict +from typing import NotRequired, Required, TypedDict from django.conf import settings @@ -33,6 +33,8 @@ class AIModelMetadata(TypedDict, total=False): class AIModelMetadataConfig(TypedDict): version: Required[int] models: Required[dict[ModelId, AIModelMetadata]] + # Provider IDs from models.dev, with full model IDs and provider-local aliases. + providers: NotRequired[dict[str, dict[ModelId, AIModelMetadata]]] def ai_model_metadata_config() -> AIModelMetadataConfig | None: diff --git a/src/sentry/tasks/ai_agent_monitoring.py b/src/sentry/tasks/ai_agent_monitoring.py index a5b75ff309ff..08c6990217e5 100644 --- a/src/sentry/tasks/ai_agent_monitoring.py +++ b/src/sentry/tasks/ai_agent_monitoring.py @@ -29,6 +29,21 @@ MODELS_DEV_TOKEN_PRICED_TYPES = frozenset({"decision"}) +def _add_provider_model_aliases(models: dict[ModelId, AIModelMetadata]) -> None: + """Keep full IDs and add unambiguous short names within a single provider.""" + aliases: dict[str, AIModelMetadata] = {} + ambiguous: set[str] = set() + for model_id, metadata in models.items(): + for alias in (model_id.split("/", 1)[-1], model_id.rsplit("/", 1)[-1]): + if alias in aliases and aliases[alias] != metadata: + ambiguous.add(alias) + aliases[alias] = metadata + for alias, metadata in aliases.items(): + if alias not in ambiguous: + models.setdefault(alias, metadata) + _add_glob_model_names(models) + + def _add_glob_model_names(models_dict: dict[ModelId, AIModelMetadata]) -> None: """ Add glob versions of model names to the models dictionary. @@ -69,17 +84,22 @@ def fetch_ai_model_metadata() -> None: This task fetches model pricing and context size data from both sources and converts it to the AIModelMetadata format. - OpenRouter data takes precedence over models.dev data. + OpenRouter data takes precedence within the OpenRouter provider and the + legacy model-only table, without overriding other providers' prices. """ if settings.SENTRY_AIR_GAP: return + # Keep the legacy model-only table for older Relays and spans without a provider. models_dict: dict[ModelId, AIModelMetadata] = {} # Fetch from OpenRouter API (takes precedence) try: openrouter_models = _fetch_openrouter_models() - models_dict.update(openrouter_models) + models_dict.update( + (model_id.split("/", 1)[-1], metadata) + for model_id, metadata in openrouter_models.items() + ) except Exception as e: logger.warning( "Failed to fetch AI model metadata from OpenRouter API", extra={"error": str(e)} @@ -88,9 +108,14 @@ def fetch_ai_model_metadata() -> None: # Fetch from models.dev API (only add models not already present) try: - models_dev_models = _fetch_models_dev_models() - for model_id, model_metadata in models_dev_models.items(): - models_dict.setdefault(model_id, model_metadata) + providers = _fetch_models_dev_models() + legacy_models = { + model_id.split("/", 1)[-1]: metadata + for models in providers.values() + for model_id, metadata in models.items() + } + for model_id, metadata in legacy_models.items(): + models_dict.setdefault(model_id, metadata) except Exception as e: logger.warning( "Failed to fetch AI model metadata from models.dev API", extra={"error": str(e)} @@ -100,7 +125,16 @@ def fetch_ai_model_metadata() -> None: # Add glob versions of model names for flexible matching _add_glob_model_names(models_dict) - metadata_config: AIModelMetadataConfig = {"version": 1, "models": models_dict} + # OpenRouter prices belong only to OpenRouter, not to the model author. + providers.setdefault("openrouter", {}).update(openrouter_models) + for models in providers.values(): + _add_provider_model_aliases(models) + + metadata_config: AIModelMetadataConfig = { + "version": 1, + "models": models_dict, + "providers": providers, + } cache.set(AI_MODEL_METADATA_CACHE_KEY, metadata_config, AI_MODEL_METADATA_CACHE_TTL) @@ -147,12 +181,6 @@ def _fetch_openrouter_models() -> dict[ModelId, AIModelMetadata]: if not model_id: continue - # OpenRouter includes provider name in the model ID, e.g. openai/gpt-4o-mini - # We need to extract the model name, since our SDKs only send the model name - # (e.g. gpt-4o-mini) - if "/" in model_id: - model_id = model_id.split("/", maxsplit=1)[1] - pricing = model_data.get("pricing", {}) # OpenRouter provides costs as strings, we need to convert to float @@ -183,7 +211,7 @@ def _fetch_openrouter_models() -> dict[ModelId, AIModelMetadata]: return models_dict -def _fetch_models_dev_models() -> dict[ModelId, AIModelMetadata]: +def _fetch_models_dev_models() -> dict[str, dict[ModelId, AIModelMetadata]]: """Fetch model metadata from models.dev API. Example response: @@ -214,7 +242,7 @@ def _fetch_models_dev_models() -> dict[ModelId, AIModelMetadata]: if not isinstance(data, dict): raise ValueError("Invalid models.dev response format: expected dict") - models_dict: dict[ModelId, AIModelMetadata] = {} + providers: dict[str, dict[ModelId, AIModelMetadata]] = {} for provider_name, provider_data in data.items(): if not isinstance(provider_data, dict): @@ -224,6 +252,7 @@ def _fetch_models_dev_models() -> dict[ModelId, AIModelMetadata]: if not isinstance(models, dict): continue + models_dict = providers.setdefault(provider_name, {}) for model_id, model_data in models.items(): if not isinstance(model_data, dict): continue @@ -237,12 +266,6 @@ def _fetch_models_dev_models() -> dict[ModelId, AIModelMetadata]: # Skip models with no cost data or empty cost data continue - # models.dev may include provider name in the model ID, e.g. google/gemini-2.0-flash-001 - # We need to extract the model name, since our SDKs only send the model name - # (e.g. gemini-2.0-flash-001) - if "/" in model_id: - model_id = model_id.split("/", maxsplit=1)[1] - # models.dev provides costs as numbers, but for extra safety convert to our format try: model_cost = AIModelCost( @@ -274,7 +297,7 @@ def _fetch_models_dev_models() -> dict[ModelId, AIModelMetadata]: ) continue - return models_dict + return providers def safe_float_conversion(value: Any) -> float: diff --git a/tests/sentry/tasks/test_ai_agent_monitoring.py b/tests/sentry/tasks/test_ai_agent_monitoring.py index 545c1700d018..a1f778c844f7 100644 --- a/tests/sentry/tasks/test_ai_agent_monitoring.py +++ b/tests/sentry/tasks/test_ai_agent_monitoring.py @@ -195,6 +195,82 @@ def _mock_models_dev_api_response(self, mock_response: dict): status=200, ) + @responses.activate + def test_fetch_ai_model_metadata_provider_prices(self) -> None: + self._mock_openrouter_api_response({"data": []}) + self._mock_models_dev_api_response( + { + "cloudflare-workers-ai": { + "models": {"@cf/cloudflare/clef": {"cost": {"input": 0.24}}} + }, + "openrouter": {"models": {"cloudflare/clef": {"cost": {"input": 0.042}}}}, + "anthropic": {"models": {"claude-sonnet-4": {"cost": {"input": 3.0}}}}, + "poe": {"models": {"anthropic/claude-sonnet-4": {"cost": {"input": 2.60}}}}, + } + ) + + fetch_ai_model_metadata() + + config = _get_metadata_from_cache() + assert config is not None + providers = config["providers"] + cloudflare = providers["cloudflare-workers-ai"] + assert cloudflare["clef"] == cloudflare["@cf/cloudflare/clef"] + assert cloudflare["*clef"] == cloudflare["clef"] + assert 55_132 * cloudflare["clef"]["costs"]["inputPerToken"] == pytest.approx(0.01323168) + assert providers["openrouter"]["cloudflare/clef"]["costs"][ + "inputPerToken" + ] == pytest.approx(0.042 / 1_000_000) + assert providers["anthropic"]["claude-sonnet-4"]["costs"]["inputPerToken"] == 3 / 1_000_000 + assert providers["poe"]["claude-sonnet-4"]["costs"]["inputPerToken"] == 2.6 / 1_000_000 + + @responses.activate + def test_provider_prices_openrouter_precedence(self) -> None: + self._mock_openrouter_api_response(MOCK_OPENROUTER_API_RESPONSE) + self._mock_models_dev_api_response( + { + "openai": {"models": {"gpt-4": {"cost": {"input": 30}}}}, + "openrouter": {"models": {"openai/gpt-4": {"cost": {"input": 20}}}}, + } + ) + + fetch_ai_model_metadata() + + config = _get_metadata_from_cache() + assert config is not None + providers = config["providers"] + assert providers["openai"]["gpt-4"]["costs"]["inputPerToken"] == 30 / 1_000_000 + assert providers["openrouter"]["openai/gpt-4"]["costs"]["inputPerToken"] == 0.0000003 + assert providers["openrouter"]["gpt-4"] == providers["openrouter"]["openai/gpt-4"] + assert config["models"]["gpt-4"] == providers["openrouter"]["gpt-4"] + + @responses.activate + def test_provider_model_aliases_do_not_overwrite_full_ids(self) -> None: + self._mock_openrouter_api_response({"data": []}) + self._mock_models_dev_api_response( + { + "openrouter": { + "models": { + "author-a/shared": {"cost": {"input": 1}}, + "author-b/shared": {"cost": {"input": 2}}, + "author/snapshot-20240101": {"cost": {"input": 3}}, + } + }, + } + ) + + fetch_ai_model_metadata() + + config = _get_metadata_from_cache() + assert config is not None + models = config["providers"]["openrouter"] + assert models["author-a/shared"]["costs"]["inputPerToken"] == 1 / 1_000_000 + assert models["author-b/shared"]["costs"]["inputPerToken"] == 2 / 1_000_000 + assert "shared" not in models + assert "*shared" not in models + assert models["author/snapshot"] == models["author/snapshot-20240101"] + assert models["*snapshot"] == models["author/snapshot-20240101"] + @responses.activate def test_fetch_ai_model_metadata_success_both_apis(self) -> None: """Test successful fetching and caching from both APIs"""