From 2326088a32324d6513e8e9ce57af9778ea84d3c1 Mon Sep 17 00:00:00 2001 From: halfaipg Date: Fri, 28 Aug 2026 21:12:12 -0400 Subject: [PATCH 1/8] feat(provider): add AI Power Grid --- litellm/constants.py | 2 + litellm/llms/openai_like/providers.json | 9 ++ ...odel_prices_and_context_window_backup.json | 38 ++++++++ litellm/types/utils.py | 1 + model_prices_and_context_window.json | 38 ++++++++ tests/llm_translation/test_aipg.py | 96 +++++++++++++++++++ 6 files changed, 184 insertions(+) create mode 100644 tests/llm_translation/test_aipg.py diff --git a/litellm/constants.py b/litellm/constants.py index fc88086805f..5c235910c1a 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -627,6 +627,7 @@ LITELLM_CHAT_PROVIDERS: Final = [ "lemonade", "docker_model_runner", "amazon_nova", + "aipg", ] # Resolving these providers runs an OAuth device flow (their provider info IS the login), so any @@ -808,6 +809,7 @@ openai_compatible_endpoints: Final[list] = [ openai_compatible_providers: Final[list] = [ + "aipg", # AI Power Grid - JSON-configured provider "anyscale", "groq", "nvidia_nim", diff --git a/litellm/llms/openai_like/providers.json b/litellm/llms/openai_like/providers.json index a458a209ea9..44f626fd987 100644 --- a/litellm/llms/openai_like/providers.json +++ b/litellm/llms/openai_like/providers.json @@ -200,5 +200,14 @@ "temperature_max": 1.99 }, "supported_endpoints": ["/v1/chat/completions"] + }, + "aipg": { + "base_url": "https://api.aipowergrid.io/v1", + "api_key_env": "AIPG_API_KEY", + "api_base_env": "AIPG_API_BASE", + "param_mappings": { + "max_completion_tokens": "max_tokens" + }, + "supported_endpoints": ["/v1/chat/completions", "/v1/responses"] } } diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index bebbcc32181..1c57107350b 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -54596,5 +54596,43 @@ "output_cost_per_token_above_200k_tokens": 5e-06, "cache_read_input_token_cost_above_200k_tokens": 4e-07, "supports_response_schema": true + }, + "aipg/gpt-oss-120b": { + "input_cost_per_token": 7.5e-08, + "litellm_provider": "aipg", + "max_input_tokens": 60000, + "max_tokens": 60000, + "mode": "chat", + "output_cost_per_token": 3e-07, + "source": "https://docs.aipowergrid.io/streaming-api", + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_system_messages": true, + "supports_tool_choice": true + }, + "aipg/deepseek-v4-flash-nvfp4": { + "input_cost_per_token": 7e-08, + "litellm_provider": "aipg", + "max_input_tokens": 262144, + "max_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 1.4e-07, + "source": "https://docs.aipowergrid.io/streaming-api", + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_system_messages": true, + "supports_tool_choice": true + }, + "aipg/Smollm-135m": { + "input_cost_per_token": 5e-09, + "litellm_provider": "aipg", + "max_input_tokens": 2048, + "max_tokens": 2048, + "mode": "chat", + "output_cost_per_token": 1e-08, + "source": "https://docs.aipowergrid.io/streaming-api", + "supports_function_calling": false, + "supports_system_messages": true, + "supports_tool_choice": false } } diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 737361e5413..a25c71a6a25 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -3695,6 +3695,7 @@ GenericBudgetConfigType = dict[str, BudgetConfig] class LlmProviders(str, Enum): + AIPG = "aipg" OPENAI = "openai" CHATGPT = "chatgpt" OPENAI_LIKE = "openai_like" # embedding only diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index bebbcc32181..1c57107350b 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -54596,5 +54596,43 @@ "output_cost_per_token_above_200k_tokens": 5e-06, "cache_read_input_token_cost_above_200k_tokens": 4e-07, "supports_response_schema": true + }, + "aipg/gpt-oss-120b": { + "input_cost_per_token": 7.5e-08, + "litellm_provider": "aipg", + "max_input_tokens": 60000, + "max_tokens": 60000, + "mode": "chat", + "output_cost_per_token": 3e-07, + "source": "https://docs.aipowergrid.io/streaming-api", + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_system_messages": true, + "supports_tool_choice": true + }, + "aipg/deepseek-v4-flash-nvfp4": { + "input_cost_per_token": 7e-08, + "litellm_provider": "aipg", + "max_input_tokens": 262144, + "max_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 1.4e-07, + "source": "https://docs.aipowergrid.io/streaming-api", + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_system_messages": true, + "supports_tool_choice": true + }, + "aipg/Smollm-135m": { + "input_cost_per_token": 5e-09, + "litellm_provider": "aipg", + "max_input_tokens": 2048, + "max_tokens": 2048, + "mode": "chat", + "output_cost_per_token": 1e-08, + "source": "https://docs.aipowergrid.io/streaming-api", + "supports_function_calling": false, + "supports_system_messages": true, + "supports_tool_choice": false } } diff --git a/tests/llm_translation/test_aipg.py b/tests/llm_translation/test_aipg.py new file mode 100644 index 00000000000..8054e9fb8bc --- /dev/null +++ b/tests/llm_translation/test_aipg.py @@ -0,0 +1,96 @@ +"""Tests for the AI Power Grid JSON provider integration.""" + +import os +from unittest import mock + +import litellm + +AIPG_API_BASE = "https://api.aipowergrid.io/v1" + + +def test_aipg_json_registry(): + from litellm.llms.openai_like.json_loader import JSONProviderRegistry + + config = JSONProviderRegistry.get("aipg") + assert config is not None + assert config.base_url == AIPG_API_BASE + assert config.api_key_env == "AIPG_API_KEY" + assert config.api_base_env == "AIPG_API_BASE" + assert config.supported_endpoints == ["/v1/chat/completions", "/v1/responses"] + assert JSONProviderRegistry.supports_responses_api("aipg") is True + + +def test_aipg_get_openai_compatible_provider_info(): + from litellm.llms.openai_like.dynamic_config import create_config_class + from litellm.llms.openai_like.json_loader import JSONProviderRegistry + + provider = JSONProviderRegistry.get("aipg") + assert provider is not None + config = create_config_class(provider)() + + with mock.patch.dict(os.environ, {"AIPG_API_KEY": "grid-test"}, clear=True): + api_base, api_key = config._get_openai_compatible_provider_info(None, None) + assert api_base == AIPG_API_BASE + assert api_key == "grid-test" + + with mock.patch.dict( + os.environ, + { + "AIPG_API_KEY": "env-key", + "AIPG_API_BASE": "https://operator.example/v1", + }, + clear=True, + ): + api_base, api_key = config._get_openai_compatible_provider_info( + "https://explicit.example/v1", "explicit-key" + ) + assert api_base == "https://explicit.example/v1" + assert api_key == "explicit-key" + + mapped = config.map_openai_params( + non_default_params={"max_completion_tokens": 12, "temperature": 0.2}, + optional_params={}, + model="gpt-oss-120b", + drop_params=False, + ) + assert mapped["max_tokens"] == 12 + assert mapped["temperature"] == 0.2 + + +def test_get_llm_provider_aipg(): + from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider + + with mock.patch.dict(os.environ, {"AIPG_API_KEY": "grid-test"}, clear=True): + model, provider, api_key, api_base = get_llm_provider("aipg/gpt-oss-120b") + + assert model == "gpt-oss-120b" + assert provider == "aipg" + assert api_key == "grid-test" + assert api_base == AIPG_API_BASE + + +def test_aipg_model_metadata(): + original_model_cost = litellm.model_cost + original_env = os.environ.get("LITELLM_LOCAL_MODEL_COST_MAP") + try: + os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" + litellm.model_cost = litellm.get_model_cost_map(url="") + + expected = { + "aipg/gpt-oss-120b": (60000, 7.5e-08, 3e-07), + "aipg/deepseek-v4-flash-nvfp4": (262144, 7e-08, 1.4e-07), + "aipg/Smollm-135m": (2048, 5e-09, 1e-08), + } + for model, (context, input_cost, output_cost) in expected.items(): + info = litellm.get_model_info(model) + assert info["litellm_provider"] == "aipg" + assert info["mode"] == "chat" + assert info["max_input_tokens"] == context + assert info["input_cost_per_token"] == input_cost + assert info["output_cost_per_token"] == output_cost + finally: + litellm.model_cost = original_model_cost + if original_env is None: + os.environ.pop("LITELLM_LOCAL_MODEL_COST_MAP", None) + else: + os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = original_env From c8a4ccaffbca4084c0520c4c56efe4d5f3fc23fe Mon Sep 17 00:00:00 2001 From: halfaipg Date: Fri, 28 Aug 2026 21:22:33 -0400 Subject: [PATCH 2/8] fix(provider): document AIPG endpoint support --- litellm/constants.py | 2 +- .../provider_endpoints_support_backup.json | 18 ++++++++ provider_endpoints_support.json | 18 ++++++++ tests/llm_translation/test_aipg.py | 41 +++++++------------ 4 files changed, 51 insertions(+), 28 deletions(-) diff --git a/litellm/constants.py b/litellm/constants.py index 5c235910c1a..c3b99921893 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -809,7 +809,7 @@ openai_compatible_endpoints: Final[list] = [ openai_compatible_providers: Final[list] = [ - "aipg", # AI Power Grid - JSON-configured provider + "aipg", "anyscale", "groq", "nvidia_nim", diff --git a/litellm/provider_endpoints_support_backup.json b/litellm/provider_endpoints_support_backup.json index 86c14fb4cd8..62dec9d4ac9 100644 --- a/litellm/provider_endpoints_support_backup.json +++ b/litellm/provider_endpoints_support_backup.json @@ -84,6 +84,24 @@ "interactions": true } }, + "aipg": { + "display_name": "AI Power Grid (`aipg`)", + "url": "https://docs.litellm.ai/docs/providers/aipg", + "endpoints": { + "chat_completions": true, + "messages": false, + "responses": true, + "embeddings": false, + "image_generations": false, + "audio_transcriptions": false, + "audio_speech": false, + "moderations": false, + "batches": false, + "rerank": false, + "a2a": false, + "interactions": false + } + }, "ai21": { "display_name": "AI21 (`ai21`)", "url": "https://docs.litellm.ai/docs/providers/ai21", diff --git a/provider_endpoints_support.json b/provider_endpoints_support.json index 1d8d374c2c4..251c676975a 100644 --- a/provider_endpoints_support.json +++ b/provider_endpoints_support.json @@ -84,6 +84,24 @@ "interactions": true } }, + "aipg": { + "display_name": "AI Power Grid (`aipg`)", + "url": "https://docs.litellm.ai/docs/providers/aipg", + "endpoints": { + "chat_completions": true, + "messages": false, + "responses": true, + "embeddings": false, + "image_generations": false, + "audio_transcriptions": false, + "audio_speech": false, + "moderations": false, + "batches": false, + "rerank": false, + "a2a": false, + "interactions": false + } + }, "ai21": { "display_name": "AI21 (`ai21`)", "url": "https://docs.litellm.ai/docs/providers/ai21", diff --git a/tests/llm_translation/test_aipg.py b/tests/llm_translation/test_aipg.py index 8054e9fb8bc..075284aebfd 100644 --- a/tests/llm_translation/test_aipg.py +++ b/tests/llm_translation/test_aipg.py @@ -41,9 +41,7 @@ def test_aipg_get_openai_compatible_provider_info(): }, clear=True, ): - api_base, api_key = config._get_openai_compatible_provider_info( - "https://explicit.example/v1", "explicit-key" - ) + api_base, api_key = config._get_openai_compatible_provider_info("https://explicit.example/v1", "explicit-key") assert api_base == "https://explicit.example/v1" assert api_key == "explicit-key" @@ -70,27 +68,16 @@ def test_get_llm_provider_aipg(): def test_aipg_model_metadata(): - original_model_cost = litellm.model_cost - original_env = os.environ.get("LITELLM_LOCAL_MODEL_COST_MAP") - try: - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" - litellm.model_cost = litellm.get_model_cost_map(url="") - - expected = { - "aipg/gpt-oss-120b": (60000, 7.5e-08, 3e-07), - "aipg/deepseek-v4-flash-nvfp4": (262144, 7e-08, 1.4e-07), - "aipg/Smollm-135m": (2048, 5e-09, 1e-08), - } - for model, (context, input_cost, output_cost) in expected.items(): - info = litellm.get_model_info(model) - assert info["litellm_provider"] == "aipg" - assert info["mode"] == "chat" - assert info["max_input_tokens"] == context - assert info["input_cost_per_token"] == input_cost - assert info["output_cost_per_token"] == output_cost - finally: - litellm.model_cost = original_model_cost - if original_env is None: - os.environ.pop("LITELLM_LOCAL_MODEL_COST_MAP", None) - else: - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = original_env + model_cost = litellm.get_model_cost_map(url="") + expected = { + "aipg/gpt-oss-120b": (60000, 7.5e-08, 3e-07), + "aipg/deepseek-v4-flash-nvfp4": (262144, 7e-08, 1.4e-07), + "aipg/Smollm-135m": (2048, 5e-09, 1e-08), + } + for model, (context, input_cost, output_cost) in expected.items(): + info = model_cost[model] + assert info["litellm_provider"] == "aipg" + assert info["mode"] == "chat" + assert info["max_input_tokens"] == context + assert info["input_cost_per_token"] == input_cost + assert info["output_cost_per_token"] == output_cost From 62ab9e5a5d6b8ed544f03c58ec0b47f1fb926f02 Mon Sep 17 00:00:00 2001 From: halfaipg Date: Fri, 28 Aug 2026 23:48:16 -0400 Subject: [PATCH 3/8] fix(provider): correct AIPG output limits --- litellm/model_prices_and_context_window_backup.json | 9 ++++++--- model_prices_and_context_window.json | 9 ++++++--- tests/llm_translation/test_aipg.py | 10 ++++++---- 3 files changed, 18 insertions(+), 10 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 1c57107350b..9f41cc85f90 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -54601,7 +54601,8 @@ "input_cost_per_token": 7.5e-08, "litellm_provider": "aipg", "max_input_tokens": 60000, - "max_tokens": 60000, + "max_output_tokens": 32768, + "max_tokens": 32768, "mode": "chat", "output_cost_per_token": 3e-07, "source": "https://docs.aipowergrid.io/streaming-api", @@ -54614,7 +54615,8 @@ "input_cost_per_token": 7e-08, "litellm_provider": "aipg", "max_input_tokens": 262144, - "max_tokens": 262144, + "max_output_tokens": 32768, + "max_tokens": 32768, "mode": "chat", "output_cost_per_token": 1.4e-07, "source": "https://docs.aipowergrid.io/streaming-api", @@ -54627,7 +54629,8 @@ "input_cost_per_token": 5e-09, "litellm_provider": "aipg", "max_input_tokens": 2048, - "max_tokens": 2048, + "max_output_tokens": 1024, + "max_tokens": 1024, "mode": "chat", "output_cost_per_token": 1e-08, "source": "https://docs.aipowergrid.io/streaming-api", diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 1c57107350b..9f41cc85f90 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -54601,7 +54601,8 @@ "input_cost_per_token": 7.5e-08, "litellm_provider": "aipg", "max_input_tokens": 60000, - "max_tokens": 60000, + "max_output_tokens": 32768, + "max_tokens": 32768, "mode": "chat", "output_cost_per_token": 3e-07, "source": "https://docs.aipowergrid.io/streaming-api", @@ -54614,7 +54615,8 @@ "input_cost_per_token": 7e-08, "litellm_provider": "aipg", "max_input_tokens": 262144, - "max_tokens": 262144, + "max_output_tokens": 32768, + "max_tokens": 32768, "mode": "chat", "output_cost_per_token": 1.4e-07, "source": "https://docs.aipowergrid.io/streaming-api", @@ -54627,7 +54629,8 @@ "input_cost_per_token": 5e-09, "litellm_provider": "aipg", "max_input_tokens": 2048, - "max_tokens": 2048, + "max_output_tokens": 1024, + "max_tokens": 1024, "mode": "chat", "output_cost_per_token": 1e-08, "source": "https://docs.aipowergrid.io/streaming-api", diff --git a/tests/llm_translation/test_aipg.py b/tests/llm_translation/test_aipg.py index 075284aebfd..d229ee0a0a2 100644 --- a/tests/llm_translation/test_aipg.py +++ b/tests/llm_translation/test_aipg.py @@ -70,14 +70,16 @@ def test_get_llm_provider_aipg(): def test_aipg_model_metadata(): model_cost = litellm.get_model_cost_map(url="") expected = { - "aipg/gpt-oss-120b": (60000, 7.5e-08, 3e-07), - "aipg/deepseek-v4-flash-nvfp4": (262144, 7e-08, 1.4e-07), - "aipg/Smollm-135m": (2048, 5e-09, 1e-08), + "aipg/gpt-oss-120b": (60000, 32768, 7.5e-08, 3e-07), + "aipg/deepseek-v4-flash-nvfp4": (262144, 32768, 7e-08, 1.4e-07), + "aipg/Smollm-135m": (2048, 1024, 5e-09, 1e-08), } - for model, (context, input_cost, output_cost) in expected.items(): + for model, (context, output_limit, input_cost, output_cost) in expected.items(): info = model_cost[model] assert info["litellm_provider"] == "aipg" assert info["mode"] == "chat" assert info["max_input_tokens"] == context + assert info["max_output_tokens"] == output_limit + assert info["max_tokens"] == output_limit assert info["input_cost_per_token"] == input_cost assert info["output_cost_per_token"] == output_cost From c947cbafbf366e0de02ceec4375e132ff682dbf5 Mon Sep 17 00:00:00 2001 From: halfaipg Date: Sat, 29 Aug 2026 00:11:15 -0400 Subject: [PATCH 4/8] fix(provider): prevent AIPG credential forwarding --- litellm/llms/openai_like/dynamic_config.py | 25 +++++++++-- litellm/llms/openai_like/json_loader.py | 1 + litellm/llms/openai_like/providers.json | 1 + tests/llm_translation/test_aipg.py | 52 ++++++++++++++++++++++ 4 files changed, 76 insertions(+), 3 deletions(-) diff --git a/litellm/llms/openai_like/dynamic_config.py b/litellm/llms/openai_like/dynamic_config.py index 19e29bcdcb2..71300bcc04a 100644 --- a/litellm/llms/openai_like/dynamic_config.py +++ b/litellm/llms/openai_like/dynamic_config.py @@ -17,6 +17,21 @@ from litellm.types.llms.openai import AllMessageValues from .json_loader import SimpleProviderConfig +def _resolve_api_key(provider: SimpleProviderConfig, api_base: str, api_key: str | None) -> str | None: + if api_key: + return api_key + + env_key: Final = get_secret_str(provider.api_key_env) + if not env_key or not provider.require_explicit_key_for_custom_base: + return env_key + + env_base: Final = get_secret_str(provider.api_base_env) if provider.api_base_env else None + trusted_bases: Final = {base.rstrip("/") for base in (provider.base_url, env_base) if base} + if api_base.rstrip("/") not in trusted_bases: + raise ValueError(f"api_key is required for custom api_base on provider {provider.slug}") + return env_key + + def create_config_class(provider: SimpleProviderConfig): """Generate config class dynamically from JSON configuration""" @@ -63,8 +78,7 @@ def create_config_class(provider: SimpleProviderConfig): if not resolved_base: resolved_base = provider.base_url - # Resolve API key - resolved_key: Final = api_key or get_secret_str(provider.api_key_env) + resolved_key: Final = _resolve_api_key(provider, resolved_base, api_key) return resolved_base, resolved_key @@ -200,7 +214,12 @@ def create_responses_config_class(provider: SimpleProviderConfig): litellm_params: GenericLiteLLMParams | None, ) -> dict: litellm_params = litellm_params or GenericLiteLLMParams() - api_key: Final = litellm_params.api_key or get_secret_str(provider.api_key_env) + api_base: Final = ( + litellm_params.api_base + or (get_secret_str(provider.api_base_env) if provider.api_base_env else None) + or provider.base_url + ) + api_key: Final = _resolve_api_key(provider, api_base, litellm_params.api_key) if api_key: headers["Authorization"] = f"Bearer {api_key}" return headers diff --git a/litellm/llms/openai_like/json_loader.py b/litellm/llms/openai_like/json_loader.py index 5cdaff90d24..36056b8e881 100644 --- a/litellm/llms/openai_like/json_loader.py +++ b/litellm/llms/openai_like/json_loader.py @@ -22,6 +22,7 @@ class SimpleProviderConfig: self.constraints = data.get("constraints", {}) self.special_handling = data.get("special_handling", {}) self.supported_endpoints = data.get("supported_endpoints", []) + self.require_explicit_key_for_custom_base = data.get("require_explicit_key_for_custom_base", False) class JSONProviderRegistry: diff --git a/litellm/llms/openai_like/providers.json b/litellm/llms/openai_like/providers.json index 44f626fd987..35d1e078374 100644 --- a/litellm/llms/openai_like/providers.json +++ b/litellm/llms/openai_like/providers.json @@ -205,6 +205,7 @@ "base_url": "https://api.aipowergrid.io/v1", "api_key_env": "AIPG_API_KEY", "api_base_env": "AIPG_API_BASE", + "require_explicit_key_for_custom_base": true, "param_mappings": { "max_completion_tokens": "max_tokens" }, diff --git a/tests/llm_translation/test_aipg.py b/tests/llm_translation/test_aipg.py index d229ee0a0a2..d7cba712ddf 100644 --- a/tests/llm_translation/test_aipg.py +++ b/tests/llm_translation/test_aipg.py @@ -3,6 +3,8 @@ import os from unittest import mock +import pytest + import litellm AIPG_API_BASE = "https://api.aipowergrid.io/v1" @@ -16,6 +18,7 @@ def test_aipg_json_registry(): assert config.base_url == AIPG_API_BASE assert config.api_key_env == "AIPG_API_KEY" assert config.api_base_env == "AIPG_API_BASE" + assert config.require_explicit_key_for_custom_base is True assert config.supported_endpoints == ["/v1/chat/completions", "/v1/responses"] assert JSONProviderRegistry.supports_responses_api("aipg") is True @@ -45,6 +48,24 @@ def test_aipg_get_openai_compatible_provider_info(): assert api_base == "https://explicit.example/v1" assert api_key == "explicit-key" + with ( + mock.patch.dict(os.environ, {"AIPG_API_KEY": "env-key"}, clear=True), + pytest.raises(ValueError, match="api_key is required for custom api_base"), + ): + config._get_openai_compatible_provider_info("https://attacker.example/v1", None) + + with mock.patch.dict( + os.environ, + { + "AIPG_API_KEY": "env-key", + "AIPG_API_BASE": "https://operator.example/v1", + }, + clear=True, + ): + api_base, api_key = config._get_openai_compatible_provider_info("https://operator.example/v1/", None) + assert api_base == "https://operator.example/v1/" + assert api_key == "env-key" + mapped = config.map_openai_params( non_default_params={"max_completion_tokens": 12, "temperature": 0.2}, optional_params={}, @@ -67,6 +88,37 @@ def test_get_llm_provider_aipg(): assert api_base == AIPG_API_BASE +def test_aipg_responses_rejects_env_key_with_custom_api_base(): + from litellm.llms.openai_like.dynamic_config import create_responses_config_class + from litellm.llms.openai_like.json_loader import JSONProviderRegistry + from litellm.types.router import GenericLiteLLMParams + + provider = JSONProviderRegistry.get("aipg") + assert provider is not None + config = create_responses_config_class(provider)() + + with ( + mock.patch.dict(os.environ, {"AIPG_API_KEY": "env-key"}, clear=True), + pytest.raises(ValueError, match="api_key is required for custom api_base"), + ): + config.validate_environment( + headers={}, + model="gpt-oss-120b", + litellm_params=GenericLiteLLMParams(api_base="https://attacker.example/v1"), + ) + + with mock.patch.dict(os.environ, {"AIPG_API_KEY": "env-key"}, clear=True): + headers = config.validate_environment( + headers={}, + model="gpt-oss-120b", + litellm_params=GenericLiteLLMParams( + api_base="https://attacker.example/v1", + api_key="explicit-key", + ), + ) + assert headers["Authorization"] == "Bearer explicit-key" + + def test_aipg_model_metadata(): model_cost = litellm.get_model_cost_map(url="") expected = { From 46c2b81f29790830a4bd6f7279effa3696b15e8c Mon Sep 17 00:00:00 2001 From: halfaipg Date: Sat, 29 Aug 2026 00:15:44 -0400 Subject: [PATCH 5/8] refactor(provider): keep trusted AIPG bases immutable --- litellm/llms/openai_like/dynamic_config.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/litellm/llms/openai_like/dynamic_config.py b/litellm/llms/openai_like/dynamic_config.py index 71300bcc04a..f4ade1e7543 100644 --- a/litellm/llms/openai_like/dynamic_config.py +++ b/litellm/llms/openai_like/dynamic_config.py @@ -26,7 +26,7 @@ def _resolve_api_key(provider: SimpleProviderConfig, api_base: str, api_key: str return env_key env_base: Final = get_secret_str(provider.api_base_env) if provider.api_base_env else None - trusted_bases: Final = {base.rstrip("/") for base in (provider.base_url, env_base) if base} + trusted_bases: Final = frozenset(base.rstrip("/") for base in (provider.base_url, env_base) if base) if api_base.rstrip("/") not in trusted_bases: raise ValueError(f"api_key is required for custom api_base on provider {provider.slug}") return env_key From 0e1430595fa7269c37fd84c28533d3531486b773 Mon Sep 17 00:00:00 2001 From: halfaipg Date: Sat, 29 Aug 2026 00:32:49 -0400 Subject: [PATCH 6/8] feat(provider): add AIPG image generation --- litellm/llms/openai_like/providers.json | 2 +- ...odel_prices_and_context_window_backup.json | 24 ++++++++ .../provider_endpoints_support_backup.json | 2 +- model_prices_and_context_window.json | 24 ++++++++ provider_endpoints_support.json | 2 +- tests/llm_translation/test_aipg.py | 61 ++++++++++++++++++- 6 files changed, 111 insertions(+), 4 deletions(-) diff --git a/litellm/llms/openai_like/providers.json b/litellm/llms/openai_like/providers.json index 35d1e078374..82585d9ed4f 100644 --- a/litellm/llms/openai_like/providers.json +++ b/litellm/llms/openai_like/providers.json @@ -209,6 +209,6 @@ "param_mappings": { "max_completion_tokens": "max_tokens" }, - "supported_endpoints": ["/v1/chat/completions", "/v1/responses"] + "supported_endpoints": ["/v1/chat/completions", "/v1/responses", "/v1/images/generations"] } } diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 9f41cc85f90..32f578712ac 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -54637,5 +54637,29 @@ "supports_function_calling": false, "supports_system_messages": true, "supports_tool_choice": false + }, + "aipg/z-image-turbo": { + "input_cost_per_image": 0.003, + "litellm_provider": "aipg", + "mode": "image_generation", + "supported_endpoints": [ + "/v1/images/generations" + ] + }, + "aipg/Krea 2 Turbo": { + "input_cost_per_image": 0.005, + "litellm_provider": "aipg", + "mode": "image_generation", + "supported_endpoints": [ + "/v1/images/generations" + ] + }, + "aipg/FLUX.2 Klein 4B FP8": { + "input_cost_per_image": 0.01, + "litellm_provider": "aipg", + "mode": "image_generation", + "supported_endpoints": [ + "/v1/images/generations" + ] } } diff --git a/litellm/provider_endpoints_support_backup.json b/litellm/provider_endpoints_support_backup.json index 62dec9d4ac9..600d11254c3 100644 --- a/litellm/provider_endpoints_support_backup.json +++ b/litellm/provider_endpoints_support_backup.json @@ -92,7 +92,7 @@ "messages": false, "responses": true, "embeddings": false, - "image_generations": false, + "image_generations": true, "audio_transcriptions": false, "audio_speech": false, "moderations": false, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 9f41cc85f90..32f578712ac 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -54637,5 +54637,29 @@ "supports_function_calling": false, "supports_system_messages": true, "supports_tool_choice": false + }, + "aipg/z-image-turbo": { + "input_cost_per_image": 0.003, + "litellm_provider": "aipg", + "mode": "image_generation", + "supported_endpoints": [ + "/v1/images/generations" + ] + }, + "aipg/Krea 2 Turbo": { + "input_cost_per_image": 0.005, + "litellm_provider": "aipg", + "mode": "image_generation", + "supported_endpoints": [ + "/v1/images/generations" + ] + }, + "aipg/FLUX.2 Klein 4B FP8": { + "input_cost_per_image": 0.01, + "litellm_provider": "aipg", + "mode": "image_generation", + "supported_endpoints": [ + "/v1/images/generations" + ] } } diff --git a/provider_endpoints_support.json b/provider_endpoints_support.json index 251c676975a..b9f88ccbc9d 100644 --- a/provider_endpoints_support.json +++ b/provider_endpoints_support.json @@ -92,7 +92,7 @@ "messages": false, "responses": true, "embeddings": false, - "image_generations": false, + "image_generations": true, "audio_transcriptions": false, "audio_speech": false, "moderations": false, diff --git a/tests/llm_translation/test_aipg.py b/tests/llm_translation/test_aipg.py index d7cba712ddf..cedc15a4232 100644 --- a/tests/llm_translation/test_aipg.py +++ b/tests/llm_translation/test_aipg.py @@ -19,7 +19,11 @@ def test_aipg_json_registry(): assert config.api_key_env == "AIPG_API_KEY" assert config.api_base_env == "AIPG_API_BASE" assert config.require_explicit_key_for_custom_base is True - assert config.supported_endpoints == ["/v1/chat/completions", "/v1/responses"] + assert config.supported_endpoints == [ + "/v1/chat/completions", + "/v1/responses", + "/v1/images/generations", + ] assert JSONProviderRegistry.supports_responses_api("aipg") is True @@ -119,6 +123,49 @@ def test_aipg_responses_rejects_env_key_with_custom_api_base(): assert headers["Authorization"] == "Bearer explicit-key" +def test_aipg_image_generation_uses_native_endpoint(): + mock_response = mock.MagicMock() + mock_response.model_dump.return_value = { + "created": 1, + "data": [{"url": "https://images.example/aipg.webp"}], + } + mock_client = mock.MagicMock() + mock_client.images.generate.return_value = mock_response + + with ( + mock.patch.dict(os.environ, {"AIPG_API_KEY": "grid-test"}, clear=True), + mock.patch("litellm.llms.openai.openai.OpenAI", return_value=mock_client) as constructor, + ): + response = litellm.image_generation( + model="aipg/z-image-turbo", + prompt="An amber square on black.", + n=1, + size="512x512", + ) + + constructor.assert_called_once() + assert constructor.call_args.kwargs["api_key"] == "grid-test" + assert str(constructor.call_args.kwargs["base_url"]).rstrip("/") == AIPG_API_BASE + request = mock_client.images.generate.call_args.kwargs + assert request["model"] == "z-image-turbo" + assert request["prompt"] == "An amber square on black." + assert request["n"] == 1 + assert request["size"] == "512x512" + assert response.data[0]["url"] == "https://images.example/aipg.webp" + + +def test_aipg_image_generation_rejects_env_key_with_custom_api_base(): + with ( + mock.patch.dict(os.environ, {"AIPG_API_KEY": "env-key"}, clear=True), + pytest.raises(litellm.BadRequestError, match="api_key is required for custom api_base"), + ): + litellm.image_generation( + model="aipg/z-image-turbo", + prompt="An amber square on black.", + api_base="https://attacker.example/v1", + ) + + def test_aipg_model_metadata(): model_cost = litellm.get_model_cost_map(url="") expected = { @@ -135,3 +182,15 @@ def test_aipg_model_metadata(): assert info["max_tokens"] == output_limit assert info["input_cost_per_token"] == input_cost assert info["output_cost_per_token"] == output_cost + + image_prices = { + "aipg/z-image-turbo": 0.003, + "aipg/Krea 2 Turbo": 0.005, + "aipg/FLUX.2 Klein 4B FP8": 0.01, + } + for model, input_cost_per_image in image_prices.items(): + info = model_cost[model] + assert info["litellm_provider"] == "aipg" + assert info["mode"] == "image_generation" + assert info["input_cost_per_image"] == input_cost_per_image + assert info["supported_endpoints"] == ["/v1/images/generations"] From f492ae7f283710cda2d42ac15fe91c4b2af5c83f Mon Sep 17 00:00:00 2001 From: halfaipg Date: Sat, 29 Aug 2026 00:39:36 -0400 Subject: [PATCH 7/8] test(provider): justify AIPG client interception --- tests/llm_translation/test_aipg.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/tests/llm_translation/test_aipg.py b/tests/llm_translation/test_aipg.py index cedc15a4232..6f4d8111ade 100644 --- a/tests/llm_translation/test_aipg.py +++ b/tests/llm_translation/test_aipg.py @@ -134,7 +134,9 @@ def test_aipg_image_generation_uses_native_endpoint(): with ( mock.patch.dict(os.environ, {"AIPG_API_KEY": "grid-test"}, clear=True), - mock.patch("litellm.llms.openai.openai.OpenAI", return_value=mock_client) as constructor, + mock.patch( # test-quality-ok: capture the SDK-created client to verify the trusted AIPG base and env key + "litellm.llms.openai.openai.OpenAI", return_value=mock_client + ) as constructor, ): response = litellm.image_generation( model="aipg/z-image-turbo", From 5cbe37e18f76e72057bfb89803af3ddf208ce2fb Mon Sep 17 00:00:00 2001 From: halfaipg Date: Sat, 29 Aug 2026 00:55:47 -0400 Subject: [PATCH 8/8] test(provider): run AIPG coverage in provider shard --- .../llms/openai_like}/test_aipg.py | 0 1 file changed, 0 insertions(+), 0 deletions(-) rename tests/{llm_translation => test_litellm/llms/openai_like}/test_aipg.py (100%) diff --git a/tests/llm_translation/test_aipg.py b/tests/test_litellm/llms/openai_like/test_aipg.py similarity index 100% rename from tests/llm_translation/test_aipg.py rename to tests/test_litellm/llms/openai_like/test_aipg.py