opengradient 0.3.16__py3-none-any.whl → 0.3.18__py3-none-any.whl

opengradient/__init__.py

@@ -1,11 +1,15 @@
+"""
+OpenGradient Python SDK for interacting with AI models and infrastructure.
+"""
+
 from typing import Dict, List, Optional, Tuple
 
 from .client import Client
 from .defaults import DEFAULT_INFERENCE_CONTRACT_ADDRESS, DEFAULT_RPC_URL
-from .types import InferenceMode, LLM
+from .types import InferenceMode, LlmInferenceMode, LLM, TEE_LLM
 from . import llm
 
-__version__ = "0.3.16"
+__version__ = "0.3.18"
 
 _client = None
 
@@ -14,15 +18,50 @@ def init(email: str,
          private_key: str,
          rpc_url=DEFAULT_RPC_URL,
          contract_address=DEFAULT_INFERENCE_CONTRACT_ADDRESS):
+    """Initialize the OpenGradient SDK with authentication and network settings.
+
+    Args:
+        email: User's email address for authentication
+        password: User's password for authentication
+        private_key: Ethereum private key for blockchain transactions
+        rpc_url: Optional RPC URL for the blockchain network, defaults to mainnet
+        contract_address: Optional inference contract address
+    """
     global _client
     _client = Client(private_key=private_key, rpc_url=rpc_url, contract_address=contract_address, email=email, password=password)
 
 def upload(model_path, model_name, version):
+    """Upload a model file to OpenGradient.
+
+    Args:
+        model_path: Path to the model file on local filesystem
+        model_name: Name of the model repository
+        version: Version string for this model upload
+
+    Returns:
+        dict: Upload response containing file metadata
+
+    Raises:
+        RuntimeError: If SDK is not initialized
+    """
     if _client is None:
         raise RuntimeError("OpenGradient client not initialized. Call og.init() first.")
     return _client.upload(model_path, model_name, version)
 
 def create_model(model_name: str, model_desc: str, model_path: str = None):
+    """Create a new model repository.
+
+    Args:
+        model_name: Name for the new model repository
+        model_desc: Description of the model
+        model_path: Optional path to model file to upload immediately
+
+    Returns:
+        dict: Creation response with model metadata and optional upload results
+
+    Raises:
+        RuntimeError: If SDK is not initialized
+    """
     if _client is None:
         raise RuntimeError("OpenGradient client not initialized. Call og.init() first.")
     
@@ -36,62 +75,193 @@ def create_model(model_name: str, model_desc: str, model_path: str = None):
     return result
 
 def create_version(model_name, notes=None, is_major=False):
+    """Create a new version for an existing model.
+
+    Args:
+        model_name: Name of the model repository
+        notes: Optional release notes for this version
+        is_major: If True, creates a major version bump instead of minor
+
+    Returns:
+        dict: Version creation response with version metadata
+
+    Raises:
+        RuntimeError: If SDK is not initialized
+    """
     if _client is None:
         raise RuntimeError("OpenGradient client not initialized. Call og.init() first.")
     return _client.create_version(model_name, notes, is_major)
 
-def infer(model_cid, inference_mode, model_input):
+def infer(model_cid, inference_mode, model_input, max_retries: Optional[int] = None):
+    """Run inference on a model.
+
+    Args:
+        model_cid: CID of the model to use
+        inference_mode: Mode of inference (e.g. VANILLA)
+        model_input: Input data for the model
+        max_retries: Maximum number of retries for failed transactions
+
+    Returns:
+        Tuple[str, Any]: Transaction hash and model output
+
+    Raises:
+        RuntimeError: If SDK is not initialized
+    """
     if _client is None:
         raise RuntimeError("OpenGradient client not initialized. Call og.init() first.")
-    return _client.infer(model_cid, inference_mode, model_input)
+    return _client.infer(model_cid, inference_mode, model_input, max_retries=max_retries)
 
 def llm_completion(model_cid: LLM, 
-                         prompt: str, 
-                         max_tokens: int = 100, 
-                         stop_sequence: Optional[List[str]] = None, 
-                         temperature: float = 0.0) -> Tuple[str, str]:
+                  prompt: str, 
+                  inference_mode: str = LlmInferenceMode.VANILLA,
+                  max_tokens: int = 100, 
+                  stop_sequence: Optional[List[str]] = None, 
+                  temperature: float = 0.0,
+                  max_retries: Optional[int] = None) -> Tuple[str, str]:
+    """Generate text completion using an LLM.
+
+    Args:
+        model_cid: CID of the LLM model to use
+        prompt: Text prompt for completion
+        inference_mode: Mode of inference, defaults to VANILLA
+        max_tokens: Maximum tokens to generate
+        stop_sequence: Optional list of sequences where generation should stop
+        temperature: Sampling temperature (0.0 = deterministic, 1.0 = creative)
+        max_retries: Maximum number of retries for failed transactions
+
+    Returns:
+        Tuple[str, str]: Transaction hash and generated text
+
+    Raises:
+        RuntimeError: If SDK is not initialized
+    """
     if _client is None:
         raise RuntimeError("OpenGradient client not initialized. Call og.init() first.")
-    return _client.llm_completion(model_cid, prompt, max_tokens, stop_sequence, temperature)
+    return _client.llm_completion(model_cid=model_cid, 
+                                inference_mode=inference_mode, 
+                                prompt=prompt,
+                                max_tokens=max_tokens, 
+                                stop_sequence=stop_sequence, 
+                                temperature=temperature,
+                                max_retries=max_retries)
 
 def llm_chat(model_cid: LLM,
-                   messages: List[Dict],
-                   max_tokens: int = 100,
-                   stop_sequence: Optional[List[str]] = None,
-                   temperature: float = 0.0,
-                   tools: Optional[List[Dict]] = None,
-                   tool_choice: Optional[str] = None):
+             messages: List[Dict],
+             inference_mode: str = LlmInferenceMode.VANILLA,
+             max_tokens: int = 100,
+             stop_sequence: Optional[List[str]] = None,
+             temperature: float = 0.0,
+             tools: Optional[List[Dict]] = None,
+             tool_choice: Optional[str] = None,
+             max_retries: Optional[int] = None) -> Tuple[str, str, Dict]:
+    """Have a chat conversation with an LLM.
+
+    Args:
+        model_cid: CID of the LLM model to use
+        messages: List of chat messages, each with 'role' and 'content'
+        inference_mode: Mode of inference, defaults to VANILLA
+        max_tokens: Maximum tokens to generate
+        stop_sequence: Optional list of sequences where generation should stop
+        temperature: Sampling temperature (0.0 = deterministic, 1.0 = creative)
+        tools: Optional list of tools the model can use
+        tool_choice: Optional specific tool to use
+        max_retries: Maximum number of retries for failed transactions
+
+    Returns:
+        Tuple[str, str, Dict]: Transaction hash, model response, and metadata
+
+    Raises:
+        RuntimeError: If SDK is not initialized
+    """
     if _client is None:
         raise RuntimeError("OpenGradient client not initialized. Call og.init() first.")
