From 53d918c47bd624efc8981dc7f3c28d9a215d6d85 Mon Sep 17 00:00:00 2001 From: Claude Date: Tue, 23 Jun 2026 21:07:20 +0000 Subject: [PATCH 1/3] Fetch provider image URLs and convert to data: URIs Providers like ByteDance Seedance return pre-signed CDN URLs rather than base64 blobs. The gateway now fetches those URLs inline and converts them to data: URIs so the SDK always receives self-contained base64 images. Without this the SDK example saved garbage by trying to base64-decode the URL string itself. Falls back to returning the raw URL if the fetch fails. Co-Authored-By: Claude Sonnet 4.6 Claude-Session: https://claude.ai/code/session_01XxsSjnuypVUhpDPiEDA1Md --- tee_gateway/llm_backend.py | 13 ++++++++- tee_gateway/test/test_image_generation.py | 33 ++++++++++++++++++----- 2 files changed, 38 insertions(+), 8 deletions(-) diff --git a/tee_gateway/llm_backend.py b/tee_gateway/llm_backend.py index 3a34e84..8d35207 100644 --- a/tee_gateway/llm_backend.py +++ b/tee_gateway/llm_backend.py @@ -7,6 +7,7 @@ """ import json +import base64 import logging from typing import List, Dict, Optional, Any, Generator from functools import lru_cache @@ -404,7 +405,17 @@ def generate_images(model: str, prompt: str, n: int = 1) -> tuple[list[str], int continue url = item.get("url") if url: - images.append(url) + # Fetch the CDN URL and convert to a data: URI so the SDK always + # receives inline base64 rather than an expiring pre-signed link. + try: + img_resp = httpx.get(url, timeout=60, follow_redirects=True) + img_resp.raise_for_status() + mime = img_resp.headers.get("content-type", "image/jpeg").split(";")[0].strip() + b64_data = base64.b64encode(img_resp.content).decode("ascii") + images.append(f"data:{mime};base64,{b64_data}") + except Exception: + logger.warning("Failed to fetch image URL %s, returning URL as-is", url) + images.append(url) return images, len(images) diff --git a/tee_gateway/test/test_image_generation.py b/tee_gateway/test/test_image_generation.py index 2dc5353..6e2e50a 100644 --- a/tee_gateway/test/test_image_generation.py +++ b/tee_gateway/test/test_image_generation.py @@ -13,6 +13,7 @@ price feed is injected. """ +import base64 import unittest from decimal import Decimal from unittest.mock import MagicMock, patch @@ -64,23 +65,38 @@ def test_b64_json_becomes_data_uri(self): self.assertEqual(payload["n"], 2) self.assertEqual(payload["response_format"], "b64_json") - def test_url_fallback_when_no_b64(self): + def _mock_httpx_get(self, content: bytes = b"\xff\xd8\xff", content_type: str = "image/jpeg"): + """Return a mock for httpx.get that yields fake image bytes.""" + img_resp = MagicMock() + img_resp.raise_for_status.return_value = None + img_resp.content = content + img_resp.headers = {"content-type": content_type} + return img_resp + + def test_url_fallback_fetches_and_converts_to_data_uri(self): client = MagicMock() client.post.return_value = _mock_response([{"url": "https://img/1.jpg"}]) - with patch.object(llm_backend, "bytedance_http_client", client): + fake_bytes = b"\xff\xd8\xff" + with patch.object(llm_backend, "bytedance_http_client", client), \ + patch("httpx.get", return_value=self._mock_httpx_get(fake_bytes)) as mock_get: images, count = generate_images(SEEDREAM, "a blue sphere", n=1) + mock_get.assert_called_once_with("https://img/1.jpg", timeout=60, follow_redirects=True) self.assertEqual(count, 1) - self.assertEqual(images, ["https://img/1.jpg"]) + expected = f"data:image/jpeg;base64,{base64.b64encode(fake_bytes).decode()}" + self.assertEqual(images, [expected]) def test_zai_glm_image_uses_documented_payload_and_url_response(self): client = MagicMock() client.post.return_value = _mock_response([{"url": "https://z.ai/img.png"}]) - with patch.object(llm_backend, "zai_http_client", client): + fake_bytes = b"\x89PNG" + with patch.object(llm_backend, "zai_http_client", client), \ + patch("httpx.get", return_value=self._mock_httpx_get(fake_bytes, "image/png")): images, count = generate_images(GLM_IMAGE, "a poster", n=3) self.assertEqual(count, 1) - self.assertEqual(images, ["https://z.ai/img.png"]) + expected = f"data:image/png;base64,{base64.b64encode(fake_bytes).decode()}" + self.assertEqual(images, [expected]) _, kwargs = client.post.call_args payload = kwargs["json"] @@ -93,11 +109,14 @@ def test_zai_glm_image_uses_documented_payload_and_url_response(self): def test_seedance_uses_url_format_and_extra_params(self): client = MagicMock() client.post.return_value = _mock_response([{"url": "https://cdn/img.jpg"}]) - with patch.object(llm_backend, "bytedance_http_client", client): + fake_bytes = b"\xff\xd8\xff" + with patch.object(llm_backend, "bytedance_http_client", client), \ + patch("httpx.get", return_value=self._mock_httpx_get(fake_bytes)): images, count = generate_images(SEEDANCE, "a black hole", n=1) self.assertEqual(count, 1) - self.assertEqual(images, ["https://cdn/img.jpg"]) + expected = f"data:image/jpeg;base64,{base64.b64encode(fake_bytes).decode()}" + self.assertEqual(images, [expected]) _, kwargs = client.post.call_args payload = kwargs["json"] From c404fc3d420078564ec773c4ba3e568f5cd27278 Mon Sep 17 00:00:00 2001 From: Claude Date: Tue, 23 Jun 2026 21:07:38 +0000 Subject: [PATCH 2/3] Add pytest as dev dependency Co-Authored-By: Claude Sonnet 4.6 Claude-Session: https://claude.ai/code/session_01XxsSjnuypVUhpDPiEDA1Md --- pyproject.toml | 1 + uv.lock | 4 +++- 2 files changed, 4 insertions(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index ac9452d..198587e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -40,6 +40,7 @@ dependencies = [ [dependency-groups] dev = [ "mypy>=1.13.0", + "pytest>=9.0.2", "ruff>=0.9.1", ] test = [ diff --git a/uv.lock b/uv.lock index 91e5521..a749268 100644 --- a/uv.lock +++ b/uv.lock @@ -1,5 +1,5 @@ version = 1 -revision = 2 +revision = 3 requires-python = "==3.12.*" [options] @@ -1875,6 +1875,7 @@ dependencies = [ [package.dev-dependencies] dev = [ { name = "mypy" }, + { name = "pytest" }, { name = "ruff" }, ] test = [ @@ -1922,6 +1923,7 @@ requires-dist = [ [package.metadata.requires-dev] dev = [ { name = "mypy", specifier = ">=1.13.0" }, + { name = "pytest", specifier = ">=9.0.2" }, { name = "ruff", specifier = ">=0.9.1" }, ] test = [ From 659e00b7c654fba072be60016f9bff0d1d65f40c Mon Sep 17 00:00:00 2001 From: Claude Date: Tue, 23 Jun 2026 21:10:18 +0000 Subject: [PATCH 3/3] =?UTF-8?q?Revert=20gateway=20URL=20fetch=20=E2=80=94?= =?UTF-8?q?=20let=20SDK=20client=20resolve=20image=20URLs?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Co-Authored-By: Claude Sonnet 4.6 Claude-Session: https://claude.ai/code/session_01XxsSjnuypVUhpDPiEDA1Md --- tee_gateway/llm_backend.py | 13 +-------- tee_gateway/test/test_image_generation.py | 33 +++++------------------ 2 files changed, 8 insertions(+), 38 deletions(-) diff --git a/tee_gateway/llm_backend.py b/tee_gateway/llm_backend.py index 8d35207..3a34e84 100644 --- a/tee_gateway/llm_backend.py +++ b/tee_gateway/llm_backend.py @@ -7,7 +7,6 @@ """ import json -import base64 import logging from typing import List, Dict, Optional, Any, Generator from functools import lru_cache @@ -405,17 +404,7 @@ def generate_images(model: str, prompt: str, n: int = 1) -> tuple[list[str], int continue url = item.get("url") if url: - # Fetch the CDN URL and convert to a data: URI so the SDK always - # receives inline base64 rather than an expiring pre-signed link. - try: - img_resp = httpx.get(url, timeout=60, follow_redirects=True) - img_resp.raise_for_status() - mime = img_resp.headers.get("content-type", "image/jpeg").split(";")[0].strip() - b64_data = base64.b64encode(img_resp.content).decode("ascii") - images.append(f"data:{mime};base64,{b64_data}") - except Exception: - logger.warning("Failed to fetch image URL %s, returning URL as-is", url) - images.append(url) + images.append(url) return images, len(images) diff --git a/tee_gateway/test/test_image_generation.py b/tee_gateway/test/test_image_generation.py index 6e2e50a..2dc5353 100644 --- a/tee_gateway/test/test_image_generation.py +++ b/tee_gateway/test/test_image_generation.py @@ -13,7 +13,6 @@ price feed is injected. """ -import base64 import unittest from decimal import Decimal from unittest.mock import MagicMock, patch @@ -65,38 +64,23 @@ def test_b64_json_becomes_data_uri(self): self.assertEqual(payload["n"], 2) self.assertEqual(payload["response_format"], "b64_json") - def _mock_httpx_get(self, content: bytes = b"\xff\xd8\xff", content_type: str = "image/jpeg"): - """Return a mock for httpx.get that yields fake image bytes.""" - img_resp = MagicMock() - img_resp.raise_for_status.return_value = None - img_resp.content = content - img_resp.headers = {"content-type": content_type} - return img_resp - - def test_url_fallback_fetches_and_converts_to_data_uri(self): + def test_url_fallback_when_no_b64(self): client = MagicMock() client.post.return_value = _mock_response([{"url": "https://img/1.jpg"}]) - fake_bytes = b"\xff\xd8\xff" - with patch.object(llm_backend, "bytedance_http_client", client), \ - patch("httpx.get", return_value=self._mock_httpx_get(fake_bytes)) as mock_get: + with patch.object(llm_backend, "bytedance_http_client", client): images, count = generate_images(SEEDREAM, "a blue sphere", n=1) - mock_get.assert_called_once_with("https://img/1.jpg", timeout=60, follow_redirects=True) self.assertEqual(count, 1) - expected = f"data:image/jpeg;base64,{base64.b64encode(fake_bytes).decode()}" - self.assertEqual(images, [expected]) + self.assertEqual(images, ["https://img/1.jpg"]) def test_zai_glm_image_uses_documented_payload_and_url_response(self): client = MagicMock() client.post.return_value = _mock_response([{"url": "https://z.ai/img.png"}]) - fake_bytes = b"\x89PNG" - with patch.object(llm_backend, "zai_http_client", client), \ - patch("httpx.get", return_value=self._mock_httpx_get(fake_bytes, "image/png")): + with patch.object(llm_backend, "zai_http_client", client): images, count = generate_images(GLM_IMAGE, "a poster", n=3) self.assertEqual(count, 1) - expected = f"data:image/png;base64,{base64.b64encode(fake_bytes).decode()}" - self.assertEqual(images, [expected]) + self.assertEqual(images, ["https://z.ai/img.png"]) _, kwargs = client.post.call_args payload = kwargs["json"] @@ -109,14 +93,11 @@ def test_zai_glm_image_uses_documented_payload_and_url_response(self): def test_seedance_uses_url_format_and_extra_params(self): client = MagicMock() client.post.return_value = _mock_response([{"url": "https://cdn/img.jpg"}]) - fake_bytes = b"\xff\xd8\xff" - with patch.object(llm_backend, "bytedance_http_client", client), \ - patch("httpx.get", return_value=self._mock_httpx_get(fake_bytes)): + with patch.object(llm_backend, "bytedance_http_client", client): images, count = generate_images(SEEDANCE, "a black hole", n=1) self.assertEqual(count, 1) - expected = f"data:image/jpeg;base64,{base64.b64encode(fake_bytes).decode()}" - self.assertEqual(images, [expected]) + self.assertEqual(images, ["https://cdn/img.jpg"]) _, kwargs = client.post.call_args payload = kwargs["json"]