cleanup

Author: dixitaniket

examples/llm_chat.py

@@ -10,7 +10,7 @@
 
 async def main():
     llm = og.LLM(private_key=os.environ.get("OG_PRIVATE_KEY"))
-    llm.ensure_opg_approval(opg_amount=0.1)
+    llm.ensure_opg_approval(opg_amount=1)
 
     messages = [
         {"role": "user", "content": "What is the capital of France?"},

examples/llm_chat_streaming.py

@@ -6,7 +6,7 @@
 
 async def main():
     llm = og.LLM(private_key=os.environ.get("OG_PRIVATE_KEY"))
-    llm.ensure_opg_approval(opg_amount=0.1)
+    llm.ensure_opg_approval(opg_amount=1)
 
     messages = [
         {"role": "user", "content": "What is Python?"},

src/opengradient/client/llm.py

@@ -3,8 +3,9 @@
 import json
 import logging
 import ssl
+import threading
 from dataclasses import dataclass
-from typing import AsyncGenerator, Dict, List, Optional, Union
+from typing import AsyncGenerator, Awaitable, Callable, Dict, List, Optional, TypeVar, Union
 
 from eth_account import Account
 from eth_account.account import LocalAccount
@@ -19,6 +20,7 @@
 from .tee_registry import TEERegistry, build_ssl_context_from_der
 
 logger = logging.getLogger(__name__)
+T = TypeVar("T")
 
 DEFAULT_RPC_URL = "https://ogevmdevnet.opengradient.ai"
 DEFAULT_TEE_REGISTRY_ADDRESS = "0x4e72238852f3c918f4E4e57AeC9280dDB0c80248"
@@ -94,31 +96,76 @@ def __init__(
         llm_server_url: Optional[str] = None,
     ):
         self._wallet_account: LocalAccount = Account.from_key(private_key)
-
-        endpoint, tls_cert_der, tee_id, tee_payment_address = self._resolve_tee(
-            llm_server_url,
-            rpc_url,
-            tee_registry_address,
+        self._rpc_url = rpc_url
+        self._tee_registry_address = tee_registry_address
+        self._llm_server_url = llm_server_url
+        self._reset_lock = threading.Lock()
+
+        self._refresh_tee_config()
+        self._init_x402_stack()
+
+    def _refresh_tee_config(self) -> None:
+        """Resolve TEE metadata from the registry and update TLS config."""
+        endpoint, tls_cert_der, tee_id, payment_addr = self._resolve_tee(
+            self._llm_server_url, self._rpc_url, self._tee_registry_address,
         )
-
+        ssl_ctx = build_ssl_context_from_der(tls_cert_der) if tls_cert_der else None
         self._tee_id = tee_id
         self._tee_endpoint = endpoint
-        self._tee_payment_address = tee_payment_address
-
-        ssl_ctx = build_ssl_context_from_der(tls_cert_der) if tls_cert_der else None
-        # When connecting directly via llm_server_url, skip cert verification —
-        # self-hosted TEE servers commonly use self-signed certificates.
-        verify_ssl = llm_server_url is None
-        self._tls_verify: Union[ssl.SSLContext, bool] = ssl_ctx if ssl_ctx else verify_ssl
+        self._tee_payment_address = payment_addr
+        self._tls_verify: Union[ssl.SSLContext, bool] = ssl_ctx if ssl_ctx else (self._llm_server_url is None)
 
-        # x402 client and signer
+    def _init_x402_stack(self) -> None:
+        """Initialize x402 signer/client/http stack."""
         signer = EthAccountSigner(self._wallet_account)
         self._x402_client = x402Client()
         register_exact_evm_client(self._x402_client, signer, networks=[BASE_TESTNET_NETWORK])
         register_upto_evm_client(self._x402_client, signer, networks=[BASE_TESTNET_NETWORK])
-        # httpx.AsyncClient subclass - construction is sync, connections open lazily
         self._http_client = x402HttpxClient(self._x402_client, verify=self._tls_verify)
 
+    @staticmethod
+    def _has_ssl_cause(exc: BaseException) -> bool:
+        """Return true when the exception chain contains an SSL error."""
+        visited: set[int] = set()
+        current: Optional[BaseException] = exc
+        while current is not None and id(current) not in visited:
+            visited.add(id(current))
+            if isinstance(current, ssl.SSLError):
+                return True
+            current = current.__cause__ or current.__context__
+        return False
+
+    async def _refresh_tee_and_reset(self) -> None:
+        """Re-resolve TEE and rebuild the HTTP client with fresh TLS config."""
+        with self._reset_lock:
+            old_http_client = self._http_client
+            self._refresh_tee_config()
+            self._init_x402_stack()
+
+        try:
+            await old_http_client.aclose()
+        except Exception:
+            logger.debug("Failed to close previous HTTP client during TEE refresh.", exc_info=True)
+
+    async def _call_with_ssl_retry(
+        self,
+        operation_name: str,
+        call: Callable[[], Awaitable[T]],
+    ) -> T:
+        """Retry once with fresh TEE/TLS state when the failure is SSL-related."""
+        try:
+            return await call()
+        except Exception as exc:
+            if not self._has_ssl_cause(exc):
+                raise
+            logger.warning(
+                "SSL failure during %s; refreshing TEE and retrying once: %s",
+                operation_name,
+                exc,
+            )
+            await self._refresh_tee_and_reset()
+            return await call()
+
     # ── TEE resolution ──────────────────────────────────────────────────
 
     @staticmethod
@@ -248,7 +295,6 @@ async def completion(
             RuntimeError: If the inference fails.
         """
         model_id = model.split("/")[1]
-        headers = self._headers(x402_settlement_mode)
         payload: Dict = {
             "model": model_id,
             "prompt": prompt,
@@ -258,11 +304,11 @@ async def completion(
         if stop_sequence:
             payload["stop"] = stop_sequence
 
-        try:
+        async def _request() -> TextGenerationOutput:
             response = await self._http_client.post(
                 self._tee_endpoint + _COMPLETION_ENDPOINT,
                 json=payload,
-                headers=headers,
+                headers=self._headers(x402_settlement_mode),
                 timeout=_REQUEST_TIMEOUT,
             )
             response.raise_for_status()
@@ -275,6 +321,9 @@ async def completion(
                 tee_timestamp=result.get("tee_timestamp"),
                 **self._tee_metadata(),
             )
+
+        try:
+            return await self._call_with_ssl_retry("completion", _request)
         except RuntimeError:
             raise
         except Exception as e:
@@ -342,14 +391,13 @@ async def chat(
 
     async def _chat_request(self, params: _ChatParams, messages: List[Dict]) -> TextGenerationOutput:
         """Non-streaming chat request."""
-        headers = self._headers(params.x402_settlement_mode)
         payload = self._chat_payload(params, messages)
 
-        try:
+        async def _request() -> TextGenerationOutput:
             response = await self._http_client.post(
                 self._tee_endpoint + _CHAT_ENDPOINT,
                 json=payload,
-                headers=headers,
+                headers=self._headers(params.x402_settlement_mode),
                 timeout=_REQUEST_TIMEOUT,
             )
             response.raise_for_status()
@@ -375,6 +423,9 @@ async def _chat_request(self, params: _ChatParams, messages: List[Dict]) -> Text
                 tee_timestamp=result.get("tee_timestamp"),
                 **self._tee_metadata(),
             )
+
+        try:
+            return await self._call_with_ssl_retry("chat", _request)
         except RuntimeError:
             raise
         except Exception as e:
@@ -410,6 +461,29 @@ async def _chat_stream(self, params: _ChatParams, messages: List[Dict]) -> Async
         headers = self._headers(params.x402_settlement_mode)
         payload = self._chat_payload(params, messages, stream=True)
 
+        chunks_yielded = False
+        try:
+            async with self._http_client.stream(
+                "POST",
+                self._tee_endpoint + _CHAT_ENDPOINT,
+                json=payload,
+                headers=headers,
+                timeout=_REQUEST_TIMEOUT,
+            ) as response:
+                async for chunk in self._parse_sse_response(response):
+                    chunks_yielded = True
+                    yield chunk
+            return
+        except Exception as exc:
+            if chunks_yielded or not self._has_ssl_cause(exc):
+                raise
+            logger.warning(
+                "SSL failure during stream setup; refreshing TEE and retrying once: %s",
+                exc,
+            )
+
+        await self._refresh_tee_and_reset()
+        headers = self._headers(params.x402_settlement_mode)
         async with self._http_client.stream(
             "POST",
             self._tee_endpoint + _CHAT_ENDPOINT,