From 38003c870f2fd8fa5fec76790537756790456ad7 Mon Sep 17 00:00:00 2001 From: shanchunhua Date: Tue, 8 Sep 2026 23:39:23 +0800 Subject: [PATCH] feat: support model fallback endpoints --- .../differential/test_runtime_parity.py | 18 ++ tests/test_agent.py | 135 ++++++++++++++- veadk/__init__.py | 8 +- veadk/agent.py | 161 +++++++++++++++++- veadk/models/retrying_lite_llm.py | 14 ++ veadk/runtime/compat.py | 19 +++ 6 files changed, 345 insertions(+), 10 deletions(-) diff --git a/tests/runtime/differential/test_runtime_parity.py b/tests/runtime/differential/test_runtime_parity.py index 251b63e07..085cefed7 100644 --- a/tests/runtime/differential/test_runtime_parity.py +++ b/tests/runtime/differential/test_runtime_parity.py @@ -667,3 +667,21 @@ def test_explicit_field_snapshot_survives_clone() -> None: assert explicit_fields(clone) == explicit_fields(agent) # The real contract: the clone still validates. check_agent_runtime_support(clone, "codex") + + +def test_model_fallbacks_warn_for_external_runtime(caplog) -> None: + from veadk import Agent + from veadk.runtime.compat import reset_warning_state + + reset_warning_state() + agent = Agent( + name="codex_agent_with_fallbacks", + model_name="scripted-model", + model_api_base="https://backend.invalid/v1", + model_api_key="backend-key", + model_fallbacks=["backup-model"], + runtime="codex", + ) + + assert agent.model_fallbacks == ["backup-model"] + assert "drops Agent(model_fallbacks=...)" in caplog.text diff --git a/tests/test_agent.py b/tests/test_agent.py index ea16a18e6..5f9c7b113 100644 --- a/tests/test_agent.py +++ b/tests/test_agent.py @@ -21,7 +21,7 @@ from google.adk.models.lite_llm import LiteLlm from google.adk.tools import load_memory -from veadk import Agent +from veadk import Agent, ModelFallbackEndpoint from veadk.consts import ( DEFAULT_AGENT_NAME, DEFAULT_MODEL_AGENT_API_BASE, @@ -31,6 +31,7 @@ ) from veadk.knowledgebase import KnowledgeBase from veadk.memory.long_term_memory import LongTermMemory +from veadk.models.retrying_lite_llm import RetryingLiteLlm from veadk.tools import load_knowledgebase_tool from veadk.tracing.telemetry.opentelemetry_tracer import OpentelemetryTracer @@ -190,6 +191,138 @@ def test_agent_configures_responses_model_fallbacks(mock_ark_llm): ] +@patch("veadk.agent.RetryingLiteLlm") +def test_agent_configures_cross_provider_litellm_fallbacks(mock_lite_llm, monkeypatch): + monkeypatch.setenv("BACKUP_MODEL_API_KEY", "backup-key") + + Agent( + model_name="primary-model", + model_provider="ark", + model_api_key="primary-key", + model_api_base="https://ark.example.com/api/v3", + model_fallbacks=[ + { + "model_provider": "openai", + "model_name": "gpt-4o-mini", + "model_api_base": "https://api.openai.com/v1", + "model_api_key_env": "BACKUP_MODEL_API_KEY", + "model_extra_config": { + "extra_headers": {"x-fallback": "1"}, + "temperature": 0.1, + }, + } + ], + ) + + assert mock_lite_llm.call_args.kwargs["model"] == "ark/primary-model" + assert mock_lite_llm.call_args.kwargs["fallbacks"] == [ + { + "model": "openai/gpt-4o-mini", + "api_key": "backup-key", + "api_base": "https://api.openai.com/v1", + "extra_headers": { + **DEFAULT_MODEL_EXTRA_CONFIG["extra_headers"], + "x-fallback": "1", + }, + "temperature": 0.1, + } + ] + + +@patch("veadk.agent.RetryingLiteLlm") +def test_agent_combines_legacy_and_explicit_litellm_fallbacks(mock_lite_llm): + Agent( + model_name=["primary-model", "same-provider-a"], + model_provider="ark", + model_api_key="primary-key", + model_api_base="https://ark.example.com/api/v3", + model_fallbacks=[ + "same-provider-b", + ModelFallbackEndpoint( + model_provider="anthropic", + model_name="claude-3-5-haiku-latest", + model_api_key="anthropic-key", + ), + ], + ) + + assert mock_lite_llm.call_args.kwargs["fallbacks"] == [ + "ark/same-provider-a", + "ark/same-provider-b", + { + "model": "anthropic/claude-3-5-haiku-latest", + "api_key": "anthropic-key", + "api_base": None, + }, + ] + + +@patch("veadk.agent.RetryingLiteLlm") +def test_agent_accepts_litellm_style_fallback_dict(mock_lite_llm): + Agent( + model_name="primary-model", + model_provider="ark", + model_api_key="primary-key", + model_api_base="https://ark.example.com/api/v3", + model_fallbacks=[ + { + "model": "openai/gpt-4o-mini", + "api_key": "openai-key", + "api_base": "https://api.openai.com/v1", + } + ], + ) + + assert mock_lite_llm.call_args.kwargs["fallbacks"] == [ + { + "model": "openai/gpt-4o-mini", + "api_key": "openai-key", + "api_base": "https://api.openai.com/v1", + } + ] + + +def test_agent_rejects_endpoint_fallbacks_for_responses_model(): + with pytest.raises(ValueError, match="Endpoint model_fallbacks"): + Agent( + model_name="primary-model", + model_provider="ark", + model_api_key="primary-key", + model_api_base="https://ark.example.com/api/v3", + enable_responses=True, + model_fallbacks=[ + { + "model_provider": "openai", + "model_name": "gpt-4o-mini", + } + ], + ) + + +def test_retrying_litellm_refreshes_mutable_fallbacks_between_calls(): + model = RetryingLiteLlm( + model="ark/primary", + fallbacks=[ + { + "model": "openai/fallback", + "api_key": "fallback-key", + "api_base": "https://fallback.example.com/v1", + } + ], + ) + + model._additional_args["fallbacks"][0].pop("model") + model._refresh_fallbacks() + + assert model._additional_args["fallbacks"] == [ + { + "model": "openai/fallback", + "api_key": "fallback-key", + "api_base": "https://fallback.example.com/v1", + } + ] + + @patch.dict("os.environ", {"MODEL_AGENT_API_KEY": "mock_api_key"}) def test_agent_with_existing_model(): existing_model = LiteLlm(model="test_model") diff --git a/veadk/__init__.py b/veadk/__init__.py index 7891d6c47..517d2d3b0 100644 --- a/veadk/__init__.py +++ b/veadk/__init__.py @@ -17,7 +17,7 @@ from veadk.version import VERSION if TYPE_CHECKING: - from veadk.agent import Agent + from veadk.agent import Agent, ModelFallbackEndpoint from veadk.runner import Runner @@ -27,6 +27,10 @@ def __getattr__(name): from veadk.agent import Agent return Agent + if name == "ModelFallbackEndpoint": + from veadk.agent import ModelFallbackEndpoint + + return ModelFallbackEndpoint if name == "Runner": from veadk.runner import Runner @@ -34,4 +38,4 @@ def __getattr__(name): raise AttributeError(f"module 'veadk' has no attribute '{name}'") -__all__ = ["Agent", "Runner", "VERSION"] +__all__ = ["Agent", "ModelFallbackEndpoint", "Runner", "VERSION"] diff --git a/veadk/agent.py b/veadk/agent.py index 84cf765c2..9b58c386c 100644 --- a/veadk/agent.py +++ b/veadk/agent.py @@ -35,7 +35,7 @@ from google.adk.agents.llm_agent import InstructionProvider, ToolUnion from google.adk.agents.run_config import ToolThreadPoolConfig from google.adk.examples.base_example_provider import BaseExampleProvider -from pydantic import ConfigDict, Field +from pydantic import AliasChoices, BaseModel, ConfigDict, Field from typing_extensions import Any from veadk.config import settings @@ -75,6 +75,117 @@ logger = get_logger(__name__) +class ModelFallbackEndpoint(BaseModel): + """A LiteLLM fallback endpoint with independent provider credentials.""" + + model_config = ConfigDict(extra="forbid", populate_by_name=True) + + model_name: str = Field(validation_alias=AliasChoices("model_name", "model")) + model_provider: Optional[str] = Field( + default=None, + validation_alias=AliasChoices("model_provider", "provider"), + ) + model_api_base: Optional[str] = Field( + default=None, + validation_alias=AliasChoices("model_api_base", "api_base", "base_url"), + ) + model_api_key: Optional[str] = Field( + default=None, + validation_alias=AliasChoices("model_api_key", "api_key"), + ) + model_api_key_env: Optional[str] = Field( + default=None, + validation_alias=AliasChoices("model_api_key_env", "api_key_env"), + ) + model_extra_config: dict[str, Any] = Field( + default_factory=dict, + validation_alias=AliasChoices("model_extra_config", "extra_config"), + ) + + +ModelFallbackConfig = Union[str, ModelFallbackEndpoint] + + +def _qualified_model_name(provider: str | None, model_name: str) -> str: + model = model_name.strip() + normalized_provider = (provider or "").strip() + if not normalized_provider: + return model + prefix = f"{normalized_provider}/" + if model.startswith(prefix): + return model + return f"{prefix}{model}" + + +def _endpoint_model_name( + endpoint: ModelFallbackEndpoint, + *, + default_provider: str, +) -> str: + if endpoint.model_provider: + return _qualified_model_name(endpoint.model_provider, endpoint.model_name) + model = endpoint.model_name.strip() + if "/" in model: + return model + return _qualified_model_name(default_provider, model) + + +def _resolve_model_api_key(endpoint: ModelFallbackEndpoint) -> str | None: + if endpoint.model_api_key: + return endpoint.model_api_key + if endpoint.model_api_key_env: + value = os.getenv(endpoint.model_api_key_env) + if value: + return value + logger.warning( + "Model fallback api key env `%s` is not set; LiteLLM will use its " + "provider defaults if available.", + endpoint.model_api_key_env, + ) + return None + + +def _merged_fallback_extra_config( + base_extra_config: dict[str, Any], + endpoint_extra_config: dict[str, Any], +) -> dict[str, Any]: + extra = dict(endpoint_extra_config) + for key in ("extra_headers", "extra_body"): + value = extra.get(key) + base_value = base_extra_config.get(key) + if isinstance(value, dict) and isinstance(base_value, dict): + extra[key] = {**base_value, **value} + return extra + + +def _build_litellm_fallback( + fallback: ModelFallbackConfig, + *, + default_provider: str, + base_extra_config: dict[str, Any], +) -> str | dict[str, Any]: + if isinstance(fallback, str): + return _qualified_model_name(default_provider, fallback) + + values = _merged_fallback_extra_config( + base_extra_config=base_extra_config, + endpoint_extra_config=fallback.model_extra_config, + ) + fallback_provider = (fallback.model_provider or "").strip() + is_cross_provider = bool( + fallback_provider and fallback_provider != default_provider + ) + values["model"] = _endpoint_model_name( + fallback, + default_provider=default_provider, + ) + if fallback.model_api_key or fallback.model_api_key_env or is_cross_provider: + values["api_key"] = _resolve_model_api_key(fallback) + if fallback.model_api_base or is_cross_provider: + values["api_base"] = fallback.model_api_base + return values + + class Agent(LlmAgent): """LLM-based Agent with Volcengine capabilities. @@ -91,6 +202,8 @@ class Agent(LlmAgent): model_provider (str): Provider of the model (e.g., openai). model_api_base (str): The base URL of the model API. model_api_key (str): The API key for accessing the model. + model_fallbacks (list): LiteLLM fallback models or endpoints tried + after the primary model fails. model_extra_config (dict): Extra configurations to include in model requests. tool_thread_pool_config (Optional[ToolThreadPoolConfig]): Default thread pool config for synchronous tool execution. @@ -131,6 +244,13 @@ class Agent(LlmAgent): """Name of the ARK API key to resolve the value from (defaults to env MODEL_AGENT_API_KEY_NAME). A key value always wins over a key name, so this is ignored when `model_api_key` or the MODEL_AGENT_API_KEY env is set.""" + model_fallbacks: list[ModelFallbackConfig] = Field(default_factory=list) + """Fallback models passed to LiteLLM. + + Strings are interpreted as same-provider model names. Use + ``ModelFallbackEndpoint`` or a matching dict when a fallback needs its own + provider, API base, API key, or LiteLLM parameters. + """ model_extra_config: dict = Field(default_factory=dict) tool_thread_pool_config: Optional[ToolThreadPoolConfig] = None @@ -290,13 +410,14 @@ def model_post_init(self, __context: Any) -> None: logger.info(f"Model extra config: {self.model_extra_config}") if not self.model: - fallbacks = None + fallbacks: list[str | dict[str, Any]] = [] if isinstance(self.model_name, list): if self.model_name: model_name = self.model_name[0] - fallbacks = [ - f"{self.model_provider}/{m}" for m in self.model_name[1:] - ] + fallbacks.extend( + _qualified_model_name(self.model_provider, m) + for m in self.model_name[1:] + ) logger.info( f"Using primary model: {model_name}, with fallbacks: {self.model_name[1:]}" ) @@ -308,14 +429,35 @@ def model_post_init(self, __context: Any) -> None: else: model_name = self.model_name + if self.model_fallbacks: + fallbacks.extend( + _build_litellm_fallback( + fallback, + default_provider=self.model_provider, + base_extra_config=self.model_extra_config, + ) + for fallback in self.model_fallbacks + ) + + litellm_fallbacks = fallbacks or None + if self.enable_responses: + unsupported_fallbacks = [ + fallback for fallback in fallbacks if not isinstance(fallback, str) + ] + if unsupported_fallbacks: + raise ValueError( + "Endpoint model_fallbacks are only supported when " + "enable_responses=False. Ark Responses fallbacks must be " + "same-provider model names." + ) from veadk.models.ark_llm import ArkLlm self.model = ArkLlm( model=f"{self.model_provider}/{model_name}", api_key=self.model_api_key, api_base=self.model_api_base, - fallbacks=fallbacks, + fallbacks=litellm_fallbacks, enable_responses_cache=self.enable_responses_cache, **self.model_extra_config, ) @@ -324,13 +466,18 @@ def model_post_init(self, __context: Any) -> None: model=f"{self.model_provider}/{model_name}", api_key=self.model_api_key, api_base=self.model_api_base, - fallbacks=fallbacks, + fallbacks=litellm_fallbacks, **self.model_extra_config, ) logger.debug( f"LiteLLM client created with config: {self.model_extra_config}" ) else: + if self.model_fallbacks: + logger.warning( + "Agent(model_fallbacks=...) is ignored when Agent(model=...) " + "is provided. Configure fallbacks on the custom model object." + ) logger.warning( "You are trying to use your own LiteLLM client, some default request headers may be missing." ) diff --git a/veadk/models/retrying_lite_llm.py b/veadk/models/retrying_lite_llm.py index ddca0d8ff..cec14d234 100644 --- a/veadk/models/retrying_lite_llm.py +++ b/veadk/models/retrying_lite_llm.py @@ -78,6 +78,18 @@ class RetryingLiteLlm(LiteLlm): def __init__(self, *, model: str, **kwargs: Any) -> None: super().__init__(model=model, **kwargs) + self._fallbacks_template = copy.deepcopy( + getattr(self, "_additional_args", {}).get("fallbacks") + ) + + def _refresh_fallbacks(self) -> None: + """Give LiteLLM a fresh fallback list for each call. + + LiteLLM's lightweight fallback helper mutates dict fallback entries when + selecting their model. Keep VeADK's model object reusable across turns. + """ + if self._fallbacks_template is not None: + self._additional_args["fallbacks"] = copy.deepcopy(self._fallbacks_template) @override async def generate_content_async( @@ -88,6 +100,7 @@ async def generate_content_async( retry_request = copy.deepcopy(llm_request) emitted = False try: + self._refresh_fallbacks() async for response in super().generate_content_async( llm_request, stream=stream, @@ -106,6 +119,7 @@ async def generate_content_async( ) await asyncio.sleep(delay) + self._refresh_fallbacks() async for response in super().generate_content_async( retry_request, stream=stream, diff --git a/veadk/runtime/compat.py b/veadk/runtime/compat.py index faa23b3f0..61b2ec506 100644 --- a/veadk/runtime/compat.py +++ b/veadk/runtime/compat.py @@ -147,6 +147,14 @@ def _model_name_fallbacks(agent: Any) -> list[str]: return [] +def _model_fallbacks(agent: Any) -> list[Any]: + """Return explicit ``model_fallbacks`` entries.""" + model_fallbacks = getattr(agent, "model_fallbacks", None) + if isinstance(model_fallbacks, list): + return list(model_fallbacks) + return [] + + SUPPORT_RULES: tuple[SupportRule, ...] = ( # --- error: silently wrong results ------------------------------------- SupportRule( @@ -242,6 +250,17 @@ def _model_name_fallbacks(agent: Any) -> list[str]: "you need fallbacks." ), ), + SupportRule( + field="model_fallbacks", + policy="warn", + predicate=lambda agent, _explicit: bool(_model_fallbacks(agent)), + message=lambda agent, rt: ( + f"{rt} runtime drops Agent(model_fallbacks=...), because the " + "fallback chain lives on the LiteLLM client this runtime never " + "builds; a backend failure will surface as an error instead of " + "failing over. Use runtime='adk' if you need fallbacks." + ), + ), SupportRule( field="model_provider", policy="warn",