Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
110 changes: 110 additions & 0 deletions contributing/samples/models/hello_world_aimlapi/README.md
Original file line number Diff line number Diff line change
@@ -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/<vendor>/<model-name>` 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.
16 changes: 16 additions & 0 deletions contributing/samples/models/hello_world_aimlapi/__init__.py
Original file line number Diff line number Diff line change
@@ -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
92 changes: 92 additions & 0 deletions contributing/samples/models/hello_world_aimlapi/agent.py
Original file line number Diff line number Diff line change
@@ -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,
],
)
90 changes: 90 additions & 0 deletions contributing/samples/models/hello_world_aimlapi/main.py
Original file line number Diff line number Diff line change
@@ -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())
1 change: 1 addition & 0 deletions src/google/adk/models/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -81,6 +81,7 @@
),
'LiteLlm': (
[
r'aiml/.*',
r'openai/.*',
r'azure/.*',
r'azure_ai/.*',
Expand Down
2 changes: 2 additions & 0 deletions src/google/adk/models/lite_llm.py
Original file line number Diff line number Diff line change
Expand Up @@ -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")
Expand Down
13 changes: 13 additions & 0 deletions tests/unittests/models/test_models.py
Original file line number Diff line number Diff line change
Expand Up @@ -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',
Expand Down Expand Up @@ -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'
Expand Down