mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-29 01:42:19 +00:00
feat(provider): add AI Power Grid
This commit is contained in:
parent
fef5d3d0f9
commit
2326088a32
6 changed files with 184 additions and 0 deletions
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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"]
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -3695,6 +3695,7 @@ GenericBudgetConfigType = dict[str, BudgetConfig]
|
|||
|
||||
|
||||
class LlmProviders(str, Enum):
|
||||
AIPG = "aipg"
|
||||
OPENAI = "openai"
|
||||
CHATGPT = "chatgpt"
|
||||
OPENAI_LIKE = "openai_like" # embedding only
|
||||
|
|
|
|||
|
|
@ -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
|
||||
}
|
||||
}
|
||||
|
|
|
|||
96
tests/llm_translation/test_aipg.py
Normal file
96
tests/llm_translation/test_aipg.py
Normal 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
|
||||
Loading…
Add table
Reference in a new issue