Skip to content

Commit dc08d86

Browse files
GWealegithub-actions[bot]
authored andcommitted
fix: Refactor LiteLlm check to avoid ImportError
The `can_use_output_schema_with_tools` function now checks if a model is a LiteLlm instance by inspecting its type's Method Resolution Order, rather than directly importing `LiteLlm` Co-authored-by: George Weale <gweale@google.com> PiperOrigin-RevId: 882253446
1 parent 066fcec commit dc08d86

File tree

2 files changed

+54
-12
lines changed

2 files changed

+54
-12
lines changed

src/google/adk/utils/output_schema_utils.py

Lines changed: 8 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -36,10 +36,14 @@ def can_use_output_schema_with_tools(model: Union[str, BaseLlm]) -> bool:
3636
# tool_choice enforcement
3737
# This is strictly more reliable than the SetModelResponseTool
3838
# prompt-based workaround.
39-
from ..models.lite_llm import LiteLlm
40-
41-
if isinstance(model, LiteLlm):
42-
return True
39+
if not isinstance(model, str):
40+
try:
41+
from ..models.lite_llm import LiteLlm
42+
except ImportError:
43+
LiteLlm = None
44+
45+
if LiteLlm is not None and isinstance(model, LiteLlm):
46+
return True
4347

4448
model_string = model if isinstance(model, str) else model.model
4549

tests/unittests/utils/test_output_schema_utils.py

Lines changed: 46 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -12,10 +12,16 @@
1212
# See the License for the specific language governing permissions and
1313
# limitations under the License.
1414

15+
from __future__ import annotations
16+
17+
import builtins
18+
import importlib.util
19+
import sys
20+
from typing import Any
1521

1622
from google.adk.models.anthropic_llm import Claude
23+
from google.adk.models.base_llm import BaseLlm
1724
from google.adk.models.google_llm import Gemini
18-
from google.adk.models.lite_llm import LiteLlm
1925
from google.adk.utils.output_schema_utils import can_use_output_schema_with_tools
2026
import pytest
2127

@@ -38,19 +44,51 @@
3844
(Claude(model="claude-3.7-sonnet"), "1", False),
3945
(Claude(model="claude-3.7-sonnet"), "0", False),
4046
(Claude(model="claude-3.7-sonnet"), None, False),
41-
(LiteLlm(model="openai/gpt-4o"), "1", True),
42-
(LiteLlm(model="openai/gpt-4o"), "0", True),
43-
(LiteLlm(model="openai/gpt-4o"), None, True),
44-
(LiteLlm(model="anthropic/claude-3.7-sonnet"), None, True),
45-
(LiteLlm(model="fireworks_ai/llama-v3p1-70b"), None, True),
4647
],
4748
)
4849
def test_can_use_output_schema_with_tools(
49-
monkeypatch, model, env_value, expected
50-
):
50+
monkeypatch: pytest.MonkeyPatch,
51+
model: str | BaseLlm,
52+
env_value: str | None,
53+
expected: bool,
54+
) -> None:
5155
"""Test can_use_output_schema_with_tools."""
5256
if env_value is not None:
5357
monkeypatch.setenv("GOOGLE_GENAI_USE_VERTEXAI", env_value)
5458
else:
5559
monkeypatch.delenv("GOOGLE_GENAI_USE_VERTEXAI", raising=False)
5660
assert can_use_output_schema_with_tools(model) == expected
61+
62+
63+
def test_can_use_output_schema_with_tools_with_litellm_model() -> None:
64+
"""Test LiteLlm detection when the optional module is available."""
65+
if importlib.util.find_spec("litellm") is None:
66+
pytest.skip("litellm is not installed")
67+
68+
from google.adk.models.lite_llm import LiteLlm
69+
70+
assert can_use_output_schema_with_tools(LiteLlm(model="openai/gpt-4o"))
71+
72+
73+
def test_can_use_output_schema_with_tools_without_litellm_module(
74+
monkeypatch: pytest.MonkeyPatch,
75+
) -> None:
76+
"""Test optional LiteLlm import failures do not affect other models."""
77+
original_import = builtins.__import__
78+
79+
def _failing_import(
80+
name: str,
81+
globals_dict: dict[str, Any] | None = None,
82+
locals_dict: dict[str, Any] | None = None,
83+
fromlist: tuple[str, ...] = (),
84+
level: int = 0,
85+
) -> Any:
86+
if name.endswith("lite_llm"):
87+
raise ImportError("litellm not installed")
88+
return original_import(name, globals_dict, locals_dict, fromlist, level)
89+
90+
monkeypatch.delenv("GOOGLE_GENAI_USE_VERTEXAI", raising=False)
91+
monkeypatch.delitem(sys.modules, "google.adk.models.lite_llm", raising=False)
92+
monkeypatch.setattr(builtins, "__import__", _failing_import)
93+
94+
assert not can_use_output_schema_with_tools(Claude(model="claude-3.7-sonnet"))

0 commit comments

Comments
 (0)