Skip to content

Commit 646c38b

Browse files
kylexqianclaudebalogh.adam@icloud.com
authored
Resolve TEE endpoint and TLS cert from on-chain registry (#175)
* feat: resolve TEE endpoint and TLS cert from on-chain registry Instead of blindly trusting the TLS certificate presented by the TEE server (TOFU), the SDK now queries the on-chain TEERegistry contract to discover active LLM proxy endpoints and their verified certificates. Key changes: - Add TEERegistry.abi and tee_registry.py to query the registry contract - Replace TOFU cert fetch in llm.py with registry-verified DER cert - Client.init queries the registry by default; og_llm_server_url still works as an explicit override (falls back to system CA verification) - Add DEFAULT_TEE_REGISTRY_ADDRESS and DEFAULT_TEE_REGISTRY_RPC_URL to defaults.py (OG EVM chain at http://13.59.43.94:8545) - Surface tee_id, tee_endpoint, tee_payment_address on every TextGenerationOutput and on the final StreamChunk so callers can audit which enclave served their request - Print TEE node info (endpoint, TEE ID, payment address) in all three CLI print helpers (completion, chat, streaming) Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> * fix: fall back to non-streaming endpoint when tool calls are requested with stream=True The TEE streaming endpoint returns an empty delta ("delta": {}) and no tool call content in SSE events when tools is supplied — the server-side streaming path simply does not emit tool call data. Introduce _tee_llm_chat_tools_as_stream which transparently calls the non-streaming /v1/chat/completions endpoint and wraps the complete TextGenerationOutput as a single final StreamChunk (with tool_calls populated in delta). chat() now routes stream=True + tools to this method, preserving the streaming iterator interface for callers and the CLI while returning correct results. Also removes the temporary [SSE RAW] debug print added during diagnosis, and fixes from_sse_data to accept "message" as a fallback for "delta" when the proxy sends a non-streaming format in SSE events. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> * fix(tests): update client_test fixtures to mock TEERegistry and fix TEE_LLM references Two issues caused CI failures: 1. The mock_web3 fixture didn't patch TEERegistry, so Client.__init__ tried to instantiate a real TEERegistry (with a live Web3 connection) even in unit tests. The mock_abi_files fallback returned {} for TEERegistry.abi but web3.eth.contract() requires a list, causing ValueError. Fix: patch src.opengradient.client.client.TEERegistry inside mock_web3 and return a fake TEEEndpoint with a stub endpoint/tee_id/payment_address. 2. Three LLM tests referenced TEE_LLM.GPT_4O which no longer exists in the enum. Updated to TEE_LLM.GPT_5. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> * fix(cli): guard against empty choices in streaming loop and format tool calls consistently Two bugs caused tool call results to be invisible in the CLI: 1. print_streaming_chat_result accessed chunk.choices[0] unconditionally. Usage-only SSE chunks carry an empty choices list, which caused an IndexError that was silently swallowed by the outer except block, truncating output before tool calls or finish_reason were printed. Fixed by guarding all choices[0] accesses with `if chunk.choices:`. 2. print_llm_chat_result (non-streaming) printed tool_calls as a raw Python dict repr. Updated to use the same formatted output as the streaming path: "Tool Calls: / Function: ... / Arguments: ...". Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> * Update pyproject.toml * fix(makefile): add system prompt to chat-tool targets to reliably trigger tool calls With tool_choice="auto", models like GPT-5 require a system prompt instructing them to use tools, otherwise they respond with finish_reason "stop" and empty content. Updated chat-tool and chat-stream-tool to include a system message matching the pattern that works in the SDK. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> * fix(makefile): use proven Dallas/Texas payload to reliably trigger GPT-5 tool calls The previous Tokyo/get_weather(location) payload still failed with finish_reason: stop on GPT-5. Switch chat-tool and chat-stream-tool to use get_current_weather(city, state, unit) with Dallas, Texas which consistently triggers tool_calls finish_reason. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> * in memory cert + registry test * use new abi * fix types * rm unused * rm tee rpc and update contract * rm _og_llm_streaming_server_url --------- Co-authored-by: Claude Sonnet 4.6 <noreply@anthropic.com> Co-authored-by: balogh.adam@icloud.com <adambalogh@mac.mynetworksettings.com>
1 parent 19adce8 commit 646c38b

11 files changed

Lines changed: 697 additions & 103 deletions

File tree

‎Makefile‎

Lines changed: 10 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -31,7 +31,7 @@ docs:
3131
# Testing
3232
# ============================================================================
3333

34-
test: utils_test client_test langchain_adapter_test opg_token_test
34+
test: utils_test client_test langchain_adapter_test opg_token_test tee_registry_test
3535

3636
utils_test:
3737
pytest tests/utils_test.py -v
@@ -45,6 +45,9 @@ langchain_adapter_test:
4545
opg_token_test:
4646
pytest tests/opg_token_test.py -v
4747

48+
tee_registry_test:
49+
pytest tests/tee_registry_test.py -v
50+
4851
integrationtest:
4952
python integrationtest/agent/test_agent.py
5053
python integrationtest/workflow_models/test_workflow_models.py
@@ -87,16 +90,16 @@ chat-stream:
8790
chat-tool:
8891
python -m opengradient.cli chat \
8992
--model $(MODEL) \
90-
--messages '[{"role":"user","content":"What is the weather in Tokyo?"}]' \
91-
--tools '[{"type":"function","function":{"name":"get_weather","description":"Get weather for a location","parameters":{"type":"object","properties":{"location":{"type":"string"},"unit":{"type":"string","enum":["celsius","fahrenheit"]}},"required":["location"]}}}]' \
92-
--max-tokens 100
93+
--messages '[{"role":"system","content":"You are a helpful assistant. Use tools when needed."},{"role":"user","content":"What'\''s the weather like in Dallas, Texas? Give me the temperature in fahrenheit."}]' \
94+
--tools '[{"type":"function","function":{"name":"get_current_weather","description":"Get the current weather in a given location","parameters":{"type":"object","properties":{"city":{"type":"string"},"state":{"type":"string"},"unit":{"type":"string","enum":["fahrenheit","celsius"]}},"required":["city","state","unit"]}}}]' \
95+
--max-tokens 200
9396

9497
chat-stream-tool:
9598
python -m opengradient.cli chat \
9699
--model $(MODEL) \
97-
--messages '[{"role":"user","content":"What is the weather in Tokyo?"}]' \
98-
--tools '[{"type":"function","function":{"name":"get_weather","description":"Get weather for a location","parameters":{"type":"object","properties":{"location":{"type":"string"},"unit":{"type":"string","enum":["celsius","fahrenheit"]}},"required":["location"]}}}]' \
99-
--max-tokens 100 \
100+
--messages '[{"role":"system","content":"You are a helpful assistant. Use tools when needed."},{"role":"user","content":"What'\''s the weather like in Dallas, Texas? Give me the temperature in fahrenheit."}]' \
101+
--tools '[{"type":"function","function":{"name":"get_current_weather","description":"Get the current weather in a given location","parameters":{"type":"object","properties":{"city":{"type":"string"},"state":{"type":"string"},"unit":{"type":"string","enum":["fahrenheit","celsius"]}},"required":["city","state","unit"]}}}]' \
102+
--max-tokens 200 \
100103
--stream
101104

102105
.PHONY: install build publish check docs test utils_test client_test langchain_adapter_test opg_token_test integrationtest examples \

‎pyproject.toml‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
44

55
[project]
66
name = "opengradient"
7-
version = "0.7.5"
7+
version = "0.7.6"
88
description = "Python SDK for OpenGradient decentralized model management & inference services"
99
authors = [{name = "OpenGradient", email = "adam@vannalabs.ai"}]
1010
readme = "README.md"

‎src/opengradient/abi/TEERegistry.abi‎

Lines changed: 80 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,80 @@
1+
[
2+
{
3+
"inputs": [{"internalType": "uint8", "name": "teeType", "type": "uint8"}],
4+
"name": "getActiveTEEs",
5+
"outputs": [
6+
{
7+
"components": [
8+
{"internalType": "address", "name": "owner", "type": "address"},
9+
{"internalType": "address", "name": "paymentAddress", "type": "address"},
10+
{"internalType": "string", "name": "endpoint", "type": "string"},
11+
{"internalType": "bytes", "name": "publicKey", "type": "bytes"},
12+
{"internalType": "bytes", "name": "tlsCertificate", "type": "bytes"},
13+
{"internalType": "bytes32", "name": "pcrHash", "type": "bytes32"},
14+
{"internalType": "uint8", "name": "teeType", "type": "uint8"},
15+
{"internalType": "bool", "name": "enabled", "type": "bool"},
16+
{"internalType": "uint256", "name": "registeredAt", "type": "uint256"},
17+
{"internalType": "uint256", "name": "lastHeartbeatAt", "type": "uint256"}
18+
],
19+
"internalType": "struct TEERegistry.TEEInfo[]",
20+
"name": "",
21+
"type": "tuple[]"
22+
}
23+
],
24+
"stateMutability": "view",
25+
"type": "function"
26+
},
27+
{
28+
"inputs": [{"internalType": "uint8", "name": "teeType", "type": "uint8"}],
29+
"name": "getEnabledTEEs",
30+
"outputs": [{"internalType": "bytes32[]", "name": "", "type": "bytes32[]"}],
31+
"stateMutability": "view",
32+
"type": "function"
33+
},
34+
{
35+
"inputs": [{"internalType": "uint8", "name": "teeType", "type": "uint8"}],
36+
"name": "getTEEsByType",
37+
"outputs": [{"internalType": "bytes32[]", "name": "", "type": "bytes32[]"}],
38+
"stateMutability": "view",
39+
"type": "function"
40+
},
41+
{
42+
"inputs": [{"internalType": "bytes32", "name": "teeId", "type": "bytes32"}],
43+
"name": "getTEE",
44+
"outputs": [
45+
{
46+
"components": [
47+
{"internalType": "address", "name": "owner", "type": "address"},
48+
{"internalType": "address", "name": "paymentAddress", "type": "address"},
49+
{"internalType": "string", "name": "endpoint", "type": "string"},
50+
{"internalType": "bytes", "name": "publicKey", "type": "bytes"},
51+
{"internalType": "bytes", "name": "tlsCertificate", "type": "bytes"},
52+
{"internalType": "bytes32", "name": "pcrHash", "type": "bytes32"},
53+
{"internalType": "uint8", "name": "teeType", "type": "uint8"},
54+
{"internalType": "bool", "name": "enabled", "type": "bool"},
55+
{"internalType": "uint256", "name": "registeredAt", "type": "uint256"},
56+
{"internalType": "uint256", "name": "lastHeartbeatAt", "type": "uint256"}
57+
],
58+
"internalType": "struct TEERegistry.TEEInfo",
59+
"name": "",
60+
"type": "tuple"
61+
}
62+
],
63+
"stateMutability": "view",
64+
"type": "function"
65+
},
66+
{
67+
"inputs": [{"internalType": "bytes32", "name": "teeId", "type": "bytes32"}],
68+
"name": "isTEEActive",
69+
"outputs": [{"internalType": "bool", "name": "", "type": "bool"}],
70+
"stateMutability": "view",
71+
"type": "function"
72+
},
73+
{
74+
"inputs": [{"internalType": "bytes32", "name": "teeId", "type": "bytes32"}],
75+
"name": "isTEEEnabled",
76+
"outputs": [{"internalType": "bool", "name": "", "type": "bool"}],
77+
"stateMutability": "view",
78+
"type": "function"
79+
}
80+
]

‎src/opengradient/cli.py‎

Lines changed: 58 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -413,13 +413,31 @@ def completion(
413413
x402_settlement_mode=x402SettlementModes[x402_settlement_mode],
414414
)
415415

416-
print_llm_completion_result(model_cid, completion_output.transaction_hash, completion_output.completion_output, is_vanilla=False)
416+
print_llm_completion_result(
417+
model_cid, completion_output.transaction_hash, completion_output.completion_output, is_vanilla=False, result=completion_output
418+
)
417419

418420
except Exception as e:
419421
click.echo(f"Error running LLM completion: {str(e)}")
420422

421423

422-
def print_llm_completion_result(model_cid, tx_hash, llm_output, is_vanilla=True):
424+
def _print_tee_info(tee_id, tee_endpoint, tee_payment_address):
425+
"""Print TEE node info if available."""
426+
if not any([tee_id, tee_endpoint, tee_payment_address]):
427+
return
428+
click.secho("TEE Node:", fg="magenta", bold=True)
429+
if tee_endpoint:
430+
click.echo(" Endpoint: ", nl=False)
431+
click.secho(tee_endpoint, fg="magenta")
432+
if tee_id:
433+
click.echo(" TEE ID: ", nl=False)
434+
click.secho(tee_id, fg="magenta")
435+
if tee_payment_address:
436+
click.echo(" Payment address: ", nl=False)
437+
click.secho(tee_payment_address, fg="magenta")
438+
439+
440+
def print_llm_completion_result(model_cid, tx_hash, llm_output, is_vanilla=True, result=None):
423441
click.secho("✅ LLM completion Successful", fg="green", bold=True)
424442
click.echo("──────────────────────────────────────")
425443
click.echo("Model: ", nl=False)
@@ -435,6 +453,9 @@ def print_llm_completion_result(model_cid, tx_hash, llm_output, is_vanilla=True)
435453
click.echo("Source: ", nl=False)
436454
click.secho("OpenGradient TEE", fg="cyan", bold=True)
437455

456+
if result is not None:
457+
_print_tee_info(result.tee_id, result.tee_endpoint, result.tee_payment_address)
458+
438459
click.echo("──────────────────────────────────────")
439460
click.secho("LLM Output:", fg="yellow", bold=True)
440461
click.echo()
@@ -578,13 +599,15 @@ def chat(
578599
if stream:
579600
print_streaming_chat_result(model_cid, result, is_tee=True)
580601
else:
581-
print_llm_chat_result(model_cid, result.transaction_hash, result.finish_reason, result.chat_output, is_vanilla=False)
602+
print_llm_chat_result(
603+
model_cid, result.transaction_hash, result.finish_reason, result.chat_output, is_vanilla=False, result=result
604+
)
582605

583606
except Exception as e:
584607
click.echo(f"Error running LLM chat inference: {str(e)}")
585608

586609

587-
def print_llm_chat_result(model_cid, tx_hash, finish_reason, chat_output, is_vanilla=True):
610+
def print_llm_chat_result(model_cid, tx_hash, finish_reason, chat_output, is_vanilla=True, result=None):
588611
click.secho("✅ LLM Chat Successful", fg="green", bold=True)
589612
click.echo("──────────────────────────────────────")
590613
click.echo("Model: ", nl=False)
@@ -600,6 +623,9 @@ def print_llm_chat_result(model_cid, tx_hash, finish_reason, chat_output, is_van
600623
click.echo("Source: ", nl=False)
601624
click.secho("OpenGradient TEE", fg="cyan", bold=True)
602625

626+
if result is not None:
627+
_print_tee_info(result.tee_id, result.tee_endpoint, result.tee_payment_address)
628+
603629
click.echo("──────────────────────────────────────")
604630
click.secho("Finish Reason: ", fg="yellow", bold=True)
605631
click.echo()
@@ -608,7 +634,16 @@ def print_llm_chat_result(model_cid, tx_hash, finish_reason, chat_output, is_van
608634
click.secho("Chat Output:", fg="yellow", bold=True)
609635
click.echo()
610636
for key, value in chat_output.items():
611-
if value is not None and value not in ("", "[]", []):
637+
if value is None or value in ("", "[]", []):
638+
continue
639+
if key == "tool_calls":
640+
# Format tool calls the same way as the streaming path
641+
click.secho("Tool Calls:", fg="yellow", bold=True)
642+
for tool_call in value:
643+
fn = tool_call.get("function", {})
644+
click.echo(f" Function: {fn.get('name', '')}")
645+
click.echo(f" Arguments: {fn.get('arguments', '')}")
646+
elif key == "content" and isinstance(value, list):
612647
# Normalize list-of-blocks content (e.g. Gemini 3 thought signatures)
613648
if key == "content" and isinstance(value, list):
614649
text = " ".join(block.get("text", "") for block in value if isinstance(block, dict) and block.get("type") == "text").strip()
@@ -638,20 +673,21 @@ def print_streaming_chat_result(model_cid, stream, is_tee=True):
638673
for chunk in stream:
639674
chunk_count += 1
640675

641-
if chunk.choices[0].delta.content:
642-
content = chunk.choices[0].delta.content
643-
sys.stdout.write(content)
644-
sys.stdout.flush()
645-
content_parts.append(content)
646-
647-
# Handle tool calls
648-
if chunk.choices[0].delta.tool_calls:
649-
sys.stdout.write("\n")
650-
sys.stdout.flush()
651-
click.secho("Tool Calls:", fg="yellow", bold=True)
652-
for tool_call in chunk.choices[0].delta.tool_calls:
653-
click.echo(f" Function: {tool_call['function']['name']}")
654-
click.echo(f" Arguments: {tool_call['function']['arguments']}")
676+
if chunk.choices:
677+
if chunk.choices[0].delta.content:
678+
content = chunk.choices[0].delta.content
679+
sys.stdout.write(content)
680+
sys.stdout.flush()
681+
content_parts.append(content)
682+
683+
# Handle tool calls
684+
if chunk.choices[0].delta.tool_calls:
685+
sys.stdout.write("\n")
686+
sys.stdout.flush()
687+
click.secho("Tool Calls:", fg="yellow", bold=True)
688+
for tool_call in chunk.choices[0].delta.tool_calls:
689+
click.echo(f" Function: {tool_call['function']['name']}")
690+
click.echo(f" Arguments: {tool_call['function']['arguments']}")
655691

656692
# Print final info when stream completes
657693
if chunk.is_final:
@@ -666,10 +702,12 @@ def print_streaming_chat_result(model_cid, stream, is_tee=True):
666702
click.echo(f" Total tokens: {chunk.usage.total_tokens}")
667703
click.echo()
668704

669-
if chunk.choices[0].finish_reason:
705+
if chunk.choices and chunk.choices[0].finish_reason:
670706
click.echo("Finish reason: ", nl=False)
671707
click.secho(chunk.choices[0].finish_reason, fg="green")
672708

709+
_print_tee_info(chunk.tee_id, chunk.tee_endpoint, chunk.tee_payment_address)
710+
673711
click.echo("──────────────────────────────────────")
674712
click.echo(f"Chunks received: {chunk_count}")
675713
click.echo(f"Content length: {len(''.join(content_parts))} characters")

‎src/opengradient/client/client.py‎

Lines changed: 46 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -1,21 +1,24 @@
11
"""Main Client class that unifies all OpenGradient service namespaces."""
22

3+
import logging
34
from typing import Optional
45

56
from web3 import Web3
67

78
from ..defaults import (
89
DEFAULT_API_URL,
910
DEFAULT_INFERENCE_CONTRACT_ADDRESS,
10-
DEFAULT_OPENGRADIENT_LLM_SERVER_URL,
11-
DEFAULT_OPENGRADIENT_LLM_STREAMING_SERVER_URL,
1211
DEFAULT_RPC_URL,
12+
DEFAULT_TEE_REGISTRY_ADDRESS,
1313
)
1414
from .alpha import Alpha
1515
from .llm import LLM
1616
from .model_hub import ModelHub
17+
from .tee_registry import TEERegistry
1718
from .twins import Twins
1819

20+
logger = logging.getLogger(__name__)
21+
1922

2023
class Client:
2124
"""
@@ -62,8 +65,8 @@ def __init__(
6265
rpc_url: str = DEFAULT_RPC_URL,
6366
api_url: str = DEFAULT_API_URL,
6467
contract_address: str = DEFAULT_INFERENCE_CONTRACT_ADDRESS,
65-
og_llm_server_url: Optional[str] = DEFAULT_OPENGRADIENT_LLM_SERVER_URL,
66-
og_llm_streaming_server_url: Optional[str] = DEFAULT_OPENGRADIENT_LLM_STREAMING_SERVER_URL,
68+
og_llm_server_url: Optional[str] = None,
69+
tee_registry_address: str = DEFAULT_TEE_REGISTRY_ADDRESS,
6770
):
6871
"""
6972
Initialize the OpenGradient client.
@@ -74,6 +77,11 @@ def __init__(
7477
You can supply a separate ``alpha_private_key`` so each chain uses its own
7578
funded wallet. When omitted, ``private_key`` is used for both.
7679
80+
By default the LLM server endpoint and its TLS certificate are fetched from
81+
the on-chain TEE Registry, which stores certificates that were verified during
82+
enclave attestation. You can override the endpoint by passing
83+
``og_llm_server_url`` explicitly (the system CA bundle is used for that URL).
84+
7785
Args:
7886
private_key: Private key whose wallet holds **Base Sepolia OPG tokens**
7987
for x402 LLM payments.
@@ -86,8 +94,11 @@ def __init__(
8694
rpc_url: RPC URL for the OpenGradient Alpha Testnet.
8795
api_url: API URL for the OpenGradient API.
8896
contract_address: Inference contract address.
89-
og_llm_server_url: OpenGradient LLM server URL.
90-
og_llm_streaming_server_url: OpenGradient LLM streaming server URL.
97+
og_llm_server_url: Override the LLM server URL instead of using the
98+
registry-discovered endpoint. When set, the TLS certificate is
99+
validated against the system CA bundle rather than the registry.
100+
tee_registry_address: Address of the TEERegistry contract used to
101+
discover active LLM proxy endpoints and their verified TLS certs.
91102
"""
92103
blockchain = Web3(Web3.HTTPProvider(rpc_url))
93104
wallet_account = blockchain.eth.account.from_key(private_key)
@@ -102,14 +113,42 @@ def __init__(
102113
if email is not None:
103114
hub_user = ModelHub._login_to_hub(email, password)
104115

116+
# Resolve LLM server URL and TLS certificate.
117+
# If the caller provided explicit URLs, use those with standard CA verification.
118+
# Otherwise, discover the endpoint and registry-verified cert from the TEE Registry.
119+
llm_tls_cert_der: Optional[bytes] = None
120+
tee = None
121+
if og_llm_server_url is None:
122+
try:
123+
registry = TEERegistry(
124+
rpc_url=rpc_url,
125+
registry_address=tee_registry_address,
126+
)
127+
tee = registry.get_llm_tee()
128+
if tee is not None:
129+
og_llm_server_url = tee.endpoint
130+
llm_tls_cert_der = tee.tls_cert_der
131+
logger.info("Using TEE endpoint from registry: %s (teeId=%s)", tee.endpoint, tee.tee_id)
132+
else:
133+
raise ValueError("No active LLM proxy TEE found in the registry. Pass og_llm_server_url explicitly to override.")
134+
except ValueError:
135+
raise
136+
except Exception as e:
137+
raise RuntimeError(
138+
f"Failed to fetch LLM TEE endpoint from registry ({tee_registry_address} on {rpc_url}): {e}. "
139+
"Pass og_llm_server_url explicitly to override."
140+
) from e
141+
105142
# Create namespaces
106143
self.model_hub = ModelHub(hub_user=hub_user)
107144
self.wallet_address = wallet_account.address
108145

109146
self.llm = LLM(
110147
wallet_account=wallet_account,
111148
og_llm_server_url=og_llm_server_url,
112-
og_llm_streaming_server_url=og_llm_streaming_server_url,
149+
tls_cert_der=llm_tls_cert_der,
150+
tee_id=tee.tee_id if tee is not None else None,
151+
tee_payment_address=tee.payment_address if tee is not None else None,
113152
)
114153

115154
self.alpha = Alpha(

0 commit comments

Comments
 (0)