google-genai 1.54.0__py3-none-any.whl → 1.55.0__py3-none-any.whl

google/genai/tests/chats/test_validate_response.py

@@ -0,0 +1,90 @@
+# Copyright 2025 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 ... import types
+from ...chats import _validate_response
+
+
+def test_validate_response_default_response():
+  response = types.GenerateContentResponse()
+
+  assert not _validate_response(response)
+
+
+def test_validate_response_empty_content():
+  response = types.GenerateContentResponse(candidates=[])
+
+  assert not _validate_response(response)
+
+
+def test_validate_response_empty_parts():
+  response = types.GenerateContentResponse(
+      candidates=[types.Candidate(content=types.Content(parts=[]))]
+  )
+
+  assert not _validate_response(response)
+
+
+def test_validate_response_empty_part():
+  response = types.GenerateContentResponse(
+      candidates=[types.Candidate(content=types.Content(parts=[types.Part()]))]
+  )
+
+  assert not _validate_response(response)
+
+
+def test_validate_response_part_with_empty_text():
+  response = types.GenerateContentResponse(
+      candidates=[
+          types.Candidate(content=types.Content(parts=[types.Part(text='')]))
+      ]
+  )
+
+  assert not _validate_response(response)
+
+
+def test_validate_response_part_with_text():
+  response = types.GenerateContentResponse(
+      candidates=[
+          types.Candidate(
+              content=types.Content(
+                  parts=[types.Part(text='response from model')]
+              )
+          )
+      ]
+  )
+
+  assert _validate_response(response)
+
+
+def test_validate_response_part_with_function_call():
+  response = types.GenerateContentResponse(
+      candidates=[
+          types.Candidate(
+              content=types.Content(
+                  parts=[
+                      types.Part(
+                          function_call=types.FunctionCall(
+                              name='foo', args={'bar': 'baz'}
+                          )
+                      )
+                  ]
+              )
+          )
+      ]
+  )
+
+  assert _validate_response(response)

google/genai/tests/client/__init__.py

@@ -0,0 +1,17 @@
+# Copyright 2025 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 the Google GenAI SDK."""

google/genai/tests/client/test_async_stream.py

