From 9ad69b1b600bfe89c0c47713ff8e0bc44ab29ea9 Mon Sep 17 00:00:00 2001 From: Claude Date: Tue, 23 Jun 2026 20:41:32 +0000 Subject: [PATCH 1/4] Add image generation models to TEE_LLM enum Add GROK_2_IMAGE (xAI), SEEDREAM_4_0 and SEEDANCE_4_5 (ByteDance), and GLM_5_2 / GLM_IMAGE (Z.ai) to the TEE_LLM enum so SDK users can access all image generation models already supported by the tee-gateway. Also add GLM_5_2 as the Z.ai text model which was missing entirely. Update the llm_image_generation.py example with a quick reference listing all available image generation models and their billing model. Co-Authored-By: Claude Sonnet 4.6 Claude-Session: https://claude.ai/code/session_01XxsSjnuypVUhpDPiEDA1Md --- examples/llm_image_generation.py | 12 ++++++++++++ src/opengradient/types.py | 19 +++++++++++++++++++ 2 files changed, 31 insertions(+) diff --git a/examples/llm_image_generation.py b/examples/llm_image_generation.py index 547ed84..68fb378 100644 --- a/examples/llm_image_generation.py +++ b/examples/llm_image_generation.py @@ -32,6 +32,18 @@ async def main(): # The text caption (if any) is in chat_output["content"]; the generated images # are in result.images as data: URIs. Images travel out-of-band and are not part # of the signed output hash. + # + # Available image generation models: + # Google (inline, billed per output token): + # og.TEE_LLM.GEMINI_2_5_FLASH_IMAGE + # og.TEE_LLM.GEMINI_3_1_FLASH_IMAGE + # xAI (flat per-image rate, $0.07/image): + # og.TEE_LLM.GROK_2_IMAGE + # ByteDance (flat per-image rate): + # og.TEE_LLM.SEEDREAM_4_0 ($0.03/image) + # og.TEE_LLM.SEEDANCE_4_5 ($0.05/image) + # Z.ai (flat per-image rate, $0.015/image): + # og.TEE_LLM.GLM_IMAGE result = await llm.chat( model=og.TEE_LLM.GEMINI_3_1_FLASH_IMAGE, messages=messages, diff --git a/src/opengradient/types.py b/src/opengradient/types.py index f2c6259..66c6abe 100644 --- a/src/opengradient/types.py +++ b/src/opengradient/types.py @@ -593,6 +593,11 @@ class TEE_LLM(str, Enum): GROK_4_20_NON_REASONING = "x-ai/grok-4.20-non-reasoning" GROK_CODE_FAST_1 = "x-ai/grok-code-fast-1" + # xAI image-generation models via TEE (Aurora, dedicated /images/generations endpoint). + # Billed at a flat rate per image. Images are returned on ``TextGenerationOutput.images`` + # and ``StreamChunk.images`` as data: URIs and are not part of the signed output hash. + GROK_2_IMAGE = "x-ai/grok-2-image" + # ByteDance Seed models via TEE (BytePlus ModelArk) SEED_1_6 = "bytedance/seed-1.6" SEED_1_8 = "bytedance/seed-1.8" @@ -602,10 +607,24 @@ class TEE_LLM(str, Enum): DEEPSEEK_V4_FLASH = "bytedance/deepseek-v4-flash" DEEPSEEK_V4_PRO = "bytedance/deepseek-v4-pro" + # ByteDance image-generation models via TEE (ModelArk, dedicated /images/generations + # endpoint). Billed at a flat rate per image. Images are returned on + # ``TextGenerationOutput.images`` and ``StreamChunk.images`` as data: URIs. + SEEDREAM_4_0 = "bytedance/seedream-4.0" + SEEDANCE_4_5 = "bytedance/seedance-4.5" + # Nous Research Hermes models via TEE (Nous Portal) HERMES_4_405B = "nous/hermes-4-405b" HERMES_4_70B = "nous/hermes-4-70b" + # Z.ai GLM models via TEE (Model API, OpenAI-compatible) + GLM_5_2 = "zai/glm-5.2" + + # Z.ai image-generation model via TEE (dedicated /images/generations endpoint). + # Billed at a flat rate per image. Images are returned on + # ``TextGenerationOutput.images`` and ``StreamChunk.images`` as data: URIs. + GLM_IMAGE = "zai/glm-image" + @dataclass class ResponseFormat: From c6a649ed068ee4a0bd329287aa94a3679c54bd34 Mon Sep 17 00:00:00 2001 From: Claude Date: Tue, 23 Jun 2026 20:41:57 +0000 Subject: [PATCH 2/4] Fix minor whitespace and formatting issues Co-Authored-By: Claude Sonnet 4.6 Claude-Session: https://claude.ai/code/session_01XxsSjnuypVUhpDPiEDA1Md --- examples/llm_chat.py | 1 + examples/llm_chat_streaming.py | 5 +++-- src/opengradient/client/tee_connection.py | 9 ++------- tests/tee_ohttp_client_test.py | 10 +++++----- tests/tee_verify_test.py | 8 ++------ 5 files changed, 13 insertions(+), 20 deletions(-) diff --git a/examples/llm_chat.py b/examples/llm_chat.py index cbc8137..3204dc7 100644 --- a/examples/llm_chat.py +++ b/examples/llm_chat.py @@ -34,4 +34,5 @@ async def main(): else: print("No settlement tx hash returned") + asyncio.run(main()) diff --git a/examples/llm_chat_streaming.py b/examples/llm_chat_streaming.py index 43e3837..3614d2d 100644 --- a/examples/llm_chat_streaming.py +++ b/examples/llm_chat_streaming.py @@ -14,7 +14,7 @@ async def main(): {"role": "user", "content": "What makes it good for beginners?"}, ] - settlement_mode=og.x402SettlementMode.INDIVIDUAL_FULL + settlement_mode = og.x402SettlementMode.INDIVIDUAL_FULL stream = await llm.chat( model=og.TEE_LLM.GPT_4_1_2025_04_14, messages=messages, @@ -28,9 +28,10 @@ async def main(): print(chunk.choices[0].delta.content, end="", flush=True) if settlement_mode == og.x402SettlementMode.INDIVIDUAL_FULL: - if chunk.data_settlement_blob_id: + if chunk.data_settlement_blob_id: print("\n Data Settlement Blob ID: ", chunk.data_settlement_blob_id) if chunk.data_settlement_transaction_hash: print("\n Data Settlement Transaction Hash: ", chunk.data_settlement_transaction_hash) + asyncio.run(main()) diff --git a/src/opengradient/client/tee_connection.py b/src/opengradient/client/tee_connection.py index e24a6e0..d3470e9 100644 --- a/src/opengradient/client/tee_connection.py +++ b/src/opengradient/client/tee_connection.py @@ -170,10 +170,7 @@ def resolve(self, tee_id: Optional[str] = None) -> ActiveTEE: continue cached = self._active_by_tee_id.get(normalized_tee_id) - if ( - cached is not None - and cached.endpoint.rstrip("/") == tee.endpoint.rstrip("/") - ): + if cached is not None and cached.endpoint.rstrip("/") == tee.endpoint.rstrip("/"): return cached resolved = self._connect_to_tee(tee) @@ -185,9 +182,7 @@ def resolve(self, tee_id: Optional[str] = None) -> ActiveTEE: ) return resolved - raise ValueError( - f"Selected TEE is not active in the registry: {normalized_tee_id}" - ) + raise ValueError(f"Selected TEE is not active in the registry: {normalized_tee_id}") # ── Connection management ─────────────────────────────────────────── diff --git a/tests/tee_ohttp_client_test.py b/tests/tee_ohttp_client_test.py index 5cc90cb..cde6e43 100644 --- a/tests/tee_ohttp_client_test.py +++ b/tests/tee_ohttp_client_test.py @@ -47,12 +47,12 @@ def iter_content(self, chunk_size=8192): def _make_endpoint(recipient): hpke_priv, hpke_pub = recipient.generate_keypair() rsa_priv = rsa.generate_private_key(public_exponent=65537, key_size=2048) - der = rsa_priv.public_key().public_bytes( - encoding=serialization.Encoding.DER, format=serialization.PublicFormat.SubjectPublicKeyInfo + der = rsa_priv.public_key().public_bytes(encoding=serialization.Encoding.DER, format=serialization.PublicFormat.SubjectPublicKeyInfo) + pem = ( + rsa_priv.public_key() + .public_bytes(encoding=serialization.Encoding.PEM, format=serialization.PublicFormat.SubjectPublicKeyInfo) + .decode() ) - pem = rsa_priv.public_key().public_bytes( - encoding=serialization.Encoding.PEM, format=serialization.PublicFormat.SubjectPublicKeyInfo - ).decode() endpoint = TEEEndpoint( tee_id=tee_id_for_key(pem), endpoint="https://gw.example", diff --git a/tests/tee_verify_test.py b/tests/tee_verify_test.py index 611a51d..00ea4653 100644 --- a/tests/tee_verify_test.py +++ b/tests/tee_verify_test.py @@ -53,9 +53,7 @@ def _sign(priv, canonical, output_content, timestamp): def _good_case(): priv, pem = _make_key() tee_id = verify.tee_id_for_key(pem) - _wire, canonical = build_inner_request( - {"model": "gpt-4.1", "messages": [{"role": "user", "content": "Hello!"}]} - ) + _wire, canonical = build_inner_request({"model": "gpt-4.1", "messages": [{"role": "user", "content": "Hello!"}]}) content = "Hi there!" response = { "choices": [{"index": 0, "message": {"role": "assistant", "content": content}, "finish_reason": "stop"}], @@ -160,9 +158,7 @@ def test_build_inner_request_strips_attachment_bytes_from_hash_only(): def test_non_list_tools_rejected(): with pytest.raises(verify.UnsupportedRequestError, match="tools"): - build_inner_request( - {"model": "gpt-4.1", "messages": [{"role": "user", "content": "x"}], "tools": {"a": 1}} - ) + build_inner_request({"model": "gpt-4.1", "messages": [{"role": "user", "content": "x"}], "tools": {"a": 1}}) def test_non_dict_message_rejected(): From 0981b7edc23fc2ca38a1261f422d4ed994a6ef6b Mon Sep 17 00:00:00 2001 From: Claude Date: Tue, 23 Jun 2026 20:57:39 +0000 Subject: [PATCH 3/4] Remove broken SEEDREAM_4_0 from TEE_LLM enum MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit SEEDREAM_4_0 used the bare model name "seedream-4-0-250828" which ByteDance rejects — their image generation API only accepts deployment endpoint IDs (ep-xxx). SEEDANCE_4_5 already points to the working deployment endpoint ep-20260624042612-7dxcv, which the gateway correctly calls with the Seedance-specific payload (response_format=url, size=2K, sequential_image_generation=disabled, watermark=false, stream=false). Co-Authored-By: Claude Sonnet 4.6 Claude-Session: https://claude.ai/code/session_01XxsSjnuypVUhpDPiEDA1Md --- examples/llm_image_generation.py | 1 - src/opengradient/types.py | 5 ++--- 2 files changed, 2 insertions(+), 4 deletions(-) diff --git a/examples/llm_image_generation.py b/examples/llm_image_generation.py index 68fb378..3fb947a 100644 --- a/examples/llm_image_generation.py +++ b/examples/llm_image_generation.py @@ -40,7 +40,6 @@ async def main(): # xAI (flat per-image rate, $0.07/image): # og.TEE_LLM.GROK_2_IMAGE # ByteDance (flat per-image rate): - # og.TEE_LLM.SEEDREAM_4_0 ($0.03/image) # og.TEE_LLM.SEEDANCE_4_5 ($0.05/image) # Z.ai (flat per-image rate, $0.015/image): # og.TEE_LLM.GLM_IMAGE diff --git a/src/opengradient/types.py b/src/opengradient/types.py index 66c6abe..2e4cf03 100644 --- a/src/opengradient/types.py +++ b/src/opengradient/types.py @@ -607,10 +607,9 @@ class TEE_LLM(str, Enum): DEEPSEEK_V4_FLASH = "bytedance/deepseek-v4-flash" DEEPSEEK_V4_PRO = "bytedance/deepseek-v4-pro" - # ByteDance image-generation models via TEE (ModelArk, dedicated /images/generations - # endpoint). Billed at a flat rate per image. Images are returned on + # ByteDance Seedream image-generation via a ModelArk deployment endpoint. + # Billed at a flat rate per image. Images are returned on # ``TextGenerationOutput.images`` and ``StreamChunk.images`` as data: URIs. - SEEDREAM_4_0 = "bytedance/seedream-4.0" SEEDANCE_4_5 = "bytedance/seedance-4.5" # Nous Research Hermes models via TEE (Nous Portal) From 75e8b8ee73a41101854748956de70856776c2c25 Mon Sep 17 00:00:00 2001 From: Claude Date: Tue, 23 Jun 2026 20:58:08 +0000 Subject: [PATCH 4/4] Restore SEEDREAM_4_0 to TEE_LLM enum Co-Authored-By: Claude Sonnet 4.6 Claude-Session: https://claude.ai/code/session_01XxsSjnuypVUhpDPiEDA1Md --- examples/llm_image_generation.py | 1 + src/opengradient/types.py | 5 +++-- 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/examples/llm_image_generation.py b/examples/llm_image_generation.py index 3fb947a..68fb378 100644 --- a/examples/llm_image_generation.py +++ b/examples/llm_image_generation.py @@ -40,6 +40,7 @@ async def main(): # xAI (flat per-image rate, $0.07/image): # og.TEE_LLM.GROK_2_IMAGE # ByteDance (flat per-image rate): + # og.TEE_LLM.SEEDREAM_4_0 ($0.03/image) # og.TEE_LLM.SEEDANCE_4_5 ($0.05/image) # Z.ai (flat per-image rate, $0.015/image): # og.TEE_LLM.GLM_IMAGE diff --git a/src/opengradient/types.py b/src/opengradient/types.py index 2e4cf03..66c6abe 100644 --- a/src/opengradient/types.py +++ b/src/opengradient/types.py @@ -607,9 +607,10 @@ class TEE_LLM(str, Enum): DEEPSEEK_V4_FLASH = "bytedance/deepseek-v4-flash" DEEPSEEK_V4_PRO = "bytedance/deepseek-v4-pro" - # ByteDance Seedream image-generation via a ModelArk deployment endpoint. - # Billed at a flat rate per image. Images are returned on + # ByteDance image-generation models via TEE (ModelArk, dedicated /images/generations + # endpoint). Billed at a flat rate per image. Images are returned on # ``TextGenerationOutput.images`` and ``StreamChunk.images`` as data: URIs. + SEEDREAM_4_0 = "bytedance/seedream-4.0" SEEDANCE_4_5 = "bytedance/seedance-4.5" # Nous Research Hermes models via TEE (Nous Portal)