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
This commit is contained in:
George Weale
2026-03-11 16:24:36 -07:00
committed by github-actions[bot]
parent 066fcec3e8
commit dc08d86534
2 changed files with 53 additions and 11 deletions
+7 -3
View File
@@ -36,10 +36,14 @@ def can_use_output_schema_with_tools(model: Union[str, BaseLlm]) -> bool:
# tool_choice enforcement
# This is strictly more reliable than the SetModelResponseTool
# prompt-based workaround.
from ..models.lite_llm import LiteLlm
if not isinstance(model, str):
try:
from ..models.lite_llm import LiteLlm
except ImportError:
LiteLlm = None
if isinstance(model, LiteLlm):
return True
if LiteLlm is not None and isinstance(model, LiteLlm):
return True
model_string = model if isinstance(model, str) else model.model
@@ -12,10 +12,16 @@
# See the License for the specific language governing permissions and
# limitations under the License.
from __future__ import annotations
import builtins
import importlib.util
import sys
from typing import Any
from google.adk.models.anthropic_llm import Claude
from google.adk.models.base_llm import BaseLlm
from google.adk.models.google_llm import Gemini
from google.adk.models.lite_llm import LiteLlm
from google.adk.utils.output_schema_utils import can_use_output_schema_with_tools
import pytest
@@ -38,19 +44,51 @@ import pytest
(Claude(model="claude-3.7-sonnet"), "1", False),
(Claude(model="claude-3.7-sonnet"), "0", False),
(Claude(model="claude-3.7-sonnet"), None, False),
(LiteLlm(model="openai/gpt-4o"), "1", True),
(LiteLlm(model="openai/gpt-4o"), "0", True),
(LiteLlm(model="openai/gpt-4o"), None, True),
(LiteLlm(model="anthropic/claude-3.7-sonnet"), None, True),
(LiteLlm(model="fireworks_ai/llama-v3p1-70b"), None, True),
],
)
def test_can_use_output_schema_with_tools(
monkeypatch, model, env_value, expected
):
monkeypatch: pytest.MonkeyPatch,
model: str | BaseLlm,
env_value: str | None,
expected: bool,
) -> None:
"""Test can_use_output_schema_with_tools."""
if env_value is not None:
monkeypatch.setenv("GOOGLE_GENAI_USE_VERTEXAI", env_value)
else:
monkeypatch.delenv("GOOGLE_GENAI_USE_VERTEXAI", raising=False)
assert can_use_output_schema_with_tools(model) == expected
def test_can_use_output_schema_with_tools_with_litellm_model() -> None:
"""Test LiteLlm detection when the optional module is available."""
if importlib.util.find_spec("litellm") is None:
pytest.skip("litellm is not installed")
from google.adk.models.lite_llm import LiteLlm
assert can_use_output_schema_with_tools(LiteLlm(model="openai/gpt-4o"))
def test_can_use_output_schema_with_tools_without_litellm_module(
monkeypatch: pytest.MonkeyPatch,
) -> None:
"""Test optional LiteLlm import failures do not affect other models."""
original_import = builtins.__import__
def _failing_import(
name: str,
globals_dict: dict[str, Any] | None = None,
locals_dict: dict[str, Any] | None = None,
fromlist: tuple[str, ...] = (),
level: int = 0,
) -> Any:
if name.endswith("lite_llm"):
raise ImportError("litellm not installed")
return original_import(name, globals_dict, locals_dict, fromlist, level)
monkeypatch.delenv("GOOGLE_GENAI_USE_VERTEXAI", raising=False)
monkeypatch.delitem(sys.modules, "google.adk.models.lite_llm", raising=False)
monkeypatch.setattr(builtins, "__import__", _failing_import)
assert not can_use_output_schema_with_tools(Claude(model="claude-3.7-sonnet"))