asteroid-odyssey 0.1.21__py3-none-any.whl → 1.0.0__py3-none-any.whl

asteroid_odyssey/client.py

@@ -1,394 +1,444 @@
-from typing import Optional, Dict, Any, List, Union, Callable, Tuple
-from uuid import UUID
-import logging
-import time
-from enum import Enum
-import os
+"""
+Asteroid Agents Python SDK - High-Level Client Interface
 
-from .api.generated.asteroid_agents_api_client.models import (
-    CreateWorkflowRequest,
-    WorkflowExecution,
-    ExecutionStatus,
-    Execution,
-    WorkflowExecutionRequest,
-    Agent,
-    ResultSchema,
-    CreateWorkflowRequestFields,
-    CreateWorkflowRequestProvider
-)
-from .api.generated.asteroid_agents_api_client.client import Client as ApiClient
-from .api.generated.asteroid_agents_api_client.api.execution.get_execution import sync_detailed as get_execution
-from .api.generated.asteroid_agents_api_client.api.agent.get_agents import sync_detailed as asteroid_get_agents
-from .api.generated.asteroid_agents_api_client.api.default.create_workflow import sync_detailed as create_workflow
-from .api.generated.asteroid_agents_api_client.api.workflow.get_agent_workflow_executions import sync_detailed as get_agent_workflow_executions
-from .api.generated.asteroid_agents_api_client.api.workflow.execute_workflow import sync_detailed as execute_workflow
-logger = logging.getLogger(__name__)
+Provides a clean, easy-to-use interface for interacting with the Asteroid Agents API,
+similar to the TypeScript SDK.
 
-class ExecutionTerminalState(Enum):
-    """Terminal states for an execution"""
-    COMPLETED = "completed"
-    FAILED = "failed"
-    CANCELLED = "cancelled"
-    ERROR = "error"
+This module provides a high-level client that wraps the generated OpenAPI client
+without modifying any generated files.
+"""
 
-agent_name = "iris"
+import time
+import os
+from typing import Dict, Any, Optional, List, Union, Tuple
+from openapi_client import (
+    Configuration,
+    ApiClient,
+    SDKApi,
+    ExecutionApi,
+    ExecutionStatusResponse,
+    ExecutionResultResponse,
+    BrowserSessionRecordingResponse,
+    UploadExecutionFiles200Response,
+    Status,
+    StructuredAgentExecutionRequest
+)
+from openapi_client.exceptions import ApiException
 
-class ExecutionResult:
-    """Wrapper class for execution results"""
-    def __init__(self, execution: Execution):
-        self.execution_id = execution.id
-        self.status = execution.status
-        self.result = execution.result
-        self.error = execution.error if hasattr(execution, 'error') else None
-        self.created_at = execution.created_at
-        self.completed_at = execution.completed_at if hasattr(execution, 'completed_at') else None
 
 class AsteroidClient:
     """
-    A high-level client for interacting with the Asteroid API.
+    High-level client for the Asteroid Agents API.
+    
+    This class provides a convenient interface for executing agents and managing
+    their execution lifecycle, similar to the TypeScript SDK.
     """
-    def __init__(
-        self,
-        api_key: Optional[str] = None,
-        base_url: Optional[str] = None,
-        verify_ssl: bool = True
-    ):
+    
+    def __init__(self, api_key: str, base_url: Optional[str] = None):
         """
-        Initialize the Asteroid client.
-
+        Create an API client with the provided API key.
+        
         Args:
-            api_key: API key for authentication. If not provided, will look for ASTEROID_API_KEY env var
-            base_url: Base URL for the API. Defaults to production URL if not specified
-            verify_ssl: Whether to verify SSL certificates. Defaults to True
+            api_key: Your API key for authentication
+            base_url: Optional base URL (defaults to https://odyssey.asteroid.ai/api/v1)
+        
+        Example:
+            client = AsteroidClient('your-api-key')
         """
