Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
18 changes: 18 additions & 0 deletions tests/runtime/differential/test_runtime_parity.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
135 changes: 134 additions & 1 deletion tests/test_agent.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -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

Expand Down Expand Up @@ -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")
Expand Down
8 changes: 6 additions & 2 deletions veadk/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -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


Expand All @@ -27,11 +27,15 @@ 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

return Runner
raise AttributeError(f"module 'veadk' has no attribute '{name}'")


__all__ = ["Agent", "Runner", "VERSION"]
__all__ = ["Agent", "ModelFallbackEndpoint", "Runner", "VERSION"]
Loading
Loading