fix: allow LlmAgent model to be provided via CodeConfig

LlmAgentConfig.model now accepts either a plain model string or a CodeConfig. This lets YAML configs pass a LiteLLM instance with managed API settings (e.g., api_base and fallbacks) so agents can hit KimiK2’s managed endpoint instead of only the default modelID.

Close #3579

Co-authored-by: George Weale <gweale@google.com>
PiperOrigin-RevId: 838978654

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

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'

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"] = (