chore: import upstream snapshot with attribution

tests/unittests/models/__init__.py

@@ -0,0 +1,13 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.

tests/unittests/models/test_apigee_llm.py

@@ -0,0 +1,799 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+import os
+from unittest import mock
+from unittest.mock import AsyncMock
+
+from google.adk.models.apigee_llm import ApigeeLlm
+from google.adk.models.apigee_llm import CompletionsHTTPClient
+from google.adk.models.llm_request import LlmRequest
+from google.genai import types
+from google.genai.types import Content
+from google.genai.types import Part
+import pytest
+
+BASE_MODEL_ID = 'gemini-2.5-flash'
+APIGEE_GEMINI_MODEL_ID = 'apigee/gemini/v1/' + BASE_MODEL_ID
+APIGEE_VERTEX_MODEL_ID = 'apigee/vertex_ai/v1beta/gemini-pro'
+VERTEX_BASE_MODEL_ID = 'gemini-pro'
+PROXY_URL = 'https://test.apigee.net'
+
+
+@pytest.fixture
+def llm_request():
+  """Provides a sample LlmRequest for testing."""
+  return LlmRequest(
+      model=APIGEE_GEMINI_MODEL_ID,
+      contents=[
+          types.Content(
+              role='user', parts=[types.Part.from_text(text='Test prompt')]
+          )
+      ],
+  )
+
+
+@pytest.mark.asyncio
+@mock.patch('google.genai.Client')
+async def test_generate_content_async_non_streaming(
+    mock_client_constructor, llm_request
+):
+  """Tests the generate_content_async method for non-streaming responses."""
+  apigee_llm_instance = ApigeeLlm(
+      model=APIGEE_GEMINI_MODEL_ID,
+      proxy_url=PROXY_URL,
+  )
+  mock_client_instance = mock.Mock()
+  mock_response = types.GenerateContentResponse(
+      candidates=[
+          types.Candidate(
+              content=Content(
+                  parts=[Part.from_text(text='Test response')],
+                  role='model',
+              )
+          )
+      ]
+  )
+  mock_client_instance.aio.models.generate_content = AsyncMock(
+      return_value=mock_response
+  )
+  mock_client_constructor.return_value = mock_client_instance
+
+  response_generator = apigee_llm_instance.generate_content_async(llm_request)
+  responses = [resp async for resp in response_generator]
+
+  assert len(responses) == 1
+  llm_response = responses[0]
+  assert llm_response.content.parts[0].text == 'Test response'
+  assert llm_response.content.role == 'model'
+
+  mock_client_constructor.assert_called_once()
+  _, kwargs = mock_client_constructor.call_args
+  assert not kwargs['enterprise']
+  http_options = kwargs['http_options']
+  assert http_options.base_url == PROXY_URL
+  assert http_options.api_version == 'v1'
+  assert 'user-agent' in http_options.headers
+  assert 'x-goog-api-client' in http_options.headers
+
+  mock_client_instance.aio.models.generate_content.assert_called_once_with(
+      model=BASE_MODEL_ID,
+      contents=llm_request.contents,
+      config=llm_request.config,
+  )
+
+
+@pytest.mark.asyncio
+@mock.patch('google.genai.Client')
+async def test_generate_content_async_streaming(
+    mock_client_constructor, llm_request
+):
+  """Tests the generate_content_async method for streaming responses."""
+  apigee_llm_instance = ApigeeLlm(
+      model=APIGEE_GEMINI_MODEL_ID,
+      proxy_url=PROXY_URL,
+  )
+  mock_client_instance = mock.Mock()
+  mock_responses = [
+      types.GenerateContentResponse(
+          candidates=[
+              types.Candidate(
+                  content=Content(
+                      parts=[Part.from_text(text='Hello')],
+                  )
+              )
+          ]
+      ),
+      types.GenerateContentResponse(
+          candidates=[
+              types.Candidate(
+                  content=Content(
+                      parts=[Part.from_text(text=',')],
+                  )
+              )
+          ]
+      ),
+      types.GenerateContentResponse(
+          candidates=[
+              types.Candidate(
+                  content=Content(
+                      parts=[Part.from_text(text=' world!')],
+                  )
+              )
+          ]
+      ),
+  ]
+
+  async def mock_stream_generator():
+    for r in mock_responses:
+      yield r
+
+  mock_client_instance.aio.models.generate_content_stream = AsyncMock(
+      return_value=mock_stream_generator()
+  )
+  mock_client_constructor.return_value = mock_client_instance
+
+  response_generator = apigee_llm_instance.generate_content_async(
+      llm_request, stream=True
+  )
+  responses = [resp async for resp in response_generator]
+
+  assert responses
+  full_text_parts = []
+  for r in responses:
+    for p in r.content.parts:
+      if p.text:
+        full_text_parts.append(p.text)
+  full_text = ''.join(full_text_parts)
+  assert 'Hello, world!' in full_text
+
+  mock_client_instance.aio.models.generate_content_stream.assert_called_once_with(
+      model=BASE_MODEL_ID,
+      contents=llm_request.contents,
+      config=llm_request.config,
+  )
+
+
+@pytest.mark.asyncio
+@mock.patch('google.genai.Client')
+async def test_generate_content_async_with_custom_headers(
+    mock_client_constructor, llm_request
+):
+  """Tests that custom headers are passed in the request."""
+  custom_headers = {
+      'X-Custom-Header': 'custom-value',
+  }
+  apigee_llm = ApigeeLlm(
+      model=APIGEE_GEMINI_MODEL_ID,
+      proxy_url=PROXY_URL,
+      custom_headers=custom_headers,
+  )
+  mock_client_instance = mock.Mock()
+  mock_response = types.GenerateContentResponse(
+      candidates=[
+          types.Candidate(
+              content=Content(
+                  parts=[Part.from_text(text='Test response')],
+                  role='model',
+              )
+          )
+      ]
+  )
+  mock_client_instance.aio.models.generate_content = AsyncMock(
+      return_value=mock_response
+  )
+  mock_client_constructor.return_value = mock_client_instance
+
+  response_generator = apigee_llm.generate_content_async(llm_request)
+  _ = [resp async for resp in response_generator]  # Consume generator
+
+  mock_client_constructor.assert_called_once()
+  _, kwargs = mock_client_constructor.call_args
+  http_options = kwargs['http_options']
+  assert http_options.headers['X-Custom-Header'] == 'custom-value'
+  assert 'user-agent' in http_options.headers
+
+
+@pytest.mark.asyncio
+@mock.patch('google.genai.Client')
+async def test_vertex_model_path_parsing(mock_client_constructor):
+  """Tests that Vertex AI model paths are parsed correctly."""
+  apigee_llm = ApigeeLlm(model=APIGEE_VERTEX_MODEL_ID, proxy_url=PROXY_URL)
+  llm_request = LlmRequest(
+      model=APIGEE_VERTEX_MODEL_ID,
+      contents=[
+          types.Content(
+              role='user', parts=[types.Part.from_text(text='Test prompt')]
+          )
+      ],
+  )
+  mock_client_instance = mock.Mock()
+  mock_client_instance.aio.models.generate_content = AsyncMock(
+      return_value=types.GenerateContentResponse(
+          candidates=[
+              types.Candidate(
+                  content=Content(
+                      parts=[Part.from_text(text='Test response')],
+                      role='model',
+                  )
+              )
+          ]
+      )
+  )
+  mock_client_constructor.return_value = mock_client_instance
+
+  _ = [resp async for resp in apigee_llm.generate_content_async(llm_request)]
+
+  mock_client_constructor.assert_called_once()
+  _, kwargs = mock_client_constructor.call_args
+  assert kwargs['enterprise']
+  assert kwargs['http_options'].api_version == 'v1beta'
+
+  mock_client_instance.aio.models.generate_content.assert_called_once()
+  call_kwargs = (
+      mock_client_instance.aio.models.generate_content.call_args.kwargs
+  )
+  assert call_kwargs['model'] == VERTEX_BASE_MODEL_ID
+
+
+@pytest.mark.asyncio
+@mock.patch('google.genai.Client')
+async def test_proxy_url_from_env_variable(mock_client_constructor):
+  """Tests that proxy_url is read from environment variable."""
+  with mock.patch.dict(
+      os.environ, {'APIGEE_PROXY_URL': 'https://env.proxy.url'}
+  ):
+    apigee_llm = ApigeeLlm(model=APIGEE_GEMINI_MODEL_ID)
+    llm_request = LlmRequest(
+        model=APIGEE_GEMINI_MODEL_ID,
+        contents=[
+            types.Content(
+                role='user', parts=[types.Part.from_text(text='Test prompt')]
+            )
+        ],
+    )
+    mock_client_instance = mock.Mock()
+    mock_client_instance.aio.models.generate_content = AsyncMock(
+        return_value=types.GenerateContentResponse(
+            candidates=[
+                types.Candidate(
+                    content=Content(
+                        parts=[Part.from_text(text='Test response')],
+                        role='model',
+                    )
+                )
+            ]
+        )
+    )
+    mock_client_constructor.return_value = mock_client_instance
+
+    _ = [resp async for resp in apigee_llm.generate_content_async(llm_request)]
+
+    mock_client_constructor.assert_called_once()
+    _, kwargs = mock_client_constructor.call_args
+    assert kwargs['http_options'].base_url == 'https://env.proxy.url'
+
+
+@pytest.mark.parametrize(
+    ('model_string', 'env_vars'),
+    [
+        (
+            'apigee/vertex_ai/gemini-2.5-flash',
+            {'GOOGLE_CLOUD_LOCATION': 'test-location'},
+        ),
+        (
+            'apigee/vertex_ai/gemini-2.5-flash',
+            {'GOOGLE_CLOUD_PROJECT': 'test-project'},
+        ),
+        (
+            'apigee/gemini-2.5-flash',
+            {
+                'GOOGLE_GENAI_USE_ENTERPRISE': 'true',
+                'GOOGLE_CLOUD_LOCATION': 'test-location',
+            },
+        ),
+        (
+            'apigee/gemini-2.5-flash',
+            {
+                'GOOGLE_GENAI_USE_ENTERPRISE': 'true',
+                'GOOGLE_CLOUD_PROJECT': 'test-project',
+            },
+        ),
+    ],
+)
+def test_vertex_model_missing_project_or_location_raises_error(
+    model_string, env_vars
+):
+  """Tests that ValueError is raised for Vertex models if project or location is missing."""
+  with mock.patch.dict(os.environ, env_vars, clear=True):
+    with pytest.raises(ValueError, match='environment variable must be set'):
+      ApigeeLlm(model=model_string, proxy_url=PROXY_URL)
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    (
+        'model_string',
+        'use_vertexai_env',
+        'expected_is_vertexai',
+        'expected_api_version',
+        'expected_model_id',
+    ),
+    [
+        ('apigee/gemini-2.5-flash', None, False, None, 'gemini-2.5-flash'),
+        ('apigee/gemini-2.5-flash', 'true', True, None, 'gemini-2.5-flash'),
+        ('apigee/gemini-2.5-flash', '1', True, None, 'gemini-2.5-flash'),
+        ('apigee/gemini-2.5-flash', 'false', False, None, 'gemini-2.5-flash'),
+        ('apigee/gemini-2.5-flash', '0', False, None, 'gemini-2.5-flash'),
+        (
+            'apigee/v1/gemini-2.5-flash',
+            None,
+            False,
+            'v1',
+            'gemini-2.5-flash',
+        ),
+        (
+            'apigee/v1/gemini-2.5-flash',
+            'true',
+            True,
+            'v1',
+            'gemini-2.5-flash',
+        ),
+        (
+            'apigee/vertex_ai/gemini-2.5-flash',
+            None,
+            True,
+            None,
+            'gemini-2.5-flash',
+        ),
+        (
+            'apigee/vertex_ai/gemini-2.5-flash',
+            'false',
+            True,
+            None,
+            'gemini-2.5-flash',
+        ),
+        (
+            'apigee/gemini/v1/gemini-2.5-flash',
+            'true',
+            False,
+            'v1',
+            'gemini-2.5-flash',
+        ),
+        (
+            'apigee/vertex_ai/v1beta/gemini-2.5-flash',
+            'false',
+            True,
+            'v1beta',
+            'gemini-2.5-flash',
+        ),
+    ],
+)
+@mock.patch('google.genai.Client')
+async def test_model_string_parsing_and_client_initialization(
+    mock_client_constructor,
+    model_string,
+    use_vertexai_env,
+    expected_is_vertexai,
+    expected_api_version,
+    expected_model_id,
+):
+  """Tests model string parsing and genai.Client initialization."""
+  env_vars = {}
+  if use_vertexai_env is not None:
+    env_vars['GOOGLE_GENAI_USE_ENTERPRISE'] = use_vertexai_env
+
+  if expected_is_vertexai:
+    env_vars['GOOGLE_CLOUD_PROJECT'] = 'test-project'
+    env_vars['GOOGLE_CLOUD_LOCATION'] = 'test-location'
+
+  # The ApigeeLlm is initialized in the 'with' block to make sure that the mock
+  # of the environment variable is active.
+  with mock.patch.dict(os.environ, env_vars, clear=True):
+    apigee_llm = ApigeeLlm(model=model_string, proxy_url=PROXY_URL)
+    request = LlmRequest(model=model_string, contents=[])
+
+    mock_client_instance = mock.Mock()
+    mock_client_instance.aio.models.generate_content = AsyncMock(
+        return_value=types.GenerateContentResponse(
+            candidates=[
+                types.Candidate(
+                    content=Content(parts=[Part.from_text(text='')])
+                )
+            ]
+        )
+    )
+    mock_client_constructor.return_value = mock_client_instance
+
+    _ = [resp async for resp in apigee_llm.generate_content_async(request)]
+
+    mock_client_constructor.assert_called_once()
+    _, kwargs = mock_client_constructor.call_args
+    assert kwargs['enterprise'] == expected_is_vertexai
+    if expected_is_vertexai:
+      assert kwargs['project'] == 'test-project'
+      assert kwargs['location'] == 'test-location'
+    http_options = kwargs['http_options']
+    assert http_options.api_version == expected_api_version
+
+    (
+        mock_client_instance.aio.models.generate_content.assert_called_once_with(
+            model=expected_model_id,
+            contents=request.contents,
+            config=request.config,
+        )
+    )
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    'invalid_model_string',
+    [
+        'apigee/',  # Missing model_id
+        'apigee',  # Invalid format
+        'gemini-pro',  # Invalid format
+        'apigee/vertex_ai/v1/model/extra',  # Too many components
+        'apigee/unknown/model',
+    ],
+)
+async def test_invalid_model_strings_raise_value_error(invalid_model_string):
+  """Tests that invalid model strings raise a ValueError."""
+  with pytest.raises(
+      ValueError, match=f'Invalid model string: {invalid_model_string}'
+  ):
+    ApigeeLlm(model=invalid_model_string, proxy_url=PROXY_URL)
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    'model',
+    [
+        'apigee/openai/gpt-4o',
+        'apigee/openai/v1/gpt-4o',
+        'apigee/openai/v1/gpt-3.5-turbo',
+    ],
+)
+async def test_validate_model_for_chat_completion_providers(model):
+  """Tests that new providers like OpenAI are accepted."""
+  # Should not raise ValueError
+  ApigeeLlm(model=model, proxy_url=PROXY_URL)
+
+
+@pytest.mark.parametrize(
+    ('model', 'api_type', 'expected_api_type'),
+    [
+        # Default case (input defaults to UNKNOWN)
+        (
+            'apigee/openai/gpt-4o',
+            ApigeeLlm.ApiType.UNKNOWN,
+            ApigeeLlm.ApiType.CHAT_COMPLETIONS,
+        ),
+        (
+            'apigee/openai/v1/gpt-3.5-turbo',
+            ApigeeLlm.ApiType.UNKNOWN,
+            ApigeeLlm.ApiType.CHAT_COMPLETIONS,
+        ),
+        (
+            'apigee/gemini/v1/gemini-pro',
+            ApigeeLlm.ApiType.UNKNOWN,
+            ApigeeLlm.ApiType.GENAI,
+        ),
+        (
+            'apigee/vertex_ai/gemini-pro',
+            ApigeeLlm.ApiType.UNKNOWN,
+            ApigeeLlm.ApiType.GENAI,
+        ),
+        (
+            'apigee/vertex_ai/v1beta/gemini-1.5-pro',
+            ApigeeLlm.ApiType.UNKNOWN,
+            ApigeeLlm.ApiType.GENAI,
+        ),
+        # Override by setting the ApiType
+        (
+            'apigee/gemini/pro',
+            ApigeeLlm.ApiType.CHAT_COMPLETIONS,
+            ApigeeLlm.ApiType.CHAT_COMPLETIONS,
+        ),
+        (
+            'apigee/gemini/pro',
+            ApigeeLlm.ApiType.GENAI,
+            ApigeeLlm.ApiType.GENAI,
+        ),
+        (
+            'apigee/openai/gpt-4o',
+            ApigeeLlm.ApiType.CHAT_COMPLETIONS,
+            ApigeeLlm.ApiType.CHAT_COMPLETIONS,
+        ),
+        (
+            'apigee/openai/gpt-4o',
+            ApigeeLlm.ApiType.GENAI,
+            ApigeeLlm.ApiType.GENAI,
+        ),
+        # Override by setting the ApiType as a string
+        (
+            'apigee/gemini/pro',
+            'chat_completions',
+            ApigeeLlm.ApiType.CHAT_COMPLETIONS,
+        ),
+        (
+            'apigee/gemini/pro',
+            'genai',
+            ApigeeLlm.ApiType.GENAI,
+        ),
+        (
+            'apigee/openai/gpt-4o',
+            'chat_completions',
+            ApigeeLlm.ApiType.CHAT_COMPLETIONS,
+        ),
+        (
+            'apigee/openai/gpt-4o',
+            'genai',
+            ApigeeLlm.ApiType.GENAI,
+        ),
+    ],
+)
+def test_api_type_resolution(model, api_type, expected_api_type):
+  """Tests that api_type is resolved correctly."""
+  llm = ApigeeLlm(
+      model=model,
+      proxy_url=PROXY_URL,
+      api_type=api_type,
+  )
+  assert llm._api_type == expected_api_type
+
+
+@pytest.mark.parametrize(
+    ('input_value', 'expected_type'),
+    [
+        ('chat_completions', ApigeeLlm.ApiType.CHAT_COMPLETIONS),
+        ('genai', ApigeeLlm.ApiType.GENAI),
+        ('unknown', ApigeeLlm.ApiType.UNKNOWN),
+        ('', ApigeeLlm.ApiType.UNKNOWN),
+        (None, ApigeeLlm.ApiType.UNKNOWN),
+    ],
+)
+def test_apitype_creation(input_value, expected_type):
+  """Tests the creation of ApiType enum members."""
+  assert ApigeeLlm.ApiType(input_value) == expected_type
+
+
+def test_apitype_creation_invalid():
+  """Tests that invalid ApiType raises ValueError."""
+  with pytest.raises(ValueError):
+    ApigeeLlm.ApiType('invalid')
+
+
+def test_invalid_api_type_raises_error():
+  """Tests that invalid string for api_type raises ValueError."""
+  with pytest.raises(ValueError):
+    ApigeeLlm(
+        model='apigee/gemini-pro',
+        proxy_url=PROXY_URL,
+        api_type='invalid_type',
+    )
+
+
+@pytest.mark.asyncio
+async def test_generate_content_async_dispatch_to_completions_client(
+    llm_request,
+):
+  """Tests that generate_content_async uses CompletionsHTTPClient for OpenAI models."""
+  llm_request.model = 'apigee/openai/gpt-4o'
+  with (
+      mock.patch.object(
+          CompletionsHTTPClient,
+          'generate_content_async',
+      ) as mock_completions_generate_content,
+      mock.patch('google.genai.Client') as mock_genai_client,
+  ):
+    apigee_llm = ApigeeLlm(model='apigee/openai/gpt-4o', proxy_url=PROXY_URL)
+    _ = [
+        r
+        async for r in apigee_llm.generate_content_async(
+            llm_request, stream=False
+        )
+    ]
+    mock_completions_generate_content.assert_called_once()
+    mock_genai_client.assert_not_called()
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    'model',
+    [
+        'apigee/openai/gpt-4o',
+        'apigee/openai/v1/gpt-3.5-turbo',
+    ],
+)
+async def test_api_key_injection_openai(model):
+  """Tests that api_key is injected for OpenAI models."""
+  apigee_llm = ApigeeLlm(
+      model=model,
+      proxy_url=PROXY_URL,
+      custom_headers={'Authorization': 'Bearer sk-test-key'},
+  )
+  client = apigee_llm._completions_http_client
+  assert client._headers['Authorization'] == 'Bearer sk-test-key'
+
+
+def test_parse_response_usage_metadata():
+  """Tests that CompletionsHTTPClient parses usage metadata correctly including reasoning tokens."""
+  client = CompletionsHTTPClient(base_url='http://test')
+  response_dict = {
+      'choices': [{
+          'message': {'role': 'assistant', 'content': 'hello'},
+          'finish_reason': 'stop',
+      }],
+      'usage': {
+          'prompt_tokens': 10,
+          'completion_tokens': 5,
+          'total_tokens': 15,
+          'completion_tokens_details': {'reasoning_tokens': 4},
+      },
+  }
+  llm_response = client._parse_response(response_dict)
+  assert llm_response.usage_metadata.prompt_token_count == 10
+  assert llm_response.usage_metadata.candidates_token_count == 5
+  assert llm_response.usage_metadata.total_token_count == 15
+  assert llm_response.usage_metadata.thoughts_token_count == 4
+
+
+@pytest.mark.asyncio
+@mock.patch('google.genai.Client')
+async def test_api_client_passes_credentials_when_provided(
+    mock_client_constructor, llm_request
+):
+  """Tests that credentials passed to __init__ are forwarded to genai.Client."""
+  mock_credentials = mock.Mock()
+
+  mock_client_instance = mock.Mock()
+  mock_client_instance.aio.models.generate_content = AsyncMock(
+      return_value=types.GenerateContentResponse(
+          candidates=[
+              types.Candidate(
+                  content=Content(
+                      parts=[Part.from_text(text='Test response')],
+                      role='model',
+                  )
+              )
+          ]
+      )
+  )
+  mock_client_constructor.return_value = mock_client_instance
+
+  apigee_llm = ApigeeLlm(
+      model=APIGEE_GEMINI_MODEL_ID,
+      proxy_url=PROXY_URL,
+      credentials=mock_credentials,
+  )
+  _ = [resp async for resp in apigee_llm.generate_content_async(llm_request)]
+
+  _, kwargs = mock_client_constructor.call_args
+  assert kwargs['credentials'] is mock_credentials
+
+
+@pytest.mark.asyncio
+@mock.patch('google.genai.Client')
+async def test_api_client_omits_credentials_when_not_provided(
+    mock_client_constructor, llm_request
+):
+  """Tests that credentials kwarg is not forwarded when not supplied."""
+  mock_client_instance = mock.Mock()
+  mock_client_instance.aio.models.generate_content = AsyncMock(
+      return_value=types.GenerateContentResponse(
+          candidates=[
+              types.Candidate(
+                  content=Content(
+                      parts=[Part.from_text(text='Test response')],
+                      role='model',
+                  )
+              )
+          ]
+      )
+  )
+  mock_client_constructor.return_value = mock_client_instance
+
+  apigee_llm = ApigeeLlm(
+      model=APIGEE_GEMINI_MODEL_ID,
+      proxy_url=PROXY_URL,
+  )
+  _ = [resp async for resp in apigee_llm.generate_content_async(llm_request)]
+
+  _, kwargs = mock_client_constructor.call_args
+  assert 'credentials' not in kwargs
+
+
+def test_parse_response_with_refusal():
+  """Tests that CompletionsHTTPClient parses refusal correctly."""
+  client = CompletionsHTTPClient(base_url='http://test')
+
+  response_dict = {
+      'choices': [{
+          'message': {
+              'role': 'assistant',
+              'refusal': 'I refuse to answer',
+          },
+          'finish_reason': 'stop',
+      }],
+  }
+  llm_response = client._parse_response(response_dict)
+  assert len(llm_response.content.parts) == 1
+  assert llm_response.content.parts[0].text == '[[REFUSAL]]: I refuse to answer'
+
+  response_dict_mixed = {
+      'choices': [{
+          'message': {
+              'role': 'assistant',
+              'content': 'Here is some content',
+              'refusal': 'But I refuse to answer the rest',
+          },
+          'finish_reason': 'stop',
+      }],
+  }
+  llm_response_mixed = client._parse_response(response_dict_mixed)
+  assert len(llm_response_mixed.content.parts) == 1
+  assert (
+      llm_response_mixed.content.parts[0].text
+      == 'Here is some content\n[[REFUSAL]]: But I refuse to answer the rest'
+  )
+
+
+@pytest.mark.parametrize(
+    ('parts', 'expected_message'),
+    [
+        (
+            [
+                types.Part.from_text(text='[[REFUSAL]]: I refuse to answer'),
+                types.Part.from_text(text='normal content'),
+            ],
+            {
+                'role': 'assistant',
+                'refusal': 'I refuse to answer',
+                'content': 'normal content',
+            },
+        ),
+        (
+            [
+                types.Part.from_text(
+                    text=(
+                        'Here is some content\n[[REFUSAL]]: But I refuse to'
+                        ' answer the rest'
+                    )
+                ),
+            ],
+            {
+                'role': 'assistant',
+                'refusal': 'But I refuse to answer the rest',
+                'content': 'Here is some content',
+            },
+        ),
+    ],
+)
+def test_construct_payload_with_refusal(parts, expected_message):
+  """Tests that CompletionsHTTPClient constructs payload with refusal correctly."""
+  client = CompletionsHTTPClient(base_url='http://test')
+  req = LlmRequest(
+      model='apigee/openai/gpt-4o',
+      contents=[
+          types.Content(
+              role='model',
+              parts=parts,
+          )
+      ],
+  )
+  payload = client._construct_payload(req, stream=False)
+  messages = payload['messages']
+  assert messages == [expected_message]

