Resolve TEE endpoint and TLS cert from on-chain registry

Makefile

@@ -87,16 +87,16 @@ chat-stream:
 chat-tool:
 	python -m opengradient.cli chat \
 		--model $(MODEL) \
-		--messages '[{"role":"user","content":"What is the weather in Tokyo?"}]' \
-		--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"]}}}]' \
-		--max-tokens 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."}]' \
+		--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"]}}}]' \
+		--max-tokens 200
 
 chat-stream-tool:
 	python -m opengradient.cli chat \
 		--model $(MODEL) \
-		--messages '[{"role":"user","content":"What is the weather in Tokyo?"}]' \
-		--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"]}}}]' \
-		--max-tokens 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."}]' \
+		--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"]}}}]' \
+		--max-tokens 200 \
 		--stream
 
 .PHONY: install build publish check docs test utils_test client_test langchain_adapter_test opg_token_test integrationtest examples \

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "opengradient"
-version = "0.7.5"
+version = "0.7.6"
 description = "Python SDK for OpenGradient decentralized model management & inference services"
 authors = [{name = "OpenGradient", email = "adam@vannalabs.ai"}]
 readme = "README.md"

src/opengradient/abi/TEERegistry.abi

@@ -0,0 +1,48 @@
+[
+  {
+    "inputs": [],
+    "name": "getActiveTEEs",
+    "outputs": [{"internalType": "bytes32[]", "name": "", "type": "bytes32[]"}],
+    "stateMutability": "view",
+    "type": "function"
+  },
+  {
+    "inputs": [{"internalType": "uint8", "name": "teeType", "type": "uint8"}],
+    "name": "getTEEsByType",
+    "outputs": [{"internalType": "bytes32[]", "name": "", "type": "bytes32[]"}],
+    "stateMutability": "view",
+    "type": "function"
+  },
+  {
+    "inputs": [{"internalType": "bytes32", "name": "teeId", "type": "bytes32"}],
+    "name": "getTEE",
+    "outputs": [
+      {
+        "components": [
+          {"internalType": "address", "name": "owner", "type": "address"},
+          {"internalType": "address", "name": "paymentAddress", "type": "address"},
+          {"internalType": "string", "name": "endpoint", "type": "string"},
+          {"internalType": "bytes", "name": "publicKey", "type": "bytes"},
+          {"internalType": "bytes", "name": "tlsCertificate", "type": "bytes"},
+          {"internalType": "bytes32", "name": "pcrHash", "type": "bytes32"},
+          {"internalType": "uint8", "name": "teeType", "type": "uint8"},
+          {"internalType": "bool", "name": "active", "type": "bool"},
+          {"internalType": "uint256", "name": "registeredAt", "type": "uint256"},
+          {"internalType": "uint256", "name": "lastUpdatedAt", "type": "uint256"}
+        ],
+        "internalType": "struct TEERegistry.TEEInfo",
+        "name": "",
+        "type": "tuple"
+      }
+    ],
+    "stateMutability": "view",
+    "type": "function"
+  },
+  {
+    "inputs": [{"internalType": "bytes32", "name": "teeId", "type": "bytes32"}],
+    "name": "isActive",
+    "outputs": [{"internalType": "bool", "name": "", "type": "bool"}],
+    "stateMutability": "view",
+    "type": "function"
+  }
+]

src/opengradient/cli.py

@@ -413,13 +413,29 @@ def completion(
             x402_settlement_mode=x402SettlementModes[x402_settlement_mode],
         )
 
-        print_llm_completion_result(model_cid, completion_output.transaction_hash, completion_output.completion_output, is_vanilla=False)
+        print_llm_completion_result(model_cid, completion_output.transaction_hash, completion_output.completion_output, is_vanilla=False, result=completion_output)
 
     except Exception as e:
         click.echo(f"Error running LLM completion: {str(e)}")
 
 
