Add new image models

examples/llm_chat.py

@@ -34,4 +34,5 @@ async def main():
     else:
         print("No settlement tx hash returned")
 
+
 asyncio.run(main())

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())

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,

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 ───────────────────────────────────────────
 

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:

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",

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():