-    return _client.llm_chat(model_cid, messages, max_tokens, stop_sequence, temperature, tools, tool_choice)
+    return _client.llm_chat(model_cid=model_cid, 
+                          inference_mode=inference_mode, 
+                          messages=messages, 
+                          max_tokens=max_tokens, 
+                          stop_sequence=stop_sequence, 
+                          temperature=temperature, 
+                          tools=tools, 
+                          tool_choice=tool_choice,
+                          max_retries=max_retries)
 
 def login(email: str, password: str):
+    """Login to OpenGradient.
+
+    Args:
+        email: User's email address
+        password: User's password
+
+    Returns:
+        dict: Login response with authentication tokens
+
+    Raises:
+        RuntimeError: If SDK is not initialized
+    """
     if _client is None:
         raise RuntimeError("OpenGradient client not initialized. Call og.init() first.")
     return _client.login(email, password)
 
 def list_files(model_name: str, version: str) -> List[Dict]:
+    """List files in a model repository version.
+
+    Args:
+        model_name: Name of the model repository
+        version: Version string to list files from
+
+    Returns:
+        List[Dict]: List of file metadata dictionaries
+
+    Raises:
+        RuntimeError: If SDK is not initialized
+    """
     if _client is None:
         raise RuntimeError("OpenGradient client not initialized. Call og.init() first.")
     return _client.list_files(model_name, version)
 
 def generate_image(model: str, prompt: str, height: Optional[int] = None, width: Optional[int] = None) -> bytes:
-    """
-    Generate an image using the specified model and prompt.
+    """Generate an image from a text prompt.
 
     Args:
-        model (str): The model identifier (e.g. "stabilityai/stable-diffusion-xl-base-1.0")
-        prompt (str): The text prompt to generate the image from
-        height (Optional[int]): Height of the generated image. Default is None.
-        width (Optional[int]): Width of the generated image. Default is None.
+        model: Model identifier (e.g. "stabilityai/stable-diffusion-xl-base-1.0")
+        prompt: Text description of the desired image
+        height: Optional height of the generated image in pixels
+        width: Optional width of the generated image in pixels
 
     Returns:
-        bytes: The raw image data bytes
+        bytes: Raw image data as bytes
 
     Raises:
-        RuntimeError: If the client is not initialized
-        OpenGradientError: If the image generation fails
+        RuntimeError: If SDK is not initialized
+        OpenGradientError: If image generation fails
     """
     if _client is None:
         raise RuntimeError("OpenGradient client not initialized. Call og.init() first.")
     return _client.generate_image(model, prompt, height=height, width=width)
+
+__all__ = [
+    'generate_image',
+    'list_files'
+    'login',
+    'llm_chat',
+    'llm_completion',
+    'infer',
+    'create_version',
+    'create_model',
+    'upload',
+    'init',
+    'LLM',
+    'TEE_LLM'
+]
+
+__pdoc__ = {
+    'account': False,
+    'cli': False,
+    'client': False,
+    'defaults': False,
+    'exceptions': False,
+    'llm': True,
+    'proto': False,
+    'types': False,
+    'utils': False
+}

opengradient/cli.py

@@ -20,7 +20,7 @@ from .defaults import (
     DEFAULT_OG_FAUCET_URL,
     DEFAULT_RPC_URL,
 )
-from .types import InferenceMode
+from .types import InferenceMode, LlmInferenceMode
 
 OG_CONFIG_FILE = Path.home() / '.opengradient_config.json'
 
@@ -65,11 +65,17 @@ InferenceModes = {
     "TEE": InferenceMode.TEE,
 }
 
+LlmInferenceModes = {
+    "VANILLA": LlmInferenceMode.VANILLA,
+    "TEE": LlmInferenceMode.TEE,
+}
+
 # Supported LLMs
 LlmModels = {
     "meta-llama/Meta-Llama-3-8B-Instruct",
     "meta-llama/Llama-3.2-3B-Instruct",
-    "mistralai/Mistral-7B-Instruct-v0.3"
+    "mistralai/Mistral-7B-Instruct-v0.3",
+    "meta-llama/Llama-3.1-70B-Instruct",
 }
 
 def initialize_config(ctx):
@@ -339,13 +345,15 @@ def infer(ctx, model_cid: str, inference_mode: str, input_data, input_file: Path
         click.echo(f"Error running inference: {str(e)}")
 
 @cli.command()
-@click.option('--model', '-m', 'model_cid', type=click.Choice(LlmModels), required=True, help='CID of the LLM model to run inference on')
+@click.option('--model', '-m', 'model_cid', type=click.Choice([e.value for e in types.LLM]), required=True, help='CID of the LLM model to run inference on')
+@click.option('--mode', 'inference_mode', type=click.Choice(LlmInferenceModes.keys()), default="VANILLA", 
+              help='Inference mode (default: VANILLA)')
 @click.option('--prompt', '-p', required=True, help='Input prompt for the LLM completion')
 @click.option('--max-tokens', type=int, default=100, help='Maximum number of tokens for LLM completion output')
 @click.option('--stop-sequence', multiple=True, help='Stop sequences for LLM')
 @click.option('--temperature', type=float, default=0.0, help='Temperature for LLM inference (0.0 to 1.0)')
 @click.pass_context
-def completion(ctx, model_cid: str, prompt: str, max_tokens: int, stop_sequence: List[str], temperature: float):
+def completion(ctx, model_cid: str, inference_mode: str,  prompt: str, max_tokens: int, stop_sequence: List[str], temperature: float):
     """
     Run completion inference on an LLM model.
 