-def print_llm_completion_result(model_cid, tx_hash, llm_output, is_vanilla=True):
+def _print_tee_info(tee_id, tee_endpoint, tee_payment_address):
+    """Print TEE node info if available."""
+    if not any([tee_id, tee_endpoint, tee_payment_address]):
+        return
+    click.secho("TEE Node:", fg="magenta", bold=True)
+    if tee_endpoint:
+        click.echo("  Endpoint:        ", nl=False)
+        click.secho(tee_endpoint, fg="magenta")
+    if tee_id:
+        click.echo("  TEE ID:          ", nl=False)
+        click.secho(tee_id, fg="magenta")
+    if tee_payment_address:
+        click.echo("  Payment address: ", nl=False)
+        click.secho(tee_payment_address, fg="magenta")
+
+
+def print_llm_completion_result(model_cid, tx_hash, llm_output, is_vanilla=True, result=None):
     click.secho("✅ LLM completion Successful", fg="green", bold=True)
     click.echo("──────────────────────────────────────")
     click.echo("Model: ", nl=False)
@@ -435,6 +451,9 @@ def print_llm_completion_result(model_cid, tx_hash, llm_output, is_vanilla=True)
         click.echo("Source: ", nl=False)
         click.secho("OpenGradient TEE", fg="cyan", bold=True)
 
+    if result is not None:
+        _print_tee_info(result.tee_id, result.tee_endpoint, result.tee_payment_address)
+
     click.echo("──────────────────────────────────────")
     click.secho("LLM Output:", fg="yellow", bold=True)
     click.echo()
