ionworks-api 0.1.0__py3-none-any.whl

ionworks/client.py

@@ -0,0 +1,126 @@
+import os
+from typing import Any, cast
+
+from dotenv import load_dotenv
+import requests
+
+from .cell_instance import CellInstanceClient
+from .cell_measurement import CellMeasurementClient
+from .cell_specification import CellSpecificationClient
+from .errors import IonworksError
+from .job import JobClient
+from .pipeline import PipelineClient
+from .simulation import SimulationClient
+
+
+class Ionworks:
+    def __init__(self, api_key: str | None = None, api_url: str | None = None):
+        """Initialize Ionworks client.
+
+        Parameters
+        ----------
+        api_key : str, optional
+            API key. If not provided, will look for IONWORKS_API_KEY env var
+        api_url : str, optional
+            API URL. If not provided, will look for IONWORKS_API_URL env var
+        """
+        load_dotenv()
+
+        self.api_key = api_key or os.getenv("IONWORKS_API_KEY")
+        if not self.api_key:
+            raise ValueError(
+                "API key must be provided either as argument or in IONWORKS_API_KEY "
+                "environment variable"
+            )
+
+        self.api_url = (
+            api_url or os.getenv("IONWORKS_API_URL") or "https://api.ionworks.com"
+        )
+
+        self.session = requests.Session()
+        self.session.headers.update(
+            {
+                "X-API-Key": self.api_key,
+                "Content-Type": "application/json",
+                "Accept": "application/json",
+            }
+        )
+
+        # Initialize components
+        self.job = JobClient(self)
+        self.pipeline = PipelineClient(self)
+        self.cell_spec = CellSpecificationClient(self)
+        self.cell_instance = CellInstanceClient(self)
+        self.cell_measurement = CellMeasurementClient(self)
+        self.simulation = SimulationClient(self)
+        # Backwards-compatible alias
+        self.measurement = self.cell_measurement
+
+    def request(
+        self, method: str, endpoint: str, json_payload: dict[str, Any] | None = None
+    ) -> Any:
+        """Make a request to the Ionworks API with standardized error handling."""
+        url = f"{self.api_url}{endpoint}"
+        try:
+            response = self.session.request(method, url, json=json_payload)
+            response.raise_for_status()
+
+            # For DELETE operations, don't try to parse JSON if response is empty
+            if method.upper() == "DELETE":
+                return None
+
+            # Return JSON response if content type is JSON and response has content
+            if (
+                response.headers.get("Content-Type") == "application/json"
+                and response.text
+            ):
+                return response.json()
+            return response
+        except requests.exceptions.HTTPError as e:
+            try:
+                error_detail = e.response.json().get("detail", str(e))
+            except requests.exceptions.JSONDecodeError:
+                error_detail = str(e)
+            # Raise the custom error with just the detail message
+            raise IonworksError(
+                error_detail, status_code=e.response.status_code
+            ) from None
+        except requests.exceptions.RequestException as e:
+            # Extract more detailed error information
+            error_msg = f"Error during request to {url}"
+            if hasattr(e, "response") and e.response is not None:
+                error_msg += f" (status code: {e.response.status_code})"
+                try:
+                    error_detail = e.response.json().get("detail", e.response.text)
+                    error_msg += f": {error_detail}"
+                except requests.exceptions.JSONDecodeError:
+                    error_msg += f": {e.response.text}"
+            else:
+                error_msg += f": {str(e)}"
+            # Raise the custom error for general request issues
+            raise IonworksError(error_msg) from None
+
+    def get(self, endpoint: str) -> Any:
+        """Make a GET request using the request helper."""
+        return self.request("GET", endpoint)
+
+    def post(self, endpoint: str, json_payload: dict[str, Any]) -> Any:
+        """Make a POST request using the request helper."""
+        return self.request("POST", endpoint, json_payload=json_payload)
+
+    def put(self, endpoint: str, json_payload: dict[str, Any]) -> Any:
+        """Make a PUT request using the request helper."""
+        return self.request("PUT", endpoint, json_payload=json_payload)
+
+    def delete(self, endpoint: str) -> None:
+        """Make a DELETE request using the request helper."""
+        self.request("DELETE", endpoint)
+
+    def health_check(self) -> dict[str, Any]:
+        """Check the health of the Ionworks API.
+
+        Returns:
+            dict[str, Any]: Health check response
+        """
+        response_data = self.get("/healthz")
+        return cast(dict[str, Any], response_data)

ionworks/errors.py