@@ -355,13 +363,14 @@ def completion(ctx, model_cid: str, prompt: str, max_tokens: int, stop_sequence:
 
     \b
     opengradient completion --model meta-llama/Meta-Llama-3-8B-Instruct --prompt "Hello, how are you?" --max-tokens 50 --temperature 0.7
-    opengradient completion -m meta-llama/Meta-Llama-3-8B-Instruct -p "Translate to French: Hello world" --stop-sequence "." --stop-sequence "\n"
+    opengradient completion -m meta-llama/Meta-Llama-3-8B-Instruct -p "Translate to French: Hello world" --stop-sequence "." --stop-sequence "\\n"
     """
     client: Client = ctx.obj['client']
     try:
         click.echo(f"Running LLM completion inference for model \"{model_cid}\"\n")
         tx_hash, llm_output = client.llm_completion(
             model_cid=model_cid,
+            inference_mode=LlmInferenceModes[inference_mode],
             prompt=prompt,
             max_tokens=max_tokens,
             stop_sequence=list(stop_sequence),
@@ -394,6 +403,9 @@ def print_llm_completion_result(model_cid, tx_hash, llm_output):
               type=click.Choice([e.value for e in types.LLM]), 
               required=True, 
               help='CID of the LLM model to run inference on')
+@click.option('--mode', 'inference_mode', type=click.Choice(LlmInferenceModes.keys()), 
+              default="VANILLA", 
+              help='Inference mode (default: VANILLA)')
 @click.option('--messages', 
               type=str,
               required=False, 
@@ -431,6 +443,7 @@ def print_llm_completion_result(model_cid, tx_hash, llm_output):
 def chat(
     ctx,
     model_cid: str,
+    inference_mode: str,
     messages: Optional[str],
     messages_file: Optional[Path],
     max_tokens: int,
@@ -444,11 +457,13 @@ def chat(
 
     This command runs a chat inference on the specified LLM model using the provided messages and parameters.
 
+    Tool call formatting is based on OpenAI documentation tool calls (see here: https://platform.openai.com/docs/guides/function-calling).
+
     Example usage:
 
     \b
     opengradient chat --model meta-llama/Meta-Llama-3-8B-Instruct --messages '[{"role":"user","content":"hello"}]' --max-tokens 50 --temperature 0.7
-    opengradient chat -m mistralai/Mistral-7B-Instruct-v0.3 --messages-file messages.json --stop-sequence "." --stop-sequence "\n"
+    opengradient chat --model mistralai/Mistral-7B-Instruct-v0.3 --messages-file messages.json --tools-file tools.json --max-tokens 200 --stop-sequence "." --stop-sequence "\\n"
     """
     client: Client = ctx.obj['client']
     try:
@@ -458,7 +473,7 @@ def chat(
             ctx.exit(1)
             return
         if messages and messages_file:
-            click.echo("Cannot have both messages and messages_file")
+            click.echo("Cannot have both messages and messages-file")
             ctx.exit(1)
             return
 
@@ -473,9 +488,9 @@ def chat(
                 messages = json.load(file)
 
         # Parse tools if provided
-        if tools is not None and tools != "[]" and tools_file:
-            click.echo("Cannot have both tools and tools_file")
-            ctx.exit(1)
+        if (tools and tools != '[]') and tools_file:
+            click.echo("Cannot have both tools and tools-file")
+            click.exit(1)
             return
         
         parsed_tools=[]
@@ -509,6 +524,7 @@ def chat(
 
         tx_hash, finish_reason, llm_chat_output = client.llm_chat(
             model_cid=model_cid,
+            inference_mode=LlmInferenceModes[inference_mode],
             messages=messages,
             max_tokens=max_tokens,
             stop_sequence=list(stop_sequence),
@@ -517,15 +533,32 @@ def chat(
             tool_choice=tool_choice,
         )
 
-        # TODO (Kyle): Make this prettier
-        print("TX Hash: ", tx_hash)
-        print("Finish reason: ", finish_reason)
-        print("Chat output: ", llm_chat_output)
+        print_llm_chat_result(model_cid, tx_hash, finish_reason, llm_chat_output)
     except Exception as e:
         click.echo(f"Error running LLM chat inference: {str(e)}")
 
-def print_llm_chat_result():
-    pass
+def print_llm_chat_result(model_cid, tx_hash, finish_reason, chat_output):
+    click.secho("✅ LLM Chat Successful", fg="green", bold=True)
+    click.echo("──────────────────────────────────────")
+    click.echo("Model CID: ", nl=False)
+    click.secho(model_cid, fg="cyan", bold=True)
+    click.echo("Transaction hash: ", nl=False)
+    click.secho(tx_hash, fg="cyan", bold=True)
+    block_explorer_link = f"{DEFAULT_BLOCKCHAIN_EXPLORER}0x{tx_hash}"
+    click.echo("Block explorer link: ", nl=False)
+    click.secho(block_explorer_link, fg="blue", underline=True)
+    click.echo("──────────────────────────────────────")
+    click.secho("Finish Reason: ", fg="yellow", bold=True)
+    click.echo()
+    click.echo(finish_reason)
+    click.echo()
+    click.secho("Chat Output:", fg="yellow", bold=True)
+    click.echo()
+    for key, value in chat_output.items():
+        # If the value doesn't give any information, don't print it
+        if value != None and value != "" and value != '[]' and value != []:
+            click.echo(f"{key}: {value}")
+    click.echo()
 
 @cli.command()
 def create_account():

opengradient/client.py

@@ -1,6 +1,7 @@
 import json
 import logging
 import os
+import random
 from typing import Dict, List, Optional, Tuple, Union
 
 import firebase
@@ -12,7 +13,7 @@ from web3.logs import DISCARD
 
 from opengradient import utils
 from opengradient.exceptions import OpenGradientError
-from opengradient.types import InferenceMode, LLM
+from opengradient.types import InferenceMode, LlmInferenceMode, LLM, TEE_LLM
 
 import grpc
 import time
@@ -23,6 +24,31 @@ from opengradient.proto import infer_pb2
 from opengradient.proto import infer_pb2_grpc
 from .defaults import DEFAULT_IMAGE_GEN_HOST, DEFAULT_IMAGE_GEN_PORT
 
+from functools import wraps
+
+def run_with_retry(txn_function, max_retries=5):
+    """
+    Execute a blockchain transaction with retry logic.
+    
+    Args:
+        txn_function: Function that executes the transaction
+        max_retries (int): Maximum number of retry attempts
+    """
+    last_error = None
+    for attempt in range(max_retries):
+        try:
+            return txn_function()
+        except Exception as e:
+            last_error = e
+            if attempt < max_retries - 1:
+                if "nonce too low" in str(e) or "nonce too high" in str(e):
+                    time.sleep(1)  # Wait before retry
+                    continue
+                # If it's not a nonce error, raise immediately
+                raise
+    # If we've exhausted all retries, raise the last error
+    raise OpenGradientError(f"Transaction failed after {max_retries} attempts: {str(last_error)}")
+
 class Client:
     FIREBASE_CONFIG = {
         "apiKey": "AIzaSyDUVckVtfl-hiteBzPopy1pDD8Uvfncs7w",
@@ -311,7 +337,8 @@ class Client:
             self, 
             model_cid: str, 
             inference_mode: InferenceMode, 
-            model_input: Dict[str, Union[str, int, float, List, np.ndarray]]
+            model_input: Dict[str, Union[str, int, float, List, np.ndarray]],
+            max_retries: Optional[int] = None
             ) -> Tuple[str, Dict[str, np.ndarray]]:
         """
         Perform inference on a model.
@@ -320,6 +347,7 @@ class Client:
             model_cid (str): The unique content identifier for the model from IPFS.
             inference_mode (InferenceMode): The inference mode.
             model_input (Dict[str, Union[str, int, float, List, np.ndarray]]): The input data for the model.
+            max_retries (int, optional): Maximum number of retry attempts. Defaults to 5.
 
         Returns:
             Tuple[str, Dict[str, np.ndarray]]: The transaction hash and the model output.
@@ -327,46 +355,22 @@ class Client:
         Raises:
             OpenGradientError: If the inference fails.
         """
-        # TODO (Kyle): Add input support for JSON tensors
-        try:
-            logging.debug("Entering infer method")
+        def execute_transaction():
             self._initialize_web3()
-            logging.debug(f"Web3 initialized. Connected: {self._w3.is_connected()}")
-
-            logging.debug(f"Creating contract instance. Address: {self.contract_address}")
             contract = self._w3.eth.contract(address=self.contract_address, abi=self.abi)
-            logging.debug("Contract instance created successfully")
-
-            logging.debug(f"Model ID: {model_cid}")
-            logging.debug(f"Inference Mode: {inference_mode}")
-            logging.debug(f"Model Input: {model_input}")
-
-            # Convert InferenceMode to uint8
+            
             inference_mode_uint8 = int(inference_mode)
-
-            # Prepare ModelInput tuple
             converted_model_input = utils.convert_to_model_input(model_input)
-            logging.debug(f"Prepared model input tuple: {converted_model_input}")
-
-            logging.debug("Preparing run function")
+            
             run_function = contract.functions.run(
                 model_cid,
                 inference_mode_uint8,
                 converted_model_input
             )
-            logging.debug("Run function prepared successfully")
 
-            # Build transaction
-            nonce = self._w3.eth.get_transaction_count(self.wallet_address)
-            logging.debug(f"Nonce: {nonce}")
-
-            # Estimate gas
+            nonce = self._w3.eth.get_transaction_count(self.wallet_address, 'pending')
             estimated_gas = run_function.estimate_gas({'from': self.wallet_address})
-            logging.debug(f"Estimated gas: {estimated_gas}")
-
-            # Increase gas limit by 20%
             gas_limit = int(estimated_gas * 3)
-            logging.debug(f"Gas limit set to: {gas_limit}")
 
             transaction = run_function.build_transaction({
                 'from': self.wallet_address,
@@ -375,62 +379,36 @@ class Client:
                 'gasPrice': self._w3.eth.gas_price,
             })
 
-            logging.debug(f"Transaction built: {transaction}")
-
-            # Sign transaction
             signed_tx = self._w3.eth.account.sign_transaction(transaction, self.private_key)
-            logging.debug("Transaction signed successfully")
-
-            # Send transaction
             tx_hash = self._w3.eth.send_raw_transaction(signed_tx.raw_transaction)
-            logging.debug(f"Transaction sent. Hash: {tx_hash.hex()}")
-
-            # Wait for transaction receipt
             tx_receipt = self._w3.eth.wait_for_transaction_receipt(tx_hash)
-            logging.debug(f"Transaction receipt received: {tx_receipt}")
 
-            # Check if the transaction was successful
             if tx_receipt['status'] == 0:
                 raise ContractLogicError(f"Transaction failed. Receipt: {tx_receipt}")
 
-            # Process the InferenceResult event
             parsed_logs = contract.events.InferenceResult().process_receipt(tx_receipt, errors=DISCARD)
-
             if len(parsed_logs) < 1:
                 raise OpenGradientError("InferenceResult event not found in transaction logs")
-            inference_result = parsed_logs[0]
-
-            # Extract the ModelOutput from the event
-            event_data = inference_result['args']
-            logging.debug(f"Raw event data: {event_data}")
-
-            try:
-                model_output = utils.convert_to_model_output(event_data)
-                logging.debug(f"Parsed ModelOutput: {model_output}")
-            except Exception as e:
-                logging.error(f"Error parsing event data: {str(e)}", exc_info=True)
-                raise OpenGradientError(f"Failed to parse event data: {str(e)}")
 
+            model_output = utils.convert_to_model_output(parsed_logs[0]['args'])
             return tx_hash.hex(), model_output
 
-        except ContractLogicError as e:
-            logging.error(f"Contract logic error: {str(e)}", exc_info=True)
-            raise OpenGradientError(f"Inference failed due to contract logic error: {str(e)}")
-        except Exception as e:
-            logging.error(f"Error in infer method: {str(e)}", exc_info=True)
-            raise OpenGradientError(f"Inference failed: {str(e)}")
-        
+        return run_with_retry(execute_transaction, max_retries or 5)
+
     def llm_completion(self, 
                        model_cid: LLM, 
+                       inference_mode: InferenceMode,
                        prompt: str, 
                        max_tokens: int = 100, 
                        stop_sequence: Optional[List[str]] = None, 
-                       temperature: float = 0.0) -> Tuple[str, str]:
+                       temperature: float = 0.0,
+                       max_retries: Optional[int] = None) -> Tuple[str, str]:
         """
         Perform inference on an LLM model using completions.
 
         Args:
             model_cid (LLM): The unique content identifier for the model.
+            inference_mode (InferenceMode): The inference mode.
             prompt (str): The input prompt for the LLM.
             max_tokens (int): Maximum number of tokens for LLM output. Default is 100.
             stop_sequence (List[str], optional): List of stop sequences for LLM. Default is None.
@@ -442,17 +420,20 @@ class Client:
         Raises:
             OpenGradientError: If the inference fails.
         """
-        try:
-            self._initialize_web3()
+        def execute_transaction():
+            # Check inference mode and supported model
+            if inference_mode != LlmInferenceMode.VANILLA and inference_mode != LlmInferenceMode.TEE:
+                raise OpenGradientError("Invalid inference mode %s: Inference mode must be VANILLA or TEE" % inference_mode)
             
-            abi_path = os.path.join(os.path.dirname(__file__), 'abi', 'inference.abi')
-            with open(abi_path, 'r') as abi_file:
-                llm_abi = json.load(abi_file)
-            contract = self._w3.eth.contract(address=self.contract_address, abi=llm_abi)
+            if inference_mode == LlmInferenceMode.TEE and model_cid not in TEE_LLM:
+                raise OpenGradientError("That model CID is not supported yet supported for TEE inference")
+
+            self._initialize_web3()
+            contract = self._w3.eth.contract(address=self.contract_address, abi=self.abi)
 
             # Prepare LLM input
             llm_request = {
-                "mode": InferenceMode.VANILLA,
+                "mode": inference_mode,
                 "modelCID": model_cid,
                 "prompt": prompt,
                 "max_tokens": max_tokens,
@@ -461,11 +442,9 @@ class Client:
             }
             logging.debug(f"Prepared LLM request: {llm_request}")
 
-            # Prepare run function
             run_function = contract.functions.runLLMCompletion(llm_request)
 
-            # Build transaction
-            nonce = self._w3.eth.get_transaction_count(self.wallet_address)
+            nonce = self._w3.eth.get_transaction_count(self.wallet_address, 'pending')
             estimated_gas = run_function.estimate_gas({'from': self.wallet_address})
             gas_limit = int(estimated_gas * 1.2)
 
@@ -476,47 +455,38 @@ class Client:
                 'gasPrice': self._w3.eth.gas_price,
             })
 
-            # Sign and send transaction
             signed_tx = self._w3.eth.account.sign_transaction(transaction, self.private_key)
             tx_hash = self._w3.eth.send_raw_transaction(signed_tx.raw_transaction)
-            logging.debug(f"Transaction sent. Hash: {tx_hash.hex()}")
-
-            # Wait for transaction receipt
             tx_receipt = self._w3.eth.wait_for_transaction_receipt(tx_hash)
 
             if tx_receipt['status'] == 0:
                 raise ContractLogicError(f"Transaction failed. Receipt: {tx_receipt}")
 
-            # Process the LLMResult event
             parsed_logs = contract.events.LLMCompletionResult().process_receipt(tx_receipt, errors=DISCARD)
-
             if len(parsed_logs) < 1:
                 raise OpenGradientError("LLM completion result event not found in transaction logs")
-            llm_result = parsed_logs[0]
 
-            llm_answer = llm_result['args']['response']['answer']
+            llm_answer = parsed_logs[0]['args']['response']['answer']
             return tx_hash.hex(), llm_answer
 
-        except ContractLogicError as e:
-            logging.error(f"Contract logic error: {str(e)}", exc_info=True)
-            raise OpenGradientError(f"LLM inference failed due to contract logic error: {str(e)}")
-        except Exception as e:
-            logging.error(f"Error in infer completion method: {str(e)}", exc_info=True)
-            raise OpenGradientError(f"LLM inference failed: {str(e)}")
-        
+        return run_with_retry(execute_transaction, max_retries or 5)
+
     def llm_chat(self,
                  model_cid: str,
+                 inference_mode: InferenceMode,
                  messages: List[Dict],
                  max_tokens: int = 100,
                  stop_sequence: Optional[List[str]] = None,
                  temperature: float = 0.0,
                  tools: Optional[List[Dict]] = [],
-                 tool_choice: Optional[str] = None) -> Tuple[str, str]:
+                 tool_choice: Optional[str] = None,
+                 max_retries: Optional[int] = None) -> Tuple[str, str]:
         """
         Perform inference on an LLM model using chat.
 
         Args:
             model_cid (LLM): The unique content identifier for the model.
+            inference_mode (InferenceMode): The inference mode.
             messages (dict): The messages that will be passed into the chat. 
                 This should be in OpenAI API format (https://platform.openai.com/docs/api-reference/chat/create)
                 Example:
@@ -567,13 +537,16 @@ class Client:
         Raises:
             OpenGradientError: If the inference fails.
         """
-        try:
-            self._initialize_web3()
+        def execute_transaction():
+            # Check inference mode and supported model
+            if inference_mode != LlmInferenceMode.VANILLA and inference_mode != LlmInferenceMode.TEE:
+                raise OpenGradientError("Invalid inference mode %s: Inference mode must be VANILLA or TEE" % inference_mode)
+            
+            if inference_mode == LlmInferenceMode.TEE and model_cid not in TEE_LLM:
+                raise OpenGradientError("That model CID is not supported yet supported for TEE inference")
             
-            abi_path = os.path.join(os.path.dirname(__file__), 'abi', 'inference.abi')
-            with open(abi_path, 'r') as abi_file:
-                llm_abi = json.load(abi_file)
-            contract = self._w3.eth.contract(address=self.contract_address, abi=llm_abi)
+            self._initialize_web3()
+            contract = self._w3.eth.contract(address=self.contract_address, abi=self.abi)
 
             # For incoming chat messages, tool_calls can be empty. Add an empty array so that it will fit the ABI.
             for message in messages:
@@ -585,17 +558,10 @@ class Client:
                     message['name'] = ""
 
             # Create simplified tool structure for smart contract
-            #
-            #   struct ToolDefinition {
-            #       string description;
-            #       string name;
-            #       string parameters; // This must be a JSON 
-            #   }
             converted_tools = []
             if tools is not None:
                 for tool in tools:
                     function = tool['function']
-
                     converted_tool = {}
                     converted_tool['name'] = function['name']
                     converted_tool['description'] = function['description']
@@ -604,12 +570,11 @@ class Client:
                             converted_tool['parameters'] = json.dumps(parameters)
                         except Exception as e:
                             raise OpenGradientError("Chat LLM failed to convert parameters into JSON: %s", e)
-                    
                     converted_tools.append(converted_tool)
 
             # Prepare LLM input
             llm_request = {
-                "mode": InferenceMode.VANILLA,
+                "mode": inference_mode,
                 "modelCID": model_cid,
                 "messages": messages,
                 "max_tokens": max_tokens,
@@ -620,11 +585,9 @@ class Client:
             }
             logging.debug(f"Prepared LLM request: {llm_request}")
 
-            # Prepare run function
             run_function = contract.functions.runLLMChat(llm_request)
 
-            # Build transaction
-            nonce = self._w3.eth.get_transaction_count(self.wallet_address)
+            nonce = self._w3.eth.get_transaction_count(self.wallet_address, 'pending')
             estimated_gas = run_function.estimate_gas({'from': self.wallet_address})
             gas_limit = int(estimated_gas * 1.2)
 
@@ -635,41 +598,25 @@ class Client:
                 'gasPrice': self._w3.eth.gas_price,
             })
 
