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:
committed by
github-actions[bot]
parent
066fcec3e8
commit
dc08d86534
@@ -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"))
|
||||
|
||||
Reference in New Issue
Block a user