@@ -0,0 +1,35 @@
+from typing import Any
+
+
+class IonworksError(Exception):
+    """Custom exception for Ionworks API errors.
+
+    Attributes
+    ----------
+    message : str
+        A string description of the error
+    data : dict[str, Any] | None
+        Structured error data if available (e.g., from API error response)
+    status_code : int | None
+        HTTP status code if applicable
+    """
+
+    def __init__(
+        self,
+        message: str | dict[str, Any],
+        status_code: int | None = None,
+    ):
+        self.status_code = status_code
+
+        # Parse message into string and optional data dict
+        if isinstance(message, dict):
+            self.message = message.get("message", str(message))
+            self.data: dict[str, Any] | None = message
+        else:
+            self.message = message
+            self.data = None
+
+        super().__init__(self.message)
+
+    def __str__(self):
+        return f"error code: {self.status_code}, message: {self.message}"

ionworks/job.py

@@ -0,0 +1,111 @@
+from typing import Any, Optional
+
+from pydantic import BaseModel, ValidationError
+
+
+# Model for the data required to create ANY job
+class JobCreationPayload(BaseModel):
+    job_type: str
+    params: dict[str, Any]
+    priority: int = 5  # Default priority
+    callback_url: Optional[str] = None  # Optional callback
+
+
+# Model for the detailed job response
+class JobResponse(BaseModel):
+    job_id: str
+    status: str
+    job_type: str
+    params: dict[str, Any]
+    priority: int
+    created_at: str
+    updated_at: str
+    metadata: Optional[dict[str, Any]]
+    error: Optional[str]
+    result: Optional[dict[str, Any]]
+
+
+class JobClient:
+    def __init__(self, client: Any):
+        self.client = client
+
+    def create(self, payload: JobCreationPayload) -> JobResponse:
+        """
+        Submit a job using the provided payload.
+
+        Parameters
+        ----------
+        payload : JobCreationPayload
+            The configuration for the job to be created.
+
+        Returns
+        -------
+        JobResponse
+            Response containing the job_id and initial status.
+
+        Raises
+        ------
+        requests.exceptions.RequestException
+            If the API request fails.
+        ValueError
+            If the response parsing fails.
+        """
+        endpoint = "/jobs/"
+        try:
+            response_data = self.client.post(
+                endpoint, json_payload=payload.model_dump(exclude_none=True)
+            )
+            # Pydantic validation is applied here
+            return JobResponse(**response_data)
+        except ValidationError as e:
+            # Catch Pydantic validation errors specifically
+            raise ValueError(
+                f"Invalid response format received from {endpoint}: {e}"
+            ) from e
+        # RequestExceptions (including HTTPError) are handled by client._post
+
+    def get(self, job_id: str) -> JobResponse:
+        """
+        Get the status and details of a specific job.
+        """
+        endpoint = f"/jobs/{job_id}"
+        try:
+            response_data = self.client.get(endpoint)
+            return JobResponse(**response_data)
+        except ValidationError as e:
+            raise ValueError(
+                f"Invalid response format received from {endpoint}: {e}"
+            ) from e
+
+    def list(self) -> list[JobResponse]:
+        """
+        List all jobs.
+        """
+        endpoint = "/jobs/"
+        try:
+            response_data = self.client.get(endpoint)
+            # Ensure response_data is a list before list comprehension
+            if not isinstance(response_data, list):
+                raise ValueError(
+                    f"Unexpected response format from {endpoint}: expected a list, "
+                    "got {type(response_data).__name__}"
+                )
+            # Apply validation within list comprehension
+            return [JobResponse(**job) for job in response_data]
+        except ValidationError as e:
+            raise ValueError(
+                f"Invalid job data format received from {endpoint}: {e}"
+            ) from e
+
+    def cancel(self, job_id: str) -> JobResponse:
+        """
+        Cancel a job.
+        """
+        endpoint = f"/jobs/{job_id}/cancel"
+        try:
+            response_data = self.client.post(endpoint, json_payload={})
+            return JobResponse(**response_data)
+        except ValidationError as e:
+            raise ValueError(
+                f"Invalid response format received from {endpoint}: {e}"
+            ) from e

ionworks/models.py

