diff --git a/src/google/adk/models/lite_llm.py b/src/google/adk/models/lite_llm.py index e3560a795d0..9742a1445cc 100644 --- a/src/google/adk/models/lite_llm.py +++ b/src/google/adk/models/lite_llm.py @@ -709,6 +709,16 @@ def _is_gemma4_model(model: str) -> bool: return bool(_GEMMA4_MODEL_PATTERN.search(model.lower())) +def _resolve_tool_result_role( + model: str, + tool_result_role: Literal["tool", "tool_responses"] | None = None, +) -> Literal["tool", "tool_responses"]: + """Returns the override, or Gemma-4 auto-detect when unset.""" + if tool_result_role is not None: + return tool_result_role + return "tool_responses" if _is_gemma4_model(model) else "tool" + + class ChatCompletionFileUrlObject(TypedDict, total=False): file_data: str file_id: str @@ -1269,6 +1279,7 @@ async def _content_to_message_param( *, provider: str = "", model: str = "", + tool_result_role: Literal["tool", "tool_responses"] | None = None, ) -> Union[Message, list[Message]] | None: """Converts a types.Content to a litellm Message or list of Messages. @@ -1279,6 +1290,8 @@ async def _content_to_message_param( content: The content to convert. provider: The LLM provider name (e.g., "openai", "azure"). model: The LiteLLM model string, used for provider-specific behavior. + tool_result_role: Override for tool-result message role. None keeps + Gemma-4 auto-detect. Returns: A litellm Message, a list of litellm Messages, or None if skipped. @@ -1305,8 +1318,8 @@ async def _content_to_message_param( # from the tool call, instead of OpenAI-compatible 'tool' role used by other models. # Earlier Gemma versions before version 4 do not support tool use, # so this check is intentionally scoped to only look for "gemma4" in the model name. - tool_role: Literal["tool", "tool_responses"] = ( - "tool_responses" if _is_gemma4_model(model) else "tool" + tool_role: Literal["tool", "tool_responses"] = _resolve_tool_result_role( + model, tool_result_role ) tool_messages.append( _tool_message( @@ -1330,6 +1343,7 @@ async def _content_to_message_param( types.Content(role=content.role, parts=non_tool_parts), provider=provider, model=model, + tool_result_role=tool_result_role, ) follow_up_messages = ( follow_up if isinstance(follow_up, list) else [follow_up] @@ -1463,7 +1477,12 @@ async def _content_to_message_param( ) -def _ensure_tool_results(messages: List[Message], model: str) -> List[Message]: +def _ensure_tool_results( + messages: List[Message], + model: str, + *, + tool_result_role: Literal["tool", "tool_responses"] | None = None, +) -> List[Message]: """Insert placeholder tool messages for missing tool results. LiteLLM-backed providers like OpenAI and Anthropic reject histories where an @@ -1482,7 +1501,7 @@ def _ensure_tool_results(messages: List[Message], model: str) -> List[Message]: healed_messages: List[Message] = [] pending_tool_call_ids: List[str] = [] expected_tool_role: Literal["tool", "tool_responses"] = ( - "tool_responses" if _is_gemma4_model(model) else "tool" + _resolve_tool_result_role(model, tool_result_role) ) for message in messages: @@ -2711,6 +2730,8 @@ def _to_litellm_response_format( async def _get_completion_inputs( llm_request: LlmRequest, model: str, + *, + tool_result_role: Literal["tool", "tool_responses"] | None = None, ) -> Tuple[ List[Message], Optional[List[Dict[str, Any]]], @@ -2737,7 +2758,10 @@ async def _get_completion_inputs( messages: List[Message] = [] for content in llm_request.contents or []: message_param_or_list = await _content_to_message_param( - content, provider=provider, model=model + content, + provider=provider, + model=model, + tool_result_role=tool_result_role, ) if isinstance(message_param_or_list, list): messages.extend(message_param_or_list) @@ -2753,7 +2777,9 @@ async def _get_completion_inputs( content=system_instruction, ), ) - messages = _ensure_tool_results(messages, model) + messages = _ensure_tool_results( + messages, model, tool_result_role=tool_result_role + ) # 2. Convert tool declarations tools: Optional[List[Dict[str, Any]]] = None @@ -3098,6 +3124,8 @@ class LiteLlm(BaseLlm): Attributes: model: The name of the LiteLlm model. llm_client: The LLM client to use for the model. + tool_result_role: Override for tool-result message role. None keeps + Gemma-4 auto-detect. """ # LiteLLMClient has no JSON serializer, so it is excluded from dumps to keep @@ -3105,6 +3133,9 @@ class LiteLlm(BaseLlm): llm_client: LiteLLMClient = Field(default_factory=LiteLLMClient, exclude=True) """The LLM client to use for the model.""" + tool_result_role: Optional[Literal["tool", "tool_responses"]] = None + """Override for tool-result message role. None keeps Gemma-4 auto-detect.""" + _additional_args: Dict[str, Any] = PrivateAttr(default_factory=dict) def __init__(self, model: str, **kwargs: Any) -> None: @@ -3122,6 +3153,7 @@ def __init__(self, model: str, **kwargs: Any) -> None: # preventing generation call with llm_client # and overriding messages, tools and stream which are managed internally self._additional_args.pop("llm_client", None) + self._additional_args.pop("tool_result_role", None) self._additional_args.pop("messages", None) self._additional_args.pop("tools", None) # public api called from runner determines to stream or not @@ -3158,7 +3190,11 @@ async def generate_content_async( effective_model = llm_request.model or self.model messages, tools, response_format, generation_params, tool_choice = ( - await _get_completion_inputs(llm_request, effective_model) + await _get_completion_inputs( + llm_request, + effective_model, + tool_result_role=self.tool_result_role, + ) ) normalized_messages = _normalize_ollama_chat_messages( messages, diff --git a/tests/unittests/models/test_lite_llm_gemma_tool_role.py b/tests/unittests/models/test_lite_llm_gemma_tool_role.py index f466d0ba17b..c6d8ae40687 100644 --- a/tests/unittests/models/test_lite_llm_gemma_tool_role.py +++ b/tests/unittests/models/test_lite_llm_gemma_tool_role.py @@ -22,7 +22,15 @@ from typing import Any from google.adk.models.lite_llm import _content_to_message_param +from google.adk.models.lite_llm import _get_completion_inputs +from google.adk.models.lite_llm import LiteLlm +from google.adk.models.lite_llm import LiteLLMClient +from google.adk.models.llm_request import LlmRequest from google.genai import types +from litellm import ChatCompletionAssistantMessage +from litellm.types.utils import Choices +from litellm.types.utils import ModelResponse +from pydantic import ValidationError import pytest @@ -189,3 +197,115 @@ async def test_non_gemma_multi_response_uses_tool_role(self): assert isinstance(result, list) for msg in result: assert _extract_role(msg) == "tool" + + +@pytest.mark.asyncio +async def test_tool_result_role_override_forces_tool_on_gemma4(): + content = _make_function_response_content() + + result = await _content_to_message_param( + content, + model="google/gemma-4-e4b", + tool_result_role="tool", + ) + + assert _extract_role(result) == "tool" + + +@pytest.mark.asyncio +async def test_tool_result_role_override_forces_tool_responses_on_openai(): + content = _make_function_response_content() + + result = await _content_to_message_param( + content, + model="openai/gpt-4o", + tool_result_role="tool_responses", + ) + + assert _extract_role(result) == "tool_responses" + + +def test_litellm_stores_tool_result_role_and_does_not_forward_it(): + lite_llm = LiteLlm(model="google/gemma-4-e4b", tool_result_role="tool") + + assert lite_llm.tool_result_role == "tool" + assert "tool_result_role" not in lite_llm._additional_args + restored = LiteLlm(**lite_llm.model_dump()) + assert restored.tool_result_role == "tool" + + +def test_litellm_rejects_invalid_tool_result_role(): + with pytest.raises(ValidationError, match="tool_result_role"): + LiteLlm(model="google/gemma-4-e4b", tool_result_role="assistant") + + +@pytest.mark.asyncio +async def test_get_completion_inputs_override_heals_with_tool_role(): + assistant_content = types.Content( + role="model", + parts=[ + types.Part.from_function_call( + name="get_weather", args={"location": "Seoul"} + ), + ], + ) + assistant_content.parts[0].function_call.id = "tool_call_1" + llm_request = LlmRequest( + contents=[ + types.Content(role="user", parts=[types.Part.from_text(text="Hi")]), + assistant_content, + types.Content(role="user", parts=[types.Part.from_text(text="Next")]), + ] + ) + + messages, _, _, _, _ = await _get_completion_inputs( + llm_request, + model="google/gemma-4-e4b", + tool_result_role="tool", + ) + + roles = [_extract_role(m) for m in messages] + assert "tool" in roles + assert "tool_responses" not in roles + + +@pytest.mark.asyncio +async def test_generate_content_uses_tool_result_role_override(): + captured = {} + + class _Client(LiteLLMClient): + + async def acompletion(self, **kwargs): + captured.update(kwargs) + return ModelResponse( + model="google/gemma-4-e4b", + choices=[ + Choices( + message=ChatCompletionAssistantMessage( + role="assistant", content="ok" + ) + ) + ], + ) + + lite_llm = LiteLlm( + model="google/gemma-4-e4b", + llm_client=_Client(), + tool_result_role="tool", + ) + llm_request = LlmRequest( + contents=[ + types.Content(role="user", parts=[types.Part.from_text(text="Hi")]), + _make_function_response_content(), + ] + ) + + _ = [ + r + async for r in lite_llm.generate_content_async(llm_request, stream=False) + ] + + roles = [_extract_role(m) for m in captured["messages"]] + assert "tool" in roles + assert "tool_responses" not in roles + assert "tool_result_role" not in captured