scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-29 01:42:19 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 100

Merge cb7295a777 into e26a6450c8

Browse source
This commit is contained in:
halfaipg 2026-09-23 14:52:04 +00:00 • committed by GitHub
commit 4e05f10113
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
10 changed files with 400 additions and 3 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
litellm/constants.py
View file

@ -758,6 +758,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
@ -942,6 +943,7 @@ openai_compatible_endpoints: Final[list] = [
openai_compatible_providers: Final[list] = [
"aipg",
"anyscale",
"groq",
"nvidia_nim",

25
litellm/llms/openai_like/dynamic_config.py
View file

@ -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 = 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
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

1
litellm/llms/openai_like/json_loader.py
View file

@ -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:

10
litellm/llms/openai_like/providers.json
View file

@ -200,5 +200,15 @@
"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",
"require_explicit_key_for_custom_base": true,
"param_mappings": {
"max_completion_tokens": "max_tokens"
},
"supported_endpoints": ["/v1/chat/completions", "/v1/responses", "/v1/images/generations"]
}
}

65
litellm/model_prices_and_context_window_backup.json
View file

@ -60892,6 +60892,71 @@
"input_cost_per_image_token": 1.25e-06,
"supports_response_schema": true
},
"aipg/gpt-oss-120b": {
"input_cost_per_token": 7.5e-08,
"litellm_provider": "aipg",
"max_input_tokens": 60000,
"max_output_tokens": 32768,
"max_tokens": 32768,
"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_output_tokens": 32768,
"max_tokens": 32768,
"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_output_tokens": 1024,
"max_tokens": 1024,
"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
},
"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"
]
},
"groq/qwen/qwen3.8-27b": {
"input_cost_per_token": 8e-07,
"litellm_provider": "groq",

18
litellm/provider_endpoints_support_backup.json
View file

@ -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": true,
"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",

1
litellm/types/utils.py
View file

@ -4143,6 +4143,7 @@ GenericBudgetConfigType = dict[str, BudgetConfig]
class LlmProviders(str, Enum):
AIPG = "aipg"
OPENAI = "openai"
CHATGPT = "chatgpt"
OPENAI_LIKE = "openai_like" # embedding only

65
model_prices_and_context_window.json
View file

@ -60892,6 +60892,71 @@
"input_cost_per_image_token": 1.25e-06,
"supports_response_schema": true
},
"aipg/gpt-oss-120b": {
"input_cost_per_token": 7.5e-08,
"litellm_provider": "aipg",
"max_input_tokens": 60000,
"max_output_tokens": 32768,
"max_tokens": 32768,
"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_output_tokens": 32768,
"max_tokens": 32768,
"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_output_tokens": 1024,
"max_tokens": 1024,
"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
},
"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"
]
},
"groq/qwen/qwen3.8-27b": {
"input_cost_per_token": 8e-07,
"litellm_provider": "groq",

18
provider_endpoints_support.json
View file

@ -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": true,
"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",

198
tests/test_litellm/llms/openai_like/test_aipg.py Normal file
View file

@ -0,0 +1,198 @@
"""Tests for the AI Power Grid JSON provider integration."""
import os
from unittest import mock
import pytest
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.require_explicit_key_for_custom_base is True
assert config.supported_endpoints == [
"/v1/chat/completions",
"/v1/responses",
"/v1/images/generations",
]
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"
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={},
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_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_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( # 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",
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 = {
"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, 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
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"]