Skip to content

Commit 87dcb3f

Browse files
google-genai-botcopybara-github
authored andcommitted
feat: Add ApigeeLlm as a model that let's ADK Agent developers to connect with an Apigee proxy
PiperOrigin-RevId: 824712152
1 parent 00d147d commit 87dcb3f

File tree

7 files changed

+960
-0
lines changed

7 files changed

+960
-0
lines changed
Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,8 @@
1+
# This is a sample .env file.
2+
# Copy this file to .env and replace the placeholder values with your actual credentials.
3+
4+
# Your Google API key for accessing Gemini models.
5+
GOOGLE_API_KEY="your-google-api-key"
6+
7+
# The URL of your Apigee proxy.
8+
APIGEE_PROXY_URL="https://your-apigee-proxy.net/basepath"
Lines changed: 84 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,84 @@
1+
# Hello World with Apigee LLM
2+
3+
This sample demonstrates how to use the Agent Development Kit (ADK) with an LLM fronted by an Apigee proxy. It showcases the flexibility of the `ApigeeLlm` class in configuring the target LLM provider (Gemini or Vertex AI) and API version through the model string.
4+
5+
## Setup
6+
7+
Before running the sample, you need to configure your environment with the necessary credentials.
8+
9+
1. **Create a `.env` file:**
10+
Copy the sample environment file to a new file named `.env` in the same directory.
11+
```bash
12+
cp .env-sample .env
13+
```
14+
15+
2. **Set Environment Variables:**
16+
Open the `.env` file and provide values for the following variables:
17+
18+
- `GOOGLE_API_KEY`: Your API key for the Google AI services (Gemini).
19+
- `APIGEE_PROXY_URL`: The full URL of your Apigee proxy endpoint.
20+
21+
Example `.env` file:
22+
```
23+
GOOGLE_API_KEY="your-google-api-key"
24+
APIGEE_PROXY_URL="https://your-apigee-proxy.net/basepath"
25+
```
26+
27+
The `main.py` script will automatically load these variables when it runs.
28+
29+
## Run the Sample
30+
31+
Once your `.env` file is configured, you can run the sample with the following command:
32+
33+
```bash
34+
python main.py
35+
```
36+
37+
## Configuring the Apigee LLM
38+
39+
The `ApigeeLlm` class is configured using a special model string format in `agent.py`. This string determines which backend provider (Vertex AI or Gemini) and which API version to use.
40+
41+
### Model String Format
42+
43+
The supported format is:
44+
45+
`apigee/[<provider>/][<version>/]<model_id>`
46+
47+
- **`provider`** (optional): Can be `vertex_ai` or `gemini`.
48+
- If specified, it forces the use of that provider.
49+
- If omitted, the provider is determined by the `GOOGLE_GENAI_USE_VERTEXAI` environment variable. If this variable is set to `true` or `1`, Vertex AI is used; otherwise, `gemini` is used by default.
50+
51+
- **`version`** (optional): The API version to use (e.g., `v1`, `v1beta`).
52+
- If omitted, the default version for the selected provider is used.
53+
54+
- **`model_id`** (required): The identifier for the model you want to use (e.g., `gemini-2.5-flash`).
55+
56+
### Configuration Examples
57+
58+
Here are some examples of how to configure the model string in `agent.py` to achieve different behaviors:
59+
60+
1. **Implicit Provider (determined by environment variable):**
61+
62+
- `model="apigee/gemini-2.5-flash"`
63+
- Uses the default API version.
64+
- Provider is Vertex AI if `GOOGLE_GENAI_USE_VERTEXAI` is true, otherwise Gemini.
65+
66+
- `model="apigee/v1/gemini-2.5-flash"`
67+
- Uses API version `v1`.
68+
- Provider is determined by the environment variable.
69+
70+
2. **Explicit Provider (ignores environment variable):**
71+
72+
- `model="apigee/vertex_ai/gemini-2.5-flash"`
73+
- Uses Vertex AI with the default API version.
74+
75+
- `model="apigee/gemini/gemini-2.5-flash"`
76+
- Uses Gemini with the default API version.
77+
78+
- `model="apigee/gemini/v1/gemini-2.5-flash"`
79+
- Uses Gemini with API version `v1`.
80+
81+
- `model="apigee/vertex_ai/v1beta/gemini-2.5-flash"`
82+
- Uses Vertex AI with API version `v1beta`.
83+
84+
By modifying the `model` string in `agent.py`, you can test various configurations without changing the core logic of the agent.
Lines changed: 108 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,108 @@
1+
# Copyright 2025 Google LLC
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
import random
16+
17+
from google.adk import Agent
18+
from google.adk.tools.tool_context import ToolContext
19+
from google.genai import types
20+
21+
22+
def roll_die(sides: int, tool_context: ToolContext) -> int:
23+
"""Roll a die and return the rolled result.
24+
25+
Args:
26+
sides: The integer number of sides the die has.
27+
28+
Returns:
29+
An integer of the result of rolling the die.
30+
"""
31+
result = random.randint(1, sides)
32+
if "rolls" not in tool_context.state:
33+
tool_context.state["rolls"] = []
34+
35+
tool_context.state["rolls"] = tool_context.state["rolls"] + [result]
36+
return result
37+
38+
39+
async def check_prime(nums: list[int]) -> str:
40+
"""Check if a given list of numbers are prime.
41+
42+
Args:
43+
nums: The list of numbers to check.
44+
45+
Returns:
46+
A str indicating which number is prime.
47+
"""
48+
primes = set()
49+
for number in nums:
50+
number = int(number)
51+
if number <= 1:
52+
continue
53+
is_prime = True
54+
for i in range(2, int(number**0.5) + 1):
55+
if number % i == 0:
56+
is_prime = False
57+
break
58+
if is_prime:
59+
primes.add(number)
60+
return (
61+
"No prime numbers found."
62+
if not primes
63+
else f"{', '.join(str(num) for num in primes)} are prime numbers."
64+
)
65+
66+
67+
root_agent = Agent(
68+
model="apigee/gemini-2.5-flash",
69+
name="hello_world_agent",
70+
description=(
71+
"hello world agent that can roll a dice of 8 sides and check prime"
72+
" numbers."
73+
),
74+
instruction="""
75+
You roll dice and answer questions about the outcome of the dice rolls.
76+
You can roll dice of different sizes.
77+
You can use multiple tools in parallel by calling functions in parallel(in one request and in one round).
78+
It is ok to discuss previous dice roles, and comment on the dice rolls.
79+
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.
80+
You should never roll a die on your own.
81+
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.
82+
You should not check prime numbers before calling the tool.
83+
When you are asked to roll a die and check prime numbers, you should always make the following two function calls:
84+
1. You should first call the roll_die tool to get a roll. Wait for the function response before calling the check_prime tool.
85+
2. After you get the function response from roll_die tool, you should call the check_prime tool with the roll_die result.
86+
2.1 If user asks you to check primes based on previous rolls, make sure you include the previous rolls in the list.
87+
3. When you respond, you must include the roll_die result from step 1.
88+
You should always perform the previous 3 steps when asking for a roll and checking prime numbers.
89+
You should not rely on the previous history on prime results.
90+
""",
91+
tools=[
92+
roll_die,
93+
check_prime,
94+
],
95+
# planner=BuiltInPlanner(
96+
# thinking_config=types.ThinkingConfig(
97+
# include_thoughts=True,
98+
# ),
99+
# ),
100+
generate_content_config=types.GenerateContentConfig(
101+
safety_settings=[
102+
types.SafetySetting( # avoid false alarm about rolling dice.
103+
category=types.HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT,
104+
threshold=types.HarmBlockThreshold.OFF,
105+
),
106+
]
107+
),
108+
)
Lines changed: 112 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,112 @@
1+
# Copyright 2025 Google LLC
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
import asyncio
16+
import os
17+
import time
18+
19+
import agent
20+
from dotenv import load_dotenv
21+
from google.adk.agents.run_config import RunConfig
22+
from google.adk.cli.utils import logs
23+
from google.adk.runners import InMemoryRunner
24+
from google.adk.sessions.session import Session
25+
from google.genai import types
26+
27+
load_dotenv(override=True)
28+
logs.log_to_tmp_folder()
29+
30+
31+
async def main():
32+
app_name = "my_app"
33+
user_id_1 = "user1"
34+
runner = InMemoryRunner(
35+
agent=agent.root_agent,
36+
app_name=app_name,
37+
)
38+
session_11 = await runner.session_service.create_session(
39+
app_name=app_name, user_id=user_id_1
40+
)
41+
42+
async def run_prompt(session: Session, new_message: str):
43+
content = types.Content(
44+
role="user", parts=[types.Part.from_text(text=new_message)]
45+
)
46+
print("** User says:", content.model_dump(exclude_none=True))
47+
async for event in runner.run_async(
48+
user_id=user_id_1,
49+
session_id=session.id,
50+
new_message=content,
51+
):
52+
if event.content.parts and event.content.parts[0].text:
53+
print(f"** {event.author}: {event.content.parts[0].text}")
54+
55+
async def run_prompt_bytes(session: Session, new_message: str):
56+
content = types.Content(
57+
role="user",
58+
parts=[
59+
types.Part.from_bytes(
60+
data=str.encode(new_message), mime_type="text/plain"
61+
)
62+
],
63+
)
64+
print("** User says:", content.model_dump(exclude_none=True))
65+
async for event in runner.run_async(
66+
user_id=user_id_1,
67+
session_id=session.id,
68+
new_message=content,
69+
run_config=RunConfig(save_input_blobs_as_artifacts=True),
70+
):
71+
if event.content.parts and event.content.parts[0].text:
72+
print(f"** {event.author}: {event.content.parts[0].text}")
73+
74+
async def check_rolls_in_state(rolls_size: int):
75+
session = await runner.session_service.get_session(
76+
app_name=app_name, user_id=user_id_1, session_id=session_11.id
77+
)
78+
assert len(session.state["rolls"]) == rolls_size
79+
for roll in session.state["rolls"]:
80+
assert roll > 0 and roll <= 100
81+
82+
start_time = time.time()
83+
print("Start time:", start_time)
84+
print("------------------------------------")
85+
await run_prompt(session_11, "Hi")
86+
await run_prompt(session_11, "Roll a die with 100 sides")
87+
await check_rolls_in_state(1)
88+
await run_prompt(session_11, "Roll a die again with 100 sides.")
89+
await check_rolls_in_state(2)
90+
await run_prompt(session_11, "What numbers did I got?")
91+
await run_prompt_bytes(session_11, "Hi bytes")
92+
print(
93+
await runner.artifact_service.list_artifact_keys(
94+
app_name=app_name, user_id=user_id_1, session_id=session_11.id
95+
)
96+
)
97+
end_time = time.time()
98+
print("------------------------------------")
99+
print("End time:", end_time)
100+
print("Total time:", end_time - start_time)
101+
102+
103+
if __name__ == "__main__":
104+
# The API key can be set in a .env file.
105+
# For example, create a .env file with the following content:
106+
# GOOGLE_API_KEY="your-api-key"
107+
# APIGEE_PROXY_URL="your-proxy-url"
108+
if not os.getenv("GOOGLE_API_KEY"):
109+
raise ValueError("GOOGLE_API_KEY environment variable is not set.")
110+
if not os.getenv("APIGEE_PROXY_URL"):
111+
raise ValueError("APIGEE_PROXY_URL environment variable is not set.")
112+
asyncio.run(main())

src/google/adk/models/__init__.py

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@
1414

1515
"""Defines the interface to support a model."""
1616

17+
from .apigee_llm import ApigeeLlm
1718
from .base_llm import BaseLlm
1819
from .gemma_llm import Gemma
1920
from .google_llm import Gemini
@@ -31,3 +32,4 @@
3132

3233
LLMRegistry.register(Gemini)
3334
LLMRegistry.register(Gemma)
35+
LLMRegistry.register(ApigeeLlm)

0 commit comments

Comments
 (0)