LLM client code refactor (#178)

* llm refactor

* refac

* async client

* format

* format

* simplify

* llm test

* format

* tests

* docs

* rm exception

* readmes

* tool

* fix ci test

* rm overlapping test

* cleanup

* docs

* c

---------

Co-authored-by: balogh.adam@icloud.com <adambalogh@mac.mynetworksettings.com>

Author: adambalogh

.github/workflows/test.yml

@@ -22,7 +22,7 @@ jobs:
         run: |
           python -m pip install --upgrade pip
           pip install -e .
-          pip install pytest
+          pip install pytest pytest-asyncio
 
       - name: Run tests
         run: make test

Makefile

@@ -31,22 +31,8 @@ docs:
 # Testing
 # ============================================================================
 
-test: utils_test client_test langchain_adapter_test opg_token_test tee_registry_test
-
-utils_test:
-	pytest tests/utils_test.py -v
-
-client_test:
-	pytest tests/client_test.py -v
-
-langchain_adapter_test:
-	pytest tests/langchain_adapter_test.py -v
-
-opg_token_test:
-	pytest tests/opg_token_test.py -v
-
-tee_registry_test:
-	pytest tests/tee_registry_test.py -v
+test:
+	pytest tests/ -v
 
 integrationtest:
 	python integrationtest/agent/test_agent.py
@@ -102,5 +88,5 @@ chat-stream-tool:
 		--max-tokens 200 \
 		--stream
 
-.PHONY: install build publish check docs test utils_test client_test langchain_adapter_test opg_token_test integrationtest examples \
+.PHONY: install build publish check docs test integrationtest examples \
 	infer completion chat chat-stream chat-tool chat-stream-tool

README.md

@@ -97,9 +97,9 @@ See [Payment Settlement](#payment-settlement) for details on settlement modes.
 
 ### TEE-Secured LLM Chat
 
-OpenGradient provides secure, verifiable inference through Trusted Execution Environments. All supported models include cryptographic attestation verified by the OpenGradient network:
+OpenGradient provides secure, verifiable inference through Trusted Execution Environments. All supported models include cryptographic attestation verified by the OpenGradient network. LLM methods are async:
 ```python
-completion = client.llm.chat(
+completion = await client.llm.chat(
     model=og.TEE_LLM.GPT_5,
     messages=[{"role": "user", "content": "Hello!"}],
 )
@@ -111,14 +111,14 @@ print(f"Transaction hash: {completion.transaction_hash}")
 
 For real-time generation, enable streaming:
 ```python
-stream = client.llm.chat(
+stream = await client.llm.chat(
     model=og.TEE_LLM.CLAUDE_SONNET_4_6,
     messages=[{"role": "user", "content": "Explain quantum computing"}],
     max_tokens=500,
     stream=True,
 )
 
-for chunk in stream:
+async for chunk in stream:
     if chunk.choices[0].delta.content:
         print(chunk.choices[0].delta.content, end="")
 ```
@@ -299,7 +299,7 @@ OpenGradient supports multiple settlement modes through the x402 payment protoco
 
 Specify settlement mode in your requests:
 ```python
-result = client.llm.chat(
+result = await client.llm.chat(
     model=og.TEE_LLM.GPT_5,
     messages=[{"role": "user", "content": "Hello"}],
     x402_settlement_mode=og.x402SettlementMode.BATCH_HASHED,

dev-requirements.txt

@@ -2,3 +2,4 @@ langgraph
 pdoc3==0.10.0
 mypy
 types-protobuf
+pytest-asyncio

docs/CLAUDE_SDK_USERS.md

@@ -23,8 +23,8 @@ client = og.Client(
     private_key=os.environ["OG_PRIVATE_KEY"],  # Required: Ethereum private key
 )
 
-# LLM Chat (TEE-verified with x402 payments)
-result = client.llm.chat(
+# LLM Chat (TEE-verified with x402 payments, async)
+result = await client.llm.chat(
     model=og.TEE_LLM.CLAUDE_HAIKU_4_5,
     messages=[{"role": "user", "content": "Hello!"}],
     max_tokens=100,
@@ -51,7 +51,7 @@ og.init(private_key="0x...", email="...", password="...")
 ### LLM Chat
 
 ```python
-result = client.llm.chat(
+result = await client.llm.chat(
     model: TEE_LLM,                    # og.TEE_LLM enum value
     messages: List[Dict],              # [{"role": "user", "content": "..."}]
     max_tokens: int = 100,
@@ -62,7 +62,7 @@ result = client.llm.chat(
     x402_settlement_mode: x402SettlementMode = x402SettlementMode.BATCH_HASHED,
     stream: bool = False,              # Enable streaming responses
 )
-# Returns: TextGenerationOutput (or TextGenerationStream if stream=True)
+# Returns: TextGenerationOutput (or AsyncGenerator[StreamChunk] if stream=True)
 #   - chat_output: Dict with role, content, tool_calls
 #   - transaction_hash: str
 #   - finish_reason: str ("stop", "tool_call")
@@ -72,7 +72,7 @@ result = client.llm.chat(
 ### LLM Completion
 
 ```python
-result = client.llm.completion(
+result = await client.llm.completion(
     model: TEE_LLM,
     prompt: str,
     max_tokens: int = 100,
@@ -139,7 +139,7 @@ og.TEE_LLM.GROK_4_1_FAST_NON_REASONING
 All models are accessed through the OpenGradient TEE infrastructure with x402 payments:
 
 ```python
-result = client.llm.chat(
+result = await client.llm.chat(
     model=og.TEE_LLM.GPT_5,
     messages=[{"role": "user", "content": "Hello"}],
 )
@@ -165,7 +165,7 @@ tools = [{
     }
 }]
 
-result = client.llm.chat(
+result = await client.llm.chat(
     model=og.TEE_LLM.CLAUDE_SONNET_4_6,
     messages=[{"role": "user", "content": "What's the weather in NYC?"}],
     tools=tools,
@@ -180,13 +180,13 @@ if result.chat_output.get("tool_calls"):
 ### Streaming
 
 ```python
-stream = client.llm.chat(
+stream = await client.llm.chat(
     model=og.TEE_LLM.CLAUDE_SONNET_4_6,
     messages=[{"role": "user", "content": "Tell me a story"}],
     stream=True,
 )
 
-for chunk in stream:
+async for chunk in stream:
     for choice in chunk.choices:
         if choice.delta.content:
             print(choice.delta.content, end="")
@@ -283,28 +283,15 @@ og.CandleOrder.ASCENDING, .DESCENDING
 
 ## Error Handling
 
-```python
-from opengradient.client.exceptions import (
-    OpenGradientError,      # Base exception
-    AuthenticationError,
-    InferenceError,
-    InvalidInputError,
-    NetworkError,
-    RateLimitError,
-    TimeoutError,
-    ServerError,
-    UnsupportedModelError,
-    InsufficientCreditsError,
-)
+LLM methods raise `RuntimeError` on failure and `ValueError` for invalid arguments:
 
+```python
 try:
-    result = client.llm.chat(...)
-except RateLimitError:
-    # Retry with backoff
-except InferenceError as e:
-    print(f"Inference failed: {e.message}")
-except OpenGradientError as e:
-    print(f"Error {e.status_code}: {e.message}")
+    result = await client.llm.chat(...)
+except RuntimeError as e:
+    print(f"Inference failed: {e}")
+except ValueError as e:
+    print(f"Invalid input: {e}")
 ```
 
 ## Environment Variables

docs/opengradient/client/alpha.md

@@ -63,7 +63,7 @@ InferenceResult (InferenceResult): A dataclass object containing the transaction
 
 **Raises**
 
-* **`OpenGradientError`**: If the inference fails.
+* **`RuntimeError`**: If the inference fails.
 
 ---
 

docs/opengradient/client/client.md

@@ -36,9 +36,9 @@ def __init__(
     twins_api_key: Optional[str] = None,
     rpc_url: str = 'https://ogevmdevnet.opengradient.ai',
     api_url: str = 'https://sdk-devnet.opengradient.ai',
-    contract_address: str = '0x8383C9bD7462F12Eb996DD02F78234C0421A6FaE',
-    og_llm_server_url: Optional[str] = 'https://3.15.214.21:443',
-    og_llm_streaming_server_url: Optional[str] = 'https://3.15.214.21:443'
+    inference_contract_address: str = '0x8383C9bD7462F12Eb996DD02F78234C0421A6FaE',
+    llm_server_url: Optional[str] = None,
+    tee_registry_address: str = '0x4e72238852f3c918f4E4e57AeC9280dDB0c80248'
 )
 ```
 
@@ -49,14 +49,20 @@ def __init__(
 * **`alpha_private_key`**: Private key whose wallet holds **OpenGradient testnet
         gas tokens** for on-chain inference. Optional -- falls back to
         ``private_key`` for backward compatibility.
-* **`email`**: Email for Model Hub authentication. Optional.
-* **`password`**: Password for Model Hub authentication. Optional.
+* **`email`**: Email for Model Hub authentication. Must be provided together
+        with ``password``.
+* **`password`**: Password for Model Hub authentication. Must be provided
+        together with ``email``.
 * **`twins_api_key`**: API key for digital twins chat (twin.fun). Optional.
 * **`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.
+* **`inference_contract_address`**: Inference contract address on the
+        OpenGradient Alpha Testnet.
+* **`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.
+* **`tee_registry_address`**: Address of the TEERegistry contract used to
+        discover active LLM proxy endpoints and their verified TLS certs.
 
 #### Methods
 
@@ -65,7 +71,7 @@ def __init__(
 #### `close()`
 
 ```python
-def close(self) ‑> None
+async def close(self) ‑> None
 ```
 Close underlying SDK resources.
 

docs/opengradient/client/index.md

@@ -39,12 +39,13 @@ client = og.init(private_key="0xLLM_KEY...", alpha_private_key="0xALPHA_KEY...")
 client.llm.ensure_opg_approval(opg_amount=5)
 
 # LLM chat (TEE-verified, streamed)
-for chunk in client.llm.chat(
+stream = await client.llm.chat(
     model=og.TEE_LLM.CLAUDE_HAIKU_4_5,
     messages=[{"role": "user", "content": "Hello!"}],
     max_tokens=200,
     stream=True,
-):
+)
+async for chunk in stream:
     if chunk.choices[0].delta.content:
         print(chunk.choices[0].delta.content, end="")
 
@@ -64,10 +65,10 @@ repo = client.model_hub.create_model("my-model", "A price prediction model")
 
 * [alpha](./alpha): Alpha Testnet features for OpenGradient SDK.
 * [client](./client): Main Client class that unifies all OpenGradient service namespaces.
-* [exceptions](./exceptions): Exception types for OpenGradient SDK errors.
 * [llm](./llm): LLM chat and completion via TEE-verified execution with x402 payments.
 * [model_hub](./model_hub): Model Hub for creating, versioning, and uploading ML models.
 * [opg_token](./opg_token): OPG token Permit2 approval utilities for x402 payments.
+* [tee_registry](./tee_registry): TEE Registry client for fetching verified TEE endpoints and TLS certificates.
 * [twins](./twins): Digital twins chat via OpenGradient verifiable inference.
 
 ## Classes
@@ -98,9 +99,9 @@ def __init__(
     twins_api_key: Optional[str] = None,
     rpc_url: str = 'https://ogevmdevnet.opengradient.ai',
     api_url: str = 'https://sdk-devnet.opengradient.ai',
-    contract_address: str = '0x8383C9bD7462F12Eb996DD02F78234C0421A6FaE',
-    og_llm_server_url: Optional[str] = 'https://3.15.214.21:443',
-    og_llm_streaming_server_url: Optional[str] = 'https://3.15.214.21:443'
+    inference_contract_address: str = '0x8383C9bD7462F12Eb996DD02F78234C0421A6FaE',
+    llm_server_url: Optional[str] = None,
+    tee_registry_address: str = '0x4e72238852f3c918f4E4e57AeC9280dDB0c80248'
 )
 ```
 
@@ -111,14 +112,20 @@ def __init__(
 * **`alpha_private_key`**: Private key whose wallet holds **OpenGradient testnet
         gas tokens** for on-chain inference. Optional -- falls back to
         ``private_key`` for backward compatibility.
-* **`email`**: Email for Model Hub authentication. Optional.
-* **`password`**: Password for Model Hub authentication. Optional.
+* **`email`**: Email for Model Hub authentication. Must be provided together
+        with ``password``.
+* **`password`**: Password for Model Hub authentication. Must be provided
+        together with ``email``.
 * **`twins_api_key`**: API key for digital twins chat (twin.fun). Optional.
 * **`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.
+* **`inference_contract_address`**: Inference contract address on the
+        OpenGradient Alpha Testnet.
+* **`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.
+* **`tee_registry_address`**: Address of the TEERegistry contract used to
+        discover active LLM proxy endpoints and their verified TLS certs.
 
 #### Methods
 
@@ -127,7 +134,7 @@ def __init__(
 #### `close()`
 
 ```python
-def close(self) ‑> None
+async def close(self) ‑> None
 ```
 Close underlying SDK resources.