-            # Sign and send transaction
             signed_tx = self._w3.eth.account.sign_transaction(transaction, self.private_key)
             tx_hash = self._w3.eth.send_raw_transaction(signed_tx.raw_transaction)
-            logging.debug(f"Transaction sent. Hash: {tx_hash.hex()}")
-
-            # Wait for transaction receipt
             tx_receipt = self._w3.eth.wait_for_transaction_receipt(tx_hash)
 
             if tx_receipt['status'] == 0:
                 raise ContractLogicError(f"Transaction failed. Receipt: {tx_receipt}")
 
-            # Process the LLMResult event
             parsed_logs = contract.events.LLMChatResult().process_receipt(tx_receipt, errors=DISCARD)
-
             if len(parsed_logs) < 1:
                 raise OpenGradientError("LLM chat result event not found in transaction logs")
-            llm_result = parsed_logs[0]['args']['response']
 
-            # Turn tool calls into normal dicts
+            llm_result = parsed_logs[0]['args']['response']
             message = dict(llm_result['message'])
-            if (tool_calls := message.get('tool_calls')) != None:
-                new_tool_calls = []
-                for tool_call in tool_calls:
-                    new_tool_calls.append(dict(tool_call))
-                message['tool_calls'] = new_tool_calls
-
-            return (tx_hash.hex(), llm_result['finish_reason'], message)
+            if (tool_calls := message.get('tool_calls')) is not None:
+                message['tool_calls'] = [dict(tool_call) for tool_call in tool_calls]
 
