From d54f8c8ea0b1f11b1697c8c1535599903641750a Mon Sep 17 00:00:00 2001 From: Stan Date: Thu, 3 Sep 2026 13:46:39 +0500 Subject: [PATCH 1/2] docs: add aimlapi.com model sample A user reaching for an aggregator has no way to learn from this repo that one works here. Nothing in ADK names aimlapi.com, so the natural reading of the registry is that only the sixteen spelled-out providers are usable, and the aggregator gets ruled out before it is ever tried. Nothing is broken: `LLMRegistry.resolve` already defers any `provider/model` name to `litellm.provider_list`, and LiteLLM has carried the `aiml` provider (base `https://api.aimlapi.com/v1`, key `AIML_API_KEY`) since well before the `litellm>=1.84` floor this project pins. The gap is documentation, so this change is documentation, following the shape the NVIDIA NIM sample set for a LiteLLM-backed provider: README, agent, runner. The two tests guard the one thing that is easy to get wrong. An aggregator id carries its own slash, so `aiml/openai/gpt-4o-mini` has two of them, and only the first is the provider. Resolution has to hand the whole string to LiteLLM or the request routes to a model nobody asked for. Verified against the live API rather than a mock: `main.py` completes both turns, including a tool call and its follow-up, on `aiml/openai/gpt-4o-mini`. --- .../models/hello_world_aimlapi/README.md | 110 ++++++++++++++++++ .../models/hello_world_aimlapi/__init__.py | 16 +++ .../models/hello_world_aimlapi/agent.py | 92 +++++++++++++++ .../models/hello_world_aimlapi/main.py | 90 ++++++++++++++ tests/unittests/models/test_models.py | 13 +++ 5 files changed, 321 insertions(+) create mode 100644 contributing/samples/models/hello_world_aimlapi/README.md create mode 100644 contributing/samples/models/hello_world_aimlapi/__init__.py create mode 100644 contributing/samples/models/hello_world_aimlapi/agent.py create mode 100644 contributing/samples/models/hello_world_aimlapi/main.py diff --git a/contributing/samples/models/hello_world_aimlapi/README.md b/contributing/samples/models/hello_world_aimlapi/README.md new file mode 100644 index 00000000000..5fe68825b5d --- /dev/null +++ b/contributing/samples/models/hello_world_aimlapi/README.md @@ -0,0 +1,110 @@ +# Using aimlapi.com Models with ADK and LiteLLM + +This example demonstrates how to use models served by +[aimlapi.com](https://aimlapi.com) with ADK through LiteLLM integration. + +aimlapi.com is an OpenAI-compatible aggregator that serves models from several +vendors behind one key. For comprehensive information about using it with +LiteLLM, refer to the +[official LiteLLM documentation](https://docs.litellm.ai/docs/providers/aiml). + +## Setup + +### 1. Get an aimlapi.com API Key + +Use the following procedure to get an aimlapi.com API key. + +1. Sign up at [aimlapi.com](https://aimlapi.com) +1. Navigate to the API Keys section and generate a key +1. Copy your API key + +### 2. Install LiteLLM + +Install LiteLLM by running the following code. + +```bash +pip install litellm +``` + +## Using aimlapi.com Models in ADK + +### Environment Variables + +Set the required environment variables: + +```bash +export AIML_API_KEY="your-aimlapi-key" +export AIML_API_BASE="https://api.aimlapi.com/v1" # Optional +``` + +`AIML_API_KEY` is the name LiteLLM reads; ADK passes it through untouched. + +### Code Examples + +#### Basic Agent Creation + +```python +from google.adk import Agent +from google.adk.models.lite_llm import LiteLlm + +# Create agent with an aimlapi.com model +agent = Agent( + model=LiteLlm(model="aiml/openai/gpt-4o-mini"), + name="aimlapi_agent", + instruction="You are a helpful assistant.", + description="Agent using a model served by aimlapi.com", +) +``` + +A plain model string works too, because `LLMRegistry` hands any `provider/model` +name LiteLLM knows about to `LiteLlm`: + +```python +agent = Agent(model="aiml/openai/gpt-4o-mini", name="aimlapi_agent") +``` + +#### Available Models + +Use the `aiml/` prefix in front of the catalog id. The catalog id keeps its own +slash — LiteLLM splits the provider off the *first* one only, so +`aiml/openai/gpt-4o-mini` reaches the API as `openai/gpt-4o-mini`. + +```python +# Examples of available models +models = [ + "aiml/openai/gpt-4o-mini", + "aiml/openai/gpt-5-5", + "aiml/anthropic/claude-sonnet-4.5", + "aiml/google/gemini-2.5-flash", + # ... and many more +] +``` + +The full catalog is `GET https://api.aimlapi.com/v1/models`, whose chat entries +are the ones with `"type": "openai/chat-completions"`. An id is valid if it +appears either as an `id` or in another entry's `aliases`. + +## Integration Details + +ADK uses LiteLLM as a wrapper to access aimlapi.com models. The integration: + +1. **Model Format**: Uses `aiml//` format +1. **Authentication**: Requires `AIML_API_KEY` environment variable +1. **Base URL**: Optional `AIML_API_BASE`, defaulting to + `https://api.aimlapi.com/v1` +1. **Compatibility**: Works with all ADK features (tools, sessions, etc.) +1. **Supported Endpoints**: `/chat/completions`. LiteLLM's `aiml` provider also + carries image generation; it has no embeddings route, so + `aiml/`-prefixed names are for chat models only. + +## Notes + +- **Do not validate a key by fetching the model catalog.** + `GET /v1/models` answers `200` with no key or a wrong one; only an actual + completion call returns `401`. +- **Catalog membership is not proof a model serves.** A few published ids + return `404` on a real call, and a few working ids are missing from the + catalog. Call an id once before depending on it. +- **`max_tokens` does not bound reasoning tokens on every model**, and some + models report `completion_tokens` that exclude reasoning tokens, so a token + count read back from a response can under-report the billed total. diff --git a/contributing/samples/models/hello_world_aimlapi/__init__.py b/contributing/samples/models/hello_world_aimlapi/__init__.py new file mode 100644 index 00000000000..044e24d3883 --- /dev/null +++ b/contributing/samples/models/hello_world_aimlapi/__init__.py @@ -0,0 +1,16 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +from . import agent diff --git a/contributing/samples/models/hello_world_aimlapi/agent.py b/contributing/samples/models/hello_world_aimlapi/agent.py new file mode 100644 index 00000000000..801166b3929 --- /dev/null +++ b/contributing/samples/models/hello_world_aimlapi/agent.py @@ -0,0 +1,92 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +import random + +from google.adk import Agent +from google.adk.models.lite_llm import LiteLlm + + +def roll_die(sides: int) -> int: + """Roll a die and return the rolled result. + + Args: + sides: The integer number of sides the die has. + + Returns: + An integer of the result of rolling the die. + """ + return random.randint(1, sides) + + +def check_prime(nums: list[int]) -> str: + """Check if a given list of numbers are prime. + + Args: + nums: The list of numbers to check. + + Returns: + A str indicating which number is prime. + """ + primes = set() + for number in nums: + if number <= 1: + continue + is_prime = True + for i in range(2, int(number**0.5) + 1): + if number % i == 0: + is_prime = False + break + if is_prime: + primes.add(number) + return ( + "No prime numbers found." + if not primes + else f"{', '.join(str(num) for num in primes)} are prime numbers." + ) + + +# The model id keeps its own slash: LiteLLM splits the provider off the first +# one, so "aiml/openai/gpt-4o-mini" reaches the API as "openai/gpt-4o-mini". +root_agent = Agent( + model=LiteLlm(model="aiml/openai/gpt-4o-mini"), + name="hello_world_aimlapi_agent", + description=( + "hello world agent powered by aimlapi.com that can roll a dice of 8" + " sides and check prime numbers." + ), + instruction=""" + You are a helpful assistant powered by models served through aimlapi.com. + You roll dice and answer questions about the outcome of the dice rolls. + You can roll dice of different sizes. + You can use multiple tools in parallel by calling functions in parallel(in one request and in one round). + It is ok to discuss previous dice rolls, and comment on the dice rolls. + When you are asked to roll a die, you must call the roll_die tool with the number of sides. Be sure to pass in an integer. Do not pass in a string. + You should never roll a die on your own. + When checking prime numbers, call the check_prime tool with a list of integers. Be sure to pass in a list of integers. You should never pass in a string. + You should not check prime numbers before calling the tool. + When you are asked to roll a die and check prime numbers, you should always make the following two function calls: + 1. You should first call the roll_die tool to get a roll. Wait for the function response before calling the check_prime tool. + 2. After you get the function response from roll_die tool, you should call the check_prime tool with the roll_die result. + 2.1 If user asks you to check primes based on previous rolls, make sure you include the previous rolls in the list. + 3. When you respond, you must include the roll_die result from step 1. + You should always perform the previous 3 steps when asking for a roll and checking prime numbers. + You should not rely on the previous history on prime results. + """, + tools=[ + roll_die, + check_prime, + ], +) diff --git a/contributing/samples/models/hello_world_aimlapi/main.py b/contributing/samples/models/hello_world_aimlapi/main.py new file mode 100644 index 00000000000..60568ac3c46 --- /dev/null +++ b/contributing/samples/models/hello_world_aimlapi/main.py @@ -0,0 +1,90 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +import asyncio +import os +import time + +import agent +from dotenv import load_dotenv +from google.adk import Runner +from google.adk.artifacts import InMemoryArtifactService +from google.adk.cli.utils import logs +from google.adk.sessions import InMemorySessionService +from google.adk.sessions import Session +from google.genai import types + +load_dotenv(override=True) +logs.log_to_tmp_folder() + + +async def main(): + # Check for the aimlapi.com API key. LiteLLM reads AIML_API_KEY. + if not os.environ.get('AIML_API_KEY'): + print( + ' Your AIML_API_KEY is not set. Set it before you continue by' + ' running the following code:' + ) + print(" export AIML_API_KEY='your-aimlapi-key'") + return + + app_name = 'aimlapi_app' + user_id_1 = 'user1' + session_service = InMemorySessionService() + artifact_service = InMemoryArtifactService() + runner = Runner( + app_name=app_name, + agent=agent.root_agent, + artifact_service=artifact_service, + session_service=session_service, + ) + session_11 = await session_service.create_session( + app_name=app_name, user_id=user_id_1 + ) + + async def run_prompt(session: Session, new_message: str): + content = types.Content( + role='user', parts=[types.Part.from_text(text=new_message)] + ) + print('** User says:', content.model_dump(exclude_none=True)) + async for event in runner.run_async( + user_id=user_id_1, + session_id=session.id, + new_message=content, + ): + if event.content.parts and event.content.parts[0].text: + print(f'** {event.author}: {event.content.parts[0].text}') + + start_time = time.time() + print('Start time:', start_time) + print('------------------------------------') + print('Testing aimlapi.com integration with ADK') + print('Model: aiml/openai/gpt-4o-mini') + print('------------------------------------') + await run_prompt(session_11, 'Hi, introduce yourself.') + # A second turn on the same session is the one that matters: it replays the + # tool call and its result back to the model. + await run_prompt( + session_11, + 'Roll a die with 20 sides and check if the result is prime.', + ) + end_time = time.time() + print('------------------------------------') + print('End time:', end_time) + print('Total time:', end_time - start_time) + + +if __name__ == '__main__': + asyncio.run(main()) diff --git a/tests/unittests/models/test_models.py b/tests/unittests/models/test_models.py index 3fa4d57de1b..0b45b653e27 100644 --- a/tests/unittests/models/test_models.py +++ b/tests/unittests/models/test_models.py @@ -101,6 +101,9 @@ def test_match_litellm_family(model_name): 'gemini/gemini-3.5-flash', 'openrouter/anthropic/claude-opus-4', 'cerebras/llama-3.3-70b', + # An aggregator id keeps its own slash: only the first segment is the + # LiteLLM provider, so the rest must survive resolution untouched. + 'aiml/openai/gpt-4o-mini', ], ) def test_match_litellm_provider_not_spelled_out_in_registry(model_name): @@ -191,6 +194,16 @@ def test_resolve_with_prefix(): assert models.LLMRegistry.resolve('LiteLlm:openai/gpt-4o') is LiteLlm +def test_new_llm_keeps_a_multi_segment_provider_model_name(): + """Test that only the class prefix, never a provider prefix, is stripped.""" + llm = models.LLMRegistry.new_llm('aiml/openai/gpt-4o-mini') + + assert isinstance(llm, LiteLlm) + # LiteLLM splits the provider off the first slash itself, so handing it a + # shortened name here would route to the wrong model. + assert llm.model == 'aiml/openai/gpt-4o-mini' + + def test_register_after_resolve_returns_the_new_class(): """Test that registering over an already-resolved name takes effect.""" model_name = 'test-registry-override-model' From 77c3a326c3fccb6a5c08bab430222f6ce1e459a2 Mon Sep 17 00:00:00 2001 From: Stan Date: Thu, 3 Sep 2026 13:47:45 +0500 Subject: [PATCH 2/2] =?UTF-8?q?chore(aimlapi):=20fork-only=20placement=20?= =?UTF-8?q?=E2=80=94=20do=20not=20send=20upstream?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Spells `aiml/.*` out at the head of both hand-ordered LiteLLM pattern lists instead of leaving it to the provider-list fallback. This is placement, not function. `LLMRegistry.resolve` already returns `LiteLlm` for these names, so the resolution result is identical either way and no behaviour changes; the entry only moves the decision from the fallback into the explicit list. It is separated into its own commit so it can be dropped before anything goes upstream, where the maintainers' stated rule is that the list spells out the common providers and defers the rest. --- src/google/adk/models/__init__.py | 1 + src/google/adk/models/lite_llm.py | 2 ++ tests/unittests/models/test_models.py | 6 +++--- 3 files changed, 6 insertions(+), 3 deletions(-) diff --git a/src/google/adk/models/__init__.py b/src/google/adk/models/__init__.py index cb1b9b06b9a..4e0a2898da8 100644 --- a/src/google/adk/models/__init__.py +++ b/src/google/adk/models/__init__.py @@ -81,6 +81,7 @@ ), 'LiteLlm': ( [ + r'aiml/.*', r'openai/.*', r'azure/.*', r'azure_ai/.*', diff --git a/src/google/adk/models/lite_llm.py b/src/google/adk/models/lite_llm.py index 5bf71169d9a..b23ea11ad9c 100644 --- a/src/google/adk/models/lite_llm.py +++ b/src/google/adk/models/lite_llm.py @@ -3497,6 +3497,8 @@ def supported_models(cls) -> list[str]: """ return [ + # For aimlapi.com models (e.g., "aiml/openai/gpt-4o-mini") + r"aiml/.*", # For OpenAI models (e.g., "openai/gpt-4o") r"openai/.*", # For Azure OpenAI models (e.g., "azure/gpt-4o") diff --git a/tests/unittests/models/test_models.py b/tests/unittests/models/test_models.py index 0b45b653e27..ea7d25f694a 100644 --- a/tests/unittests/models/test_models.py +++ b/tests/unittests/models/test_models.py @@ -81,6 +81,9 @@ def test_match_claude_family(model_name): @pytest.mark.parametrize( 'model_name', [ + # An aggregator id keeps its own slash: only the first segment is the + # LiteLLM provider, so the rest must survive resolution untouched. + 'aiml/openai/gpt-4o-mini', 'openai/gpt-4o', 'openai/gpt-4o-mini', 'groq/llama3-70b-8192', @@ -101,9 +104,6 @@ def test_match_litellm_family(model_name): 'gemini/gemini-3.5-flash', 'openrouter/anthropic/claude-opus-4', 'cerebras/llama-3.3-70b', - # An aggregator id keeps its own slash: only the first segment is the - # LiteLLM provider, so the rest must survive resolution untouched. - 'aiml/openai/gpt-4o-mini', ], ) def test_match_litellm_provider_not_spelled_out_in_registry(model_name):