Skip to content

Commit ef010fd

Browse files
authored
Fix e2e tests by switching to native Gemini integration (#1375)
This PR fixes the failed e2e tests for restaurant_finder in CI. As suggested by the warning in the logs, I switched from using LiteLlm to native Gemini integration in RestaurantAgent (samples/agent/adk/restaurant_finder/agent.py). This avoids the issue with LiteLLM trying to use an unavailable API version.
1 parent 4d62cfa commit ef010fd

1 file changed

Lines changed: 8 additions & 3 deletions

File tree

  • samples/agent/adk/restaurant_finder

samples/agent/adk/restaurant_finder/agent.py

Lines changed: 8 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -32,7 +32,7 @@
3232
from google.adk.agents.llm_agent import LlmAgent
3333
from google.adk.artifacts import InMemoryArtifactService
3434
from google.adk.memory.in_memory_memory_service import InMemoryMemoryService
35-
from google.adk.models.lite_llm import LiteLlm
35+
from google.adk.models import Gemini
3636
from google.adk.runners import Runner
3737
from google.adk.sessions import InMemorySessionService
3838
from google.genai import types
@@ -149,7 +149,12 @@ def _build_llm_agent(
149149
self, schema_manager: Optional[A2uiSchemaManager] = None
150150
) -> LlmAgent:
151151
"""Builds the LLM agent for the restaurant agent."""
152-
LITELLM_MODEL = os.getenv("LITELLM_MODEL", "gemini/gemini-2.5-flash")
152+
model_env = (
153+
os.getenv("MODEL_NAME")
154+
or os.getenv("LITELLM_MODEL")
155+
or "gemini-3-flash-preview"
156+
)
157+
model_name = model_env.split("/")[-1]
153158

154159
instruction = (
155160
schema_manager.generate_system_prompt(
@@ -164,7 +169,7 @@ def _build_llm_agent(
164169
)
165170

166171
return LlmAgent(
167-
model=LiteLlm(model=LITELLM_MODEL),
172+
model=Gemini(model=model_name),
168173
name="restaurant_agent",
169174
description="An agent that finds restaurants and helps book tables.",
170175
instruction=instruction,

0 commit comments

Comments
 (0)