chore: import upstream snapshot with attribution

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'
+  )