diff --git a/src/google/adk/agents/llm_agent.py b/src/google/adk/agents/llm_agent.py index 71a07488..005d073c 100644 --- a/src/google/adk/agents/llm_agent.py +++ b/src/google/adk/agents/llm_agent.py @@ -469,7 +469,10 @@ class LlmAgent(BaseAgent): if ctx.is_resumable: events = ctx._get_events(current_invocation=True, current_branch=True) - if any(ctx.should_pause_invocation(e) for e in events[-2:]): + if events and ( + ctx.should_pause_invocation(events[-1]) + or ctx.should_pause_invocation(events[-2]) + ): return # Only yield an end state if the last event is no longer a long running # tool call. @@ -907,7 +910,9 @@ class LlmAgent(BaseAgent): from .config_agent_utils import resolve_callbacks from .config_agent_utils import resolve_code_reference - if config.model: + if config.model_code: + kwargs['model'] = resolve_code_reference(config.model_code) + elif config.model: kwargs['model'] = config.model if config.instruction: kwargs['instruction'] = config.instruction diff --git a/src/google/adk/agents/llm_agent_config.py b/src/google/adk/agents/llm_agent_config.py index 7d249359..59c6d588 100644 --- a/src/google/adk/agents/llm_agent_config.py +++ b/src/google/adk/agents/llm_agent_config.py @@ -15,6 +15,7 @@ from __future__ import annotations import logging +from typing import Any from typing import List from typing import Literal from typing import Optional @@ -22,6 +23,7 @@ from typing import Optional from google.genai import types from pydantic import ConfigDict from pydantic import Field +from pydantic import model_validator from ..tools.tool_configs import ToolConfig from .base_agent_config import BaseAgentConfig @@ -52,11 +54,48 @@ class LlmAgentConfig(BaseAgentConfig): model: Optional[str] = Field( default=None, description=( - 'Optional. LlmAgent.model. If not set, the model will be inherited' - ' from the ancestor.' + 'Optional. LlmAgent.model. Provide a model name string (e.g.' + ' "gemini-2.0-flash"). If not set, the model will be inherited from' + ' the ancestor. To construct a model instance from code, use' + ' model_code.' ), ) + model_code: Optional[CodeConfig] = Field( + default=None, + description=( + 'Optional. A CodeConfig that instantiates a BaseLlm implementation' + ' such as LiteLlm with custom arguments (API base, fallbacks,' + ' etc.). Cannot be set together with `model`.' + ), + ) + + @model_validator(mode='before') + @classmethod + def _normalize_model_code(cls, data: Any) -> dict[str, Any] | Any: + if not isinstance(data, dict): + return data + + model_value = data.get('model') + model_code = data.get('model_code') + if isinstance(model_value, dict) and model_code is None: + logger.warning( + 'Detected legacy `model` mapping. Use `model_code` to provide a' + ' CodeConfig for custom model construction.' + ) + data = dict(data) + data['model_code'] = model_value + data['model'] = None + + return data + + @model_validator(mode='after') + def _validate_model_sources(self) -> LlmAgentConfig: + if self.model and self.model_code: + raise ValueError('Only one of `model` or `model_code` should be set.') + + return self + instruction: str = Field( description=( 'Required. LlmAgent.instruction. Dynamic instructions with' diff --git a/tests/unittests/agents/test_agent_config.py b/tests/unittests/agents/test_agent_config.py index 3d8e9209..86fda7fc 100644 --- a/tests/unittests/agents/test_agent_config.py +++ b/tests/unittests/agents/test_agent_config.py @@ -29,6 +29,7 @@ from google.adk.agents.llm_agent import LlmAgent from google.adk.agents.loop_agent import LoopAgent from google.adk.agents.parallel_agent import ParallelAgent from google.adk.agents.sequential_agent import SequentialAgent +from google.adk.models.lite_llm import LiteLlm import pytest import yaml @@ -259,6 +260,53 @@ sub_agents: assert config.root.agent_class == agent_class_value +def test_agent_config_litellm_model_with_custom_args(tmp_path: Path): + yaml_content = """\ +name: managed_api_agent +description: Agent using LiteLLM managed endpoint +instruction: Respond concisely. +model_code: + name: google.adk.models.lite_llm.LiteLlm + args: + - name: model + value: kimi/k2 + - name: api_base + value: https://proxy.litellm.ai/v1 +""" + config_file = tmp_path / "litellm_agent.yaml" + config_file.write_text(yaml_content) + + agent = config_agent_utils.from_config(str(config_file)) + + assert isinstance(agent, LlmAgent) + assert isinstance(agent.model, LiteLlm) + assert agent.model.model == "kimi/k2" + assert agent.model._additional_args.get("api_base") == ( + "https://proxy.litellm.ai/v1" + ) + + +def test_agent_config_legacy_model_mapping_still_supported(tmp_path: Path): + yaml_content = """\ +name: managed_api_agent +description: Agent using LiteLLM managed endpoint +instruction: Respond concisely. +model: + name: google.adk.models.lite_llm.LiteLlm + args: + - name: model + value: kimi/k2 +""" + config_file = tmp_path / "legacy_litellm_agent.yaml" + config_file.write_text(yaml_content) + + agent = config_agent_utils.from_config(str(config_file)) + + assert isinstance(agent, LlmAgent) + assert isinstance(agent.model, LiteLlm) + assert agent.model.model == "kimi/k2" + + def test_agent_config_discriminator_custom_agent(): class MyCustomAgentConfig(BaseAgentConfig): agent_class: Literal["mylib.agents.MyCustomAgent"] = (