@@ -0,0 +1,185 @@
+"""
+Pydantic models for the Ionworks API client.
+
+These models use extra="allow" to accept any fields from the API response,
+letting the API handle validation. Required fields are kept minimal.
+"""
+
+from typing import Any
+
+import pandas as pd  # type: ignore
+from pydantic import BaseModel, ConfigDict, field_validator
+
+from .validators import dict_to_df_validator
+
+# --- Cell Specification Models --- #
+
+
+class CellSpecification(BaseModel):
+    """
+    Cell specification model - accepts any fields from the API.
+
+    The API returns nested component/material data and ratings objects.
+    This model is permissive to allow the API to define the schema.
+    """
+
+    model_config = ConfigDict(extra="allow")
+
+    id: str
+    slug: str
+    name: str
+
+
+# --- Cell Instance Models --- #
+
+
+class CellInstance(BaseModel):
+    """
+    Cell instance model - accepts any fields from the API.
+    """
+
+    model_config = ConfigDict(extra="allow")
+
+    id: str
+    slug: str
+
+
+# --- Cell Measurement Models --- #
+
+
+class CellMeasurement(BaseModel):
+    """
+    Cell measurement model - accepts any fields from the API.
+    """
+
+    model_config = ConfigDict(extra="allow")
+
+    id: str
+    slug: str
+    name: str
+
+
+# --- Bundle Models --- #
+
+
+class CellMeasurementBundleResponse(BaseModel):
+    """Response from creating a measurement bundle."""
+
+    measurement: CellMeasurement
+    steps_created: int
+    file_path: str
+
+
+class InitiateUploadResponse(BaseModel):
+    """Response from the initiate-upload endpoint for signed URL uploads."""
+
+    measurement_id: str
+    signed_url: str
+    token: str
+    path: str
+
+
+class ConfirmUploadResponse(BaseModel):
+    """Response from the confirm-upload endpoint after successful upload."""
+
+    instance: CellInstance
+    measurement: CellMeasurement
+    steps_created: int
+    file_path: str
+
+
+# --- Detail Models --- #
+
+
+class CellMeasurementDetailBase(BaseModel):
+    """Base model for measurement detail with steps and time series."""
+
+    measurement: CellMeasurement
+    steps: pd.DataFrame
+    time_series: pd.DataFrame
+
+    @field_validator("steps", "time_series", mode="before")
+    @classmethod
+    def convert_dict_to_df(cls, v: Any) -> Any:
+        return dict_to_df_validator(v)
+
+    class Config:
+        arbitrary_types_allowed = True  # Allow pandas DataFrame
+
+
+class CellInstanceDetail(BaseModel):
+    """Detail model for a cell instance with all measurements."""
+
+    specification: CellSpecification
+    instance: CellInstance
+    measurements: list[CellMeasurementDetailBase]
+
+
+class CellMeasurementDetail(CellMeasurementDetailBase):
+    """Detail model for a single measurement with context."""
+
+    specification: CellSpecification
+    instance: CellInstance
+
+
+# --- Group Response Models --- #
+
+
+class CellInstanceMeasurementsWithSteps(BaseModel):
+    """
+    Response model for all measurements + steps associated with a specific cell
+    instance.
+    Returns normalized data: separate lists for instance, measurements, and steps.
+    Includes parent information for consistency with other group endpoints.
+    """
+
+    specification: CellSpecification
+    instance: CellInstance
+    measurements: list[CellMeasurement]
+    steps: list[pd.DataFrame]
+
+    @field_validator("steps", mode="before")
+    @classmethod
+    def convert_dict_to_df(cls, v: Any) -> Any:
+        if isinstance(v, list):
+            return [dict_to_df_validator(item) for item in v]
+        return dict_to_df_validator(v)
+
+    class Config:
+        arbitrary_types_allowed = True
+
+
+class CellSpecificationInstances(BaseModel):
+    """
+    Response model for all instances associated with a cell specification.
+    Returns normalized data: separate lists for specification, instances,
+    and measurements.
+    """
+
+    specification: CellSpecification
+    instances: list[CellInstance]
+    measurements: list[CellMeasurement]
+
+
+class CellSpecificationInstancesWithSteps(BaseModel):
+    """
+    Response model for all instances + measurements + steps associated with a cell
+    specification.
+    Returns normalized data: separate lists for specification, instances, measurements,
+    and steps.
+    """
+
+    specification: CellSpecification
+    instances: list[CellInstance]
+    measurements: list[CellMeasurement]
+    steps: list[pd.DataFrame]
+
+    @field_validator("steps", mode="before")
+    @classmethod
+    def convert_dict_to_df(cls, v: Any) -> Any:
+        if isinstance(v, list):
+            return [dict_to_df_validator(item) for item in v]
+        return dict_to_df_validator(v)
+
+    class Config:
+        arbitrary_types_allowed = True