tests/unittests/models/test_cache_metadata.py

@@ -0,0 +1,346 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Tests for CacheMetadata."""
+
+import time
+
+from google.adk.models.cache_metadata import CacheMetadata
+from pydantic import ValidationError
+import pytest
+
+
+class TestCacheMetadata:
+  """Test suite for CacheMetadata."""
+
+  def test_required_fields(self):
+    """Test that all required fields must be provided."""
+    # Valid creation with all required fields
+    metadata = CacheMetadata(
+        cache_name="projects/123/locations/us-central1/cachedContents/456",
+        expire_time=time.time() + 1800,
+        fingerprint="abc123",
+        invocations_used=5,
+        contents_count=3,
+    )
+
+    assert (
+        metadata.cache_name
+        == "projects/123/locations/us-central1/cachedContents/456"
+    )
+    assert metadata.expire_time > time.time()
+    assert metadata.fingerprint == "abc123"
+    assert metadata.invocations_used == 5
+    assert metadata.contents_count == 3
+    assert metadata.created_at is None  # Optional field
+
+  def test_optional_created_at(self):
+    """Test that created_at is optional."""
+    current_time = time.time()
+
+    metadata = CacheMetadata(
+        cache_name="projects/123/locations/us-central1/cachedContents/456",
+        expire_time=time.time() + 1800,
+        fingerprint="abc123",
+        invocations_used=3,
+        contents_count=2,
+        created_at=current_time,
+    )
+
+    assert metadata.created_at == current_time
+
+  def test_invocations_used_validation(self):
+    """Test invocations_used validation constraints."""
+    # Valid: zero or positive
+    metadata = CacheMetadata(
+        cache_name="projects/123/locations/us-central1/cachedContents/456",
+        expire_time=time.time() + 1800,
+        fingerprint="abc123",
+        invocations_used=0,
+        contents_count=1,
+    )
+    assert metadata.invocations_used == 0
+
+    metadata = CacheMetadata(
+        cache_name="projects/123/locations/us-central1/cachedContents/456",
+        expire_time=time.time() + 1800,
+        fingerprint="abc123",
+        invocations_used=10,
+        contents_count=1,
+    )
+    assert metadata.invocations_used == 10
+
+    # Invalid: negative
+    with pytest.raises(ValidationError) as exc_info:
+      CacheMetadata(
+          cache_name="projects/123/locations/us-central1/cachedContents/456",
+          expire_time=time.time() + 1800,
+          fingerprint="abc123",
+          invocations_used=-1,
+          contents_count=1,
+      )
+    assert "greater than or equal to 0" in str(exc_info.value)
+
+  def test_contents_count_validation(self):
+    """Test contents_count validation constraints."""
+    # Valid: zero or positive
+    metadata = CacheMetadata(
+        cache_name="projects/123/locations/us-central1/cachedContents/456",
+        expire_time=time.time() + 1800,
+        fingerprint="abc123",
+        invocations_used=1,
+        contents_count=0,
+    )
+    assert metadata.contents_count == 0
+
+    metadata = CacheMetadata(
+        cache_name="projects/123/locations/us-central1/cachedContents/456",
+        expire_time=time.time() + 1800,
+        fingerprint="abc123",
+        invocations_used=1,
+        contents_count=10,
+    )
+    assert metadata.contents_count == 10
+
+    # Invalid: negative
+    with pytest.raises(ValidationError) as exc_info:
+      CacheMetadata(
+          cache_name="projects/123/locations/us-central1/cachedContents/456",
+          expire_time=time.time() + 1800,
+          fingerprint="abc123",
+          invocations_used=1,
+          contents_count=-1,
+      )
+    assert "greater than or equal to 0" in str(exc_info.value)
+
+  def test_expire_soon_property(self):
+    """Test expire_soon property."""
+    # Cache that expires in 10 minutes (should not expire soon)
+    future_time = time.time() + 600  # 10 minutes
+    metadata = CacheMetadata(
+        cache_name="projects/123/locations/us-central1/cachedContents/456",
+        expire_time=future_time,
+        fingerprint="abc123",
+        invocations_used=1,
+        contents_count=1,
+    )
+    assert not metadata.expire_soon
+
+    # Cache that expires in 1 minute (should expire soon)
+    soon_time = time.time() + 60  # 1 minute
+    metadata = CacheMetadata(
+        cache_name="projects/123/locations/us-central1/cachedContents/456",
+        expire_time=soon_time,
+        fingerprint="abc123",
+        invocations_used=1,
+        contents_count=1,
+    )
+    assert metadata.expire_soon
+
+  def test_str_representation(self):
+    """Test string representation."""
+    current_time = time.time()
+    expire_time = current_time + 1800  # 30 minutes
+
+    metadata = CacheMetadata(
+        cache_name="projects/123/locations/us-central1/cachedContents/test456",
+        expire_time=expire_time,
+        fingerprint="abc123",
+        invocations_used=7,
+        contents_count=4,
+    )
+
+    str_repr = str(metadata)
+    assert "test456" in str_repr  # Cache ID
+    assert "used 7 invocations" in str_repr
+    assert "cached 4 contents" in str_repr
+    assert "expires in" in str_repr
+
+  def test_immutability(self):
+    """Test that CacheMetadata is immutable (frozen)."""
+    metadata = CacheMetadata(
+        cache_name="projects/123/locations/us-central1/cachedContents/456",
+        expire_time=time.time() + 1800,
+        fingerprint="abc123",
+        invocations_used=5,
+        contents_count=3,
+    )
+
+    # Should not be able to modify fields
+    with pytest.raises(ValidationError):
+      metadata.invocations_used = 10
+
+  def test_model_config(self):
+    """Test that model config is set correctly."""
+    metadata = CacheMetadata(
+        cache_name="projects/123/locations/us-central1/cachedContents/456",
+        expire_time=time.time() + 1800,
+        fingerprint="abc123",
+        invocations_used=5,
+        contents_count=3,
+    )
+
+    assert metadata.model_config["extra"] == "forbid"
+    assert metadata.model_config["frozen"] == True
+
+  def test_field_descriptions(self):
+    """Test that fields have proper descriptions."""
+    metadata = CacheMetadata(
+        cache_name="projects/123/locations/us-central1/cachedContents/456",
+        expire_time=time.time() + 1800,
+        fingerprint="abc123",
+        invocations_used=5,
+        contents_count=3,
+    )
+    schema = metadata.model_json_schema()
+
+    assert "invocations_used" in schema["properties"]
+    assert (
+        "Number of invocations"
+        in schema["properties"]["invocations_used"]["description"]
+    )
+
+    assert "contents_count" in schema["properties"]
+    assert (
+        "Number of contents"
+        in schema["properties"]["contents_count"]["description"]
+    )
+
+  def test_realistic_cache_scenarios(self):
+    """Test realistic cache scenarios."""
+    current_time = time.time()
+
+    # Fresh cache
+    fresh_cache = CacheMetadata(
+        cache_name="projects/123/locations/us-central1/cachedContents/fresh123",
+        expire_time=current_time + 1800,
+        fingerprint="fresh_fingerprint",
+        invocations_used=1,
+        contents_count=5,
+        created_at=current_time,
+    )
+    assert fresh_cache.invocations_used == 1
+    assert not fresh_cache.expire_soon
+
+    # Well-used cache
+    used_cache = CacheMetadata(
+        cache_name="projects/123/locations/us-central1/cachedContents/used456",
+        expire_time=current_time + 600,
+        fingerprint="used_fingerprint",
+        invocations_used=8,
+        contents_count=3,
+        created_at=current_time - 1200,
+    )
+    assert used_cache.invocations_used == 8
+
+    # Expiring cache
+    expiring_cache = CacheMetadata(
+        cache_name=(
+            "projects/123/locations/us-central1/cachedContents/expiring789"
+        ),
+        expire_time=current_time + 60,  # 1 minute
+        fingerprint="expiring_fingerprint",
+        invocations_used=15,
+        contents_count=10,
+    )
+    assert expiring_cache.expire_soon
+
+  def test_cache_name_extraction(self):
+    """Test cache name ID extraction in string representation."""
+    metadata = CacheMetadata(
+        cache_name=(
+            "projects/123/locations/us-central1/cachedContents/extracted_id"
+        ),
+        expire_time=time.time() + 1800,
+        fingerprint="abc123",
+        invocations_used=1,
+        contents_count=2,
+    )
+
+    str_repr = str(metadata)
+    assert "extracted_id" in str_repr
+
+  def test_no_performance_metrics(self):
+    """Test that performance metrics are not in CacheMetadata."""
+    metadata = CacheMetadata(
+        cache_name="projects/123/locations/us-central1/cachedContents/456",
+        expire_time=time.time() + 1800,
+        fingerprint="abc123",
+        invocations_used=5,
+        contents_count=3,
+    )
+
+    # Verify that token counts are NOT in CacheMetadata
+    # (they should be in LlmResponse.usage_metadata)
+    assert not hasattr(metadata, "cached_tokens")
+    assert not hasattr(metadata, "total_tokens")
+    assert not hasattr(metadata, "prompt_tokens")
+
+  def test_missing_required_fields(self):
+    """Test validation when truly required fields are missing."""
+    # Only fingerprint and contents_count are required now
+    # Other fields are optional (for fingerprint-only state)
+    required_fields = [
+        "fingerprint",
+        "contents_count",
+    ]
+
+    base_args = {
+        "fingerprint": "abc123",
+        "contents_count": 2,
+    }
+
+    for field in required_fields:
+      args = base_args.copy()
+      del args[field]
+
+      with pytest.raises(ValidationError):
+        CacheMetadata(**args)
+
+    # Test that optional fields can be omitted (fingerprint-only state)
+    metadata = CacheMetadata(
+        fingerprint="abc123",
+        contents_count=5,
+    )
+    assert metadata.cache_name is None
+    assert metadata.expire_time is None
+    assert metadata.invocations_used is None
+    assert metadata.created_at is None
+
+  def test_partial_active_state_rejected(self):
+    """cache_name, expire_time, invocations_used must all be set or all None."""
+    # Only cache_name set.
+    with pytest.raises(ValidationError, match="must all be set"):
+      CacheMetadata(
+          cache_name="projects/123/locations/us-central1/cachedContents/x",
+          fingerprint="abc",
+          contents_count=1,
+      )
+
+    # cache_name + expire_time but no invocations_used.
+    with pytest.raises(ValidationError, match="must all be set"):
+      CacheMetadata(
+          cache_name="projects/123/locations/us-central1/cachedContents/x",
+          expire_time=time.time() + 1800,
+          fingerprint="abc",
+          contents_count=1,
+      )
+
+    # invocations_used set without cache_name (e.g. construction bug).
+    with pytest.raises(ValidationError, match="must all be set"):
+      CacheMetadata(
+          fingerprint="abc",
+          invocations_used=3,
+          contents_count=1,
+      )

tests/unittests/models/test_completions_http_client.py

@@ -0,0 +1,831 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import json
+from unittest import mock
+from unittest.mock import AsyncMock
+
+from google.adk.models.apigee_llm import ChatCompletionsResponseHandler
+from google.adk.models.apigee_llm import CompletionsHTTPClient
+from google.adk.models.llm_request import LlmRequest
+from google.genai import types
+import httpx
+import pytest
+
+
+@pytest.fixture
+def client():
+  return CompletionsHTTPClient(base_url='https://localhost')
+
+
+@pytest.fixture(name='llm_request')
+def fixture_llm_request():
+  return LlmRequest(
+      model='apigee/open_llama',
+      contents=[
+          types.Content(role='user', parts=[types.Part.from_text(text='Hello')])
+      ],
+  )
+
+
+@pytest.mark.asyncio
+async def test_construct_payload_basic_payload(client, llm_request):
+  mock_response = AsyncMock(spec=httpx.Response)
+  mock_response.json.return_value = {
+      'choices': [{'message': {'role': 'assistant', 'content': 'Hi'}}]
+  }
+  mock_response.status_code = 200
+
+  with mock.patch.object(
+      httpx.AsyncClient, 'post', return_value=mock_response
+  ) as mock_post:
+    _ = [
+        r
+        async for r in client.generate_content_async(llm_request, stream=False)
+    ]
+
+    mock_post.assert_called_once()
+    call_args = mock_post.call_args
+    url = call_args[0][0]
+    kwargs = call_args[1]
+
+    assert url == 'https://localhost/chat/completions'
+    payload = kwargs['json']
+    assert payload['model'] == 'open_llama'
+    assert payload['stream'] is False
+    assert len(payload['messages']) == 1
+    assert payload['messages'][0]['role'] == 'user'
+    assert payload['messages'][0]['content'] == 'Hello'
+
+
+@pytest.mark.asyncio
+async def test_construct_payload_with_config(client, llm_request):
+  llm_request.config = types.GenerateContentConfig(
+      temperature=0.7,
+      top_p=0.9,
+      max_output_tokens=100,
+      stop_sequences=['STOP'],
+      frequency_penalty=0.5,
+      presence_penalty=0.5,
+      seed=42,
+      candidate_count=2,
+      response_mime_type='application/json',
+  )
+
+  mock_response = AsyncMock(spec=httpx.Response)
+  mock_response.json.return_value = {
+      'choices': [{'message': {'role': 'assistant', 'content': 'Hi'}}]
+  }
+  mock_response.status_code = 200
+
+  with mock.patch.object(
+      httpx.AsyncClient, 'post', return_value=mock_response
+  ) as mock_post:
+    _ = [
+        r
+        async for r in client.generate_content_async(llm_request, stream=False)
+    ]
+
+    mock_post.assert_called_once()
+    payload = mock_post.call_args[1]['json']
+
+    assert payload['temperature'] == 0.7
+    assert payload['top_p'] == 0.9
+    assert payload['max_tokens'] == 100
+    assert payload['stop'] == ['STOP']
+    assert payload['frequency_penalty'] == 0.5
+    assert payload['presence_penalty'] == 0.5
+    assert payload['seed'] == 42
+    assert payload['n'] == 2
+    assert payload['response_format'] == {'type': 'json_object'}
+
+
+@pytest.mark.asyncio
+async def test_construct_payload_with_tools(client, llm_request):
+  tool = types.Tool(
+      function_declarations=[
+          types.FunctionDeclaration(
+              name='get_weather',
+              description='Get weather',
+              parameters=types.Schema(
+                  type=types.Type.OBJECT,
+                  properties={'location': types.Schema(type=types.Type.STRING)},
+              ),
+          )
+      ]
+  )
+  llm_request.config = types.GenerateContentConfig(tools=[tool])
+
+  mock_response = AsyncMock(spec=httpx.Response)
+  mock_response.json.return_value = {
+      'choices': [{'message': {'role': 'assistant', 'content': 'Hi'}}]
+  }
+  mock_response.status_code = 200
+
+  with mock.patch.object(
+      httpx.AsyncClient, 'post', return_value=mock_response
+  ) as mock_post:
+    _ = [
+        r
+        async for r in client.generate_content_async(llm_request, stream=False)
+    ]
+
+    mock_post.assert_called_once()
+    payload = mock_post.call_args[1]['json']
+    assert 'tools' in payload
+    assert payload['tools'][0]['function']['name'] == 'get_weather'
+
+
+@pytest.mark.asyncio
+async def test_construct_payload_system_instruction(client, llm_request):
+  llm_request.config = types.GenerateContentConfig(
+      system_instruction='You are a helpful assistant.'
+  )
+  mock_response = AsyncMock(spec=httpx.Response)
+  mock_response.json.return_value = {
+      'choices': [{'message': {'role': 'assistant', 'content': 'Hi'}}]
+  }
+  mock_response.status_code = 200
+
+  with mock.patch.object(
+      httpx.AsyncClient, 'post', return_value=mock_response
+  ) as mock_post:
+    _ = [
+        r
+        async for r in client.generate_content_async(llm_request, stream=False)
+    ]
+
+    payload = mock_post.call_args[1]['json']
+    assert payload['messages'][0]['role'] == 'system'
+    assert payload['messages'][0]['content'] == 'You are a helpful assistant.'
+    # Ensure user message follows system
+    assert payload['messages'][1]['role'] == 'user'
+
+
+@pytest.mark.asyncio
+async def test_construct_payload_multimodal_content(client):
+  # Mock inline_data for image
+  image_data = b'fake_image_bytes'
+  llm_request = LlmRequest(
+      model='apigee/open_llama',
+      contents=[
+          types.Content(
+              role='user',
+              parts=[
+                  types.Part.from_text(text='What is this?'),
+                  types.Part.from_bytes(
+                      data=image_data, mime_type='image/jpeg'
+                  ),
+              ],
+          )
+      ],
+  )
+
+  mock_response = AsyncMock(spec=httpx.Response)
+  mock_response.json.return_value = {
+      'choices': [
+          {'message': {'role': 'assistant', 'content': 'It is an image'}}
+      ]
+  }
+
+  mock_response.status_code = 200
+
+  with mock.patch.object(
+      httpx.AsyncClient, 'post', return_value=mock_response
+  ) as mock_post:
+    _ = [
+        r
+        async for r in client.generate_content_async(llm_request, stream=False)
+    ]
+
+    mock_post.assert_called_once()
+    payload = mock_post.call_args[1]['json']
+    assert len(payload['messages']) == 1
+    message = payload['messages'][0]
+    assert message['role'] == 'user'
+    assert isinstance(message['content'], list)
+    assert len(message['content']) == 2
+    assert message['content'][0] == {'type': 'text', 'text': 'What is this?'}
+    assert message['content'][1]['type'] == 'image_url'
+    # Base64 encoding of b'fake_image_bytes' is 'ZmFrZV9pbWFnZV9ieXRlcw=='
+    assert message['content'][1]['image_url']['url'] == (
+        'data:image/jpeg;base64,ZmFrZV9pbWFnZV9ieXRlcw=='
+    )
+
+
+@pytest.mark.asyncio
+async def test_construct_payload_image_file_uri(client):
+  llm_request = LlmRequest(
+      model='apigee/open_llama',
+      contents=[
+          types.Content(
+              role='user',
+              parts=[
+                  types.Part.from_uri(
+                      file_uri='https://localhost/image.jpg',
+                      mime_type='image/jpeg',
+                  )
+              ],
+          )
+      ],
+  )
+
+  mock_response = AsyncMock(spec=httpx.Response)
+  mock_response.json.return_value = {
+      'choices': [
+          {'message': {'role': 'assistant', 'content': 'It is an image'}}
+      ]
+  }
+  mock_response.status_code = 200
+
+  with mock.patch.object(
+      httpx.AsyncClient, 'post', return_value=mock_response
+  ) as mock_post:
+    _ = [
+        r
+        async for r in client.generate_content_async(llm_request, stream=False)
+    ]
+
+    mock_post.assert_called_once()
+    payload = mock_post.call_args[1]['json']
+    assert len(payload['messages']) == 1
+    message = payload['messages'][0]
+    assert message['role'] == 'user'
+    assert isinstance(message['content'], list)
+    assert message['content'][0] == {
+        'type': 'image_url',
+        'image_url': {'url': 'https://localhost/image.jpg'},
+    }
+
+
+@pytest.mark.asyncio
+async def test_generate_content_async_function_call_response(
+    client, llm_request
+):
+  # Mock response with tool call
+  mock_response = AsyncMock(spec=httpx.Response)
+  mock_response.json.return_value = {
+      'choices': [{
+          'message': {
+              'role': 'assistant',
+              'content': None,
+              'tool_calls': [{
+                  'id': 'call_123',
+                  'type': 'function',
+                  'function': {
+                      'name': 'get_weather',
+                      'arguments': '{"location": "London"}',
+                  },
+              }],
+          }
+      }]
+  }
+  mock_response.status_code = 200
+
+  with mock.patch.object(httpx.AsyncClient, 'post', return_value=mock_response):
+    responses = [
+        r
+        async for r in client.generate_content_async(llm_request, stream=False)
+    ]
+
+    assert len(responses) == 1
+    part = responses[0].content.parts[0]
+    assert part.function_call
+    assert part.function_call.name == 'get_weather'
+    assert part.function_call.args == {'location': 'London'}
+    assert part.function_call.id == 'call_123'
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    ('response_json_schema', 'response_mime_type', 'expected_response_format'),
+    [
+        # Case 1: Only response_json_schema is provided
+        (
+            {'type': 'object', 'properties': {'name': {'type': 'string'}}},
+            None,
+            {
+                'type': 'json_schema',
+                'json_schema': {
+                    'type': 'object',
+                    'properties': {'name': {'type': 'string'}},
+                },
+            },
+        ),
+        # Case 2: Both provided, schema takes precedence
+        (
+            {'type': 'object', 'properties': {'name': {'type': 'string'}}},
+            'application/json',
+            {
+                'type': 'json_schema',
+                'json_schema': {
+                    'type': 'object',
+                    'properties': {'name': {'type': 'string'}},
+                },
+            },
+        ),
+        # Case 3: Only response_mime_type is provided
+        (
+            None,
+            'application/json',
+            {'type': 'json_object'},
+        ),
+    ],
+)
+async def test_construct_payload_response_format(
+    client,
+    llm_request,
+    response_json_schema,
+    response_mime_type,
+    expected_response_format,
+):
+  llm_request.config = types.GenerateContentConfig(
+      response_json_schema=response_json_schema,
+      response_mime_type=response_mime_type,
+  )
+  mock_response = AsyncMock(spec=httpx.Response)
+  mock_response.json.return_value = {
+      'choices': [{'message': {'role': 'assistant', 'content': '{}'}}]
+  }
+  mock_response.status_code = 200
+
+  with mock.patch.object(
+      httpx.AsyncClient, 'post', return_value=mock_response
+  ) as mock_post:
+    _ = [
+        r
+        async for r in client.generate_content_async(llm_request, stream=False)
+    ]
+
+    mock_post.assert_called_once()
+    payload = mock_post.call_args[1]['json']
+
+    assert payload['response_format'] == expected_response_format
+
+
+@pytest.mark.asyncio
+async def test_generate_content_async_invalid_tool_call_type_raises_error(
+    client, llm_request
+):
+  # Mock response with invalid tool call type
+  mock_response = AsyncMock(spec=httpx.Response)
+  mock_response.json.return_value = {
+      'choices': [{
+          'message': {
+              'role': 'assistant',
+              'content': None,
+              'tool_calls': [{
+                  'id': 'call_123',
+                  # Invalid type
+                  'type': 'custom',
+                  'custom': {
+                      'name': 'read_string',
+                      'input': 'Hi! The this is a custom tool call!',
+                  },
+              }],
+          }
+      }]
+  }
+  mock_response.status_code = 200
+
+  with mock.patch.object(httpx.AsyncClient, 'post', return_value=mock_response):
+    with pytest.raises(ValueError, match='Unsupported tool_call type: custom'):
+      _ = [
+          r
+          async for r in client.generate_content_async(
+              llm_request, stream=False
+          )
+      ]
+
+
+@pytest.mark.asyncio
+async def test_generate_content_async_function_call_response(
+    client, llm_request
+):
+  # Mock response with deprecated function call
+  mock_response = AsyncMock(spec=httpx.Response)
+  mock_response.json.return_value = {
+      'choices': [{
+          'message': {
+              'role': 'assistant',
+              'content': None,
+              'function_call': {
+                  'name': 'get_weather',
+                  'arguments': '{"location": "London"}',
+              },
+          }
+      }]
+  }
+  mock_response.status_code = 200
+
+  with mock.patch.object(httpx.AsyncClient, 'post', return_value=mock_response):
+    responses = [
+        r
+        async for r in client.generate_content_async(llm_request, stream=False)
+    ]
+
+    assert len(responses) == 1
+    part = responses[0].content.parts[0]
+    assert part.function_call
+    assert part.function_call.name == 'get_weather'
+    assert part.function_call.args == {'location': 'London'}
+    assert part.function_call.id is None
+
+
+@pytest.mark.asyncio
+async def test_generate_content_async_streaming_function_call():
+  local_client = CompletionsHTTPClient(base_url='https://localhost')
+  llm_request = LlmRequest(
+      model='apigee/test',
+      contents=[
+          types.Content(role='user', parts=[types.Part.from_text(text='hi')])
+      ],
+  )
+
+  # Mock chunks simulating split arguments
+  chunk_data_0 = {
+      'id': 'chatcmpl-123',
+      'object': 'chat.completion.chunk',
+      'created': 1234567890,
+      'model': 'gpt-3.5-turbo',
+      'service_tier': 'default',
+      'choices': [{
+          'index': 0,
+          'delta': {
+              'tool_calls': [{
+                  'index': 0,
+                  'id': 'call_123',
+                  'type': 'function',
+                  'function': {'name': 'get_weather', 'arguments': ''},
+              }]
+          },
+          'finish_reason': None,
+      }],
+  }
+  chunk_data_1 = {
+      'id': 'chatcmpl-123',
+      'object': 'chat.completion.chunk',
+      'created': 1234567890,
+      'model': 'gpt-3.5-turbo',
+      'service_tier': 'default',
+      'choices': [{
+          'index': 0,
+          'delta': {
+              'tool_calls': [{
+                  'index': 0,
+                  'function': {'arguments': '{"location": "London"}'},
+              }]
+          },
+          'finish_reason': None,
+      }],
+  }
+  chunk_data_2 = {
+      'id': 'chatcmpl-123',
+      'object': 'chat.completion.chunk',
+      'created': 1234567890,
+      'model': 'gpt-3.5-turbo',
+      'service_tier': 'default',
+      'choices': [{
+          'index': 0,
+          'delta': {
+              'tool_calls': [{
+                  'index': 0,
+                  'function': {'arguments': '{"country": "UK"}'},
+              }]
+          },
+          'finish_reason': None,
+      }],
+  }
+  chunk_data_3 = {
+      'id': 'chatcmpl-123',
+      'object': 'chat.completion.chunk',
+      'created': 1234567890,
+      'model': 'gpt-3.5-turbo',
+      'service_tier': 'default',
+      'choices': [{'index': 0, 'delta': {}, 'finish_reason': 'tool_calls'}],
+      'usage': {
+          'prompt_tokens': 10,
+          'completion_tokens': 20,
+          'total_tokens': 30,
+      },
+  }
+
+  chunks = [
+      f'{json.dumps(chunk_data_0)}\n',
+      f'{json.dumps(chunk_data_1)}\n',
+      f'{json.dumps(chunk_data_2)}\n',
+      f'{json.dumps(chunk_data_3)}\n',
+  ]
+
+  async def mock_aiter_lines():
+    for chunk in chunks:
+      yield chunk
+
+  mock_response = AsyncMock(spec=httpx.Response)
+  mock_response.aiter_lines.return_value = mock_aiter_lines()
+  mock_response.status_code = 200
+
+  mock_stream_ctx = mock.AsyncMock()
+  mock_stream_ctx.__aenter__.return_value = mock_response
+
+  with mock.patch.object(
+      httpx.AsyncClient, 'stream', return_value=mock_stream_ctx
+  ):
+    responses = [
+        r
+        async for r in local_client.generate_content_async(
+            llm_request, stream=True
+        )
+    ]
+    # Check that we get 5 responses (one per chunk + extra final accumulated)
+    assert len(responses) == 5
+
+    # Check 1st response: partial tool call, empty args
+    assert responses[0].partial is True
+    assert responses[0].content.parts[0].function_call.name == 'get_weather'
+    assert responses[0].content.parts[0].function_call.id == 'call_123'
+
+    # Check 2nd response: full args for first update
+    assert responses[1].partial is True
+    assert responses[1].content.parts[0].function_call.args == {
+        'location': 'London'
+    }
+
+    # Check 3rd response: full args for second update (merged)
+    assert responses[2].partial is True
+    assert responses[2].content.parts[0].function_call.args == {'country': 'UK'}
+
+    # Check 4th response: last delta (empty)
+    assert responses[3].partial is True
+    assert responses[3].content.parts == []
+
+    # Check 5th response: final accumulated
+    assert responses[4].finish_reason == types.FinishReason.STOP
+    # Full accumulated args
+    assert responses[4].content.parts[0].function_call.args == {
+        'location': 'London',
+        'country': 'UK',
+    }
+
+    # Check metadata and usage
+    assert responses[4].model_version == 'gpt-3.5-turbo'
+    assert responses[4].custom_metadata['id'] == 'chatcmpl-123'
+    assert responses[4].custom_metadata['created'], 1234567890
+    assert responses[4].custom_metadata['object'], 'chat.completion.chunk'
+    assert responses[4].custom_metadata['service_tier'], 'default'
+    assert responses[4].usage_metadata is not None
+    assert responses[4].usage_metadata.prompt_token_count == 10
+    assert responses[4].usage_metadata.candidates_token_count == 20
+    assert responses[4].usage_metadata.total_token_count == 30
+
+
+@pytest.mark.asyncio
+async def test_generate_content_async_streaming_multiple_function_calls():
+  # Mock streaming response with multiple tool calls
+  local_client = CompletionsHTTPClient(base_url='https://localhost')
+  llm_request = LlmRequest(
+      model='apigee/test',
+      contents=[
+          types.Content(role='user', parts=[types.Part.from_text(text='hi')])
+      ],
+  )
+  chunk_data_1 = {
+      'choices': [{
+          'index': 0,
+          'delta': {
+              'tool_calls': [
+                  {
+                      'index': 0,
+                      'id': 'call_1',
+                      'type': 'function',
+                      'function': {'name': 'func_1', 'arguments': ''},
+                  },
+                  {
+                      'index': 1,
+                      'id': 'call_2',
+                      'type': 'function',
+                      'function': {'name': 'func_2', 'arguments': ''},
+                  },
+              ]
+          },
+          'finish_reason': None,
+      }]
+  }
+  # the tool_call type is optional in chunk responses.
+  chunk_data_2 = {
+      'choices': [{
+          'index': 0,
+          'delta': {
+              'tool_calls': [
+                  {'index': 0, 'function': {'arguments': '{"arg": 1}'}},
+                  {'index': 1, 'function': {'arguments': '{"arg": 2}'}},
+              ]
+          },
+          'finish_reason': None,
+      }]
+  }
+  chunk_data_3 = {
+      'choices': [{'index': 0, 'delta': {}, 'finish_reason': 'tool_calls'}]
+  }
+
+  chunks = [
+      f'{json.dumps(chunk_data_1)}\n',
+      f'{json.dumps(chunk_data_2)}\n',
+      f'{json.dumps(chunk_data_3)}\n',
+  ]
+
+  async def mock_aiter_lines():
+    for chunk in chunks:
+      yield chunk
+
+  mock_response = AsyncMock(spec=httpx.Response)
+  mock_response.aiter_lines.return_value = mock_aiter_lines()
+  mock_response.status_code = 200
+
+  mock_stream_ctx = mock.AsyncMock()
+  mock_stream_ctx.__aenter__.return_value = mock_response
+
+  with mock.patch.object(
+      httpx.AsyncClient, 'stream', return_value=mock_stream_ctx
+  ):
+    responses = [
+        r
+        async for r in local_client.generate_content_async(
+            llm_request, stream=True
+        )
+    ]
+
+    assert len(responses) == 4
+    parts = responses[-1].content.parts
+    assert len(parts) == 2
+
+    assert parts[0].function_call.name == 'func_1'
+    assert parts[0].function_call.args == {'arg': 1}
+    assert parts[0].function_call.id == 'call_1'
+
+    assert parts[1].function_call.name == 'func_2'
+    assert parts[1].function_call.args == {'arg': 2}
+
+    assert parts[1].function_call.id == 'call_2'
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    ('chunks', 'expected_response_count'),
+    [
+        (
+            [
+                '\n',
+                '   \n',
+                (
+                    'data: {"choices": [{"index": 0, "delta": {"content":'
+                    ' "Hello"}, "finish_reason": null}]}\n'
+                ),
+            ],
+            1,
+        ),
+        (
+            [
+                (
+                    'data: {"choices": [{"index": 0, "delta": {"content":'
+                    ' "Hello"}, "finish_reason": null}]}\n'
+                ),
+                '[DONE]\n',
+                (
+                    'data: {"choices": [{"index": 0, "delta": {"content":'
+                    ' "World"}, "finish_reason": "stop"}]}\n'
+                ),
+            ],
+            1,  # Should stop after [DONE]
+        ),
+        (
+            [
+                (
+                    'data: {"choices": [{"index": 0, "delta": {"content":'
+                    ' "Hello"}, "finish_reason": null}]}\n'
+                ),
+                '   [DONE]   \n',
+                (
+                    'data: {"choices": [{"index": 0, "delta": {"content":'
+                    ' "World"}, "finish_reason": "stop"}]}\n'
+                ),
+            ],
+            1,  # Should stop after [DONE]
+        ),
+        (
+            [
+                (
+                    'data: {"choices": [{"index": 0, "delta": {"content":'
+                    ' "Hello"}, "finish_reason": null}]}\n'
+                ),
+                'data: [DONE]\n',
+                (
+                    'data: {"choices": [{"index": 0, "delta": {"content":'
+                    ' "World"}, "finish_reason": "stop"}]}\n'
+                ),
+            ],
+            1,  # Should stop after [DONE]
+        ),
+    ],
+)
+async def test_generate_content_async_streaming_parse_lines(
+    chunks, expected_response_count
+):
+  local_client = CompletionsHTTPClient(base_url='https://localhost')
+  llm_request = LlmRequest(
+      model='apigee/test',
+      contents=[
+          types.Content(role='user', parts=[types.Part.from_text(text='hi')])
+      ],
+  )
+
+  async def mock_aiter_lines():
+    for chunk in chunks:
+      yield chunk
+
+  mock_response = AsyncMock(spec=httpx.Response)
+  mock_response.aiter_lines.return_value = mock_aiter_lines()
+  mock_response.status_code = 200
+
+  mock_stream_ctx = mock.AsyncMock()
+  mock_stream_ctx.__aenter__.return_value = mock_response
+
+  with mock.patch.object(
+      httpx.AsyncClient, 'stream', return_value=mock_stream_ctx
+  ):
+    responses = [
+        r
+        async for r in local_client.generate_content_async(
+            llm_request, stream=True
+        )
+    ]
+    assert len(responses) == expected_response_count
+    assert responses[0].content.parts[0].text == 'Hello'
+
+
+def test_process_chunk_with_refusal_streaming():
+  handler = ChatCompletionsResponseHandler()
+
+  chunk1 = {
+      'choices': [{
+          'delta': {
+              'role': 'assistant',
+              'content': 'Hello',
+          },
+          'index': 0,
+      }]
+  }
+  responses1 = list(handler.process_chunk(chunk1))
+  assert len(responses1) == 1
+  assert responses1[0].content.parts[0].text == 'Hello'
+
+  chunk2 = {
+      'choices': [{
+          'delta': {
+              'refusal': 'I refuse',
+          },
+          'index': 0,
+      }]
+  }
+  responses2 = list(handler.process_chunk(chunk2))
+  assert len(responses2) == 1
+  assert responses2[0].content.parts[0].text == '\n[[REFUSAL]]: I refuse'
+
+  chunk3 = {
+      'choices': [{
+          'delta': {
+              'refusal': ' to answer',
+          },
+          'index': 0,
+      }]
+  }
+  responses3 = list(handler.process_chunk(chunk3))
+  assert len(responses3) == 1
+  assert responses3[0].content.parts[0].text == ' to answer'
+
+  chunk4 = {
+      'choices': [{
+          'delta': {},
+          'finish_reason': 'stop',
+          'index': 0,
+      }]
+  }
+  responses4 = list(handler.process_chunk(chunk4))
+  assert len(responses4) == 2
+  final_response = responses4[1]
+  assert final_response.finish_reason == types.FinishReason.STOP
+  assert (
+      final_response.content.parts[0].text
+      == 'Hello\n[[REFUSAL]]: I refuse to answer'
+  )

tests/unittests/models/test_gemma_llm.py

@@ -0,0 +1,558 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from google.adk import models
+from google.adk.models.gemma_llm import Gemma
+from google.adk.models.google_llm import Gemini
+from google.adk.models.llm_request import LlmRequest
+from google.adk.models.llm_response import LlmResponse
+from google.genai import types
+from google.genai.types import Content
+from google.genai.types import Part
+import pytest
+
+
+@pytest.fixture
+def llm_request():
+  return LlmRequest(
+      model="gemma-3-4b-it",
+      contents=[Content(role="user", parts=[Part.from_text(text="Hello")])],
+      config=types.GenerateContentConfig(
+          temperature=0.1,
+          response_modalities=[types.Modality.TEXT],
+          system_instruction="You are a helpful assistant",
+      ),
+  )
+
+
+@pytest.fixture
+def llm_request_with_duplicate_instruction():
+  return LlmRequest(
+      model="gemma-3-1b-it",
+      contents=[
+          Content(
+              role="user",
+              parts=[Part.from_text(text="Talk like a pirate.")],
+          ),
+          Content(role="user", parts=[Part.from_text(text="Hello")]),
+      ],
+      config=types.GenerateContentConfig(
+          response_modalities=[types.Modality.TEXT],
+          system_instruction="Talk like a pirate.",
+      ),
+  )
+
+
+@pytest.fixture
+def llm_request_with_tools():
+  return LlmRequest(
+      model="gemma-3-1b-it",
+      contents=[Content(role="user", parts=[Part.from_text(text="Hello")])],
+      config=types.GenerateContentConfig(
+          tools=[
+              types.Tool(
+                  function_declarations=[
+                      types.FunctionDeclaration(
+                          name="search_web",
+                          description="Search the web for a query.",
+                          parameters=types.Schema(
+                              type=types.Type.OBJECT,
+                              properties={
+                                  "query": types.Schema(type=types.Type.STRING)
+                              },
+                              required=["query"],
+                          ),
+                      ),
+                      types.FunctionDeclaration(
+                          name="get_current_time",
+                          description="Gets the current time.",
+                          parameters=types.Schema(
+                              type=types.Type.OBJECT, properties={}
+                          ),
+                      ),
+                  ]
+              )
+          ],
+      ),
+  )
+
+
+def test_supported_models_matches_gemma4():
+  """Gemma 4 model strings must resolve to the Gemini class via the registry."""
+  assert models.LLMRegistry.resolve("gemma-4-31b-it") is Gemini
+
+
+def test_supported_models_matches_gemma3():
+  """Gemma 3 model strings must continue to resolve to the Gemma class."""
+  assert models.LLMRegistry.resolve("gemma-3-27b-it") is Gemma
+
+
+@pytest.mark.asyncio
+async def test_not_gemma_model():
+  llm = Gemma()
+  llm_request_bad_model = LlmRequest(
+      model="not-a-gemma-model",
+  )
+  with pytest.raises(AssertionError, match=r".*model.*"):
+    async for _ in llm.generate_content_async(llm_request_bad_model):
+      pass
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    "llm_request",
+    ["llm_request", "llm_request_with_duplicate_instruction"],
+    indirect=True,
+)
+async def test_preprocess_request(llm_request):
+  llm = Gemma()
+  want_content_text = llm_request.config.system_instruction
+
+  await llm._preprocess_request(llm_request)
+
+  # system instruction should be cleared
+  assert not llm_request.config.system_instruction
+  # should be two content bits now (deduped, if needed)
+  assert len(llm_request.contents) == 2
+  # first message in contents should be "user": <original sys instruction>
+  assert llm_request.contents[0].role == "user"
+  assert llm_request.contents[0].parts[0].text == want_content_text
+
+
+@pytest.mark.asyncio
+async def test_preprocess_request_with_tools(llm_request_with_tools):
+
+  gemma = Gemma()
+  await gemma._preprocess_request(llm_request_with_tools)
+
+  assert not llm_request_with_tools.config.tools
+
+  # The original user content should now be the second item
+  assert llm_request_with_tools.contents[1].role == "user"
+  assert llm_request_with_tools.contents[1].parts[0].text == "Hello"
+
+  sys_instruct_text = llm_request_with_tools.contents[0].parts[0].text
+  assert sys_instruct_text is not None
+  assert "You have access to the following functions" in sys_instruct_text
+  assert (
+      """{"description":"Search the web for a query.","name":"search_web","""
+      in sys_instruct_text
+  )
+  assert (
+      """{"description":"Gets the current time.","name":"get_current_time","parameters":{"properties":{}"""
+      in sys_instruct_text
+  )
+
+
+@pytest.mark.asyncio
+async def test_preprocess_request_with_function_response():
+  # Simulate an LlmRequest with a function response
+  func_response_data = types.FunctionResponse(
+      name="search_web", response={"results": [{"title": "ADK"}]}
+  )
+  llm_request = LlmRequest(
+      model="gemma-3-1b-it",
+      contents=[
+          types.Content(
+              role="model",
+              parts=[types.Part(function_response=func_response_data)],
+          )
+      ],
+      config=types.GenerateContentConfig(),
+  )
+
+  gemma = Gemma()
+  await gemma._preprocess_request(llm_request)
+
+  # Assertions: function response converted to user role text content
+  assert llm_request.contents
+  assert len(llm_request.contents) == 1
+  assert llm_request.contents[0].role == "user"
+  assert llm_request.contents[0].parts
+  assert (
+      llm_request.contents[0].parts[0].text
+      == 'Invoking tool `search_web` produced: `{"results": [{"title":'
+      ' "ADK"}]}`.'
+  )
+  assert llm_request.contents[0].parts[0].function_response is None
+  assert llm_request.contents[0].parts[0].function_call is None
+
+
+@pytest.mark.asyncio
+async def test_preprocess_request_with_function_call():
+  func_call_data = types.FunctionCall(name="get_current_time", args={})
+  llm_request = LlmRequest(
+      model="gemma-3-1b-it",
+      contents=[
+          types.Content(
+              role="user", parts=[types.Part(function_call=func_call_data)]
+          )
+      ],
+  )
+
+  gemma = Gemma()
+  await gemma._preprocess_request(llm_request)
+
+  assert len(llm_request.contents) == 1
+  assert llm_request.contents[0].role == "model"
+  expected_text = func_call_data.model_dump_json(exclude_none=True)
+  assert llm_request.contents[0].parts
+  got_part = llm_request.contents[0].parts[0]
+  assert got_part.text == expected_text
+  assert got_part.function_call is None
+  assert got_part.function_response is None
+
+
+@pytest.mark.asyncio
+async def test_preprocess_request_with_mixed_content():
+  func_call = types.FunctionCall(name="get_weather", args={"city": "London"})
+  func_response = types.FunctionResponse(
+      name="get_weather", response={"temp": "15C"}
+  )
+
+  llm_request = LlmRequest(
+      model="gemma-3-1b-it",
+      contents=[
+          types.Content(
+              role="user", parts=[types.Part.from_text(text="Hello!")]
+          ),
+          types.Content(
+              role="model", parts=[types.Part(function_call=func_call)]
+          ),
+          types.Content(
+              role="some_function",
+              parts=[types.Part(function_response=func_response)],
+          ),
+          types.Content(
+              role="user", parts=[types.Part.from_text(text="How are you?")]
+          ),
+      ],
+  )
+
+  gemma = Gemma()
+  await gemma._preprocess_request(llm_request)
+
+  # Assertions
+  assert len(llm_request.contents) == 4
+
+  # First part: original user text
+  assert llm_request.contents[0].role == "user"
+  assert llm_request.contents[0].parts
+  assert llm_request.contents[0].parts[0].text == "Hello!"
+
+  # Second part: function call converted to model text
+  assert llm_request.contents[1].role == "model"
+  assert llm_request.contents[1].parts
+  assert llm_request.contents[1].parts[0].text == func_call.model_dump_json(
+      exclude_none=True
+  )
+
+  # Third part: function response converted to user text
+  assert llm_request.contents[2].role == "user"
+  assert llm_request.contents[2].parts
+  assert (
+      llm_request.contents[2].parts[0].text
+      == 'Invoking tool `get_weather` produced: `{"temp": "15C"}`.'
+  )
+
+  # Fourth part: original user text
+  assert llm_request.contents[3].role == "user"
+  assert llm_request.contents[3].parts
+  assert llm_request.contents[3].parts[0].text == "How are you?"
+
+
+def test_process_response():
+  # Simulate a response from Gemma that should be converted to a FunctionCall
+  json_function_call_str = (
+      '{"name": "search_web", "parameters": {"query": "latest news"}}'
+  )
+  llm_response = LlmResponse(
+      content=Content(
+          role="model", parts=[Part.from_text(text=json_function_call_str)]
+      )
+  )
+
+  gemma = Gemma()
+  gemma._extract_function_calls_from_response(llm_response=llm_response)
+
+  # Assert that the content was transformed into a FunctionCall
+  assert llm_response.content
+  assert llm_response.content.parts
+  assert len(llm_response.content.parts) == 1
+  part = llm_response.content.parts[0]
+  assert part.function_call is not None
+  assert part.function_call.name == "search_web"
+  assert part.function_call.args == {"query": "latest news"}
+  # Assert that the entire part matches the expected structure
+  expected_function_call = types.FunctionCall(
+      name="search_web", args={"query": "latest news"}
+  )
+  expected_part = Part(function_call=expected_function_call)
+  assert part == expected_part
+  assert part.text is None  # Ensure text part is cleared
+
+
+def test_process_response_invalid_json_text():
+  # Simulate a response with plain text that is not JSON
+  original_text = "This is a regular text response."
+  llm_response = LlmResponse(
+      content=Content(role="model", parts=[Part.from_text(text=original_text)])
+  )
+
+  gemma = Gemma()
+  gemma._extract_function_calls_from_response(llm_response=llm_response)
+
+  # Assert that the content remains unchanged
+  assert llm_response.content
+  assert llm_response.content.parts
+  assert len(llm_response.content.parts) == 1
+  assert llm_response.content.parts[0].text == original_text
+  assert llm_response.content.parts[0].function_call is None
+
+
+def test_process_response_malformed_json():
+  # Simulate a response with valid JSON but not in the function call format
+  malformed_json_str = '{"not_a_function": "value", "another_field": 123}'
+  llm_response = LlmResponse(
+      content=Content(
+          role="model", parts=[Part.from_text(text=malformed_json_str)]
+      )
+  )
+  gemma = Gemma()
+  gemma._extract_function_calls_from_response(llm_response=llm_response)
+
+  # Assert that the content remains unchanged because it doesn't match the expected schema
+  assert llm_response.content
+  assert llm_response.content.parts
+  assert len(llm_response.content.parts) == 1
+  assert llm_response.content.parts[0].text == malformed_json_str
+  assert llm_response.content.parts[0].function_call is None
+
+
+def test_process_response_empty_content_or_multiple_parts():
+  gemma = Gemma()
+
+  # Test case 1: LlmResponse with no content
+  llm_response_no_content = LlmResponse(content=None)
+  gemma._extract_function_calls_from_response(
+      llm_response=llm_response_no_content
+  )
+  assert llm_response_no_content.content is None
+
+  # Test case 2: LlmResponse with empty parts list
+  llm_response_empty_parts = LlmResponse(
+      content=Content(role="model", parts=[])
+  )
+  gemma._extract_function_calls_from_response(
+      llm_response=llm_response_empty_parts
+  )
+  assert llm_response_empty_parts.content
+  assert not llm_response_empty_parts.content.parts
+
+  # Test case 3: LlmResponse with multiple parts
+  llm_response_multiple_parts = LlmResponse(
+      content=Content(
+          role="model",
+          parts=[
+              Part.from_text(text="part one"),
+              Part.from_text(text="part two"),
+          ],
+      )
+  )
+  original_parts = list(
+      llm_response_multiple_parts.content.parts
+  )  # Copy for comparison
+  gemma._extract_function_calls_from_response(
+      llm_response=llm_response_multiple_parts
+  )
+  assert llm_response_multiple_parts.content
+  assert (
+      llm_response_multiple_parts.content.parts == original_parts
+  )  # Should remain unchanged
+
+  # Test case 4: LlmResponse with one part, but empty text
+  llm_response_empty_text_part = LlmResponse(
+      content=Content(role="model", parts=[Part.from_text(text="")])
+  )
+  gemma._extract_function_calls_from_response(
+      llm_response=llm_response_empty_text_part
+  )
+  assert llm_response_empty_text_part.content
+  assert llm_response_empty_text_part.content.parts
+  assert llm_response_empty_text_part.content.parts[0].text == ""
+  assert llm_response_empty_text_part.content.parts[0].function_call is None
+
+
+def test_process_response_with_markdown_json_block():
+  # Simulate a response from Gemma with a JSON function call in a markdown block
+  json_function_call_str = """
+```json
+{"name": "search_web", "parameters": {"query": "latest news"}}
+```"""
+  llm_response = LlmResponse(
+      content=Content(
+          role="model", parts=[Part.from_text(text=json_function_call_str)]
+      )
+  )
+
+  gemma = Gemma()
+  gemma._extract_function_calls_from_response(llm_response)
+
+  assert llm_response.content
+  assert llm_response.content.parts
+  assert len(llm_response.content.parts) == 1
+  part = llm_response.content.parts[0]
+  assert part.function_call is not None
+  assert part.function_call.name == "search_web"
+  assert part.function_call.args == {"query": "latest news"}
+  assert part.text is None
+
+
+def test_process_response_with_markdown_tool_code_block():
+  # Simulate a response from Gemma with a JSON function call in a 'tool_code' markdown block
+  json_function_call_str = """
+Some text before.
+```tool_code
+{"name": "get_current_time", "parameters": {}}
+```
+And some text after."""
+  llm_response = LlmResponse(
+      content=Content(
+          role="model", parts=[Part.from_text(text=json_function_call_str)]
+      )
+  )
+
+  gemma = Gemma()
+  gemma._extract_function_calls_from_response(llm_response)
+
+  assert llm_response.content
+  assert llm_response.content.parts
+  assert len(llm_response.content.parts) == 1
+  part = llm_response.content.parts[0]
+  assert part.function_call is not None
+  assert part.function_call.name == "get_current_time"
+  assert part.function_call.args == {}
+  assert part.text is None
+
+
+def test_process_response_with_embedded_json():
+  # Simulate a response with valid JSON embedded in text
+  embedded_json_str = (
+      'Please call the tool: {"name": "search_web", "parameters": {"query":'
+      ' "new features"}} thanks!'
+  )
+  llm_response = LlmResponse(
+      content=Content(
+          role="model", parts=[Part.from_text(text=embedded_json_str)]
+      )
+  )
+
+  gemma = Gemma()
+  gemma._extract_function_calls_from_response(llm_response)
+
+  assert llm_response.content
+  assert llm_response.content.parts
+  assert len(llm_response.content.parts) == 1
+  part = llm_response.content.parts[0]
+  assert part.function_call is not None
+  assert part.function_call.name == "search_web"
+  assert part.function_call.args == {"query": "new features"}
+  assert part.text is None
+
+
+def test_process_response_flexible_parsing():
+  # Test with "function" and "args" keys as supported by GemmaFunctionCallModel
+  flexible_json_str = '{"function": "do_something", "args": {"value": 123}}'
+  llm_response = LlmResponse(
+      content=Content(
+          role="model", parts=[Part.from_text(text=flexible_json_str)]
+      )
+  )
+
+  gemma = Gemma()
+  gemma._extract_function_calls_from_response(llm_response)
+
+  assert llm_response.content
+  assert llm_response.content.parts
+  assert len(llm_response.content.parts) == 1
+  part = llm_response.content.parts[0]
+  assert part.function_call is not None
+  assert part.function_call.name == "do_something"
+  assert part.function_call.args == {"value": 123}
+  assert part.text is None
+
+
+def test_process_response_last_json_object():
+  # Simulate a response with multiple JSON objects, ensuring the last valid one is picked
+  multiple_json_str = (
+      'I thought about {"name": "first_call", "parameters": {"a": 1}} but then'
+      ' decided to call: {"name": "second_call", "parameters": {"b": 2}}'
+  )
+  llm_response = LlmResponse(
+      content=Content(
+          role="model", parts=[Part.from_text(text=multiple_json_str)]
+      )
+  )
+
+  gemma = Gemma()
+  gemma._extract_function_calls_from_response(llm_response)
+
+  assert llm_response.content
+  assert llm_response.content.parts
+  assert len(llm_response.content.parts) == 1
+  part = llm_response.content.parts[0]
+  assert part.function_call is not None
+  assert part.function_call.name == "second_call"
+  assert part.function_call.args == {"b": 2}
+  assert part.text is None
+
+
+# Tests for Gemma 4 registry routing
+def test_gemma4_resolves_to_gemini_not_gemma():
+  """Gemma 4 models should resolve to Gemini, not the Gemma workaround class."""
+  resolved = models.LLMRegistry.resolve("gemma-4-31b-it")
+  assert resolved is not Gemma
+  assert resolved is Gemini
+
+
+# Tests for Gemma3Ollama (only run when LiteLLM is installed)
+try:
+  from google.adk.models.gemma_llm import Gemma3Ollama
+  from google.adk.models.lite_llm import LiteLlm
+
+  def test_gemma3_ollama_supported_models():
+    assert Gemma3Ollama.supported_models() == [r"ollama/gemma3.*"]
+
+  def test_gemma3_ollama_registry_resolution():
+    assert models.LLMRegistry.resolve("ollama/gemma3:12b") is Gemma3Ollama
+
+  def test_non_gemma_ollama_registry_resolution():
+    assert models.LLMRegistry.resolve("ollama/llama3.2") is LiteLlm
+
+  @pytest.mark.parametrize(
+      "model_arg,expected_model",
+      [
+          (None, "ollama/gemma3:12b"),
+          ("ollama/gemma3:27b", "ollama/gemma3:27b"),
+      ],
+  )
+  def test_gemma3_ollama_model(model_arg, expected_model):
+    model = (
+        Gemma3Ollama() if model_arg is None else Gemma3Ollama(model=model_arg)
+    )
+    assert model.model == expected_model
+
+except ImportError:
+  # LiteLLM not installed, skip Gemma3Ollama tests
+  pass

tests/unittests/models/test_lite_llm_gemma_tool_role.py

@@ -0,0 +1,177 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Tests for Gemma-specific tool role handling in _content_to_message_param.
+
+Gemma's chat template expects role='tool_responses' for tool result messages,
+while the OpenAI-compatible default is role='tool'. This module verifies that
+_content_to_message_param sets the correct role based on the model name.
+"""
+
+from typing import Any
+
+from google.adk.models.lite_llm import _content_to_message_param
+from google.genai import types
+import pytest
+
+
+def _make_function_response_content(
+    function_name: str = "get_weather",
+    response_data: dict[str, Any] | None = None,
+    call_id: str = "call_001",
+) -> types.Content:
+  """Builds a types.Content with a single function_response part."""
+  if response_data is None:
+    response_data = {"city": "Santiago de Cuba", "condition": "sunny"}
+  return types.Content(
+      role="user",
+      parts=[
+          types.Part(
+              function_response=types.FunctionResponse(
+                  name=function_name,
+                  response=response_data,
+                  id=call_id,
+              )
+          )
+      ],
+  )
+
+
+def _make_multi_function_response_content(
+    call_ids: list[str] | None = None,
+) -> types.Content:
+  """Builds a types.Content with multiple function_response parts."""
+  if call_ids is None:
+    call_ids = ["call_001", "call_002"]
+  return types.Content(
+      role="user",
+      parts=[
+          types.Part(
+              function_response=types.FunctionResponse(
+                  name=f"tool_{i}",
+                  response={"result": f"value_{i}"},
+                  id=call_id,
+              )
+          )
+          for i, call_id in enumerate(call_ids)
+      ],
+  )
+
+
+def _extract_role(msg) -> str:
+  """Extracts role from a litellm message, whether dict or object."""
+  if isinstance(msg, dict):
+    return msg["role"]
+  return msg.role
+
+
+class TestToolRoleSingleResponse:
+  """_content_to_message_param with a single function_response part."""
+
+  @pytest.mark.asyncio
+  async def test_gemma4_model_uses_tool_responses_role(self):
+    """Models containing 'gemma4' should get role='tool_responses'."""
+    content = _make_function_response_content()
+
+    result = await _content_to_message_param(content, model="ollama/gemma4:e2b")
+
+    assert _extract_role(result) == "tool_responses", (
+        "Gemma models require role='tool_responses' to match their chat "
+        "template; role='tool' causes infinite tool-calling loops."
+    )
+
+  @pytest.mark.asyncio
+  async def test_gemma4_uppercase_model_name(self):
+    """Model name matching should be case-insensitive."""
+    content = _make_function_response_content()
+
+    result = await _content_to_message_param(content, model="ollama/Gemma4:31b")
+
+    assert _extract_role(result) == "tool_responses"
+
+  @pytest.mark.asyncio
+  async def test_tool_call_id_and_content_preserved(self):
+    """Fix must not alter tool_call_id or content — only role changes."""
+    content = _make_function_response_content(
+        response_data={"status": "ok"}, call_id="my_call_123"
+    )
+
+    result = await _content_to_message_param(content, model="ollama/gemma4:e2b")
+
+    if isinstance(result, dict):
+      assert result["tool_call_id"] == "my_call_123"
+      assert "ok" in result["content"]
+    else:
+      assert result.tool_call_id == "my_call_123"
+      assert "ok" in result.content
+
+  @pytest.mark.asyncio
+  async def test_empty_model_string_uses_tool_role(self):
+    """Empty model string should fall back to default role='tool'."""
+    content = _make_function_response_content()
+
+    result = await _content_to_message_param(content, model="")
+
+    assert _extract_role(result) == "tool"
+
+  @pytest.mark.asyncio
+  async def test_unrelated_models_use_tool_role(self):
+    """Models that do not contain 'gemma4' must not be affected."""
+    unaffected_models = [
+        "ollama/llama3:8b",
+        "ollama/qwen2.5-coder:3b",
+        "anthropic/claude-3-opus",
+        "openai/gpt-4o",
+        "ollama/gemma3:4b",  # gemma3 != gemma4
+    ]
+    for model in unaffected_models:
+      content = _make_function_response_content()
+      result = await _content_to_message_param(content, model=model)
+      assert (
+          _extract_role(result) == "tool"
+      ), f"Model '{model}' should not be affected by the Gemma4 fix."
+
+
+class TestToolRoleMultipleResponses:
+  """_content_to_message_param with multiple function_response parts."""
+
+  @pytest.mark.asyncio
+  async def test_gemma4_all_messages_use_tool_responses_role(self):
+    """All messages in a multi-response must have role='tool_responses'."""
+    content = _make_multi_function_response_content(
+        call_ids=["call_a", "call_b", "call_c"]
+    )
+
+    result = await _content_to_message_param(content, model="ollama/gemma4:4b")
+
+    assert isinstance(result, list)
+    assert len(result) == 3
+    for msg in result:
+      assert _extract_role(msg) == "tool_responses", (
+          "Every tool message in a multi-response must use 'tool_responses' "
+          "for Gemma4 models."
+      )
+
+  @pytest.mark.asyncio
+  async def test_non_gemma_multi_response_uses_tool_role(self):
+    """Non-Gemma multi-response messages should all have role='tool'."""
+    content = _make_multi_function_response_content(
+        call_ids=["call_a", "call_b"]
+    )
+
+    result = await _content_to_message_param(content, model="openai/gpt-4o")
+
+    assert isinstance(result, list)
+    for msg in result:
+      assert _extract_role(msg) == "tool"

tests/unittests/models/test_litellm_import.py

@@ -0,0 +1,108 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import importlib.util
+import os
+import subprocess
+import sys
+
+import pytest
+
+
+def _subprocess_env() -> dict[str, str]:
+  env = dict(os.environ)
+  src_path = os.path.join(os.getcwd(), "src")
+  pythonpath = env.get("PYTHONPATH", "")
+  env["PYTHONPATH"] = (
+      f"{src_path}{os.pathsep}{pythonpath}" if pythonpath else src_path
+  )
+  return env
+
+
+def test_importing_models_does_not_import_litellm_or_set_mode():
+  env = _subprocess_env()
+  env.pop("LITELLM_MODE", None)
+
+  result = subprocess.run(
+      [
+          sys.executable,
+          "-c",
+          (
+              "import os, sys\n"
+              "import google.adk.models\n"
+              "print('litellm' in sys.modules)\n"
+              "print(os.environ.get('LITELLM_MODE'))\n"
+          ),
+      ],
+      check=True,
+      capture_output=True,
+      text=True,
+      env=env,
+  )
+  stdout_lines = result.stdout.strip().splitlines()
+  assert stdout_lines == ["False", "None"]
+
+
+def test_ensure_litellm_imported_defaults_to_production():
+  if importlib.util.find_spec("litellm") is None:
+    pytest.skip("litellm is not installed")
+
+  env = _subprocess_env()
+  env.pop("LITELLM_MODE", None)
+
+  result = subprocess.run(
+      [
+          sys.executable,
+          "-c",
+          (
+              "import os\n"
+              "from google.adk.models.lite_llm import"
+              " _ensure_litellm_imported\n"
+              "_ensure_litellm_imported()\n"
+              "print(os.environ.get('LITELLM_MODE'))\n"
+          ),
+      ],
+      check=True,
+      capture_output=True,
+      text=True,
+      env=env,
+  )
+  assert result.stdout.strip() == "PRODUCTION"
+
+
+def test_ensure_litellm_imported_does_not_override():
+  if importlib.util.find_spec("litellm") is None:
+    pytest.skip("litellm is not installed")
+
+  env = _subprocess_env()
+  env["LITELLM_MODE"] = "DEV"
+
+  result = subprocess.run(
+      [
+          sys.executable,
+          "-c",
+          (
+              "import os\n"
+              "from google.adk.models.lite_llm import"
+              " _ensure_litellm_imported\n"
+              "_ensure_litellm_imported()\n"
+              "print(os.environ.get('LITELLM_MODE'))\n"
+          ),
+      ],
+      check=True,
+      capture_output=True,
+      text=True,
+      env=env,
+  )
+  assert result.stdout.strip() == "DEV"

tests/unittests/models/test_llm_request.py

@@ -0,0 +1,849 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Tests for LlmRequest functionality."""
+
+import asyncio
+from typing import Optional
+
+from google.adk.agents.invocation_context import InvocationContext
+from google.adk.agents.sequential_agent import SequentialAgent
+from google.adk.models.llm_request import LlmRequest
+from google.adk.sessions.in_memory_session_service import InMemorySessionService
+from google.adk.tools.base_tool import BaseTool
+from google.adk.tools.function_tool import FunctionTool
+from google.adk.tools.tool_context import ToolContext
+from google.genai import types
+import pytest
+
+
+def dummy_tool(query: str) -> str:
+  """A dummy tool for testing."""
+  return f'Searched for: {query}'
+
+
+def test_append_tools_with_none_config_tools():
+  """Test that append_tools initializes config.tools when it's None."""
+  request = LlmRequest()
+
+  # Initially config.tools should be None
+  assert request.config.tools is None
+
+  # Create a tool to append
+  tool = FunctionTool(func=dummy_tool)
+
+  # This should not raise an AttributeError
+  request.append_tools([tool])
+
+  # Now config.tools should be initialized and contain the tool
+  assert request.config.tools is not None
+  assert len(request.config.tools) == 1
+  assert len(request.config.tools[0].function_declarations) == 1
+  assert request.config.tools[0].function_declarations[0].name == 'dummy_tool'
+
+  # Tool should also be in tools_dict
+  assert 'dummy_tool' in request.tools_dict
+  assert request.tools_dict['dummy_tool'] == tool
+
+
+def test_append_tools_with_existing_tools():
+  """Test that append_tools works correctly when config.tools already exists."""
+  request = LlmRequest()
+
+  # Pre-initialize config.tools with an existing tool
+  existing_declaration = types.FunctionDeclaration(
+      name='existing_tool', description='An existing tool', parameters={}
+  )
+  request.config.tools = [
+      types.Tool(function_declarations=[existing_declaration])
+  ]
+
+  # Create a new tool to append
+  tool = FunctionTool(func=dummy_tool)
+
+  # Append the new tool
+  request.append_tools([tool])
+
+  # Should still have 1 tool but now with 2 function declarations
+  assert len(request.config.tools) == 1
+  assert len(request.config.tools[0].function_declarations) == 2
+
+  # Verify both declarations are present
+  decl_names = {
+      decl.name for decl in request.config.tools[0].function_declarations
+  }
+  assert decl_names == {'existing_tool', 'dummy_tool'}
+
+
+def test_append_tools_empty_list():
+  """Test that append_tools handles empty list correctly."""
+  request = LlmRequest()
+
+  # This should not modify anything
+  request.append_tools([])
+
+  # config.tools should still be None
+  assert request.config.tools is None
+  assert len(request.tools_dict) == 0
+
+
+def test_append_tools_tool_with_no_declaration():
+  """Test append_tools with a BaseTool that returns None from _get_declaration."""
+  from google.adk.tools.base_tool import BaseTool
+
+  request = LlmRequest()
+
+  # Create a mock tool that inherits from BaseTool and returns None for declaration
+  class NoDeclarationTool(BaseTool):
+
+    def __init__(self):
+      super().__init__(
+          name='no_decl_tool', description='A tool with no declaration'
+      )
+
+    def _get_declaration(self):
+      return None
+
+  tool = NoDeclarationTool()
+
+  # This should not add anything to config.tools but should handle gracefully
+  request.append_tools([tool])
+
+  # config.tools should still be None since no declarations were added
+  assert request.config.tools is None
+  # tools_dict should be empty since no valid declaration
+  assert len(request.tools_dict) == 0
+
+
+def test_append_tools_consolidates_declarations_in_single_tool():
+  """Test that append_tools puts all function declarations in a single Tool."""
+  request = LlmRequest()
+
+  # Create multiple tools
+  tool1 = FunctionTool(func=dummy_tool)
+
+  def another_tool(param: str) -> str:
+    return f'Another: {param}'
+
+  def third_tool(value: int) -> int:
+    return value * 2
+
+  tool2 = FunctionTool(func=another_tool)
+  tool3 = FunctionTool(func=third_tool)
+
+  # Append all tools at once
+  request.append_tools([tool1, tool2, tool3])
+
+  # Should have exactly 1 Tool with 3 function declarations
+  assert len(request.config.tools) == 1
+  assert len(request.config.tools[0].function_declarations) == 3
+
+  # Verify all tools are in tools_dict
+  assert len(request.tools_dict) == 3
+  assert 'dummy_tool' in request.tools_dict
+  assert 'another_tool' in request.tools_dict
+  assert 'third_tool' in request.tools_dict
+
+
+def test_append_instructions_with_string_list():
+  """Test that append_instructions works with list of strings (existing behavior)."""
+  request = LlmRequest()
+
+  # Initially system_instruction should be None
+  assert request.config.system_instruction is None
+
+  # Append first set of instructions
+  request.append_instructions(['First instruction', 'Second instruction'])
+
+  # Should be joined with double newlines
+  expected = 'First instruction\n\nSecond instruction'
+  assert request.config.system_instruction == expected
+  assert len(request.contents) == 0
+
+
+def test_append_instructions_with_string_list_multiple_calls():
+  """Test multiple calls to append_instructions with string lists."""
+  request = LlmRequest()
+
+  # First call
+  request.append_instructions(['First instruction'])
+  assert request.config.system_instruction == 'First instruction'
+
+  # Second call should append with double newlines
+  request.append_instructions(['Second instruction', 'Third instruction'])
+  expected = 'First instruction\n\nSecond instruction\n\nThird instruction'
+  assert request.config.system_instruction == expected
+
+
+def test_append_instructions_with_content():
+  """Test that append_instructions works with types.Content (new behavior)."""
+  request = LlmRequest()
+
+  # Create a Content object
+  content = types.Content(
+      role='user', parts=[types.Part(text='This is content-based instruction')]
+  )
+
+  # Append content
+  request.append_instructions(content)
+
+  # Should be set as system_instruction
+  assert len(request.contents) == 0
+  assert request.config.system_instruction == content
+
+
+def test_append_instructions_with_content_multiple_calls():
+  """Test multiple calls to append_instructions with Content objects."""
+  request = LlmRequest()
+
+  # Add some existing content first
+  existing_content = types.Content(
+      role='user', parts=[types.Part(text='Existing content')]
+  )
+  request.contents.append(existing_content)
+
+  # First Content instruction
+  content1 = types.Content(
+      role='user', parts=[types.Part(text='First instruction')]
+  )
+  request.append_instructions(content1)
+
+  # Should be set as system_instruction, existing content unchanged
+  assert len(request.contents) == 1
+  assert request.contents[0] == existing_content
+  assert request.config.system_instruction == content1
+
+  # Second Content instruction
+  content2 = types.Content(
+      role='user', parts=[types.Part(text='Second instruction')]
+  )
+  request.append_instructions(content2)
+
+  # Second Content should be merged with first in system_instruction
+  assert len(request.contents) == 1
+  assert request.contents[0] == existing_content
+  assert isinstance(request.config.system_instruction, types.Content)
+  assert len(request.config.system_instruction.parts) == 2
+  assert request.config.system_instruction.parts[0].text == 'First instruction'
+  assert request.config.system_instruction.parts[1].text == 'Second instruction'
+
+
+def test_append_instructions_with_content_multipart():
+  """Test append_instructions with Content containing multiple parts."""
+  request = LlmRequest()
+
+  # Create Content with multiple parts (text and potentially files)
+  content = types.Content(
+      role='user',
+      parts=[
+          types.Part(text='Text instruction'),
+          types.Part(text='Additional text part'),
+      ],
+  )
+
+  request.append_instructions(content)
+
+  assert len(request.contents) == 0
+  assert request.config.system_instruction == content
+  assert len(request.config.system_instruction.parts) == 2
+  assert request.config.system_instruction.parts[0].text == 'Text instruction'
+  assert (
+      request.config.system_instruction.parts[1].text == 'Additional text part'
+  )
+
+
+def test_append_instructions_mixed_string_and_content():
+  """Test mixing string list and Content instructions."""
+  request = LlmRequest()
+
+  # First add string instructions
+  request.append_instructions(['String instruction'])
+  assert request.config.system_instruction == 'String instruction'
+
+  # Then add Content instruction
+  content = types.Content(
+      role='user', parts=[types.Part(text='Content instruction')]
+  )
+  request.append_instructions(content)
+
+  # String and Content should be merged in system_instruction
+  assert len(request.contents) == 0
+  assert isinstance(request.config.system_instruction, types.Content)
+  assert len(request.config.system_instruction.parts) == 2
+  assert request.config.system_instruction.parts[0].text == 'String instruction'
+  assert (
+      request.config.system_instruction.parts[1].text == 'Content instruction'
+  )
+
+
+def test_append_instructions_empty_string_list():
+  """Test append_instructions with empty list of strings."""
+  request = LlmRequest()
+
+  # Empty list should not modify anything
+  request.append_instructions([])
+
+  assert request.config.system_instruction is None
+  assert len(request.contents) == 0
+
+
+def test_append_instructions_invalid_input():
+  """Test append_instructions with invalid input types."""
+  request = LlmRequest()
+
+  # Test with invalid types
+  with pytest.raises(
+      TypeError, match='instructions must be list\\[str\\] or types.Content'
+  ):
+    request.append_instructions('single string')  # Should be list[str]
+
+  with pytest.raises(
+      TypeError, match='instructions must be list\\[str\\] or types.Content'
+  ):
+    request.append_instructions(123)  # Invalid type
+
+  with pytest.raises(
+      TypeError, match='instructions must be list\\[str\\] or types.Content'
+  ):
+    request.append_instructions(
+        ['valid string', 123]
+    )  # Mixed valid/invalid in list
+
+
+def test_append_instructions_content_preserves_role_and_parts():
+  """Test that Content objects have text extracted regardless of role or parts."""
+  request = LlmRequest()
+
+  # Create Content with specific role and parts
+  content = types.Content(
+      role='system',  # Different role
+      parts=[
+          types.Part(text='System instruction'),
+          types.Part(text='Additional system part'),
+      ],
+  )
+
+  request.append_instructions(content)
+
+  # Text should be extracted and concatenated to system_instruction string
+  assert len(request.contents) == 0
+  assert (
+      request.config.system_instruction
+      == 'System instruction\n\nAdditional system part'
+  )
+
+
+async def _create_tool_context() -> ToolContext:
+  """Helper to create a ToolContext for testing."""
+  session_service = InMemorySessionService()
+  session = await session_service.create_session(
+      app_name='test_app', user_id='test_user'
+  )
+  agent = SequentialAgent(name='test_agent')
+  invocation_context = InvocationContext(
+      invocation_id='invocation_id',
+      agent=agent,
+      session=session,
+      session_service=session_service,
+  )
+  return ToolContext(invocation_context)
+
+
+class _MockTool(BaseTool):
+  """Mock tool for testing process_llm_request behavior."""
+
+  def __init__(self, name: str):
+    super().__init__(name=name, description=f'Mock tool {name}')
+
+  def _get_declaration(self) -> Optional[types.FunctionDeclaration]:
+    return types.FunctionDeclaration(
+        name=self.name,
+        description=self.description,
+        parameters=types.Schema(type=types.Type.STRING, title='param'),
+    )
+
+
+@pytest.mark.asyncio
+async def test_process_llm_request_consolidates_declarations_in_single_tool():
+  """Test that multiple process_llm_request calls consolidate in single Tool."""
+  request = LlmRequest()
+  tool_context = await _create_tool_context()
+
+  # Create multiple tools
+  tool1 = _MockTool('tool1')
+  tool2 = _MockTool('tool2')
+  tool3 = _MockTool('tool3')
+
+  # Process each tool individually (simulating what happens in real usage)
+  await tool1.process_llm_request(
+      tool_context=tool_context, llm_request=request
+  )
+  await tool2.process_llm_request(
+      tool_context=tool_context, llm_request=request
+  )
+  await tool3.process_llm_request(
+      tool_context=tool_context, llm_request=request
+  )
+
+  # Should have exactly 1 Tool with 3 function declarations
+  assert len(request.config.tools) == 1
+  assert len(request.config.tools[0].function_declarations) == 3
+
+  # Verify all function declaration names
+  decl_names = [
+      decl.name for decl in request.config.tools[0].function_declarations
+  ]
+  assert 'tool1' in decl_names
+  assert 'tool2' in decl_names
+  assert 'tool3' in decl_names
+
+  # Verify all tools are in tools_dict
+  assert len(request.tools_dict) == 3
+  assert 'tool1' in request.tools_dict
+  assert 'tool2' in request.tools_dict
+  assert 'tool3' in request.tools_dict
+
+
+@pytest.mark.asyncio
+async def test_append_tools_and_process_llm_request_consistent_behavior():
+  """Test that append_tools and process_llm_request produce same structure."""
+  tool_context = await _create_tool_context()
+
+  # Test 1: Using append_tools
+  request1 = LlmRequest()
+  tool1 = _MockTool('tool1')
+  tool2 = _MockTool('tool2')
+  tool3 = _MockTool('tool3')
+  request1.append_tools([tool1, tool2, tool3])
+
+  # Test 2: Using process_llm_request
+  request2 = LlmRequest()
+  tool4 = _MockTool('tool1')  # Same names for comparison
+  tool5 = _MockTool('tool2')
+  tool6 = _MockTool('tool3')
+  await tool4.process_llm_request(
+      tool_context=tool_context, llm_request=request2
+  )
+  await tool5.process_llm_request(
+      tool_context=tool_context, llm_request=request2
+  )
+  await tool6.process_llm_request(
+      tool_context=tool_context, llm_request=request2
+  )
+
+  # Both approaches should produce identical structure
+  assert len(request1.config.tools) == len(request2.config.tools) == 1
+  assert len(request1.config.tools[0].function_declarations) == 3
+  assert len(request2.config.tools[0].function_declarations) == 3
+
+  # Function declaration names should match
+  decl_names1 = {
+      decl.name for decl in request1.config.tools[0].function_declarations
+  }
+  decl_names2 = {
+      decl.name for decl in request2.config.tools[0].function_declarations
+  }
+  assert decl_names1 == decl_names2 == {'tool1', 'tool2', 'tool3'}
+
+
+def test_multiple_append_tools_calls_consolidate():
+  """Test that multiple append_tools calls add to the same Tool."""
+  request = LlmRequest()
+
+  # First call to append_tools
+  tool1 = FunctionTool(func=dummy_tool)
+  request.append_tools([tool1])
+
+  # Should have 1 tool with 1 declaration
+  assert len(request.config.tools) == 1
+  assert len(request.config.tools[0].function_declarations) == 1
+  assert request.config.tools[0].function_declarations[0].name == 'dummy_tool'
+
+  # Second call to append_tools with different tools
+  def another_tool(param: str) -> str:
+    return f'Another: {param}'
+
+  def third_tool(value: int) -> int:
+    return value * 2
+
+  tool2 = FunctionTool(func=another_tool)
+  tool3 = FunctionTool(func=third_tool)
+  request.append_tools([tool2, tool3])
+
+  # Should still have 1 tool but now with 3 declarations
+  assert len(request.config.tools) == 1
+  assert len(request.config.tools[0].function_declarations) == 3
+
+  # Verify all declaration names are present
+  decl_names = {
+      decl.name for decl in request.config.tools[0].function_declarations
+  }
+  assert decl_names == {'dummy_tool', 'another_tool', 'third_tool'}
+
+  # Verify all tools are in tools_dict
+  assert len(request.tools_dict) == 3
+  assert 'dummy_tool' in request.tools_dict
+  assert 'another_tool' in request.tools_dict
+  assert 'third_tool' in request.tools_dict
+
+
+# Updated tests for simplified string-only append_instructions behavior
+
+
+def test_append_instructions_with_content():
+  """Test that append_instructions extracts text from types.Content."""
+  request = LlmRequest()
+
+  # Create a Content object
+  content = types.Content(
+      role='user', parts=[types.Part(text='This is content-based instruction')]
+  )
+
+  # Append content
+  request.append_instructions(content)
+
+  # Should extract text and set as system_instruction string
+  assert len(request.contents) == 0
+  assert (
+      request.config.system_instruction == 'This is content-based instruction'
+  )
+
+
+def test_append_instructions_with_content_multiple_calls():
+  """Test multiple calls to append_instructions with Content objects."""
+  request = LlmRequest()
+
+  # Add some existing content first
+  existing_content = types.Content(
+      role='user', parts=[types.Part(text='Existing content')]
+  )
+  request.contents.append(existing_content)
+
+  # First Content instruction
+  content1 = types.Content(
+      role='user', parts=[types.Part(text='First instruction')]
+  )
+  request.append_instructions(content1)
+
+  # Should extract text and set as system_instruction, existing content unchanged
+  assert len(request.contents) == 1
+  assert request.contents[0] == existing_content
+  assert request.config.system_instruction == 'First instruction'
+
+  # Second Content instruction
+  content2 = types.Content(
+      role='user', parts=[types.Part(text='Second instruction')]
+  )
+  request.append_instructions(content2)
+
+  # Second Content text should be appended to existing string
+  assert len(request.contents) == 1
+  assert request.contents[0] == existing_content
+  assert (
+      request.config.system_instruction
+      == 'First instruction\n\nSecond instruction'
+  )
+
+
+def test_append_instructions_with_content_multipart():
+  """Test append_instructions with Content containing multiple text parts."""
+  request = LlmRequest()
+
+  # Create Content with multiple text parts
+  content = types.Content(
+      role='user',
+      parts=[
+          types.Part(text='Text instruction'),
+          types.Part(text='Additional text part'),
+      ],
+  )
+
+  request.append_instructions(content)
+
+  # Should extract and join all text parts
+  assert len(request.contents) == 0
+  assert (
+      request.config.system_instruction
+      == 'Text instruction\n\nAdditional text part'
+  )
+
+
+def test_append_instructions_mixed_string_and_content():
+  """Test mixing string list and Content instructions."""
+  request = LlmRequest()
+
+  # First add string instructions
+  request.append_instructions(['String instruction'])
+  assert request.config.system_instruction == 'String instruction'
+
+  # Then add Content instruction
+  content = types.Content(
+      role='user', parts=[types.Part(text='Content instruction')]
+  )
+  request.append_instructions(content)
+
+  # Content text should be appended to existing string
+  assert len(request.contents) == 0
+  assert (
+      request.config.system_instruction
+      == 'String instruction\n\nContent instruction'
+  )
+
+
+def test_append_instructions_content_extracts_text_only():
+  """Test that Content objects have text extracted regardless of role."""
+  request = LlmRequest()
+
+  # Create Content with specific role and parts
+  content = types.Content(
+      role='system',  # Different role
+      parts=[
+          types.Part(text='System instruction'),
+          types.Part(text='Additional system part'),
+      ],
+  )
+
+  request.append_instructions(content)
+
+  # Only text should be extracted and concatenated
+  assert len(request.contents) == 0
+  assert (
+      request.config.system_instruction
+      == 'System instruction\n\nAdditional system part'
+  )
+
+
+def test_append_instructions_content_with_non_text_parts():
+  """Test that non-text parts in Content are processed with references."""
+  request = LlmRequest()
+
+  # Create Content with text and non-text parts
+  content = types.Content(
+      role='user',
+      parts=[
+          types.Part(text='Text instruction'),
+          types.Part(
+              inline_data=types.Blob(data=b'file_data', mime_type='text/plain')
+          ),
+          types.Part(text='More text'),
+      ],
+  )
+
+  user_contents = request.append_instructions(content)
+
+  # Text parts should be extracted with references to non-text parts
+  expected_system = (
+      'Text instruction\n\n'
+      '[Reference to inline binary data: inline_data_0 (type: text/plain)]\n\n'
+      'More text'
+  )
+  assert request.config.system_instruction == expected_system
+
+  # Should return user content for the non-text part
+  assert len(user_contents) == 1
+  assert user_contents[0].role == 'user'
+  assert len(user_contents[0].parts) == 2
+  assert (
+      user_contents[0].parts[0].text == 'Referenced inline data: inline_data_0'
+  )
+  assert user_contents[0].parts[1].inline_data.data == b'file_data'
+
+
+def test_append_instructions_content_no_text_parts():
+  """Test that Content with no text parts processes non-text parts with references."""
+  request = LlmRequest()
+
+  # Set initial system instruction
+  request.config.system_instruction = 'Initial'
+
+  # Create Content with only non-text parts
+  content = types.Content(
+      role='user',
+      parts=[
+          types.Part(
+              inline_data=types.Blob(data=b'file_data', mime_type='text/plain')
+          ),
+      ],
+  )
+
+  user_contents = request.append_instructions(content)
+
+  # Should add reference to non-text part to system instruction
+  expected_system = (
+      'Initial\n\n[Reference to inline binary data: inline_data_0 (type:'
+      ' text/plain)]'
+  )
+  assert request.config.system_instruction == expected_system
+
+  # Should return user content for the non-text part
+  assert len(user_contents) == 1
+  assert user_contents[0].role == 'user'
+  assert len(user_contents[0].parts) == 2
+  assert (
+      user_contents[0].parts[0].text == 'Referenced inline data: inline_data_0'
+  )
+  assert user_contents[0].parts[1].inline_data.data == b'file_data'
+
+
+def test_append_instructions_content_empty_text_parts():
+  """Test that Content with empty text parts are skipped."""
+  request = LlmRequest()
+
+  # Create Content with empty and non-empty text parts
+  content = types.Content(
+      role='user',
+      parts=[
+          types.Part(text='Valid text'),
+          types.Part(text=''),  # Empty text
+          types.Part(text=None),  # None text
+          types.Part(text='More valid text'),
+      ],
+  )
+
+  request.append_instructions(content)
+
+  # Only non-empty text should be extracted
+  assert request.config.system_instruction == 'Valid text\n\nMore valid text'
+
+
+def test_append_instructions_warning_unsupported_system_instruction_type(
+    caplog,
+):
+  """Test that warnings are logged for unsupported system_instruction types."""
+  import logging
+
+  request = LlmRequest()
+
+  # Set unsupported type as system_instruction
+  request.config.system_instruction = {'unsupported': 'dict'}
+
+  with caplog.at_level(logging.WARNING):
+    # Try appending Content - should log warning and skip
+    content = types.Content(role='user', parts=[types.Part(text='Test')])
+    request.append_instructions(content)
+
+    # Should remain unchanged
+    assert request.config.system_instruction == {'unsupported': 'dict'}
+
+    # Try appending strings - should also log warning and skip
+    request.append_instructions(['Test string'])
+
+    # Should remain unchanged
+    assert request.config.system_instruction == {'unsupported': 'dict'}
+
+  # Check that warnings were logged
+  assert (
+      len(
+          [record for record in caplog.records if record.levelname == 'WARNING']
+      )
+      >= 1
+  )
+  assert (
+      'Cannot append to system_instruction of unsupported type' in caplog.text
+  )
+
+
+def test_append_instructions_with_mixed_content():
+  """Test append_instructions with mixed text and non-text content."""
+  request = LlmRequest()
+
+  # Create static instruction with mixed content
+  static_content = types.Content(
+      role='user',
+      parts=[
+          types.Part(text='Analyze this:'),
+          types.Part(
+              inline_data=types.Blob(
+                  data=b'test_data',
+                  mime_type='image/png',
+                  display_name='test.png',
+              )
+          ),
+          types.Part(text='Focus on details.'),
+          types.Part(
+              file_data=types.FileData(
+                  file_uri='files/doc123',
+                  mime_type='text/plain',
+                  display_name='document.txt',
+              )
+          ),
+      ],
+  )
+
+  user_contents = request.append_instructions(static_content)
+
+  # System instruction should contain text with references
+  expected_system = (
+      'Analyze this:\n\n[Reference to inline binary data: inline_data_0'
+      " ('test.png', type: image/png)]\n\nFocus on details.\n\n[Reference to"
+      " file data: file_data_1 ('document.txt', URI: files/doc123, type:"
+      ' text/plain)]'
+  )
+  assert request.config.system_instruction == expected_system
+
+  # Should return user contents for non-text parts
+  assert len(user_contents) == 2
+
+  # Check inline_data content
+  assert user_contents[0].role == 'user'
+  assert len(user_contents[0].parts) == 2
+  assert (
+      user_contents[0].parts[0].text == 'Referenced inline data: inline_data_0'
+  )
+  assert user_contents[0].parts[1].inline_data.data == b'test_data'
+  assert user_contents[0].parts[1].inline_data.display_name == 'test.png'
+
+  # Check file_data content
+  assert user_contents[1].role == 'user'
+  assert len(user_contents[1].parts) == 2
+  assert user_contents[1].parts[0].text == 'Referenced file data: file_data_1'
+  assert user_contents[1].parts[1].file_data.file_uri == 'files/doc123'
+  assert user_contents[1].parts[1].file_data.display_name == 'document.txt'
+
+
+def test_append_instructions_with_only_text_parts():
+  """Test append_instructions with only text parts."""
+  request = LlmRequest()
+
+  static_content = types.Content(
+      role='user',
+      parts=[
+          types.Part(text='First instruction'),
+          types.Part(text='Second instruction'),
+      ],
+  )
+
+  user_contents = request.append_instructions(static_content)
+
+  # Should only have text in system instruction
+  assert (
+      request.config.system_instruction
+      == 'First instruction\n\nSecond instruction'
+  )
+
+  # Should return empty list since no non-text parts
+  assert user_contents == []
+
+
+def test_is_managed_agent_defaults_false():
+  """_is_managed_agent defaults to False for ordinary requests."""
+  request = LlmRequest()
+  assert request._is_managed_agent is False
+
+
+def test_is_managed_agent_can_be_set_true():
+  """_is_managed_agent is an internal flag set after construction."""
+  request = LlmRequest()
+  request._is_managed_agent = True
+  assert request._is_managed_agent is True

tests/unittests/models/test_llm_response.py

@@ -0,0 +1,459 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Tests for LlmResponse, including log probabilities feature."""
+
+from google.adk.models.llm_response import LlmResponse
+from google.genai import types
+
+
+def test_llm_response_create_with_logprobs():
+  """Test LlmResponse.create() extracts logprobs from candidate."""
+  avg_logprobs = -0.75
+  logprobs_result = types.LogprobsResult(
+      chosen_candidates=[], top_candidates=[]
+  )
+
+  generate_content_response = types.GenerateContentResponse(
+      candidates=[
+          types.Candidate(
+              content=types.Content(parts=[types.Part(text='Response text')]),
+              finish_reason=types.FinishReason.STOP,
+              avg_logprobs=avg_logprobs,
+              logprobs_result=logprobs_result,
+          )
+      ]
+  )
+
+  response = LlmResponse.create(generate_content_response)
+
+  assert response.avg_logprobs == avg_logprobs
+  assert response.logprobs_result == logprobs_result
+  assert response.content.parts[0].text == 'Response text'
+  assert response.finish_reason == types.FinishReason.STOP
+
+
+def test_llm_response_create_without_logprobs():
+  """Test LlmResponse.create() handles missing logprobs gracefully."""
+  generate_content_response = types.GenerateContentResponse(
+      candidates=[
+          types.Candidate(
+              content=types.Content(parts=[types.Part(text='Response text')]),
+              finish_reason=types.FinishReason.STOP,
+              avg_logprobs=None,
+              logprobs_result=None,
+          )
+      ]
+  )
+
+  response = LlmResponse.create(generate_content_response)
+
+  assert response.avg_logprobs is None
+  assert response.logprobs_result is None
+  assert response.content.parts[0].text == 'Response text'
+
+
+def test_llm_response_create_error_case_with_logprobs():
+  """Test LlmResponse.create() includes logprobs in error cases."""
+  avg_logprobs = -2.1
+
+  generate_content_response = types.GenerateContentResponse(
+      candidates=[
+          types.Candidate(
+              content=None,  # No content - error case
+              finish_reason=types.FinishReason.SAFETY,
+              finish_message='Safety filter triggered',
+              avg_logprobs=avg_logprobs,
+              logprobs_result=None,
+          )
+      ]
+  )
+
+  response = LlmResponse.create(generate_content_response)
+
+  assert response.avg_logprobs == avg_logprobs
+  assert response.logprobs_result is None
+  assert response.error_code == types.FinishReason.SAFETY
+  assert response.error_message == 'Safety filter triggered'
+
+
+def test_llm_response_create_no_candidates():
+  """Test LlmResponse.create() with no candidates."""
+  generate_content_response = types.GenerateContentResponse(
+      candidates=[],
+      prompt_feedback=types.GenerateContentResponsePromptFeedback(
+          block_reason=types.BlockedReason.SAFETY,
+          block_reason_message='Prompt blocked for safety',
+      ),
+  )
+
+  response = LlmResponse.create(generate_content_response)
+
+  # No candidates means no logprobs
+  assert response.avg_logprobs is None
+  assert response.logprobs_result is None
+  assert response.error_code == types.BlockedReason.SAFETY
+  assert response.error_message == 'Prompt blocked for safety'
+
+
+def test_llm_response_create_no_candidates_without_prompt_feedback():
+  """Test LlmResponse.create() for empty successful model responses."""
+  usage_metadata = types.GenerateContentResponseUsageMetadata(
+      prompt_token_count=10,
+      candidates_token_count=0,
+      total_token_count=10,
+  )
+  generate_content_response = types.GenerateContentResponse(
+      candidates=[],
+      usage_metadata=usage_metadata,
+      model_version='gemini-2.5-flash',
+  )
+
+  response = LlmResponse.create(generate_content_response)
+
+  assert response.error_code is None
+  assert response.error_message is None
+  assert response.finish_reason is None
+  assert response.content is not None
+  assert response.content.role == 'model'
+  assert not response.content.parts
+  assert response.usage_metadata == usage_metadata
+  assert response.model_version == 'gemini-2.5-flash'
+
+
+def test_llm_response_create_with_concrete_logprobs_result():
+  """Test LlmResponse.create() with detailed logprobs_result containing actual token data."""
+  # Create realistic logprobs data
+  chosen_candidates = [
+      types.LogprobsResultCandidate(
+          token='The', log_probability=-0.1, token_id=123
+      ),
+      types.LogprobsResultCandidate(
+          token=' capital', log_probability=-0.5, token_id=456
+      ),
+      types.LogprobsResultCandidate(
+          token=' of', log_probability=-0.2, token_id=789
+      ),
+  ]
+
+  top_candidates = [
+      types.LogprobsResultTopCandidates(
+          candidates=[
+              types.LogprobsResultCandidate(
+                  token='The', log_probability=-0.1, token_id=123
+              ),
+              types.LogprobsResultCandidate(
+                  token='A', log_probability=-2.3, token_id=124
+              ),
+              types.LogprobsResultCandidate(
+                  token='This', log_probability=-3.1, token_id=125
+              ),
+          ]
+      ),
+      types.LogprobsResultTopCandidates(
+          candidates=[
+              types.LogprobsResultCandidate(
+                  token=' capital', log_probability=-0.5, token_id=456
+              ),
+              types.LogprobsResultCandidate(
+                  token=' city', log_probability=-1.2, token_id=457
+              ),
+              types.LogprobsResultCandidate(
+                  token=' main', log_probability=-2.8, token_id=458
+              ),
+          ]
+      ),
+  ]
+
+  avg_logprobs = -0.27  # Average of -0.1, -0.5, -0.2
+  logprobs_result = types.LogprobsResult(
+      chosen_candidates=chosen_candidates, top_candidates=top_candidates
+  )
+
+  generate_content_response = types.GenerateContentResponse(
+      candidates=[
+          types.Candidate(
+              content=types.Content(
+                  parts=[types.Part(text='The capital of France is Paris.')]
+              ),
+              finish_reason=types.FinishReason.STOP,
+              avg_logprobs=avg_logprobs,
+              logprobs_result=logprobs_result,
+          )
+      ]
+  )
+
+  response = LlmResponse.create(generate_content_response)
+
+  assert response.avg_logprobs == avg_logprobs
+  assert response.logprobs_result is not None
+
+  # Test chosen candidates
+  assert len(response.logprobs_result.chosen_candidates) == 3
+  assert response.logprobs_result.chosen_candidates[0].token == 'The'
+  assert response.logprobs_result.chosen_candidates[0].log_probability == -0.1
+  assert response.logprobs_result.chosen_candidates[0].token_id == 123
+  assert response.logprobs_result.chosen_candidates[1].token == ' capital'
+  assert response.logprobs_result.chosen_candidates[1].log_probability == -0.5
+  assert response.logprobs_result.chosen_candidates[1].token_id == 456
+
+  # Test top candidates
+  assert len(response.logprobs_result.top_candidates) == 2
+  assert (
+      len(response.logprobs_result.top_candidates[0].candidates) == 3
+  )  # 3 alternatives for first token
+  assert response.logprobs_result.top_candidates[0].candidates[0].token == 'The'
+  assert (
+      response.logprobs_result.top_candidates[0].candidates[0].token_id == 123
+  )
+  assert response.logprobs_result.top_candidates[0].candidates[1].token == 'A'
+  assert (
+      response.logprobs_result.top_candidates[0].candidates[1].token_id == 124
+  )
+  assert (
+      response.logprobs_result.top_candidates[0].candidates[2].token == 'This'
+  )
+  assert (
+      response.logprobs_result.top_candidates[0].candidates[2].token_id == 125
+  )
+
+
+def test_llm_response_create_with_partial_logprobs_result():
+  """Test LlmResponse.create() with logprobs_result having only chosen_candidates."""
+  chosen_candidates = [
+      types.LogprobsResultCandidate(
+          token='Hello', log_probability=-0.05, token_id=111
+      ),
+      types.LogprobsResultCandidate(
+          token=' world', log_probability=-0.8, token_id=222
+      ),
+  ]
+
+  logprobs_result = types.LogprobsResult(
+      chosen_candidates=chosen_candidates,
+      top_candidates=[],  # Empty top candidates
+  )
+
+  generate_content_response = types.GenerateContentResponse(
+      candidates=[
+          types.Candidate(
+              content=types.Content(parts=[types.Part(text='Hello world')]),
+              finish_reason=types.FinishReason.STOP,
+              avg_logprobs=-0.425,  # Average of -0.05 and -0.8
+              logprobs_result=logprobs_result,
+          )
+      ]
+  )
+
+  response = LlmResponse.create(generate_content_response)
+
+  assert response.avg_logprobs == -0.425
+  assert response.logprobs_result is not None
+  assert len(response.logprobs_result.chosen_candidates) == 2
+  assert len(response.logprobs_result.top_candidates) == 0
+  assert response.logprobs_result.chosen_candidates[0].token == 'Hello'
+  assert response.logprobs_result.chosen_candidates[1].token == ' world'
+
+
+def test_llm_response_create_with_citation_metadata():
+  """Test LlmResponse.create() extracts citation_metadata from candidate."""
+  citation_metadata = types.CitationMetadata(
+      citations=[
+          types.Citation(
+              start_index=0,
+              end_index=10,
+              uri='https://example.com',
+          )
+      ]
+  )
+
+  generate_content_response = types.GenerateContentResponse(
+      candidates=[
+          types.Candidate(
+              content=types.Content(parts=[types.Part(text='Response text')]),
+              finish_reason=types.FinishReason.STOP,
+              citation_metadata=citation_metadata,
+          )
+      ]
+  )
+
+  response = LlmResponse.create(generate_content_response)
+
+  assert response.citation_metadata == citation_metadata
+  assert response.content.parts[0].text == 'Response text'
+
+
+def test_llm_response_create_without_citation_metadata():
+  """Test LlmResponse.create() handles missing citation_metadata gracefully."""
+  generate_content_response = types.GenerateContentResponse(
+      candidates=[
+          types.Candidate(
+              content=types.Content(parts=[types.Part(text='Response text')]),
+              finish_reason=types.FinishReason.STOP,
+              citation_metadata=None,
+          )
+      ]
+  )
+
+  response = LlmResponse.create(generate_content_response)
+
+  assert response.citation_metadata is None
+  assert response.content.parts[0].text == 'Response text'
+
+
+def test_llm_response_create_error_case_with_citation_metadata():
+  """Test LlmResponse.create() includes citation_metadata in error cases."""
+  citation_metadata = types.CitationMetadata(
+      citations=[
+          types.Citation(
+              start_index=0,
+              end_index=10,
+              uri='https://example.com',
+          )
+      ]
+  )
+
+  generate_content_response = types.GenerateContentResponse(
+      candidates=[
+          types.Candidate(
+              content=None,  # No content - blocked case
+              finish_reason=types.FinishReason.RECITATION,
+              finish_message='Response blocked due to recitation triggered',
+              citation_metadata=citation_metadata,
+          )
+      ]
+  )
+
+  response = LlmResponse.create(generate_content_response)
+
+  assert response.citation_metadata == citation_metadata
+  assert response.error_code == types.FinishReason.RECITATION
+  assert (
+      response.error_message == 'Response blocked due to recitation triggered'
+  )
+
+
+def test_llm_response_create_empty_content_with_stop_reason():
+  """Empty content + STOP stays a successful response at the model layer.
+
+  Surfacing the empty turn as an error is the flow's job (non-streaming only);
+  the model/streaming layer must not classify a terminal finish-only chunk as
+  an error or it breaks streaming consumers that batch parts across chunks.
+  """
+  generate_content_response = types.GenerateContentResponse(
+      candidates=[
+          types.Candidate(
+              content=types.Content(parts=[]),
+              finish_reason=types.FinishReason.STOP,
+          )
+      ]
+  )
+
+  response = LlmResponse.create(generate_content_response)
+
+  assert response.error_code is None
+  assert response.content is not None
+  assert response.finish_reason == types.FinishReason.STOP
+
+
+def test_llm_response_create_non_empty_parts_with_stop_is_success():
+  """Regression guard: real text + STOP must remain a successful response."""
+  generate_content_response = types.GenerateContentResponse(
+      candidates=[
+          types.Candidate(
+              content=types.Content(
+                  role='model', parts=[types.Part(text='ok')]
+              ),
+              finish_reason=types.FinishReason.STOP,
+          )
+      ]
+  )
+
+  response = LlmResponse.create(generate_content_response)
+
+  assert response.error_code is None
+  assert response.content is not None
+
+
+def test_llm_response_create_includes_model_version():
+  """Test LlmResponse.create() includes model version."""
+  generate_content_response = types.GenerateContentResponse(
+      model_version='gemini-2.5-flash',
+      candidates=[
+          types.Candidate(
+              content=types.Content(parts=[types.Part(text='Response text')]),
+              finish_reason=types.FinishReason.STOP,
+          )
+      ],
+  )
+  response = LlmResponse.create(generate_content_response)
+  assert response.model_version == 'gemini-2.5-flash'
+
+
+def test_get_function_calls_returns_calls_in_order():
+  fc1 = types.FunctionCall(name='a', args={})
+  fc2 = types.FunctionCall(name='b', args={'x': 1})
+  response = LlmResponse(
+      content=types.Content(
+          parts=[
+              types.Part(function_call=fc1),
+              types.Part(text='ignored'),
+              types.Part(function_call=fc2),
+          ]
+      )
+  )
+  assert response.get_function_calls() == [fc1, fc2]
+
+
+def test_get_function_calls_empty_when_no_content():
+  assert LlmResponse().get_function_calls() == []
+
+
+def test_get_function_calls_empty_when_no_parts():
+  response = LlmResponse(content=types.Content(parts=None))
+  assert response.get_function_calls() == []
+
+
+def test_get_function_responses_returns_responses_in_order():
+  fr1 = types.FunctionResponse(name='a', response={'r': 1})
+  fr2 = types.FunctionResponse(name='b', response={'r': 2})
+  response = LlmResponse(
+      content=types.Content(
+          parts=[
+              types.Part(function_response=fr1),
+              types.Part(text='ignored'),
+              types.Part(function_response=fr2),
+          ]
+      )
+  )
+  assert response.get_function_responses() == [fr1, fr2]
+
+
+def test_get_function_responses_empty_when_no_content():
+  assert LlmResponse().get_function_responses() == []
+
+
+def test_get_function_responses_empty_when_no_parts():
+  response = LlmResponse(content=types.Content(parts=None))
+  assert response.get_function_responses() == []
+
+
+def test_environment_id_defaults_to_none_and_roundtrips():
+  resp = LlmResponse()
+  assert resp.environment_id is None
+
+  resp.environment_id = 'env_abc'
+  dumped = resp.model_dump(exclude_none=True)
+  assert dumped['environment_id'] == 'env_abc'
+  assert LlmResponse.model_validate(dumped).environment_id == 'env_abc'

tests/unittests/models/test_models.py

@@ -0,0 +1,143 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from google.adk import models
+from google.adk.models.anthropic_llm import Claude
+from google.adk.models.google_llm import Gemini
+from google.adk.models.lite_llm import LiteLlm
+import pytest
+
+
+@pytest.mark.parametrize(
+    'model_name',
+    [
+        'gemini-1.5-pro',
+        'gemini-1.5-pro-001',
+        'gemini-1.5-pro-002',
+        'gemini-2.5-flash',
+        'projects/123456/locations/us-central1/endpoints/123456',  # finetuned vertex gemini endpoint
+        'projects/123456/locations/us-central1/publishers/google/models/gemini-2.5-flash',  # vertex gemini long name
+    ],
+)
+def test_match_gemini_family(model_name):
+  """Test that Gemini models are resolved correctly."""
+  assert models.LLMRegistry.resolve(model_name) is Gemini
+
+
+@pytest.mark.parametrize(
+    'model_name',
+    [
+        'claude-3-5-haiku@20241022',
+        'claude-3-5-sonnet-v2@20241022',
+        'claude-3-5-sonnet@20240620',
+        'claude-3-haiku@20240307',
+        'claude-3-opus@20240229',
+        'claude-3-sonnet@20240229',
+        'claude-sonnet-4@20250514',
+        'claude-opus-4@20250514',
+    ],
+)
+def test_match_claude_family(model_name):
+  """Test that Claude models are resolved correctly."""
+  assert models.LLMRegistry.resolve(model_name) is Claude
+
+
+@pytest.mark.parametrize(
+    'model_name',
+    [
+        'openai/gpt-4o',
+        'openai/gpt-4o-mini',
+        'groq/llama3-70b-8192',
+        'groq/mixtral-8x7b-32768',
+        'anthropic/claude-3-opus-20240229',
+        'anthropic/claude-3-5-sonnet-20241022',
+    ],
+)
+def test_match_litellm_family(model_name):
+  """Test that LiteLLM models are resolved correctly."""
+  assert models.LLMRegistry.resolve(model_name) is LiteLlm
+
+
+def test_non_exist_model():
+  with pytest.raises(ValueError) as e_info:
+    models.LLMRegistry.resolve('non-exist-model')
+  assert 'Model non-exist-model not found.' in str(e_info.value)
+
+
+def test_helpful_error_for_claude_without_extensions():
+  """Test that missing Claude models show helpful install instructions.
+
+  Note: This test may pass even when anthropic IS installed, because it
+  only checks the error message format when a model is not found.
+  """
+  # Use a non-existent Claude model variant to trigger error
+  with pytest.raises(ValueError) as e_info:
+    models.LLMRegistry.resolve('claude-nonexistent-model-xyz')
+
+  error_msg = str(e_info.value)
+  # The error should mention anthropic package and installation instructions
+  # These checks work whether or not anthropic is actually installed
+  assert 'Model claude-nonexistent-model-xyz not found' in error_msg
+  assert 'anthropic package' in error_msg
+  assert 'pip install' in error_msg
+
+
+def test_helpful_error_for_litellm_without_extensions():
+  """Test that missing LiteLLM models show helpful install instructions.
+
+  Note: This test may pass even when litellm IS installed, because it
+  only checks the error message format when a model is not found.
+  """
+  # Use a non-existent provider to trigger error
+  with pytest.raises(ValueError) as e_info:
+    models.LLMRegistry.resolve('unknown-provider/gpt-4o')
+
+  error_msg = str(e_info.value)
+  # The error should mention litellm package for provider-style models
+  assert 'Model unknown-provider/gpt-4o not found' in error_msg
+  assert 'litellm package' in error_msg
+  assert 'pip install' in error_msg
+  assert 'Provider-style models' in error_msg
+
+
+def test_resolve_with_prefix():
+  """Test that model resolution can be overridden with a prefix."""
+  assert models.LLMRegistry.resolve('gemini:gemini-1.5-flash') is Gemini
+  assert models.LLMRegistry.resolve('Claude:claude-3-opus@20240229') is Claude
+  assert models.LLMRegistry.resolve('lite:openai/gpt-4o') is LiteLlm
+  assert models.LLMRegistry.resolve('LiteLlm:openai/gpt-4o') is LiteLlm
+
+
+def test_new_llm_with_prefix(mocker):
+  """Test that new_llm strips prefix when creating instance if it matches class."""
+  mock_class = mocker.MagicMock()
+  mock_class.__name__ = 'MockLlm'
+  mocker.patch.object(models.LLMRegistry, 'resolve', return_value=mock_class)
+
+  models.LLMRegistry.new_llm('mock:gpt-4')
+  mock_class.assert_called_once_with(model='gpt-4')
+
+  mock_class.reset_mock()
+  models.LLMRegistry.new_llm('MockLlm:gpt-4')
+  mock_class.assert_called_once_with(model='gpt-4')
+
+
+def test_new_llm_with_non_matching_prefix(mocker):
+  """Test that new_llm keeps prefix if it does not match class."""
+  mock_class = mocker.MagicMock()
+  mock_class.__name__ = 'MockLlm'
+  mocker.patch.object(models.LLMRegistry, 'resolve', return_value=mock_class)
+
+  models.LLMRegistry.new_llm('custom:gpt-4')
+  mock_class.assert_called_once_with(model='custom:gpt-4')