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/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 3fa4d57de1b..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', @@ -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'