feat(provider): add AI Power Grid

This commit is contained in:
halfaipg 2026-08-28 21:12:12 -04:00
parent fef5d3d0f9
commit 2326088a32
6 changed files with 184 additions and 0 deletions

View file

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

View file

@ -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"]
}
}

View file

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

View file

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

View file

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

View file

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