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