@@ -0,0 +1,427 @@
+# Copyright 2025 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 async stream."""
+
+import asyncio
+from typing import List
+from unittest import mock
+from unittest.mock import AsyncMock
+from unittest.mock import MagicMock
+from unittest.mock import patch
+import pytest
+
+
+try:
+  import aiohttp
+
+  AIOHTTP_NOT_INSTALLED = False
+except ImportError:
+  AIOHTTP_NOT_INSTALLED = True
+  aiohttp = mock.MagicMock()
+
+import httpx
+
+from ... import _api_client as api_client
+from ... import Client
+from ... import errors
+from ... import types
+
+
+EVENT_STREAM_DATA_WITH_ERROR = [
+    b'{"candidates":[{"content":{"parts":[{"text":"test"}],"role":"model"}}]}',
+    b'\n',
+    b'{"error":{"code":500,"message":"Error","status":"INTERNAL"}}',
+]
+
+
+class MockHTTPXResponse(httpx.Response):
+  """Mock httpx.Response class for testing."""
+
+  def __init__(self, lines: List[str]):
+    self.aiter_lines = MagicMock()
+    self.aiter_lines.return_value.__aiter__ = MagicMock(
+        return_value=self._async_line_iterator(lines)
+    )
+    self.aclose = AsyncMock()
+
+  async def _async_line_iterator(self, lines: List[str]):
+    for line in lines:
+      yield line
+
+
+class MockAIOHTTPResponse(aiohttp.ClientResponse):
+
+  def __init__(self, lines: List[str]):
+    self.content = MagicMock()
+    self.content.readline = AsyncMock()
+    # Simulate reading lines, each ending with newline bytes for readline behavior
+    self._read_data = b"\n".join(line.encode("utf-8") for line in lines) + b"\n"
+    self._read_pos = 0
+    self.content.readline.side_effect = self._async_read_line
+    self.release = MagicMock()
+
+  async def _async_read_line(self) -> bytes:
+    if self._read_pos >= len(self._read_data):
+      return b""  # End of stream
+
+    newline_pos = self._read_data.find(b"\n", self._read_pos)
+    if newline_pos == -1:  # Should not happen with the appended '\n'
+      line = self._read_data[self._read_pos :]
+      self._read_pos = len(self._read_data)
+      return line
+    else:
+      line = self._read_data[self._read_pos : newline_pos + 1]
+      self._read_pos = newline_pos + 1
+      return line
+
+
+@pytest.fixture
+def responses() -> api_client.HttpResponse:
+  return api_client.HttpResponse(headers={})
+
+
+requires_aiohttp = pytest.mark.skipif(
+    AIOHTTP_NOT_INSTALLED, reason="aiohttp is not installed, skipping test."
+)
+
+
+@pytest.fixture(autouse=True)
+def reset_has_aiohttp():
+  yield
+  api_client.has_aiohttp = False
+
+
+def test_invalid_response_stream_type(responses: api_client.HttpResponse):
+  """Tests that an invalid response stream type raises an error."""
+  api_client.has_aiohttp = False
+  with pytest.raises(
+      TypeError,
+      match=(
+          "Expected self.response_stream to be an httpx.Response or"
+          " aiohttp.ClientResponse object"
+      ),
+  ):
+
+    async def run():
+      async for _ in responses._aiter_response_stream():
+        pass
+
+    asyncio.run(run())
+
+
+@pytest.mark.asyncio
+async def test_httpx_simple_lines(responses: api_client.HttpResponse):
+  lines = ["hello", "world", "testing"]
+  mock_response = MockHTTPXResponse(lines)
+  responses.response_stream = mock_response
+
+  results = [line async for line in responses._aiter_response_stream()]
+
+  assert results == lines
+  mock_response.aiter_lines.assert_called_once()
+  mock_response.aclose.assert_called_once()
+
+
+@pytest.mark.asyncio
+async def test_httpx_data_prefix(responses: api_client.HttpResponse):
+  lines = ["data: { 'message': 'hello' }", "data: { 'status': 'ok' }"]
+  mock_response = MockHTTPXResponse(lines)
+  responses.response_stream = mock_response
+
+  results = [line async for line in responses._aiter_response_stream()]
+
+  assert results == ["{ 'message': 'hello' }", "{ 'status': 'ok' }"]
+  mock_response.aiter_lines.assert_called_once()
+  mock_response.aclose.assert_called_once()
+
+
+@pytest.mark.asyncio
+async def test_httpx_multiple_json_chunk(responses: api_client.HttpResponse):
+  lines = [
+      '{ "id": 1 }',
+      "",
+      'data: { "id": 2 }',
+      'data: { "id": 3 }',
+  ]
+  mock_response = MockHTTPXResponse(lines)
+  responses.response_stream = mock_response
+
+  results = [line async for line in responses._aiter_response_stream()]
+
+  assert results == ['{ "id": 1 }', '{ "id": 2 }', '{ "id": 3 }']
+  mock_response.aiter_lines.assert_called_once()
+  mock_response.aclose.assert_called_once()
+
+
+@pytest.mark.asyncio
+async def test_httpx_incomplete_json_at_end(responses: api_client.HttpResponse):
+  lines = ['{ "partial": "data"']  # Missing closing brace
+  mock_response = MockHTTPXResponse(lines)
+  responses.response_stream = mock_response
+
+  results = [line async for line in responses._aiter_response_stream()]
+
+  # The remaining chunk is yielded
+  assert results == ['{ "partial": "data"']
+  mock_response.aiter_lines.assert_called_once()
+  mock_response.aclose.assert_called_once()
+
+
+@pytest.mark.asyncio
+async def test_httpx_empty_stream(responses: api_client.HttpResponse):
+  lines: List[str] = []
+  mock_response = MockHTTPXResponse(lines)
+  responses.response_stream = mock_response
+
+  results = [line async for line in responses._aiter_response_stream()]
+
+  assert results == []
+  mock_response.aiter_lines.assert_called_once()
+  mock_response.aclose.assert_called_once()
+
+
+# Async aiohttp
+@requires_aiohttp
+@pytest.mark.asyncio
+async def test_aiohttp_simple_lines(responses: api_client.HttpResponse):
+  api_client.has_aiohttp = True  # Force aiohttp
+  lines = ["hello", "world", "testing"]
+  # Use the mock class that pretends to be aiohttp.ClientResponse
+  mock_response = MockAIOHTTPResponse(lines)
+  responses.response_stream = mock_response
+
+  results = [line async for line in responses._aiter_response_stream()]
+
+  assert results == lines
+  mock_response.content.readline.assert_any_call()
+  mock_response.release.assert_called_once()
+
+
+@requires_aiohttp
+@pytest.mark.asyncio
+async def test_aiohttp_data_prefix(responses: api_client.HttpResponse):
+  api_client.has_aiohttp = True  # Force aiohttp
+  lines = ["data: { 'message': 'hello' }", "data: { 'status': 'ok' }"]
+  # Use the mock class that pretends to be aiohttp.ClientResponse
+  mock_response = MockAIOHTTPResponse(lines)
+  responses.response_stream = mock_response
+
+  results = [line async for line in responses._aiter_response_stream()]
+
+  assert results == ["{ 'message': 'hello' }", "{ 'status': 'ok' }"]
+  mock_response.content.readline.assert_any_call()
+  mock_response.release.assert_called_once()
+
+
+@requires_aiohttp
+@pytest.mark.asyncio
+async def test_aiohttp_multiple_json_chunks(responses: api_client.HttpResponse):
+  api_client.has_aiohttp = True  # Force aiohttp
+  lines = [
+      '{ "id": 1 }',
+      "",  # empty line to check robustness
+      'data: { "id": 2 }',
+      'data: { "id": 3 }',
+  ]
+  # Use the mock class that pretends to be aiohttp.ClientResponse
+  mock_response = MockAIOHTTPResponse(lines)
+  responses.response_stream = mock_response
+
+  results = [line async for line in responses._aiter_response_stream()]
+
+  assert results == ['{ "id": 1 }', '{ "id": 2 }', '{ "id": 3 }']
+  mock_response.content.readline.assert_any_call()
+  mock_response.release.assert_called_once()
+
+
+@requires_aiohttp
+@pytest.mark.asyncio
+async def test_aiohttp_incomplete_json_at_end(
+    responses: api_client.HttpResponse,
+):
+  api_client.has_aiohttp = True  # Force aiohttp
+  lines = ['{ "partial": "data"']  # Missing closing brace
+  # Use the mock class that pretends to be aiohttp.ClientResponse
+  mock_response = MockAIOHTTPResponse(lines)
+  responses.response_stream = mock_response
+
+  results = [line async for line in responses._aiter_response_stream()]
+
+  assert results == ['{ "partial": "data"']
+  mock_response.content.readline.assert_any_call()
+  mock_response.release.assert_called_once()
+
+
+def mock_response(chunks):
+  mock_stream = MagicMock(spec=httpx.SyncByteStream)
+  mock_stream.__iter__.return_value = chunks
+  return httpx.Response(
+      status_code=200,
+      stream=mock_stream,
+  )
+
+
+@patch('httpx.Client.send')
+def test_error_event_in_streamed_responses_bad_json(mock_send_method):
+  with_bad_json = [
+      b'{"candidates":[{"content":{"parts":[{"text":"test"}],"role":"model"}}]}',
+      b'\n',
+      b'{"error": bad_json}',
+  ]
+  mock_send_method.return_value = mock_response(with_bad_json)
+
+  client = api_client.BaseApiClient(api_key='test_api_key')
+  stream = client.request_streamed('POST', 'models/gemini-2.5-flash', {})
+
+  chunk = next(stream)
+  assert chunk == types.HttpResponse(
+      headers={},
+      body=(
+          '{"candidates": [{"content": {"parts": [{"text": "test"}], "role":'
+          ' "model"}}]}'
+      ),
+  )
+
+  with pytest.raises(errors.UnknownApiResponseError):
+    next(stream)
+
+
+@patch('httpx.Client.send')
+def test_error_event_in_streamed_responses(mock_send_method):
+  mock_send_method.return_value = mock_response(EVENT_STREAM_DATA_WITH_ERROR)
+
+  client = api_client.BaseApiClient(api_key='test_api_key')
+  stream = client.request_streamed('POST', 'models/gemini-2.5-flash', {})
+
+  chunk = next(stream)
+  assert chunk == types.HttpResponse(
+      body=(
+          '{"candidates": [{"content": {"parts": [{"text": "test"}], "role":'
+          ' "model"}}]}'
+      ),
+      headers={},
+  )
+
+  with pytest.raises(errors.ServerError):
+    next(stream)
+
+
+@patch('httpx.Client.send')
+def test_error_event_in_generate_content_stream(mock_send_method):
+  mock_send_method.return_value = mock_response(EVENT_STREAM_DATA_WITH_ERROR)
+
+  client = Client(api_key='test_api_key')
+  generated_response = client.models.generate_content_stream(
+      model='gemini-2.5-flash',
+      contents='Tell me a story in 300 words.',
+  )
+
+  chunk = next(generated_response)
+  assert chunk == types.GenerateContentResponse(
+      candidates=[
+          types.Candidate(
+              content=types.Content(
+                  parts=[
+                      types.Part(
+                          text='test'
+                      ),
+                  ],
+                  role='model'
+              )
+          ),
+      ],
+      sdk_http_response=types.HttpResponse(headers={})
+  )
+
+  with pytest.raises(errors.ServerError):
+    next(generated_response)
+
+
+async def _async_httpx_response(_):
+  mock_stream = MagicMock(spec=httpx.AsyncByteStream)
+  mock_stream.__aiter__.return_value = EVENT_STREAM_DATA_WITH_ERROR
+  mock_stream.aclose = AsyncMock()
+  return httpx.Response(
+      status_code=200,
+      stream=mock_stream,
+  )
+
+
+@patch('httpx.AsyncBaseTransport')
+@pytest.mark.asyncio
+async def test_error_event_in_streamed_responses_async(mock_transport):
+  client = api_client.BaseApiClient(
+      api_key='test_api_key',
+      http_options=types.HttpOptions(
+          async_client_args={'transport': mock_transport}
+          ),
+      )
+  mock_transport.handle_async_request = _async_httpx_response
+  mock_transport.aclose = AsyncMock()
+
+  resp = await client.async_request_streamed(
+      'POST', 'models/gemini-2.5-flash', {'key': 'value'}
+  )
+
+  chunk = await anext(resp)
+  assert chunk == types.HttpResponse(
+      headers={},
+      body=(
+          '{"candidates": [{"content": {"parts": [{"text": "test"}], "role":'
+          ' "model"}}]}'
+      ),
+  )
+
+  with pytest.raises(errors.ServerError):
+    await anext(resp)
+
+
+@patch('httpx.AsyncBaseTransport')
+@pytest.mark.asyncio
+async def test_error_event_in_generate_content_stream_async(mock_transport):
+  client = Client(
+      api_key='test_api_key',
+      http_options=types.HttpOptions(
+          async_client_args={'transport': mock_transport}
+      ),
+  )
+  mock_transport.handle_async_request = _async_httpx_response
+  mock_transport.aclose = AsyncMock()
+
+  generated_response = await client.aio.models.generate_content_stream(
+      model='gemini-2.5-flash',
+      contents='Tell me a story in 300 words.',
+  )
+
+  chunk = await anext(generated_response)
+  assert chunk == types.GenerateContentResponse(
+      candidates=[
+          types.Candidate(
+              content=types.Content(
+                  parts=[
+                      types.Part(
+                          text='test'
+                      ),
+                  ],
+                  role='model'
+              )
+          ),
+      ],
+      sdk_http_response=types.HttpResponse(headers={})
+  )
+
+  with pytest.raises(errors.ServerError):
+    await anext(generated_response)