-        except ContractLogicError as e:
-            logging.error(f"Contract logic error: {str(e)}", exc_info=True)
-            raise OpenGradientError(f"LLM inference failed due to contract logic error: {str(e)}")
-        except Exception as e:
-            logging.error(f"Error in infer chat method: {str(e)}", exc_info=True)
-            raise OpenGradientError(f"LLM inference failed: {str(e)}")
+            return tx_hash.hex(), llm_result['finish_reason'], message
 
+        return run_with_retry(execute_transaction, max_retries or 5)
 
     def list_files(self, model_name: str, version: str) -> List[Dict]:
         """

opengradient/llm/__init__.py

@@ -1,5 +1,38 @@
-from .chat import *
+"""
+OpenGradient LLM Adapters
+
+This module provides adapter interfaces to use OpenGradient LLMs with popular AI frameworks
+like LangChain and OpenAI. These adapters allow seamless integration of OpenGradient models
+into existing applications and agent frameworks.
+"""
+
+from .og_langchain import *
+from .og_openai import *
+
+def langchain_adapter(private_key: str, model_cid: str, max_tokens: int = 300) -> OpenGradientChatModel:
+    """
+    Returns an OpenGradient LLM that implements LangChain's LLM interface
+    and can be plugged into LangChain agents.
+    """
+    return OpenGradientChatModel(
+        private_key=private_key,
+        model_cid=model_cid,
+        max_tokens=max_tokens)
+
+def openai_adapter(private_key: str) -> OpenGradientOpenAIClient:
+    """
+    Returns an generic OpenAI LLM client that can be plugged into Swarm and can
+    be used with any LLM model on OpenGradient. The LLM is usually defined in the
+    agent.
+    """
+    return OpenGradientOpenAIClient(private_key=private_key)
 
 __all__ = [
-    'OpenGradientChatModel'
-]
+    'langchain_adapter',
+    'openai_adapter',
+]
+
+__pdoc__ = {
+    'og_langchain': False,
+    'og_openai': False
+}

opengradient/llm/og_openai.py

@@ -0,0 +1,121 @@
+from openai.types.chat import ChatCompletion
+import opengradient as og
+from opengradient.defaults import DEFAULT_RPC_URL, DEFAULT_INFERENCE_CONTRACT_ADDRESS
+
+from typing import List
+import time
+import json
+import uuid
+
+class OGCompletions(object):
+    client: og.Client
+
+    def __init__(self, client: og.Client):
+        self.client = client
+
+    def create(
+        self, 
+        model: str, 
+        messages: List[object], 
+        tools: List[object], 
+        tool_choice: str, 
+        stream: bool = False, 
+        parallel_tool_calls: bool = False) -> ChatCompletion:
+
+        # convert OpenAI message format so it's compatible with the SDK
+        sdk_messages = OGCompletions.convert_to_abi_compatible(messages)
+
+        _, finish_reason, chat_completion = self.client.llm_chat(
+            model_cid=model,
+            messages=sdk_messages,
+            max_tokens=200,
+            tools=tools,
+            tool_choice=tool_choice, 
+            temperature=0.25,
+            inference_mode=og.LlmInferenceMode.VANILLA
+        )
+
+        choice = {
+            'index': 0,  # Add missing index field
+            'finish_reason': finish_reason,
+            'message': {
+                'role': chat_completion['role'],
+                'content': chat_completion['content'],
+                'tool_calls': [
+                    {
+                        'id': tool_call['id'],
+                        'type': 'function',  # Add missing type field
+                        'function': {        # Add missing function field
+                            'name': tool_call['name'],
+                            'arguments': tool_call['arguments']
+                        }
+                    }
+                    for tool_call in chat_completion.get('tool_calls', [])
+                ] 
+            }
+        }
+
+        return ChatCompletion(
+            id=str(uuid.uuid4()),
+            created=int(time.time()),
+            model=model,
+            object='chat.completion',
+            choices=[choice]
+        )
+
+    @staticmethod
+    @staticmethod
+    def convert_to_abi_compatible(messages):
+        sdk_messages = []
+
+        for message in messages:
+            role = message['role']
+            sdk_message = {
+                'role': role
+            }
+
+            if role == 'system':
+                sdk_message['content'] = message['content']
+            elif role == 'user':
+                sdk_message['content'] = message['content']
+            elif role == 'tool':
+                sdk_message['content'] = message['content']
+                sdk_message['tool_call_id'] = message['tool_call_id']
+            elif role == 'assistant':
+                flattened_calls = []
+                for tool_call in message['tool_calls']:
+                    # OpenAI format
+                    flattened_call = {
+                        'id': tool_call['id'],
+                        'name': tool_call['function']['name'],
+                        'arguments': tool_call['function']['arguments']
+                    }
+                    flattened_calls.append(flattened_call)
+
+                sdk_message['tool_calls'] = flattened_calls
+                sdk_message['content'] = message['content']
+
+            sdk_messages.append(sdk_message)
+
+        return sdk_messages
+
+class OGChat(object):
+    completions: OGCompletions
+
+    def __init__(self, client: og.Client):
+        self.completions = OGCompletions(client)
+
+class OpenGradientOpenAIClient(object):
+    """OpenAI client implementation"""
+    client: og.Client
+    chat: OGChat
+
+    def __init__(self, private_key: str):
+        self.client = og.Client(
+            private_key=private_key,
+            rpc_url=DEFAULT_RPC_URL,
+            contract_address=DEFAULT_INFERENCE_CONTRACT_ADDRESS,
+            email=None,
+            password=None
+        )
+        self.chat = OGChat(self.client)

opengradient/types.py

@@ -27,6 +27,10 @@ class InferenceMode:
     ZKML = 1
     TEE = 2
 
+class LlmInferenceMode:
+    VANILLA = 0
+    TEE = 1
+
 @dataclass
 class ModelOutput:
     numbers: List[NumberTensor]
@@ -75,8 +79,15 @@ class Abi:
         return result
     
 class LLM(str, Enum):
+    """Enum for available LLM models"""
+
     META_LLAMA_3_8B_INSTRUCT = "meta-llama/Meta-Llama-3-8B-Instruct"
     LLAMA_3_2_3B_INSTRUCT = "meta-llama/Llama-3.2-3B-Instruct"
     MISTRAL_7B_INSTRUCT_V3 = "mistralai/Mistral-7B-Instruct-v0.3"
     HERMES_3_LLAMA_3_1_70B = "NousResearch/Hermes-3-Llama-3.1-70B"
+    META_LLAMA_3_1_70B_INSTRUCT = "meta-llama/Llama-3.1-70B-Instruct"
+
+class TEE_LLM(str, Enum):
+    """Enum for LLM models available for TEE execution"""
+
     META_LLAMA_3_1_70B_INSTRUCT = "meta-llama/Llama-3.1-70B-Instruct"

{opengradient-0.3.16.dist-info → opengradient-0.3.18.dist-info}/METADATA

@@ -1,6 +1,6 @@
-Metadata-Version: 2.3
+Metadata-Version: 2.4
 Name: opengradient
-Version: 0.3.16
+Version: 0.3.18
 Summary: Python SDK for OpenGradient decentralized model management & inference services
 Project-URL: Homepage, https://opengradient.ai
 Author-email: OpenGradient <oliver@opengradient.ai>
@@ -25,6 +25,7 @@ License: MIT License
         LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
         OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
         SOFTWARE.
+License-File: LICENSE
 Classifier: Development Status :: 3 - Alpha
 Classifier: Intended Audience :: Developers
 Classifier: License :: OSI Approved :: MIT License
@@ -85,6 +86,7 @@ Requires-Dist: langchain==0.3.7
 Requires-Dist: more-itertools==10.5.0
 Requires-Dist: msgpack==1.1.0
 Requires-Dist: multidict==6.1.0
+Requires-Dist: openai==1.58.1
 Requires-Dist: packaging==24.1
 Requires-Dist: pandas==2.2.3
 Requires-Dist: parsimonious==0.10.0
@@ -215,18 +217,7 @@ tx_hash, finish_reason, message = og.llm_chat(
 )
 ```
 
-### Image Generation
-```python
-tx_hash, image_data = og.generate_image(
-    model="stabilityai/stable-diffusion-xl-base-1.0",
-    prompt="A beautiful sunset over mountains",
-    width=1024,
-    height=1024
-)
 
-with open("generated_image.png", "wb") as f:
-    f.write(image_data)
-```
 
 ## Using the CLI
 
@@ -281,21 +272,4 @@ Or you can use files instead of text input in order to simplify your command:
 opengradient chat --model "mistralai/Mistral-7B-Instruct-v0.3" --messages-file messages.json --tools-file tools.json --max-tokens 200
 ```
 
-### Image Generation
-```bash
-opengradient generate-image \
---model "stabilityai/stable-diffusion-xl-base-1.0" \
---prompt "A beautiful sunset over mountains" \
---output-path sunset.png \
---width 1024 \
---height 1024
-```
-
-Options:
-- `--model`, `-m`: Model identifier for image generation (required)
-- `--prompt`, `-p`: Text prompt for generating the image (required)  
-- `--output-path`, `-o`: Output file path for the generated image (required)
-- `--width`: Output image width in pixels (default: 1024)
-- `--height`: Output image height in pixels (default: 1024)
-
 For more information read the OpenGradient [documentation](https://docs.opengradient.ai/).

opengradient-0.3.18.dist-info/RECORD

@@ -0,0 +1,21 @@
+opengradient/__init__.py,sha256=CRII5thG_vRVtEqyhc48IWgYwjif2S4kOBbBf6MgH8I,9262
+opengradient/account.py,sha256=2B7rtCXQDX-yF4U69h8B9-OUreJU4IqoGXG_1Hn9nWs,1150
+opengradient/cli.py,sha256=niN8tlLaiVEpdtkdWEUbxidG75nxrlb6mMUfUAIjiVw,26400
+opengradient/client.py,sha256=axMbfqAQ6OWq3sM_D4bGVWrpAd-15Ru-bOHJ6R6GruA,35197
+opengradient/defaults.py,sha256=6tsW9Z84z6YtITCsULTTgnN0KRUZjfSoeWJZqdWkYCo,384
+opengradient/exceptions.py,sha256=v4VmUGTvvtjhCZAhR24Ga42z3q-DzR1Y5zSqP_yn2Xk,3366
+opengradient/types.py,sha256=-lGWv_yfXMN48bbvARKIFrj1L0AotIwr2c7GOv1JBcI,2464
+opengradient/utils.py,sha256=lUDPmyPqLwpZI-owyN6Rm3QvUjOn5pLN5G1QyriVm-E,6994
+opengradient/abi/inference.abi,sha256=MR5u9npZ-Yx2EqRW17_M-UnGgFF3mMEMepOwaZ-Bkgc,7040
+opengradient/llm/__init__.py,sha256=J1W_AKPntqlDqLeflhn2x7A0i-dkMT-ol3jlEdFgMWU,1135
+opengradient/llm/og_langchain.py,sha256=F32yN1o8EvRbzZSJkUwI0-FmSVAWMuMTI9ho7wgW5hk,4470
+opengradient/llm/og_openai.py,sha256=GilEkIVDac5Cacennb7YNXD6BKwUj69mfnMkvxkiW5Y,3865
+opengradient/proto/__init__.py,sha256=AhaSmrqV0TXGzCKaoPV8-XUvqs2fGAJBM2aOmDpkNbE,55
+opengradient/proto/infer.proto,sha256=13eaEMcppxkBF8yChptsX9HooWFwJKze7oLZNl-LEb8,1217
+opengradient/proto/infer_pb2.py,sha256=wg2vjLQCNv6HRhYuIqgj9xivi3nO4IPz6E5wh2dhDqY,3446
+opengradient/proto/infer_pb2_grpc.py,sha256=y5GYwD1EdNs892xx58jdfyA0fO5QC7k3uZOtImTHMiE,6891
+opengradient-0.3.18.dist-info/METADATA,sha256=eH1-0CRuheW8Ho8mMbiivn2WIrgAKZf37c4zlFMloes,8744
+opengradient-0.3.18.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+opengradient-0.3.18.dist-info/entry_points.txt,sha256=yUKTaJx8RXnybkob0J62wVBiCp_1agVbgw9uzsmaeJc,54
+opengradient-0.3.18.dist-info/licenses/LICENSE,sha256=xEcvQ3AxZOtDkrqkys2Mm6Y9diEnaSeQRKvxi-JGnNA,1069
+opengradient-0.3.18.dist-info/RECORD,,

{opengradient-0.3.16.dist-info → opengradient-0.3.18.dist-info}/WHEEL

@@ -1,4 +1,4 @@
 Wheel-Version: 1.0
-Generator: hatchling 1.26.3
+Generator: hatchling 1.27.0
 Root-Is-Purelib: true
 Tag: py3-none-any

opengradient-0.3.16.dist-info/RECORD

@@ -1,20 +0,0 @@
-opengradient/__init__.py,sha256=IlVjsDPLaC8Zqaj8_QZCzi6-F20Ta-wyefd1mh4k53Q,4015
-opengradient/account.py,sha256=2B7rtCXQDX-yF4U69h8B9-OUreJU4IqoGXG_1Hn9nWs,1150
-opengradient/cli.py,sha256=kdYR_AFKHV99HtO_son7vHpM5jWVZe8FO0iMWxJ7pJE,24444
-opengradient/client.py,sha256=RdlTz60NJKVJihYY6oVYLfNOg6RGnJbfG-2UIxUk-ws,37069
-opengradient/defaults.py,sha256=6tsW9Z84z6YtITCsULTTgnN0KRUZjfSoeWJZqdWkYCo,384
-opengradient/exceptions.py,sha256=v4VmUGTvvtjhCZAhR24Ga42z3q-DzR1Y5zSqP_yn2Xk,3366
-opengradient/types.py,sha256=QTEsygwT5AnIf8Dg9mexvVUe49nCo9N0pgZOOIp3trc,2214
-opengradient/utils.py,sha256=lUDPmyPqLwpZI-owyN6Rm3QvUjOn5pLN5G1QyriVm-E,6994
-opengradient/abi/inference.abi,sha256=MR5u9npZ-Yx2EqRW17_M-UnGgFF3mMEMepOwaZ-Bkgc,7040
-opengradient/llm/__init__.py,sha256=n_11WFPoU8YtGc6wg9cK6gEy9zBISf1183Loip3dAbI,62
-opengradient/llm/chat.py,sha256=F32yN1o8EvRbzZSJkUwI0-FmSVAWMuMTI9ho7wgW5hk,4470
-opengradient/proto/__init__.py,sha256=AhaSmrqV0TXGzCKaoPV8-XUvqs2fGAJBM2aOmDpkNbE,55
-opengradient/proto/infer.proto,sha256=13eaEMcppxkBF8yChptsX9HooWFwJKze7oLZNl-LEb8,1217
-opengradient/proto/infer_pb2.py,sha256=wg2vjLQCNv6HRhYuIqgj9xivi3nO4IPz6E5wh2dhDqY,3446
-opengradient/proto/infer_pb2_grpc.py,sha256=y5GYwD1EdNs892xx58jdfyA0fO5QC7k3uZOtImTHMiE,6891
-opengradient-0.3.16.dist-info/METADATA,sha256=FrkTwhsj8sHnpaDxKVwCJosvfAwljdJ9yFE9KKHkLMk,9536
-opengradient-0.3.16.dist-info/WHEEL,sha256=C2FUgwZgiLbznR-k0b_5k3Ai_1aASOXDss3lzCUsUug,87
-opengradient-0.3.16.dist-info/entry_points.txt,sha256=yUKTaJx8RXnybkob0J62wVBiCp_1agVbgw9uzsmaeJc,54
-opengradient-0.3.16.dist-info/licenses/LICENSE,sha256=xEcvQ3AxZOtDkrqkys2Mm6Y9diEnaSeQRKvxi-JGnNA,1069
-opengradient-0.3.16.dist-info/RECORD,,