From 6d21fbec0016a279d1f54cdc39df956d82420429 Mon Sep 17 00:00:00 2001 From: Aniket Dixit Date: Fri, 10 Jul 2026 00:51:05 +0530 Subject: [PATCH 1/2] model updates --- tee_gateway/model_registry.py | 42 ++-------------------- tee_gateway/test/test_tee_core.py | 6 ++-- tee_gateway/test/test_web_search.py | 2 +- tests/test_pricing.py | 56 +++++++++++------------------ tests/test_server.py | 20 +++++------ tests/test_structured_outputs.py | 6 ++-- 6 files changed, 40 insertions(+), 92 deletions(-) diff --git a/tee_gateway/model_registry.py b/tee_gateway/model_registry.py index ae6c51f..e6463c4 100644 --- a/tee_gateway/model_registry.py +++ b/tee_gateway/model_registry.py @@ -287,9 +287,6 @@ class SupportedModel(Enum): ) # ── Google Gemini ─────────────────────────────────────────────────── - # Note: gemini-2.5-flash, gemini-2.5-pro, and gemini-2.5-flash-lite are scheduled - # for deprecation on June 17, 2026 (flash-lite: July 22, 2026). Use the Gemini 3 - # replacements below for new integrations. GEMINI_2_5_FLASH = ModelConfig( provider="google", api_name="gemini-2.5-flash", @@ -324,9 +321,9 @@ class SupportedModel(Enum): output_price_usd=Decimal("0.000012"), thinking_budget=128, ) - GEMINI_3_1_FLASH_LITE_PREVIEW = ModelConfig( + GEMINI_3_1_FLASH_LITE = ModelConfig( provider="google", - api_name="gemini-3.1-flash-lite-preview", + api_name="gemini-3.1-flash-lite", input_price_usd=Decimal("0.00000025"), output_price_usd=Decimal("0.0000015"), thinking_budget=0, @@ -392,12 +389,6 @@ class SupportedModel(Enum): input_price_usd=Decimal("0.0000002"), output_price_usd=Decimal("0.0000005"), ) - GROK_4_1_FAST_NON_REASONING = ModelConfig( - provider="x-ai", - api_name="grok-4-1-fast-non-reasoning", - input_price_usd=Decimal("0.0000002"), - output_price_usd=Decimal("0.0000005"), - ) GROK_4_20_REASONING = ModelConfig( provider="x-ai", api_name="grok-4.20-reasoning", @@ -410,12 +401,6 @@ class SupportedModel(Enum): input_price_usd=Decimal("0.000002"), output_price_usd=Decimal("0.000006"), ) - GROK_CODE_FAST_1 = ModelConfig( - provider="x-ai", - api_name="grok-code-fast-1", - input_price_usd=Decimal("0.0000002"), - output_price_usd=Decimal("0.0000015"), - ) # Image generation via xAI's OpenAI-compatible /images/generations endpoint # (Aurora). Billed at a flat $0.07 per generated image; token prices unused. GROK_2_IMAGE = ModelConfig( @@ -554,20 +539,6 @@ class SupportedModel(Enum): image_extra_params={"size": "1280x1280"}, ) - # ── Legacy models (not in current SDK — retained for older SDK versions) ── - GROK_3_MINI = ModelConfig( - provider="x-ai", - api_name="grok-3-mini", - input_price_usd=Decimal("0.0000003"), - output_price_usd=Decimal("0.0000005"), - ) - GROK_3 = ModelConfig( - provider="x-ai", - api_name="grok-3-latest", - input_price_usd=Decimal("0.000003"), - output_price_usd=Decimal("0.000015"), - ) - # Canonical lookup: user-facing model name → SupportedModel # The "user-facing name" is what callers pass in the `model` field of requests. @@ -610,7 +581,7 @@ class SupportedModel(Enum): "gemini-2.5-flash-lite": SupportedModel.GEMINI_2_5_FLASH_LITE, "gemini-3-flash-preview": SupportedModel.GEMINI_3_FLASH_PREVIEW, "gemini-3.1-pro-preview": SupportedModel.GEMINI_3_1_PRO_PREVIEW, - "gemini-3.1-flash-lite-preview": SupportedModel.GEMINI_3_1_FLASH_LITE_PREVIEW, + "gemini-3.1-flash-lite": SupportedModel.GEMINI_3_1_FLASH_LITE, "gemini-2.5-flash-image": SupportedModel.GEMINI_2_5_FLASH_IMAGE, "gemini-3.1-flash-image": SupportedModel.GEMINI_3_1_FLASH_IMAGE, "gemini-3.5-flash": SupportedModel.GEMINI_3_5_FLASH, @@ -622,10 +593,8 @@ class SupportedModel(Enum): "grok-4-fast": SupportedModel.GROK_4_FAST, "grok-4-1-fast": SupportedModel.GROK_4_1_FAST, "grok-4.1-fast": SupportedModel.GROK_4_1_FAST, - "grok-4-1-fast-non-reasoning": SupportedModel.GROK_4_1_FAST_NON_REASONING, "grok-4.20-reasoning": SupportedModel.GROK_4_20_REASONING, "grok-4.20-non-reasoning": SupportedModel.GROK_4_20_NON_REASONING, - "grok-code-fast-1": SupportedModel.GROK_CODE_FAST_1, "grok-2-image": SupportedModel.GROK_2_IMAGE, "grok-2-image-1212": SupportedModel.GROK_2_IMAGE, "grok-2-image-latest": SupportedModel.GROK_2_IMAGE, @@ -656,11 +625,6 @@ class SupportedModel(Enum): # Z.ai "glm-5.2": SupportedModel.GLM_5_2, "glm-image": SupportedModel.GLM_IMAGE, - # Legacy — not in current SDK, retained for older SDK versions - "grok-3-mini-beta": SupportedModel.GROK_3_MINI, # old beta alias - "grok-3-mini": SupportedModel.GROK_3_MINI, - "grok-3-beta": SupportedModel.GROK_3, # old beta alias - "grok-3": SupportedModel.GROK_3, } # Build the rate card automatically from the enum (for backward compat with util.py) diff --git a/tee_gateway/test/test_tee_core.py b/tee_gateway/test/test_tee_core.py index 55f87aa..2a3f0da 100644 --- a/tee_gateway/test/test_tee_core.py +++ b/tee_gateway/test/test_tee_core.py @@ -445,7 +445,7 @@ def test_all_four_providers_reachable(self): providers = { get_model_config("gpt-4.1").provider, get_model_config("claude-sonnet-4-5").provider, - get_model_config("gemini-2.5-flash").provider, + get_model_config("gemini-3.5-flash").provider, get_model_config("grok-4").provider, } self.assertEqual(providers, {"openai", "anthropic", "google", "x-ai"}) @@ -460,7 +460,7 @@ def test_anthropic_model_lookup(self): self.assertEqual(cfg.provider, "anthropic") def test_google_model_lookup(self): - cfg = get_model_config("gemini-2.5-flash") + cfg = get_model_config("gemini-3.5-flash") self.assertEqual(cfg.provider, "google") def test_xai_model_lookup(self): @@ -493,7 +493,7 @@ def test_model_alias_and_dated_name_resolve_identically(self): def test_models_with_force_temperature_have_positive_value(self): """Any model that forces a temperature must set it to a positive float. We don't pin the exact value — just confirm the field is plausible if set.""" - cfg = get_model_config("o4-mini") + cfg = get_model_config("o3") if cfg.force_temperature is not None: self.assertGreater(cfg.force_temperature, 0) diff --git a/tee_gateway/test/test_web_search.py b/tee_gateway/test/test_web_search.py index 277dcd3..c983b7c 100644 --- a/tee_gateway/test/test_web_search.py +++ b/tee_gateway/test/test_web_search.py @@ -46,7 +46,7 @@ def test_price_uses_provider_default(self): # grok-4 -> xAI default ($0.025/search unit) self.assertEqual(get_web_search_price_usd("grok-4"), Decimal("0.025")) # gemini -> google default ($0.035/grounded request) - self.assertEqual(get_web_search_price_usd("gemini-2.5-flash"), Decimal("0.035")) + self.assertEqual(get_web_search_price_usd("gemini-3.5-flash"), Decimal("0.035")) def test_unsupported_provider_is_free(self): # ByteDance has no native web search -> no charge diff --git a/tests/test_pricing.py b/tests/test_pricing.py index 89d9136..347274a 100644 --- a/tests/test_pricing.py +++ b/tests/test_pricing.py @@ -76,6 +76,21 @@ def test_all_lookup_keys_resolve(self): self.assertGreater(cfg.input_price_usd, 0) self.assertGreater(cfg.output_price_usd, 0) + def test_removed_upstream_models_are_unsupported(self): + removed_models = [ + "gemini-3.1-flash-lite-preview", + "grok-4-1-fast-non-reasoning", + "grok-code-fast-1", + "grok-3", + "grok-3-beta", + "grok-3-mini", + "grok-3-mini-beta", + ] + for model in removed_models: + with self.subTest(model=model): + with self.assertRaises(ValueError): + get_model_config(model) + # ── Anthropic Sonnet ──────────────────────────────────────────────────── def test_claude_sonnet_4_5_resolves(self): @@ -262,8 +277,8 @@ def test_gemini_3_1_pro_preview_resolves(self): self.assertEqual(cfg.output_price_usd, Decimal("0.000012")) self.assertEqual(cfg.thinking_budget, 128) - def test_gemini_3_1_flash_lite_preview_resolves(self): - cfg = get_model_config("gemini-3.1-flash-lite-preview") + def test_gemini_3_1_flash_lite_resolves(self): + cfg = get_model_config("gemini-3.1-flash-lite") self.assertEqual(cfg.provider, "google") self.assertEqual(cfg.input_price_usd, Decimal("0.00000025")) self.assertEqual(cfg.output_price_usd, Decimal("0.0000015")) @@ -314,14 +329,6 @@ def test_grok_4_1_fast_dot_notation_resolves(self): cfg = get_model_config("grok-4.1-fast") self.assertEqual(cfg, get_model_config("grok-4-1-fast")) - def test_grok_3_mini_resolves(self): - cfg = get_model_config("grok-3-mini") - self.assertEqual(cfg.provider, "x-ai") - - def test_grok_3_resolves(self): - cfg = get_model_config("grok-3") - self.assertEqual(cfg.provider, "x-ai") - def test_grok_4_20_reasoning_resolves(self): cfg = get_model_config("grok-4.20-reasoning") self.assertEqual(cfg.provider, "x-ai") @@ -334,12 +341,6 @@ def test_grok_4_20_non_reasoning_resolves(self): self.assertEqual(cfg.input_price_usd, Decimal("0.000002")) self.assertEqual(cfg.output_price_usd, Decimal("0.000006")) - def test_grok_code_fast_1_resolves(self): - cfg = get_model_config("grok-code-fast-1") - self.assertEqual(cfg.provider, "x-ai") - self.assertEqual(cfg.input_price_usd, Decimal("0.0000002")) - self.assertEqual(cfg.output_price_usd, Decimal("0.0000015")) - def test_claude_opus_4_7_resolves(self): cfg = get_model_config("claude-opus-4-7") self.assertEqual(cfg.provider, "anthropic") @@ -657,9 +658,9 @@ def test_gemini_3_1_pro_preview_cost(self): # 1000*0.000002 + 500*0.000012 = 0.002 + 0.006 = 0.008 USD = 8e15 wei self.assertEqual(cost, 8_000_000_000_000_000) - def test_gemini_3_1_flash_lite_preview_cost(self): - cost = self._calc("gemini-3.1-flash-lite-preview", 1000, 500) - expected = _expected_cost_opg("gemini-3.1-flash-lite-preview", 1000, 500) + def test_gemini_3_1_flash_lite_cost(self): + cost = self._calc("gemini-3.1-flash-lite", 1000, 500) + expected = _expected_cost_opg("gemini-3.1-flash-lite", 1000, 500) self.assertEqual(cost, expected) # 1000*0.00000025 + 500*0.0000015 = 0.00025 + 0.00075 = 0.001 USD = 1e15 wei self.assertEqual(cost, 1_000_000_000_000_000) @@ -708,23 +709,6 @@ def test_grok_4_20_non_reasoning_cost(self): cost = self._calc("grok-4.20-non-reasoning", 1000, 500) self.assertEqual(cost, self._calc("grok-4.20-reasoning", 1000, 500)) - def test_grok_code_fast_1_cost(self): - cost = self._calc("grok-code-fast-1", 1000, 500) - expected = _expected_cost_opg("grok-code-fast-1", 1000, 500) - self.assertEqual(cost, expected) - # 1000*0.0000002 + 500*0.0000015 = 0.0002 + 0.00075 = 0.00095 USD = 9.5e14 wei - self.assertEqual(cost, 950_000_000_000_000) - - def test_grok_3_mini_cost(self): - cost = self._calc("grok-3-mini", 1000, 500) - expected = _expected_cost_opg("grok-3-mini", 1000, 500) - self.assertEqual(cost, expected) - - def test_grok_3_cost(self): - cost = self._calc("grok-3", 1000, 500) - expected = _expected_cost_opg("grok-3", 1000, 500) - self.assertEqual(cost, expected) - # ── ByteDance (BytePlus ModelArk) ─────────────────────────────────────── def test_seed_1_6_cost(self): diff --git a/tests/test_server.py b/tests/test_server.py index 07bdee4..b3356b6 100644 --- a/tests/test_server.py +++ b/tests/test_server.py @@ -346,7 +346,7 @@ def test_chat_completion_google(mock_get_model, mock_google_model): response = client.post( "/v1/chat/completions", json={ - "model": "gemini-2.5-flash-preview", + "model": "gemini-3.5-flash", "messages": [{"role": "user", "content": "Say hello"}], "temperature": 0.7, "max_tokens": 100, @@ -358,7 +358,7 @@ def test_chat_completion_google(mock_get_model, mock_google_model): assert "message" in data assert "Hello from Google!" in data["message"]["content"] assert data["finish_reason"] == "stop" - assert data["model"] == "gemini-2.5-flash-preview" + assert data["model"] == "gemini-3.5-flash" assert data["usage"]["total_tokens"] == 17 @@ -370,7 +370,7 @@ def test_chat_completion_xai(mock_get_model, mock_xai_model): response = client.post( "/v1/chat/completions", json={ - "model": "grok-3-beta", + "model": "grok-4.3", "messages": [{"role": "user", "content": "Say hello"}], "temperature": 0.7, "max_tokens": 100, @@ -382,7 +382,7 @@ def test_chat_completion_xai(mock_get_model, mock_xai_model): assert "message" in data assert "Hello from xAI!" in data["message"]["content"] assert data["finish_reason"] == "stop" - assert data["model"] == "grok-3-beta" + assert data["model"] == "grok-4.3" assert data["usage"]["total_tokens"] == 20 @@ -488,7 +488,7 @@ def test_chat_tool_calls_google(mock_get_model, mock_tool_call_model): response = client.post( "/v1/chat/completions", json={ - "model": "gemini-2.5-flash-preview", + "model": "gemini-3.5-flash", "messages": [{"role": "user", "content": "What's the weather?"}], "tools": [ { @@ -516,7 +516,7 @@ def test_chat_tool_calls_xai(mock_get_model, mock_tool_call_model): response = client.post( "/v1/chat/completions", json={ - "model": "grok-3-beta", + "model": "grok-4.3", "messages": [{"role": "user", "content": "What's the weather?"}], "tools": [ { @@ -628,7 +628,7 @@ def test_chat_streaming_google(mock_get_model, mock_google_model): response = client.post( "/v1/chat/completions/stream", json={ - "model": "gemini-2.5-flash-preview", + "model": "gemini-3.5-flash", "messages": [{"role": "user", "content": "Say hello"}], }, ) @@ -662,7 +662,7 @@ def test_chat_streaming_xai(mock_get_model, mock_xai_model): response = client.post( "/v1/chat/completions/stream", json={ - "model": "grok-3-beta", + "model": "grok-4.3", "messages": [{"role": "user", "content": "Say hello"}], }, ) @@ -799,7 +799,7 @@ def test_chat_streaming_tool_calls_google(mock_get_model, mock_tool_call_model): response = client.post( "/v1/chat/completions/stream", json={ - "model": "gemini-2.5-flash-preview", + "model": "gemini-3.5-flash", "messages": [{"role": "user", "content": "What's the weather?"}], "tools": [ { @@ -848,7 +848,7 @@ def test_chat_streaming_tool_calls_xai(mock_get_model, mock_tool_call_model): response = client.post( "/v1/chat/completions/stream", json={ - "model": "grok-3-beta", + "model": "grok-4.3", "messages": [{"role": "user", "content": "What's the weather?"}], "tools": [ { diff --git a/tests/test_structured_outputs.py b/tests/test_structured_outputs.py index 71fbd20..a7c1d17 100644 --- a/tests/test_structured_outputs.py +++ b/tests/test_structured_outputs.py @@ -1252,7 +1252,7 @@ def test_prompt_tokens_preserved_from_first_chunk( mock_tee_keys.return_value = mock_keys req = CreateChatCompletionRequest( - model="gemini-2.5-flash", + model="gemini-3.5-flash", messages=[], temperature=1.0, stream=True, @@ -1299,7 +1299,7 @@ def test_replace_behaviour_would_lose_prompt_tokens( mock_tee_keys.return_value = mock_keys req = CreateChatCompletionRequest( - model="gemini-2.5-flash", + model="gemini-3.5-flash", messages=[], temperature=1.0, stream=True, @@ -1340,7 +1340,7 @@ def test_no_usage_chunks_gives_no_usage_in_final( mock_tee_keys.return_value = mock_keys req = CreateChatCompletionRequest( - model="gemini-2.5-flash", + model="gemini-3.5-flash", messages=[], temperature=1.0, stream=True, From 8af7ce725f91bd37bedd9ef19d55be39b88bfe50 Mon Sep 17 00:00:00 2001 From: Aniket Dixit Date: Fri, 10 Jul 2026 01:52:40 +0530 Subject: [PATCH 2/2] 404 error for model shutdown --- tee_gateway/controllers/chat_controller.py | 21 +++- .../controllers/completions_controller.py | 6 ++ tee_gateway/image_generation.py | 15 ++- tee_gateway/model_errors.py | 81 ++++++++++++++++ tee_gateway/model_registry.py | 55 +---------- tee_gateway/test/test_tee_core.py | 8 +- tee_gateway/test/test_web_search.py | 4 +- tests/test_pricing.py | 96 ++++--------------- 8 files changed, 146 insertions(+), 140 deletions(-) create mode 100644 tee_gateway/model_errors.py diff --git a/tee_gateway/controllers/chat_controller.py b/tee_gateway/controllers/chat_controller.py index 1392b6a..3a4dbe4 100644 --- a/tee_gateway/controllers/chat_controller.py +++ b/tee_gateway/controllers/chat_controller.py @@ -39,6 +39,10 @@ create_image_generation_response, create_image_generation_streaming_response, ) +from tee_gateway.model_errors import ( + is_upstream_model_not_served, + upstream_model_not_served_payload, +) from tee_gateway.model_registry import get_model_config from tee_gateway.pricing import compute_session_cost @@ -367,6 +371,8 @@ def _create_non_streaming_response(chat_request: CreateChatCompletionRequest): except Exception as e: logger.error(f"Chat completion error: {str(e)}", exc_info=True) + if is_upstream_model_not_served(e): + return upstream_model_not_served_payload(chat_request.model, e), 404 return { "error": str(e) or "Request processing failed", "exception_type": type(e).__name__, @@ -805,10 +811,15 @@ def generate(): # event IS the terminal marker for the error path — we do NOT # emit a trailing `[DONE]`, which conventionally signals a # clean completion and would mis-signal an errored stream. - error_payload = { - "error": str(e) or "Stream processing failed", - "exception_type": type(e).__name__, - } + if is_upstream_model_not_served(e): + error_payload = upstream_model_not_served_payload( + chat_request.model, e + ) + else: + error_payload = { + "error": str(e) or "Stream processing failed", + "exception_type": type(e).__name__, + } yield f"data: {json.dumps(error_payload)}\n\n" return Response( @@ -822,6 +833,8 @@ def generate(): except Exception as e: logger.error(f"Stream setup error: {str(e)}", exc_info=True) + if is_upstream_model_not_served(e): + return upstream_model_not_served_payload(chat_request.model, e), 404 return { "error": str(e) or "Stream setup failed", "exception_type": type(e).__name__, diff --git a/tee_gateway/controllers/completions_controller.py b/tee_gateway/controllers/completions_controller.py index 773dd74..ad982d6 100644 --- a/tee_gateway/controllers/completions_controller.py +++ b/tee_gateway/controllers/completions_controller.py @@ -18,6 +18,10 @@ extract_web_search_count, extract_usage, ) +from tee_gateway.model_errors import ( + is_upstream_model_not_served, + upstream_model_not_served_payload, +) from tee_gateway.pricing import compute_session_cost logger = logging.getLogger(__name__) @@ -116,6 +120,8 @@ def create_completion(body): except Exception as e: logger.error(f"Completion error: {str(e)}", exc_info=True) + if is_upstream_model_not_served(e): + return upstream_model_not_served_payload(body.model, e), 404 return { "error": str(e) or "Request processing failed", "exception_type": type(e).__name__, diff --git a/tee_gateway/image_generation.py b/tee_gateway/image_generation.py index 07776fa..52976ca 100644 --- a/tee_gateway/image_generation.py +++ b/tee_gateway/image_generation.py @@ -41,6 +41,10 @@ from tee_gateway.models.create_chat_completion_response import ( CreateChatCompletionResponse, ) +from tee_gateway.model_errors import ( + is_upstream_model_not_served, + upstream_model_not_served_payload, +) from tee_gateway.model_registry import get_model_config from tee_gateway.pricing import compute_session_cost from tee_gateway.tee_manager import get_tee_keys, compute_tee_msg_hash @@ -498,10 +502,13 @@ def generate(): # Surface the real exception detail to the client (matching the chat # streaming path) so browser-side logs show the actual cause instead # of an opaque generic string. - error_payload = { - "error": str(e) or "Stream processing failed", - "exception_type": type(e).__name__, - } + if is_upstream_model_not_served(e): + error_payload = upstream_model_not_served_payload(chat_request.model, e) + else: + error_payload = { + "error": str(e) or "Stream processing failed", + "exception_type": type(e).__name__, + } yield f"data: {json.dumps(error_payload)}\n\n" return Response( diff --git a/tee_gateway/model_errors.py b/tee_gateway/model_errors.py new file mode 100644 index 0000000..f81d3a5 --- /dev/null +++ b/tee_gateway/model_errors.py @@ -0,0 +1,81 @@ +"""Helpers for turning provider model failures into client-facing errors.""" + +from __future__ import annotations + +from typing import Any + +from tee_gateway.model_registry import get_model_config + + +def exception_status_code(exc: Exception) -> int | None: + """Extract an HTTP/gRPC-ish status code from common provider SDK exceptions.""" + status = getattr(exc, "status_code", None) + if isinstance(status, int): + return status + + response = getattr(exc, "response", None) + response_status = getattr(response, "status_code", None) + if isinstance(response_status, int): + return response_status + + code_attr = getattr(exc, "code", None) + code_value: Any + if callable(code_attr): + try: + code_value = code_attr() + except Exception: + code_value = None + else: + code_value = code_attr + + if isinstance(code_value, int): + return code_value + if getattr(code_value, "name", None) == "NOT_FOUND": + return 404 + + return None + + +def is_upstream_model_not_served(exc: Exception) -> bool: + """Return True when a provider says a registered model is not available.""" + if exception_status_code(exc) == 404: + return True + + exc_type = type(exc).__name__.lower() + if "notfound" in exc_type or "not_found" in exc_type: + return True + + message = str(exc).lower() + model_missing_phrases = ( + "model not found", + "model is not found", + "model does not exist", + "model_not_found", + "not found for model", + "not served", + "is not supported for", + ) + return any(phrase in message for phrase in model_missing_phrases) + + +def upstream_model_not_served_payload(model: str, exc: Exception) -> dict[str, Any]: + """Build an OpenAI-style-ish payload for a provider model availability miss.""" + payload: dict[str, Any] = { + "error": ( + f"Upstream provider does not currently serve model '{model}'. " + "The model may have been removed, renamed, or disabled upstream." + ), + "code": "model_not_served_upstream", + "model": model, + "exception_type": type(exc).__name__, + "upstream_error": str(exc) or type(exc).__name__, + } + + try: + cfg = get_model_config(model) + except Exception: + return payload + + payload["provider"] = cfg.provider + payload["api_name"] = cfg.api_name + return payload diff --git a/tee_gateway/model_registry.py b/tee_gateway/model_registry.py index e6463c4..3458c18 100644 --- a/tee_gateway/model_registry.py +++ b/tee_gateway/model_registry.py @@ -287,27 +287,6 @@ class SupportedModel(Enum): ) # ── Google Gemini ─────────────────────────────────────────────────── - GEMINI_2_5_FLASH = ModelConfig( - provider="google", - api_name="gemini-2.5-flash", - input_price_usd=Decimal("0.0000003"), - output_price_usd=Decimal("0.0000025"), - thinking_budget=0, - ) - GEMINI_2_5_PRO = ModelConfig( - provider="google", - api_name="gemini-2.5-pro", - input_price_usd=Decimal("0.00000125"), - output_price_usd=Decimal("0.00001"), - thinking_budget=128, - ) - GEMINI_2_5_FLASH_LITE = ModelConfig( - provider="google", - api_name="gemini-2.5-flash-lite", - input_price_usd=Decimal("0.0000001"), - output_price_usd=Decimal("0.0000004"), - thinking_budget=0, - ) GEMINI_3_FLASH_PREVIEW = ModelConfig( provider="google", api_name="gemini-3-flash-preview", @@ -371,33 +350,15 @@ class SupportedModel(Enum): input_price_usd=Decimal("0.00000125"), output_price_usd=Decimal("0.0000025"), ) - GROK_4 = ModelConfig( - provider="x-ai", - api_name="grok-4", - input_price_usd=Decimal("0.000003"), - output_price_usd=Decimal("0.000015"), - ) - GROK_4_FAST = ModelConfig( - provider="x-ai", - api_name="grok-4-fast", - input_price_usd=Decimal("0.0000002"), - output_price_usd=Decimal("0.0000005"), - ) - GROK_4_1_FAST = ModelConfig( - provider="x-ai", - api_name="grok-4-1-fast", - input_price_usd=Decimal("0.0000002"), - output_price_usd=Decimal("0.0000005"), - ) GROK_4_20_REASONING = ModelConfig( provider="x-ai", - api_name="grok-4.20-reasoning", + api_name="grok-4.20-0309-reasoning", input_price_usd=Decimal("0.000002"), output_price_usd=Decimal("0.000006"), ) GROK_4_20_NON_REASONING = ModelConfig( provider="x-ai", - api_name="grok-4.20-non-reasoning", + api_name="grok-4.20-0309-non-reasoning", input_price_usd=Decimal("0.000002"), output_price_usd=Decimal("0.000006"), ) @@ -405,7 +366,7 @@ class SupportedModel(Enum): # (Aurora). Billed at a flat $0.07 per generated image; token prices unused. GROK_2_IMAGE = ModelConfig( provider="x-ai", - api_name="grok-2-image-1212", + api_name="grok-imagine-image", input_price_usd=Decimal("0"), output_price_usd=Decimal("0"), image_generation=True, @@ -576,9 +537,6 @@ class SupportedModel(Enum): "claude-opus-4-8": SupportedModel.CLAUDE_OPUS_4_8, "claude-fable-5": SupportedModel.CLAUDE_FABLE_5, # Google - "gemini-2.5-flash": SupportedModel.GEMINI_2_5_FLASH, - "gemini-2.5-pro": SupportedModel.GEMINI_2_5_PRO, - "gemini-2.5-flash-lite": SupportedModel.GEMINI_2_5_FLASH_LITE, "gemini-3-flash-preview": SupportedModel.GEMINI_3_FLASH_PREVIEW, "gemini-3.1-pro-preview": SupportedModel.GEMINI_3_1_PRO_PREVIEW, "gemini-3.1-flash-lite": SupportedModel.GEMINI_3_1_FLASH_LITE, @@ -589,15 +547,10 @@ class SupportedModel(Enum): "grok-4.5": SupportedModel.GROK_4_5, "grok-4.5-latest": SupportedModel.GROK_4_5, "grok-4.3": SupportedModel.GROK_4_3, - "grok-4": SupportedModel.GROK_4, - "grok-4-fast": SupportedModel.GROK_4_FAST, - "grok-4-1-fast": SupportedModel.GROK_4_1_FAST, - "grok-4.1-fast": SupportedModel.GROK_4_1_FAST, "grok-4.20-reasoning": SupportedModel.GROK_4_20_REASONING, "grok-4.20-non-reasoning": SupportedModel.GROK_4_20_NON_REASONING, + "grok-imagine-image": SupportedModel.GROK_2_IMAGE, "grok-2-image": SupportedModel.GROK_2_IMAGE, - "grok-2-image-1212": SupportedModel.GROK_2_IMAGE, - "grok-2-image-latest": SupportedModel.GROK_2_IMAGE, # ByteDance "seed-1-6-250615": SupportedModel.SEED_1_6, "seed-1.6": SupportedModel.SEED_1_6, diff --git a/tee_gateway/test/test_tee_core.py b/tee_gateway/test/test_tee_core.py index 2a3f0da..500365b 100644 --- a/tee_gateway/test/test_tee_core.py +++ b/tee_gateway/test/test_tee_core.py @@ -446,7 +446,7 @@ def test_all_four_providers_reachable(self): get_model_config("gpt-4.1").provider, get_model_config("claude-sonnet-4-5").provider, get_model_config("gemini-3.5-flash").provider, - get_model_config("grok-4").provider, + get_model_config("grok-4.3").provider, } self.assertEqual(providers, {"openai", "anthropic", "google", "x-ai"}) @@ -464,7 +464,7 @@ def test_google_model_lookup(self): self.assertEqual(cfg.provider, "google") def test_xai_model_lookup(self): - cfg = get_model_config("grok-4") + cfg = get_model_config("grok-4.3") self.assertEqual(cfg.provider, "x-ai") def test_unknown_model_raises_value_error(self): @@ -810,7 +810,7 @@ def test_plain_text_request_passes(self): def test_image_blocked_when_model_lacks_support(self): with mock.patch(self.CAPS, return_value={"image_inputs": False}): with self.assertRaises(AttachmentValidationError): - validate_attachments(self._image_msg("aGVsbG8="), "grok-4") + validate_attachments(self._image_msg("aGVsbG8="), "grok-4.3") def test_image_allowed_when_model_supports(self): with mock.patch(self.CAPS, return_value={"image_inputs": True}): @@ -826,7 +826,7 @@ def test_pdf_blocked_when_model_lacks_support(self): self.CAPS, return_value={"image_inputs": True, "pdf_inputs": False} ): with self.assertRaises(AttachmentValidationError): - validate_attachments(self._pdf_msg("JVBERi0="), "grok-4") + validate_attachments(self._pdf_msg("JVBERi0="), "grok-4.3") # --------------------------------------------------------------------------- diff --git a/tee_gateway/test/test_web_search.py b/tee_gateway/test/test_web_search.py index c983b7c..3a31ace 100644 --- a/tee_gateway/test/test_web_search.py +++ b/tee_gateway/test/test_web_search.py @@ -43,8 +43,8 @@ def test_provider_support_predicate(self): def test_price_uses_provider_default(self): # gpt-4.1 -> openai default ($0.01/search) self.assertEqual(get_web_search_price_usd("gpt-4.1"), Decimal("0.01")) - # grok-4 -> xAI default ($0.025/search unit) - self.assertEqual(get_web_search_price_usd("grok-4"), Decimal("0.025")) + # grok-4.3 -> xAI default ($0.025/search unit) + self.assertEqual(get_web_search_price_usd("grok-4.3"), Decimal("0.025")) # gemini -> google default ($0.035/grounded request) self.assertEqual(get_web_search_price_usd("gemini-3.5-flash"), Decimal("0.035")) diff --git a/tests/test_pricing.py b/tests/test_pricing.py index 347274a..6110295 100644 --- a/tests/test_pricing.py +++ b/tests/test_pricing.py @@ -79,12 +79,21 @@ def test_all_lookup_keys_resolve(self): def test_removed_upstream_models_are_unsupported(self): removed_models = [ "gemini-3.1-flash-lite-preview", + "gemini-2.5-flash", + "gemini-2.5-pro", + "gemini-2.5-flash-lite", + "grok-4", + "grok-4-fast", + "grok-4-1-fast", + "grok-4.1-fast", "grok-4-1-fast-non-reasoning", "grok-code-fast-1", "grok-3", "grok-3-beta", "grok-3-mini", "grok-3-mini-beta", + "grok-2-image-1212", + "grok-2-image-latest", ] for model in removed_models: with self.subTest(model=model): @@ -253,19 +262,6 @@ def test_gpt_5_6_luna_resolves(self): # ── Google ────────────────────────────────────────────────────────────── - def test_gemini_2_5_flash_resolves(self): - cfg = get_model_config("gemini-2.5-flash") - self.assertEqual(cfg.provider, "google") - self.assertEqual(cfg.input_price_usd, Decimal("0.0000003")) - - def test_gemini_2_5_pro_resolves(self): - cfg = get_model_config("gemini-2.5-pro") - self.assertEqual(cfg.provider, "google") - - def test_gemini_2_5_flash_lite_resolves(self): - cfg = get_model_config("gemini-2.5-flash-lite") - self.assertEqual(cfg.provider, "google") - def test_gemini_3_flash_preview_resolves(self): cfg = get_model_config("gemini-3-flash-preview") self.assertEqual(cfg.provider, "google") @@ -313,34 +309,26 @@ def test_grok_4_5_latest_resolves(self): cfg = get_model_config("grok-4.5-latest") self.assertEqual(cfg, get_model_config("grok-4.5")) - def test_grok_4_resolves(self): - cfg = get_model_config("grok-4") - self.assertEqual(cfg.provider, "x-ai") - - def test_grok_4_fast_resolves(self): - cfg = get_model_config("grok-4-fast") - self.assertEqual(cfg.provider, "x-ai") - - def test_grok_4_1_fast_resolves(self): - cfg = get_model_config("grok-4-1-fast") - self.assertEqual(cfg.provider, "x-ai") - - def test_grok_4_1_fast_dot_notation_resolves(self): - cfg = get_model_config("grok-4.1-fast") - self.assertEqual(cfg, get_model_config("grok-4-1-fast")) - def test_grok_4_20_reasoning_resolves(self): cfg = get_model_config("grok-4.20-reasoning") self.assertEqual(cfg.provider, "x-ai") + self.assertEqual(cfg.api_name, "grok-4.20-0309-reasoning") self.assertEqual(cfg.input_price_usd, Decimal("0.000002")) self.assertEqual(cfg.output_price_usd, Decimal("0.000006")) def test_grok_4_20_non_reasoning_resolves(self): cfg = get_model_config("grok-4.20-non-reasoning") self.assertEqual(cfg.provider, "x-ai") + self.assertEqual(cfg.api_name, "grok-4.20-0309-non-reasoning") self.assertEqual(cfg.input_price_usd, Decimal("0.000002")) self.assertEqual(cfg.output_price_usd, Decimal("0.000006")) + def test_grok_imagine_image_resolves(self): + cfg = get_model_config("grok-imagine-image") + self.assertEqual(cfg.provider, "x-ai") + self.assertEqual(cfg.api_name, "grok-imagine-image") + self.assertTrue(cfg.image_generation) + def test_claude_opus_4_7_resolves(self): cfg = get_model_config("claude-opus-4-7") self.assertEqual(cfg.provider, "anthropic") @@ -627,25 +615,6 @@ def test_claude_opus_4_7_cost(self): # ── Google Gemini ──────────────────────────────────────────────────────── - def test_gemini_2_5_flash_cost(self): - cost = self._calc("gemini-2.5-flash", 1000, 500) - expected = _expected_cost_opg("gemini-2.5-flash", 1000, 500) - self.assertEqual(cost, expected) - # 1000*0.0000003 + 500*0.0000025 = 0.0003 + 0.00125 = 0.00155 USD - self.assertEqual(cost, 1_550_000_000_000_000) - - def test_gemini_2_5_flash_lite_cost(self): - cost = self._calc("gemini-2.5-flash-lite", 1000, 500) - expected = _expected_cost_opg("gemini-2.5-flash-lite", 1000, 500) - self.assertEqual(cost, expected) - # 1000*0.0000001 + 500*0.0000004 = 0.0001 + 0.0002 = 0.0003 USD - self.assertEqual(cost, 300_000_000_000_000) - - def test_gemini_2_5_pro_cost(self): - cost = self._calc("gemini-2.5-pro", 1000, 500) - expected = _expected_cost_opg("gemini-2.5-pro", 1000, 500) - self.assertEqual(cost, expected) - def test_gemini_3_flash_preview_cost(self): cost = self._calc("gemini-3-flash-preview", 1000, 500) expected = _expected_cost_opg("gemini-3-flash-preview", 1000, 500) @@ -680,24 +649,6 @@ def test_grok_4_5_latest_cost(self): self._calc("grok-4.5", 1000, 500), ) - def test_grok_4_cost(self): - cost = self._calc("grok-4", 1000, 500) - expected = _expected_cost_opg("grok-4", 1000, 500) - self.assertEqual(cost, expected) - # Same pricing tier as claude-sonnet-4-5 - self.assertEqual(cost, 10_500_000_000_000_000) - - def test_grok_4_fast_cost(self): - cost = self._calc("grok-4-fast", 1000, 500) - expected = _expected_cost_opg("grok-4-fast", 1000, 500) - self.assertEqual(cost, expected) - # 1000*0.0000002 + 500*0.0000005 = 0.0002 + 0.00025 = 0.00045 USD - self.assertEqual(cost, 450_000_000_000_000) - - def test_grok_4_1_fast_cost(self): - cost = self._calc("grok-4-1-fast", 1000, 500) - self.assertEqual(cost, self._calc("grok-4-fast", 1000, 500)) - def test_grok_4_20_reasoning_cost(self): cost = self._calc("grok-4.20-reasoning", 1000, 500) expected = _expected_cost_opg("grok-4.20-reasoning", 1000, 500) @@ -757,15 +708,10 @@ def test_haiku_cheaper_than_sonnet(self): self.assertLess(haiku, sonnet) def test_gemini_flash_lite_cheaper_than_flash(self): - lite = self._calc("gemini-2.5-flash-lite", 1000, 1000) - flash = self._calc("gemini-2.5-flash", 1000, 1000) + lite = self._calc("gemini-3.1-flash-lite", 1000, 1000) + flash = self._calc("gemini-3.5-flash", 1000, 1000) self.assertLess(lite, flash) - def test_grok_4_fast_cheaper_than_grok_4(self): - fast = self._calc("grok-4-fast", 1000, 1000) - full = self._calc("grok-4", 1000, 1000) - self.assertLess(fast, full) - class TestCalculateSessionCostEdgeCases(unittest.TestCase): """Edge cases for compute_session_cost.""" @@ -782,8 +728,8 @@ def test_rounding_ceiling(self): """Fractional token costs are always rounded UP.""" # 1 output token of Haiku: 0.000005 USD = 5e12 wei — exact self.assertEqual(_calc_opg("claude-haiku-4-5", 0, 1), 5_000_000_000_000) - # 1 input token of Gemini Flash Lite: 0.0000001 USD = 1e11 wei — exact - self.assertEqual(_calc_opg("gemini-2.5-flash-lite", 1, 0), 100_000_000_000) + # 1 input token of Gemini 3.1 Flash Lite: 0.00000025 USD = 2.5e11 wei — exact + self.assertEqual(_calc_opg("gemini-3.1-flash-lite", 1, 0), 250_000_000_000) def test_model_name_case_insensitive(self): self.assertEqual(