-        self.api_key = api_key or os.getenv("ASTEROID_API_KEY")
-        if not self.api_key:
-            raise ValueError(
-                "API key is required. Either pass it directly or set ASTEROID_API_KEY environment variable"
-            )
-
-        self.base_url = base_url or "https://odyssey.asteroid.ai/api/v1"
-        # Initialize API client
-        self.client = ApiClient(
-            base_url=self.base_url,
-            verify_ssl=verify_ssl,
-            headers={"X-Asteroid-Agents-Api-Key": f"{self.api_key}"}
+        if api_key is None:
+            raise TypeError("API key cannot be None")
+        
+        # Configure the API client
+        config = Configuration(
+            host=base_url or "https://odyssey.asteroid.ai/api/v1",
+            api_key={'ApiKeyAuth': api_key}
         )
-
-    def get_agents(self) -> List["Agent"]:
-        """
-        Get list of available agents.
         
-        Returns:
-            List of agent details
-        """
-        try:
-            result = asteroid_get_agents(client=self.client)
-            if result.parsed is None:
-                raise ValueError("No agents were returned from the API")
-            if not isinstance(result.parsed, list):
-                raise ValueError(f"The result is not of type list, it is of type: {type(result.parsed)}")
-            return result.parsed
-        except Exception as e:
-            logger.error(f"Failed to get agents: {str(e)}")
-            raise
-
-    def create_workflow(
-        self, 
-        workflow_name: str,
-        start_url: str,
-        prompt: str,
-        result_schema: Optional[ResultSchema] = None
-    ) -> str:
+        self.api_client = ApiClient(config)
+        self.sdk_api = SDKApi(self.api_client)
+        self.execution_api = ExecutionApi(self.api_client)
+        
+    def execute_agent(self, agent_id: str, agent_profile_id: str, execution_data: Dict[str, Any]) -> str: 
         """
-        Create a new workflow for an agent.
-
+        Execute an agent with the provided parameters.
+        
         Args:
-            workflow_name: Name of the workflow
-            start_url: Starting URL for the workflow
-            prompt: Prompt for the workflow
-            result_schema: Optional custom result schema. Currently not fully supported.
-
+            agent_id: The ID of the agent to execute
+            agent_profile_id: The ID of the agent profile
+            execution_data: The execution parameters
+            
         Returns:
-            Workflow ID
-
-        Warning:
-            Custom result schemas are not fully supported yet and will be added in a future update.
-            Currently, only the default schema will be used regardless of input.
+            The execution ID
+            
+        Raises:
+            Exception: If the execution request fails
+            
+        Example:
+            execution_id = client.execute_structured_agent('my-agent-id', 'agent-profile-id', {'input': 'some dynamic value'})
         """
-        if result_schema is not None:
-            logger.warning("Custom result schemas are not fully supported yet and will be ignored. Using default schema.")
-
-        # Default result schema
-        default_schema = ResultSchema()
-        default_schema.additional_properties = {
-            "properties": {
-                "explanation": {
-                    "description": "Detailed explanation of the result",
-                    "type": "string"
-                },
-                "success": {
-                    "description": "Whether the operation was successful",
-                    "type": "boolean"
-                }
-            },
-            "required": [
-                "explanation",
-                "success"
-            ],
-            "type": "object"
-        }
-        
-        if result_schema is None:
-            result_schema = default_schema
-
+        req = StructuredAgentExecutionRequest(agent_profile_id=agent_profile_id, dynamic_data=execution_data)
         try:
-            fields = CreateWorkflowRequestFields()
-            fields.additional_properties = {
-                "workflow_name": workflow_name,
-                "start_url": start_url
-            }
-
-            request = CreateWorkflowRequest(
-                name=workflow_name,
-                result_schema=result_schema,
-                fields=fields,
-                prompts=[prompt],
-                provider=CreateWorkflowRequestProvider.OPENAI
-            )
-
-            result = create_workflow(
-                agent_name=agent_name,
-                body=request,
-                client=self.client
-            ).parsed
-            if not isinstance(result, str):
-                raise ValueError("The result is not of type str")
-            return result
-        except Exception as e:
-            logger.error(f"Failed to create workflow: {str(e)}")
-            raise
-
-    def execute_workflow(
-        self,
-        workflow_id: UUID,
-        execution_params: Dict[str, Any]
-    ) -> str:
+            response = self.sdk_api.execute_agent_structured(agent_id, req)
+            return response.execution_id
+        except ApiException as e:
+            raise Exception(f"Failed to execute agent: {e}")
+    
+    def get_execution_status(self, execution_id: str) -> ExecutionStatusResponse:
         """
-        Execute an existing workflow.
-
+        Get the current status for an execution.
+        
         Args:
-            workflow_id: ID of workflow to execute
-            execution_params: Parameters for workflow execution
-
-        Returns:
-            Execution ID
-        """
-        try:
-            # Convert execution_params to WorkflowExecutionRequest using from_dict
-            request_body = WorkflowExecutionRequest.from_dict(execution_params)
+            execution_id: The execution identifier
             
-            result = execute_workflow(
-                workflow_id=workflow_id,
-                body=request_body,
-                client=self.client
-            ).parsed
-            if not isinstance(result, str):
-                raise ValueError("The result is not of type str")
-            return result
-        except Exception as e:
-            logger.error(f"Failed to execute workflow: {str(e)}")
-            raise
-
-    def get_workflow_executions(self) -> List[WorkflowExecution]:
-        """
-        Get list of workflow executions.
-
         Returns:
-            List of workflow executions
+            The execution status details
+            
+        Raises:
+            Exception: If the status request fails
+            
+        Example:
+            status = client.get_execution_status(execution_id)
+            print(status.status)
         """
         try:
-            result = get_agent_workflow_executions(
-                agent_name=agent_name,
-                client=self.client
-            ).parsed
-            if not isinstance(result, List):
-                raise ValueError("The result is not of type List")
-            return result
-        except Exception as e:
-            logger.error(f"Failed to get workflow executions: {str(e)}")
-            raise
-
-    def get_execution(self, execution_id: str) -> Execution:
+            return self.sdk_api.get_execution_status(execution_id)
+        except ApiException as e:
+            raise Exception(f"Failed to get execution status: {e}")
+    
+    def get_execution_result(self, execution_id: str) -> Dict[str, Any]:
         """
-        Get the full execution details.
-
+        Get the final result of an execution.
+        
         Args:
-            execution_id: ID of the execution to retrieve
-
+            execution_id: The execution identifier
+            
         Returns:
-            Execution object with full details
+            The result object of the execution
+            
+        Raises:
+            Exception: If the result request fails or execution failed
+            
+        Example:
+            result = client.get_execution_result(execution_id)
+            print(result)
         """
         try:
-            result = get_execution(id=UUID(execution_id), client=self.client).parsed
-            if not isinstance(result, Execution):
-                raise ValueError("The result is not of type Execution")
-            return result
-        except Exception as e:
-            logger.error(f"Failed to get execution: {str(e)}")
-            raise
-
-    def get_execution_status(self, execution_id: str) -> ExecutionStatus:
+            response = self.sdk_api.get_execution_result(execution_id)
+            
+            if response.error:
+                raise Exception(response.error)
+            
+            return response.result or {}
+        except ApiException as e:
+            raise Exception(f"Failed to get execution result: {e}")
+    
+    def wait_for_execution_result(
+        self, 
+        execution_id: str, 
+        interval: float = 1.0, 
+        timeout: float = 3600.0
+    ) -> Dict[str, Any]:
         """
-        Get the current status of an execution.
-
+        Wait for an execution to reach a terminal state and return the result.
+        
+        Continuously polls the execution status until it's either "completed", 
+        "cancelled", or "failed".
+        
         Args:
-            execution_id: ID of the execution to check
-
+            execution_id: The execution identifier
+            interval: Polling interval in seconds (default is 1.0)
+            timeout: Maximum wait time in seconds (default is 3600 - 1 hour)
+            
         Returns:
-            Current execution status
+            The execution result if completed
+            
+        Raises:
+            Exception: If the execution ends as "cancelled" or "failed", or times out
+            
+        Example:
+            result = client.wait_for_execution_result(execution_id, interval=2.0)
         """
-        execution = self.get_execution(execution_id)
-        if not isinstance(execution.status, ExecutionStatus):
-            raise ValueError("The execution status is not of type ExecutionStatus")
-        return execution.status
-
-    def get_execution_result(self, execution_id: str) -> ExecutionResult:
+        start_time = time.time()
+        
+        while True:
+            elapsed_time = time.time() - start_time
+            if elapsed_time >= timeout:
+                raise Exception(f"Execution {execution_id} timed out after {timeout}s")
+            
+            status_response = self.get_execution_status(execution_id)
+            current_status = status_response.status
+            
+            if current_status == Status.COMPLETED:
+                return self.get_execution_result(execution_id)
+            elif current_status in [Status.FAILED, Status.CANCELLED]:
+                reason = f" - {status_response.reason}" if status_response.reason else ""
+                raise Exception(f"Execution {execution_id} ended with status: {current_status.value}{reason}")
+            
+            # Wait for the specified interval before polling again
+            time.sleep(interval)
+    
+    def upload_execution_files(
+        self, 
+        execution_id: str, 
+        files: List[Union[bytes, str, Tuple[str, bytes]]],
+        default_filename: str = "file.txt"
+    ) -> UploadExecutionFiles200Response:
         """
-        Get the result of an execution.
-
+        Upload files to an execution.
+        
         Args:
-            execution_id: ID of the execution to get results for
-
+            execution_id: The execution identifier
+            files: List of files to upload. Each file can be:
+                   - bytes: Raw file content (will use default_filename)
+                   - str: File path as string (will read file and use filename)
+                   - Tuple[str, bytes]: (filename, file_content) tuple
+            default_filename: Default filename to use when file is provided as bytes
+            
         Returns:
-            ExecutionResult object containing status, result, and other metadata
-
+            The upload response containing message and file IDs
+            
         Raises:
-            ValueError: If execution doesn't exist or hasn't completed
+            Exception: If the upload request fails
+            
+        Example:
+            # Upload with file content (file should be in your current working directory)
+            with open('hello.txt', 'r') as f:
+                file_content = f.read()
+            
+            response = client.upload_execution_files(execution_id, [file_content.encode()])
+            print(f"Uploaded files: {response.file_ids}")
+            
+            # Upload with filename and content
+            files = [('hello.txt', file_content.encode())]
+            response = client.upload_execution_files(execution_id, files)
+            
+            # Or create content directly
+            hello_content = "Hello World!".encode()
+            response = client.upload_execution_files(execution_id, [hello_content])
         """
-        execution = self.get_execution(execution_id)
-        return ExecutionResult(execution)
-
-    def wait_for_execution(
-        self,
-        execution_id: str,
-        polling_interval: float = 1.0,
-        timeout: Optional[float] = None,
-        status_callback: Optional[Callable[[ExecutionStatus], None]] = None
-    ) -> ExecutionStatus:
+        try:
+            # Process files to ensure proper format
+            processed_files = []
+            for file_item in files:
+                if isinstance(file_item, tuple):
+                    # Already in (filename, content) format
+                    filename, content = file_item
+                    if isinstance(content, str):
+                        content = content.encode()
+                    processed_files.append((filename, content))
+                elif isinstance(file_item, str):
+                    # Check if string is a file path that exists, otherwise treat as content
+                    if os.path.isfile(file_item):
+                        # File path - read the file
+                        filename = os.path.basename(file_item)
+                        with open(file_item, 'rb') as f:
+                            content = f.read()
+                        processed_files.append((filename, content))
+                    else:
+                        # String content - encode and use default filename
+                        content = file_item.encode()
+                        processed_files.append((default_filename, content))
+                elif isinstance(file_item, bytes):
+                    # Raw bytes - use default filename
+                    processed_files.append((default_filename, file_item))
+                else:
+                    # Other types - convert to string content and encode
+                    content = str(file_item).encode()
+                    processed_files.append((default_filename, content))
+            
+            response = self.execution_api.upload_execution_files(execution_id, files=processed_files)
+            return response
+        except ApiException as e:
+            raise Exception(f"Failed to upload execution files: {e}")
+    
+    def get_browser_session_recording(self, execution_id: str) -> str:
         """
-        Wait for an execution to reach a terminal state.
-
+        Get the browser session recording URL for a completed execution.
+        
         Args:
-            execution_id: ID of the execution to wait for
-            polling_interval: Time in seconds between status checks
-            timeout: Maximum time in seconds to wait. None means wait indefinitely
-            status_callback: Optional callback function that will be called with each status update
-
+            execution_id: The execution identifier
+            
         Returns:
-            Final execution status
-
+            The URL of the browser session recording
+            
         Raises:
-            TimeoutError: If timeout is reached before execution reaches terminal state
-            ValueError: If execution_id is invalid
+            Exception: If the recording request fails
+            
+        Example:
+            recording_url = client.get_browser_session_recording(execution_id)
+            print(f"Recording available at: {recording_url}")
         """
-        start_time = time.time()
-        last_status = None
-
-        while True:
-            # Check if we've exceeded timeout
-            if timeout and (time.time() - start_time) > timeout:
-                raise TimeoutError(f"Execution {execution_id} did not complete within {timeout} seconds")
-
-            # Get current status
-            current_status = self.get_execution_status(execution_id)
-
-            # Call status callback if status has changed
-            if status_callback and current_status != last_status:
-                status_callback(current_status)
-            last_status = current_status
-
-            # Check if we've reached a terminal state
-            if current_status.status.value in [state.value for state in ExecutionTerminalState]:
-                return current_status
+        try:
+            response = self.sdk_api.get_browser_session_recording(execution_id)
+            return response.recording_url
+        except ApiException as e:
+            raise Exception(f"Failed to get browser session recording: {e}")
+    
+    def __enter__(self):
+        """Context manager entry."""
+        return self
+    
+    def __exit__(self, exc_type, exc_value, traceback):
+        """Context manager exit."""
+        pass
 
-            # Wait before next check
-            time.sleep(polling_interval)
 
-    def wait_for_execution_result(
-        self,
-        execution_id: str,
-        polling_interval: float = 1.0,
-        timeout: Optional[float] = None,
-        status_callback: Optional[Callable[[ExecutionStatus], None]] = None
-    ) -> ExecutionResult:
-        """
-        Wait for an execution to complete and get its result.
+# Convenience functions that mirror the TypeScript SDK pattern
+def create_client(api_key: str, base_url: Optional[str] = None) -> AsteroidClient:
+    """
+    Create an API client with a provided API key.
+    
+    This is a convenience function that creates an AsteroidClient instance.
+    
+    Args:
+        api_key: Your API key
+        base_url: Optional base URL
+        
+    Returns:
+        A configured AsteroidClient instance
+        
+    Example:
+        client = create_client('your-api-key')
+    """
+    return AsteroidClient(api_key, base_url)
 
-        Args:
-            execution_id: ID of the execution to wait for
-            polling_interval: Time in seconds between status checks
-            timeout: Maximum time in seconds to wait. None means wait indefinitely
-            status_callback: Optional callback function that will be called with each status update
+def execute_agent(client: AsteroidClient, agent_id: str, agent_profile_id: str, execution_data: Dict[str, Any]) -> str:
+    """
+    Execute an agent with the provided parameters.
+    
+    Args:
+        client: The AsteroidClient instance
+        agent_id: The ID of the agent to execute
+        agent_profile_id: The ID of the agent profile
+        execution_data: The execution parameters
+        
+    Returns:
+        The execution ID
+        
+    Example:
+        execution_id = execute_agent(client, 'my-agent-id', {'input': 'some dynamic value'})
+    """
+    return client.execute_agent(agent_id, agent_profile_id, execution_data)
 
-        Returns:
-            ExecutionResult object containing final status, result, and other metadata
 
-        Raises:
-            TimeoutError: If timeout is reached before execution completes
-            ValueError: If execution_id is invalid
-        """
-        # Wait for execution to reach terminal state
-        final_status = self.wait_for_execution(
-            execution_id=execution_id,
-            polling_interval=polling_interval,
-            timeout=timeout,
-            status_callback=status_callback
-        )
 
-        # Get the final result
-        result = self.get_execution_result(execution_id)
+def get_execution_status(client: AsteroidClient, execution_id: str) -> ExecutionStatusResponse:
+    """
+    Get the current status for an execution.
+    
+    Args:
+        client: The AsteroidClient instance
+        execution_id: The execution identifier
+        
+    Returns:
+        The execution status details
+        
+    Example:
+        status = get_execution_status(client, execution_id)
+        print(status.status)
+    """
+    return client.get_execution_status(execution_id)
 
-        # If execution failed, include error information in logs
-        if final_status in [ExecutionTerminalState.FAILED, ExecutionTerminalState.ERROR]:
-            logger.error(f"Execution {execution_id} failed with error: {result.error}")
 
-        return result
+def get_execution_result(client: AsteroidClient, execution_id: str) -> Dict[str, Any]:
+    """
+    Get the final result of an execution.
+    
+    Args:
+        client: The AsteroidClient instance
+        execution_id: The execution identifier
+        
+    Returns:
+        The result object of the execution
+        
+    Example:
+        result = get_execution_result(client, execution_id)
+        print(result)
+    """
+    return client.get_execution_result(execution_id)
 
-    def execute_workflow_and_get_result(
-        self,
-        workflow_id: UUID,
-        execution_params: Dict[str, Any],
-        polling_interval: float = 1.0,
-        timeout: Optional[float] = None,
-        status_callback: Optional[Callable[[ExecutionStatus], None]] = None
-    ) -> ExecutionResult:
-        """
-        Execute a workflow and wait for its result.
 
-        Args:
-            workflow_id: ID of workflow to execute
-            execution_params: Parameters for workflow execution
-            polling_interval: Time in seconds between status checks
-            timeout: Maximum time in seconds to wait. None means wait indefinitely
-            status_callback: Optional callback function that will be called with each status update
+def wait_for_execution_result(
+    client: AsteroidClient, 
+    execution_id: str, 
+    interval: float = 1.0, 
+    timeout: float = 3600.0
+) -> Dict[str, Any]:
+    """
+    Wait for an execution to reach a terminal state and return the result.
+    
+    Args:
+        client: The AsteroidClient instance
+        execution_id: The execution identifier
+        interval: Polling interval in seconds (default is 1.0)
+        timeout: Maximum wait time in seconds (default is 3600 - 1 hour)
+        
+    Returns:
+        The execution result if completed
+        
+    Example:
+        result = wait_for_execution_result(client, execution_id, interval=2.0)
+    """
+    return client.wait_for_execution_result(execution_id, interval, timeout)
 
-        Returns:
-            ExecutionResult object containing final status, result, and other metadata
-        """
-        # Start execution
-        execution_id = self.execute_workflow(workflow_id, execution_params)
 
-        # Wait for result
-        return self.wait_for_execution_result(
-            execution_id=execution_id,
-            polling_interval=polling_interval,
-            timeout=timeout,
-            status_callback=status_callback
-        )
+def upload_execution_files(
+    client: AsteroidClient, 
+    execution_id: str, 
+    files: List[Union[bytes, str, Tuple[str, bytes]]],
+    default_filename: str = "file.txt"
+) -> UploadExecutionFiles200Response:
+    """
+    Upload files to an execution.
+    
+    Args:
+        client: The AsteroidClient instance
+        execution_id: The execution identifier
+        files: List of files to upload
+        default_filename: Default filename to use when file is provided as bytes
+        
+    Returns:
+        The upload response containing message and file IDs
+        
+    Example:
+        # Create a simple text file with "Hello World!" content
+        hello_content = "Hello World!".encode()
+        response = upload_execution_files(client, execution_id, [hello_content])
+        print(f"Uploaded files: {response.file_ids}")
+        
+        # Or specify filename with content
+        files = [('hello.txt', "Hello World!".encode())]
+        response = upload_execution_files(client, execution_id, files)
+    """
+    return client.upload_execution_files(execution_id, files, default_filename)
 
-    def __enter__(self):
-        return self
 
-    def __exit__(self, exc_type, exc_val, exc_tb):
-        pass
+def get_browser_session_recording(client: AsteroidClient, execution_id: str) -> str:
+    """
+    Get the browser session recording URL for a completed execution.
+    
+    Args:
+        client: The AsteroidClient instance
+        execution_id: The execution identifier
+        
+    Returns:
+        The URL of the browser session recording
+        
+    Example:
+        recording_url = get_browser_session_recording(client, execution_id)
+        print(f"Recording available at: {recording_url}")
+    """
+    return client.get_browser_session_recording(execution_id)
+
+
+# Re-export common types for convenience
+__all__ = [
+    'AsteroidClient',
+    'create_client',
+    'execute_agent',
+    'get_execution_status',
+    'get_execution_result',
+    'wait_for_execution_result',
+    'upload_execution_files',
+    'get_browser_session_recording'
+]