@@ -578,13 +597,13 @@ def chat(
         if stream:
             print_streaming_chat_result(model_cid, result, is_tee=True)
         else:
-            print_llm_chat_result(model_cid, result.transaction_hash, result.finish_reason, result.chat_output, is_vanilla=False)
+            print_llm_chat_result(model_cid, result.transaction_hash, result.finish_reason, result.chat_output, is_vanilla=False, result=result)
 
     except Exception as e:
         click.echo(f"Error running LLM chat inference: {str(e)}")
 
 
-def print_llm_chat_result(model_cid, tx_hash, finish_reason, chat_output, is_vanilla=True):
+def print_llm_chat_result(model_cid, tx_hash, finish_reason, chat_output, is_vanilla=True, result=None):
     click.secho("✅ LLM Chat Successful", fg="green", bold=True)
     click.echo("──────────────────────────────────────")
     click.echo("Model: ", nl=False)
@@ -600,6 +619,9 @@ def print_llm_chat_result(model_cid, tx_hash, finish_reason, chat_output, is_van
         click.echo("Source: ", nl=False)
         click.secho("OpenGradient TEE", fg="cyan", bold=True)
 
+    if result is not None:
+        _print_tee_info(result.tee_id, result.tee_endpoint, result.tee_payment_address)
+
     click.echo("──────────────────────────────────────")
     click.secho("Finish Reason: ", fg="yellow", bold=True)
     click.echo()
@@ -608,16 +630,24 @@ def print_llm_chat_result(model_cid, tx_hash, finish_reason, chat_output, is_van
     click.secho("Chat Output:", fg="yellow", bold=True)
     click.echo()
     for key, value in chat_output.items():
-        if value is not None and value not in ("", "[]", []):
+        if value is None or value in ("", "[]", []):
+            continue
+        if key == "tool_calls":
+            # Format tool calls the same way as the streaming path
+            click.secho("Tool Calls:", fg="yellow", bold=True)
+            for tool_call in value:
+                fn = tool_call.get("function", {})
+                click.echo(f"  Function: {fn.get('name', '')}")
+                click.echo(f"  Arguments: {fn.get('arguments', '')}")
+        elif key == "content" and isinstance(value, list):
             # Normalize list-of-blocks content (e.g. Gemini 3 thought signatures)
-            if key == "content" and isinstance(value, list):
-                text = " ".join(
-                    block.get("text", "") for block in value
-                    if isinstance(block, dict) and block.get("type") == "text"
-                ).strip()
-                click.echo(f"{key}: {text}")
-            else:
-                click.echo(f"{key}: {value}")
+            text = " ".join(
+                block.get("text", "") for block in value
+                if isinstance(block, dict) and block.get("type") == "text"
+            ).strip()
+            click.echo(f"{key}: {text}")
+        else:
+            click.echo(f"{key}: {value}")
     click.echo()
 
 
@@ -641,20 +671,21 @@ def print_streaming_chat_result(model_cid, stream, is_tee=True):
         for chunk in stream:
             chunk_count += 1
 
-            if chunk.choices[0].delta.content:
-                content = chunk.choices[0].delta.content
-                sys.stdout.write(content)
-                sys.stdout.flush()
-                content_parts.append(content)
-
-            # Handle tool calls
-            if chunk.choices[0].delta.tool_calls:
-                sys.stdout.write("\n")
-                sys.stdout.flush()
-                click.secho("Tool Calls:", fg="yellow", bold=True)
-                for tool_call in chunk.choices[0].delta.tool_calls:
-                    click.echo(f"  Function: {tool_call['function']['name']}")
-                    click.echo(f"  Arguments: {tool_call['function']['arguments']}")
+            if chunk.choices:
+                if chunk.choices[0].delta.content:
+                    content = chunk.choices[0].delta.content
+                    sys.stdout.write(content)
+                    sys.stdout.flush()
+                    content_parts.append(content)
+
+                # Handle tool calls
+                if chunk.choices[0].delta.tool_calls:
+                    sys.stdout.write("\n")
+                    sys.stdout.flush()
+                    click.secho("Tool Calls:", fg="yellow", bold=True)
+                    for tool_call in chunk.choices[0].delta.tool_calls:
+                        click.echo(f"  Function: {tool_call['function']['name']}")
+                        click.echo(f"  Arguments: {tool_call['function']['arguments']}")
 
             # Print final info when stream completes
             if chunk.is_final:
@@ -669,10 +700,12 @@ def print_streaming_chat_result(model_cid, stream, is_tee=True):
                     click.echo(f"  Total tokens: {chunk.usage.total_tokens}")
                     click.echo()
 
-                if chunk.choices[0].finish_reason:
+                if chunk.choices and chunk.choices[0].finish_reason:
                     click.echo("Finish reason: ", nl=False)
                     click.secho(chunk.choices[0].finish_reason, fg="green")
 
+                _print_tee_info(chunk.tee_id, chunk.tee_endpoint, chunk.tee_payment_address)
+
                 click.echo("──────────────────────────────────────")
                 click.echo(f"Chunks received: {chunk_count}")
                 click.echo(f"Content length: {len(''.join(content_parts))} characters")

src/opengradient/client/client.py

@@ -1,21 +1,25 @@
 """Main Client class that unifies all OpenGradient service namespaces."""
 
+import logging
 from typing import Optional
 
 from web3 import Web3
 
 from ..defaults import (
     DEFAULT_API_URL,
     DEFAULT_INFERENCE_CONTRACT_ADDRESS,
-    DEFAULT_OPENGRADIENT_LLM_SERVER_URL,
-    DEFAULT_OPENGRADIENT_LLM_STREAMING_SERVER_URL,
     DEFAULT_RPC_URL,
+    DEFAULT_TEE_REGISTRY_ADDRESS,
+    DEFAULT_TEE_REGISTRY_RPC_URL,
 )
 from .alpha import Alpha
 from .llm import LLM
 from .model_hub import ModelHub
+from .tee_registry import TEERegistry
 from .twins import Twins
 
+logger = logging.getLogger(__name__)
+
 
 class Client:
     """
@@ -62,8 +66,10 @@ def __init__(
         rpc_url: str = DEFAULT_RPC_URL,
         api_url: str = DEFAULT_API_URL,
         contract_address: str = DEFAULT_INFERENCE_CONTRACT_ADDRESS,
-        og_llm_server_url: Optional[str] = DEFAULT_OPENGRADIENT_LLM_SERVER_URL,
-        og_llm_streaming_server_url: Optional[str] = DEFAULT_OPENGRADIENT_LLM_STREAMING_SERVER_URL,
+        og_llm_server_url: Optional[str] = None,
+        og_llm_streaming_server_url: Optional[str] = None,
+        tee_registry_address: str = DEFAULT_TEE_REGISTRY_ADDRESS,
+        tee_registry_rpc_url: str = DEFAULT_TEE_REGISTRY_RPC_URL,
     ):
         """
         Initialize the OpenGradient client.
@@ -74,6 +80,11 @@ def __init__(
         You can supply a separate ``alpha_private_key`` so each chain uses its own
         funded wallet. When omitted, ``private_key`` is used for both.
 
+        By default the LLM server endpoint and its TLS certificate are fetched from
+        the on-chain TEE Registry, which stores certificates that were verified during
+        enclave attestation.  You can override the endpoint by passing
+        ``og_llm_server_url`` explicitly (the system CA bundle is used for that URL).
+
         Args:
             private_key: Private key whose wallet holds **Base Sepolia OPG tokens**
                 for x402 LLM payments.
@@ -86,8 +97,15 @@ def __init__(
             rpc_url: RPC URL for the OpenGradient Alpha Testnet.
             api_url: API URL for the OpenGradient API.
             contract_address: Inference contract address.
-            og_llm_server_url: OpenGradient LLM server URL.
-            og_llm_streaming_server_url: OpenGradient LLM streaming server URL.
+            og_llm_server_url: Override the LLM server URL instead of using the
+                registry-discovered endpoint. When set, the TLS certificate is
+                validated against the system CA bundle rather than the registry.
+            og_llm_streaming_server_url: Override the LLM streaming server URL.
+                Defaults to ``og_llm_server_url`` when that is provided.
+            tee_registry_address: Address of the TEERegistry contract used to
+                discover active LLM proxy endpoints and their verified TLS certs.
+            tee_registry_rpc_url: RPC endpoint for the chain that hosts the
+                TEERegistry contract.
         """
         blockchain = Web3(Web3.HTTPProvider(rpc_url))
         wallet_account = blockchain.eth.account.from_key(private_key)
@@ -102,6 +120,38 @@ def __init__(
         if email is not None:
             hub_user = ModelHub._login_to_hub(email, password)
 
+        # Resolve LLM server URL and TLS certificate.
+        # If the caller provided explicit URLs, use those with standard CA verification.
+        # Otherwise, discover the endpoint and registry-verified cert from the TEE Registry.
+        llm_tls_cert_der: Optional[bytes] = None
+        tee = None
+        if og_llm_server_url is None:
+            try:
+                registry = TEERegistry(
+                    rpc_url=tee_registry_rpc_url,
+                    registry_address=tee_registry_address,
+                )
+                tee = registry.get_llm_tee()
+                if tee is not None:
+                    og_llm_server_url = tee.endpoint
+                    og_llm_streaming_server_url = og_llm_streaming_server_url or tee.endpoint
+                    llm_tls_cert_der = tee.tls_cert_der
+                    logger.info("Using TEE endpoint from registry: %s (teeId=%s)", tee.endpoint, tee.tee_id)
+                else:
+                    raise ValueError(
+                        "No active LLM proxy TEE found in the registry. "
+                        "Pass og_llm_server_url explicitly to override."
+                    )
+            except ValueError:
+                raise
+            except Exception as e:
+                raise RuntimeError(
+                    f"Failed to fetch LLM TEE endpoint from registry ({tee_registry_address}): {e}. "
+                    "Pass og_llm_server_url explicitly to override."
+                ) from e
+        else:
+            og_llm_streaming_server_url = og_llm_streaming_server_url or og_llm_server_url
+
         # Create namespaces
         self.model_hub = ModelHub(hub_user=hub_user)
         self.wallet_address = wallet_account.address
@@ -110,6 +160,9 @@ def __init__(
             wallet_account=wallet_account,
             og_llm_server_url=og_llm_server_url,
             og_llm_streaming_server_url=og_llm_streaming_server_url,
+            tls_cert_der=llm_tls_cert_der,
+            tee_id=tee.tee_id if tee is not None else None,
+            tee_payment_address=tee.payment_address if tee is not None else None,
         )
 
         self.alpha = Alpha(