craft-ai-sdk 0.63.1rc1__tar.gz → 0.64.0rc1__tar.gz

{craft_ai_sdk-0.63.1rc1 → craft_ai_sdk-0.64.0rc1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: craft-ai-sdk
-Version: 0.63.1rc1
+Version: 0.64.0rc1
 Summary: Craft AI MLOps platform SDK
 License: Apache-2.0
 Author: Craft AI

{craft_ai_sdk-0.63.1rc1 → craft_ai_sdk-0.64.0rc1}/craft_ai_sdk/__init__.py

@@ -13,4 +13,4 @@ from .io import (  # noqa: F401
 )
 from .sdk import CraftAiSdk  # noqa: F401
 
-__version__ = "0.63.1rc1"
+__version__ = "0.64.0rc1"

{craft_ai_sdk-0.63.1rc1 → craft_ai_sdk-0.64.0rc1}/craft_ai_sdk/core/data_store.py

@@ -1,4 +1,5 @@
 import io
+from typing import TypedDict, Union
 
 import requests
 
@@ -8,7 +9,19 @@ from ..shared.request_response_handler import handle_data_store_response
 from ..utils import chunk_buffer, convert_size
 
 
-def get_data_store_object_information(sdk: BaseCraftAiSdk, object_path_in_datastore):
+class DataStoreObjectInformation(TypedDict):
+    path: str
+    last_modified: str
+    size: str
+
+
+class DataStoreDeletedObject(TypedDict):
+    path: str
+
+
+def get_data_store_object_information(
+    sdk: BaseCraftAiSdk, object_path_in_datastore: str
+) -> DataStoreObjectInformation:
     """Get information about a single object in the data store.
 
     Args:
@@ -33,7 +46,9 @@ def get_data_store_object_information(sdk: BaseCraftAiSdk, object_path_in_datast
     return result
 
 
-def list_data_store_objects(sdk: BaseCraftAiSdk, folder_path=None):
+def list_data_store_objects(
+    sdk: BaseCraftAiSdk, folder_path: Union[str, None] = None
+) -> list[DataStoreObjectInformation]:
     """Get the list of the objects stored in the data store.
 
     Args:
@@ -63,7 +78,7 @@ def list_data_store_objects(sdk: BaseCraftAiSdk, folder_path=None):
     return all_items
 
 
-def _get_upload_presigned_url(sdk: BaseCraftAiSdk, object_path_in_datastore):
+def _get_upload_presigned_url(sdk: BaseCraftAiSdk, object_path_in_datastore: str):
     url = f"{sdk.base_environment_api_url}/data-store/upload"
     params = {"path_to_object": object_path_in_datastore}
     resp = sdk._get(url, params=params)
@@ -74,7 +89,9 @@ def _get_upload_presigned_url(sdk: BaseCraftAiSdk, object_path_in_datastore):
 
 @log_func_result("Object upload")
 def upload_data_store_object(
-    sdk: BaseCraftAiSdk, filepath_or_buffer, object_path_in_datastore
+    sdk: BaseCraftAiSdk,
+    filepath_or_buffer: Union[str, io.IOBase],
+    object_path_in_datastore: str,
 ):
     """Upload a file as an object into the data store.
 
@@ -117,7 +134,7 @@ def upload_data_store_object(
 
 
 def _upload_singlepart_data_store_object(
-    sdk: BaseCraftAiSdk, buffer, object_path_in_datastore
+    sdk: BaseCraftAiSdk, buffer: io.IOBase, object_path_in_datastore: str
 ):
     files = {"file": buffer}
 
@@ -128,7 +145,7 @@ def _upload_singlepart_data_store_object(
 
 
 def _upload_multipart_data_store_object(
-    sdk: BaseCraftAiSdk, buffer, object_path_in_datastore
+    sdk: BaseCraftAiSdk, buffer: io.IOBase, object_path_in_datastore: str
 ):
     multipart_base_url = f"{sdk.base_environment_api_url}/data-store/upload/multipart"
     multipart_upload_configuration = sdk._post(
@@ -156,10 +173,10 @@ def _upload_multipart_data_store_object(
             data=chunk,
             headers=multipart_part_result["headers"],
         )
-        partData = {"number": part_idx}
+        part_data: dict = {"number": part_idx}
         if "ETag" in resp.headers:
-            partData["metadata"] = resp.headers["ETag"]
-        parts.append(partData)
+            part_data["metadata"] = resp.headers["ETag"]
+        parts.append(part_data)
         if part["lastChunk"]:
             break
 
@@ -173,7 +190,9 @@ def _upload_multipart_data_store_object(
     )
 
 
-def _get_download_presigned_url(sdk: BaseCraftAiSdk, object_path_in_datastore):
+def _get_download_presigned_url(
+    sdk: BaseCraftAiSdk, object_path_in_datastore: str
+) -> str:
     url = f"{sdk.base_environment_api_url}/data-store/download"
     data = {
         "path_to_object": object_path_in_datastore,
@@ -184,7 +203,9 @@ def _get_download_presigned_url(sdk: BaseCraftAiSdk, object_path_in_datastore):
 
 @log_func_result("Object download")
 def download_data_store_object(
-    sdk: BaseCraftAiSdk, object_path_in_datastore, filepath_or_buffer
+    sdk: BaseCraftAiSdk,
+    object_path_in_datastore: str,
+    filepath_or_buffer: Union[str, io.IOBase],
 ):
     """Download an object in the data store and save it into a file.
 
@@ -218,7 +239,9 @@ def download_data_store_object(
 
 
 @log_func_result("Object deletion")
-def delete_data_store_object(sdk: BaseCraftAiSdk, object_path_in_datastore):
+def delete_data_store_object(
+    sdk: BaseCraftAiSdk, object_path_in_datastore: str
+) -> DataStoreDeletedObject:
     """Delete an object on the datastore.
 
     Args:

{craft_ai_sdk-0.63.1rc1 → craft_ai_sdk-0.64.0rc1}/craft_ai_sdk/core/deployments.py

@@ -1,6 +1,17 @@
+from datetime import datetime
+from typing import Literal, TypedDict, Union
+from typing_extensions import NotRequired
+
+from craft_ai_sdk.shared.types import Log
 from ..constants import DEPLOYMENT_EXECUTION_RULES, DEPLOYMENT_MODES, DEPLOYMENT_STATUS
 from ..exceptions import SdkException
-from ..io import _validate_inputs_mapping, _validate_outputs_mapping
+from ..io import (
+    InputSource,
+    InputSourceDict,
+    OutputDestination,
+    _validate_inputs_mapping,
+    _validate_outputs_mapping,
+)
 from ..sdk import BaseCraftAiSdk
 from ..shared.logger import log_action, log_func_result
 from ..utils import (
@@ -10,25 +21,98 @@ from ..utils import (
 )
 
 
+class DeploymentInputMapping(InputSourceDict):
+    data_type: str
+    description: str
+
+
+class DeploymentOutputMapping(OutputDestination):
+    data_type: str
+    description: str
+
+
+class DeploymentPipeline(TypedDict):
+    name: str
+
+
+class DeploymentPod(TypedDict):
+    pod_id: str
+    status: str
+
+
+class DeploymentBase(TypedDict):
+    name: str
+    mode: str
+    pipeline: DeploymentPipeline
+    inputs_mapping: list[DeploymentInputMapping]
+    outputs_mapping: list[DeploymentOutputMapping]
+    created_at: str
+    created_by: str
+    updated_at: str
+    updated_by: str
+    last_execution_id: str
+    is_enabled: bool
+    description: str
+    status: str
+    enable_parallel_executions: bool
+    max_parallel_executions_per_pod: int
+    pods: NotRequired[list[DeploymentPod]]
+
+
+class DeploymentEndpoint(DeploymentBase):
+    endpoint_token: str
+    endpoint_url_path: str
+    execution_rule: Literal["endpoint"]
+
+
+class DeploymentPeriodic(DeploymentBase):
+    schedule: str
+    human_readable_schedule: str
+    execution_rule: Literal["periodic"]
+
+
+Deployment = Union[DeploymentEndpoint, DeploymentPeriodic]
+
+
+class DeploymentLog(Log):
+    stream: str
+    pod_id: str
+    type: str
+
+
+class DeploymentListItem(TypedDict):
+    name: str
+    pipeline_name: str
+    execution_rule: str
+    mode: str
+    is_enabled: bool
+    created_at: str
+
+
+class DeploymentDeleted(TypedDict):
+    name: str
+    execution_rule: str
+
+
 @log_func_result("Deployment creation")
 def create_deployment(
     sdk: BaseCraftAiSdk,
-    pipeline_name,
-    deployment_name,
-    execution_rule,
+    pipeline_name: str,
+    deployment_name: str,
+    execution_rule: DEPLOYMENT_EXECUTION_RULES,
     mode=DEPLOYMENT_MODES.ELASTIC,
-    schedule=None,
-    endpoint_url_path=None,
-    inputs_mapping=None,
-    outputs_mapping=None,
-    description=None,
-    enable_parallel_executions=None,
-    max_parallel_executions_per_pod=None,
-    ram_request=None,
-    gpu_request=None,
+    schedule: Union[str, None] = None,
+    endpoint_url_path: Union[str, None] = None,
+    inputs_mapping: Union[list[InputSource], None] = None,
+    outputs_mapping: Union[list[OutputDestination], None] = None,
+    description: Union[str, None] = None,
+    enable_parallel_executions: Union[bool, None] = None,
+    max_parallel_executions_per_pod: Union[int, None] = None,
+    ram_request: Union[str, None] = None,
+    gpu_request: Union[int, None] = None,
     wait_for_completion=True,
-    timeout_s=None,
-    execution_timeout_s=None,
+    timeout_s: Union[int, None] = None,
+    execution_timeout_s: Union[int, None] = None,
 ):
     """Create a deployment associated with a given pipeline.
 
@@ -397,8 +481,11 @@ FINAL_DEPLOYMENT_STATUSES = [
 
 
 def get_deployment(
-    sdk: BaseCraftAiSdk, deployment_name, wait_for_completion=False, timeout_s=None
-):
+    sdk: BaseCraftAiSdk,
+    deployment_name: str,
+    wait_for_completion=False,
+    timeout_s: Union[int, None] = None,
+) -> Deployment:
     """Get information of a deployment.
 
     Args:
@@ -563,13 +650,13 @@ wait_for_completion parameter set to false.',
 @log_func_result("Deployment update")
 def update_deployment(
     sdk: BaseCraftAiSdk,
-    deployment_name,
-    is_enabled=None,
-    inputs_mapping=None,
-    outputs_mapping=None,
-    schedule=None,
+    deployment_name: str,
+    is_enabled: Union[bool, None] = None,
+    inputs_mapping: Union[list[InputSource], None] = None,
+    outputs_mapping: Union[list[OutputDestination], None] = None,
+    schedule: Union[str, None] = None,
     wait_for_completion=True,
-    timeout_s=None,
+    timeout_s: Union[int, None] = None,
 ):
     """Update the specified properties of a deployment. The properties that can be
     updated include enabling/disabling the deployment, updating input/output values,
@@ -624,11 +711,11 @@ def update_deployment(
 
 def get_deployment_logs(
     sdk: BaseCraftAiSdk,
-    deployment_name,
-    from_datetime=None,
-    to_datetime=None,
-    limit=None,
-):
+    deployment_name: str,
+    from_datetime: Union[datetime, None] = None,
+    to_datetime: Union[datetime, None] = None,
+    limit: Union[int, None] = None,
+) -> list[DeploymentLog]:
     """Get the logs of a deployment with ``"low_latency"`` mode.
 
     Args:
@@ -667,7 +754,7 @@ def get_deployment_logs(
     return sdk._post(url, json=data)
 
 
-def list_deployments(sdk: BaseCraftAiSdk):
+def list_deployments(sdk: BaseCraftAiSdk) -> list[DeploymentListItem]:
     """Get the list of all deployments.
 
     Returns:
@@ -689,7 +776,7 @@ def list_deployments(sdk: BaseCraftAiSdk):
 
 
 @log_func_result("Deployment deletion")
-def delete_deployment(sdk: BaseCraftAiSdk, deployment_name):
+def delete_deployment(sdk: BaseCraftAiSdk, deployment_name: str) -> DeploymentDeleted:
     """Delete a deployment identified by its name.
 
     Args:

{craft_ai_sdk-0.63.1rc1 → craft_ai_sdk-0.64.0rc1}/craft_ai_sdk/core/endpoints.py

@@ -1,4 +1,5 @@
 import io
+from typing import Any, Literal, TypedDict, Union, overload
 from urllib.parse import urlencode
 
 import requests
@@ -9,7 +10,7 @@ from ..shared.request_response_handler import handle_http_response
 from .deployments import get_deployment
 
 
-def _get_endpoint_url_path(sdk: BaseCraftAiSdk, endpoint_name):
+def _get_endpoint_url_path(sdk: BaseCraftAiSdk, endpoint_name: str):
     deployment = get_deployment(sdk, endpoint_name)
 
     if deployment.get("execution_rule", "") != "endpoint":
@@ -18,14 +19,46 @@ def _get_endpoint_url_path(sdk: BaseCraftAiSdk, endpoint_name):
     return deployment.get("endpoint_url_path", "")
 
 
+class EndpointTriggerBase(TypedDict):
+    execution_id: str
+
+
+class EndpointTriggerWithOutputs(EndpointTriggerBase):
+    outputs: dict[str, Any]
+
+
+class EndpointNewToken(TypedDict):
+    endpoint_token: str
+
+
+@overload
+def trigger_endpoint(
+    sdk: BaseCraftAiSdk,
+    endpoint_name: str,
+    endpoint_token: str,
+    inputs: dict[str, Any],
+    wait_for_completion: Literal[True],
+) -> EndpointTriggerWithOutputs: ...
+
+
+@overload
+def trigger_endpoint(
+    sdk: BaseCraftAiSdk,
+    endpoint_name: str,
+    endpoint_token: str,
+    inputs: dict[str, Any],
+    wait_for_completion: Literal[False],
+) -> EndpointTriggerBase: ...
+
+
 @log_func_result("Endpoint trigger")
 def trigger_endpoint(
     sdk: BaseCraftAiSdk,
-    endpoint_name,
-    endpoint_token,
-    inputs=None,
+    endpoint_name: str,
+    endpoint_token: str,
+    inputs: Union[dict[str, Any], None] = None,
     wait_for_completion=True,
-):
+) -> Union[EndpointTriggerWithOutputs, EndpointTriggerBase]:
     """Trigger an endpoint.
 
     Args:
@@ -88,10 +121,10 @@ def trigger_endpoint(
 @log_func_result("Endpoint result retrieval")
 def retrieve_endpoint_results(
     sdk: BaseCraftAiSdk,
-    endpoint_name,
-    execution_id,
-    endpoint_token,
-):
+    endpoint_name: str,
+    execution_id: str,
+    endpoint_token: str,
+) -> EndpointTriggerWithOutputs:
     """Get the results of an endpoint execution.
 
     Args:
@@ -135,7 +168,9 @@ def retrieve_endpoint_results(
         return {**handle_http_response(response), "execution_id": execution_id}
 
 
-def generate_new_endpoint_token(sdk: BaseCraftAiSdk, endpoint_name):
+def generate_new_endpoint_token(
+    sdk: BaseCraftAiSdk, endpoint_name: str
+) -> EndpointNewToken:
     """Generate a new endpoint token for an endpoint.
 
     Args:

{craft_ai_sdk-0.63.1rc1 → craft_ai_sdk-0.64.0rc1}/craft_ai_sdk/core/environment_variables.py

@@ -1,10 +1,16 @@
+from typing import TypedDict
 from ..sdk import BaseCraftAiSdk
 from ..shared.logger import log_func_result
 
 
+class EnvironmentVariable(TypedDict):
+    name: str
+    value: str
+
+
 @log_func_result("Environment variable definition")
 def create_or_update_environment_variable(
-    sdk: BaseCraftAiSdk, environment_variable_name, environment_variable_value
+    sdk: BaseCraftAiSdk, environment_variable_name: str, environment_variable_value: str
 ):
     """Create or update an environment variable available for
     all pipelines executions.
@@ -28,7 +34,7 @@ def create_or_update_environment_variable(
     return None
 
 
-def list_environment_variables(sdk: BaseCraftAiSdk):
+def list_environment_variables(sdk: BaseCraftAiSdk) -> list[EnvironmentVariable]:
     """Get a list of all environments variables.
 
     Returns:
@@ -43,7 +49,9 @@ def list_environment_variables(sdk: BaseCraftAiSdk):
 
 
 @log_func_result("Environment variable deletion")
-def delete_environment_variable(sdk: BaseCraftAiSdk, environment_variable_name):
+def delete_environment_variable(
+    sdk: BaseCraftAiSdk, environment_variable_name: str
+) -> EnvironmentVariable:
     """Delete the specified environment variable
 
     Args:

{craft_ai_sdk-0.63.1rc1 → craft_ai_sdk-0.64.0rc1}/craft_ai_sdk/core/pipeline_metrics.py

@@ -1,3 +1,4 @@
+from typing import TypedDict, Union
 import warnings
 
 from ..sdk import BaseCraftAiSdk
@@ -6,8 +7,26 @@ from ..shared.logger import log_func_result
 from ..utils import remove_none_values
 
 
+class Metric(TypedDict):
+    name: str
+    value: float
+    created_at: str
+    execution_id: str
+    deployment_name: str
+    pipeline_name: str
+
+
+class ListMetric(TypedDict):
+    name: str
+    values: list[float]
+    created_at: str
+    execution_id: str
+    deployment_name: str
+    pipeline_name: str
+
+
 @log_func_result("Pipeline metrics definition", get_execution_id)
-def record_metric_value(sdk: BaseCraftAiSdk, name, value):
+def record_metric_value(sdk: BaseCraftAiSdk, name: str, value: float):
     """Create or update a pipeline metric. Note that this function can only be used
     inside a step code.
 
@@ -34,7 +53,7 @@ been sent",
 
 
 @log_func_result("Pipeline list metric definition", get_execution_id)
-def record_list_metric_values(sdk: BaseCraftAiSdk, name, values):
+def record_list_metric_values(sdk: BaseCraftAiSdk, name: str, values: list[float]):
     """Add values to a pipeline metric list. Note that this function can only be
     used inside a step code.
 
@@ -74,11 +93,11 @@ been sent",
 @log_func_result("Pipeline metrics listing")
 def get_metrics(
     sdk: BaseCraftAiSdk,
-    name=None,
-    pipeline_name=None,
-    deployment_name=None,
-    execution_id=None,
-):
+    name: Union[str, None] = None,
+    pipeline_name: Union[str, None] = None,
+    deployment_name: Union[str, None] = None,
+    execution_id: Union[str, None] = None,
+) -> list[Metric]:
     """Get a list of pipeline metrics. Note that only one of the
     parameters (pipeline_name, deployment_name, execution_id) can be set.
 
@@ -140,11 +159,11 @@ def get_metrics(
 @log_func_result("Pipeline list metrics listing")
 def get_list_metrics(
     sdk: BaseCraftAiSdk,
-    name=None,
-    pipeline_name=None,
-    deployment_name=None,
-    execution_id=None,
-):
+    name: Union[str, None] = None,
+    pipeline_name: Union[str, None] = None,
+    deployment_name: Union[str, None] = None,
+    execution_id: Union[str, None] = None,
+) -> list[ListMetric]:
     """Get a list of pipeline metric lists. Note that only one of the
     parameters (pipeline_name, deployment_name, execution_id) can be set.
 
@@ -162,7 +181,7 @@ def get_list_metrics(
         with the following keys:
 
           * ``name`` (:obj:`str`): Name of the metric.
-          * ``value`` (:obj:`float`): Value of the metric.
+          * ``values`` (:obj:`list[float]`): Values list of the metric.
           * ``created_at`` (:obj:`str`): Date of the metric creation.
           * ``execution_id`` (:obj:`str`): Name of the execution the metric
             belongs to.

{craft_ai_sdk-0.63.1rc1 → craft_ai_sdk-0.64.0rc1}/craft_ai_sdk/core/pipelines.py

@@ -1,21 +1,70 @@
 import os
+from typing import TypedDict, Union, cast
+from typing_extensions import NotRequired
 import warnings
 
 import requests
 
+from craft_ai_sdk.io import Input, Output
+from craft_ai_sdk.shared.types import Log
+
 from ..sdk import BaseCraftAiSdk
 from ..shared.helpers import wait_create_until_ready
 from ..shared.logger import log_action, log_func_result
 from ..shared.request_response_handler import handle_data_store_response
 from ..utils import datetime_to_timestamp_in_ms, multipartify, remove_keys_from_dict
 from .steps import (
+    ContainerConfig,
+    StepCreationInfo,
     _prepare_create_step_data,
     _prepare_create_step_files,
     _validate_create_step_parameters,
 )
 
 
-def _create_pipeline_with_step(sdk: BaseCraftAiSdk, pipeline_name: str, step_name: str):
+class PipelineParameter(TypedDict):
+    step_name: NotRequired[str]
+    pipeline_name: str
+    function_path: str
+    function_name: str
+    description: str
+    inputs: list[Input]
+    outputs: list[Output]
+    container_config: ContainerConfig
+
+
+class PipelineCreationInfo(StepCreationInfo):
+    last_execution_id: str
+
+
+class Pipeline(TypedDict):
+    parameters: PipelineParameter
+    creation_info: PipelineCreationInfo
+
+
+class PipelineListItem(TypedDict):
+    pipeline_name: str
+    created_at: str
+    status: str
+
+
+class PipelineDeletedPipeline(TypedDict):
+    name: str
+
+
+class PipelineDeletedDeployment(TypedDict):
+    name: str
+    execution_rule: str
+
+
+class PipelineDeleted(TypedDict):
+    pipeline: PipelineDeletedPipeline
+    deployments: list[PipelineDeletedDeployment]
+
+
+def _create_pipeline_with_step(
+    sdk: BaseCraftAiSdk, pipeline_name: str, step_name: str
+) -> Pipeline:
     url = f"{sdk.base_environment_api_url}/pipelines"
     body = {
         "pipeline_name": pipeline_name,
@@ -30,15 +79,15 @@ def _create_pipeline_with_step(sdk: BaseCraftAiSdk, pipeline_name: str, step_nam
 @log_func_result("Pipeline creation")
 def create_pipeline(
     sdk: BaseCraftAiSdk,
-    pipeline_name,
-    function_path=None,
-    function_name=None,
-    description=None,
-    container_config=None,
-    inputs=None,
-    outputs=None,
+    pipeline_name: str,
+    function_path: Union[str, None] = None,
+    function_name: Union[str, None] = None,
+    description: Union[str, None] = None,
+    container_config: Union[ContainerConfig, None] = None,
+    inputs: Union[list[Input], None] = None,
+    outputs: Union[list[Output], None] = None,
     wait_for_completion=True,
-    timeout_s=None,
+    timeout_s: Union[int, None] = None,
     **kwargs,
 ):
     """Create a pipeline from a function located on a remote repository or locally.
@@ -211,7 +260,11 @@ def create_pipeline(
         )
 
     # Otherwise, the pipeline is created with a "hidden" step
-    container_config = {} if container_config is None else container_config.copy()
+    container_config = (
+        cast(ContainerConfig, {})
+        if container_config is None
+        else cast(ContainerConfig, container_config.copy())
+    )
     _validate_create_step_parameters(inputs, outputs, timeout_s)
 
     url = f"{sdk.base_environment_api_url}/pipelines"
@@ -239,7 +292,7 @@ def create_pipeline(
 
 def get_pipeline(
     sdk: BaseCraftAiSdk, pipeline_name, wait_for_completion=False, timeout_s=None
-):
+) -> Pipeline:
     """Get a single pipeline if it exists.
 
     Args:
@@ -336,7 +389,7 @@ def get_pipeline(
     )
 
 
-def list_pipelines(sdk: BaseCraftAiSdk):
+def list_pipelines(sdk: BaseCraftAiSdk) -> list[PipelineListItem]:
     """Get the list of all pipelines.
 
     Returns:
@@ -354,8 +407,8 @@ def list_pipelines(sdk: BaseCraftAiSdk):
 
 @log_func_result("Pipeline deletion")
 def delete_pipeline(
-    sdk: BaseCraftAiSdk, pipeline_name, force_deployments_deletion=False
-):
+    sdk: BaseCraftAiSdk, pipeline_name: str, force_deployments_deletion=False
+) -> PipelineDeleted:
     """Delete a pipeline identified by its name.
 
     Args:
@@ -386,14 +439,16 @@ def delete_pipeline(
     return sdk._delete(url, params=params)
 
 
-def _get_download_presigned_url(sdk: BaseCraftAiSdk, pipeline_name):
+def _get_download_presigned_url(sdk: BaseCraftAiSdk, pipeline_name: str) -> str:
     url = f"{sdk.base_environment_api_url}/pipelines/{pipeline_name}/download"
     presigned_url = sdk._get(url)["signed_url"]
     return presigned_url
 
 
 @log_func_result("pipeline download")
-def download_pipeline_local_folder(sdk: BaseCraftAiSdk, pipeline_name, folder):
+def download_pipeline_local_folder(
+    sdk: BaseCraftAiSdk, pipeline_name: str, folder: str
+):
     """Download a pipeline's local folder as a `.tgz` archive.
 
     Only available if the pipeline's ``origin`` is ``"local_folder"``. This archive
@@ -427,7 +482,7 @@ def get_pipeline_logs(
     from_datetime=None,
     to_datetime=None,
     limit=None,
-):
+) -> list[Log]:
     """Get the logs of a pipeline.
 
     Args:

{craft_ai_sdk-0.63.1rc1 → craft_ai_sdk-0.64.0rc1}/craft_ai_sdk/core/resource_metrics.py

@@ -1,8 +1,58 @@
+from datetime import datetime
+from typing import Literal, TypedDict, Union, overload
+from typing_extensions import NotRequired
 from ..sdk import BaseCraftAiSdk
 from ..utils import datetime_to_timestamp_in_ms
 
 
-def get_resource_metrics(sdk: BaseCraftAiSdk, start_date, end_date, csv=False):
+class AdditionalData(TypedDict):
+    total_disk: int
+    total_ram: int
+    total_vram: NotRequired[str]
+
+
+class MetricWorker(TypedDict):
+    worker: str
+
+
+class Metric(TypedDict):
+    metric: MetricWorker
+    values: list[list[Union[int, float]]]
+
+
+class MetricsDict(TypedDict):
+    cpu_usage: list[Metric]
+    disk_usage: list[Metric]
+    ram_usage: list[Metric]
+    vram_usage: NotRequired[list[Metric]]
+    gpu_usage: NotRequired[list[Metric]]
+    network_input_usage: list[Metric]
+    network_output_usage: list[Metric]
+
+
+class ResourceMetrics(TypedDict):
+    additional_data: AdditionalData
+    metrics: MetricsDict
+
+
+@overload
+def get_resource_metrics(
+    sdk: BaseCraftAiSdk, start_date: datetime, end_date: datetime, csv: Literal[True]
+) -> bytes: ...
+
+
+@overload
+def get_resource_metrics(
+    sdk: BaseCraftAiSdk,
+    start_date: datetime,
+    end_date: datetime,
+    csv: Literal[False] = False,
+) -> ResourceMetrics: ...
+
+
+def get_resource_metrics(
+    sdk: BaseCraftAiSdk, start_date: datetime, end_date: datetime, csv=False
+) -> Union[ResourceMetrics, bytes]:
     """Get resource metrics of the environment.
 
     Args:

{craft_ai_sdk-0.63.1rc1 → craft_ai_sdk-0.64.0rc1}/craft_ai_sdk/core/steps.py

@@ -1,9 +1,14 @@
+from datetime import datetime
 import io
 import os
 import tarfile
+from typing import TypeVar, TypedDict, Union, cast
+from typing_extensions import NotRequired
 
 import requests
 
+from craft_ai_sdk.shared.types import Log
+
 from ..constants import CREATION_PARAMETER_VALUE
 from ..io import Input, Output
 from ..sdk import BaseCraftAiSdk
@@ -12,8 +17,67 @@ from ..shared.logger import log_action, log_func_result
 from ..shared.request_response_handler import handle_data_store_response
 from ..utils import datetime_to_timestamp_in_ms, multipartify, remove_none_values
 
+T = TypeVar("T")
+
+CreationParameter = Union[T, CREATION_PARAMETER_VALUE, None]
+
+
+class ContainerConfigBase(TypedDict):
+    language: NotRequired[CreationParameter[str]]
+    requirements_path: NotRequired[CreationParameter[str]]
+    included_folders: NotRequired[CreationParameter[list[str]]]
+    system_dependencies: NotRequired[CreationParameter[list[str]]]
+    dockerfile_path: NotRequired[CreationParameter[str]]
+
+
+class ContainerConfigWithGit(ContainerConfigBase):
+    repository_url: NotRequired[CreationParameter[str]]
+    repository_branch: NotRequired[CreationParameter[str]]
+    repository_deploy_key: NotRequired[CreationParameter[str]]
+
+
+class ContainerConfigWithLocalFolder(ContainerConfigBase):
+    local_folder: str
+
+
+ContainerConfig = Union[ContainerConfigWithGit, ContainerConfigWithLocalFolder]
+
+
+class StepParameter(TypedDict):
+    step_name: str
+    function_path: str
+    function_name: str
+    description: str
+    inputs: list[Input]
+    outputs: list[Output]
+    container_config: ContainerConfig
+    "This type is actually too large, you'll not get any CREATION_PARAMETER_VALUE here"
+
+
+class StepCreationInfo(TypedDict):
+    created_at: str
+    created_by: str
+    commit_id: str
+    status: str
+    origin: str
 
-def _compress_folder_to_memory(local_folder, include):
+
+class Step(TypedDict):
+    parameters: StepParameter
+    creation_info: StepCreationInfo
+
+
+class StepListItem(TypedDict):
+    step_name: str
+    status: str
+    created_at: str
+
+
+class StepDeleted(TypedDict):
+    step_name: str
+
+
+def _compress_folder_to_memory(local_folder: str, include: list[str]):
     tar_data = io.BytesIO()
     # Remove leading slashes from the paths
     include = [item.lstrip("/") for item in include]
@@ -24,7 +88,11 @@ def _compress_folder_to_memory(local_folder, include):
     return tar_data
 
 
-def _validate_create_step_parameters(inputs, outputs, timeout_s):
+def _validate_create_step_parameters(
+    inputs: Union[list[Input], None],
+    outputs: Union[list[Output], None],
+    timeout_s: Union[int, None],
+):
     if timeout_s is not None and timeout_s <= 0:
         raise ValueError("The timeout must be greater than 0 or None.")
 
@@ -37,7 +105,7 @@ def _validate_create_step_parameters(inputs, outputs, timeout_s):
             raise ValueError("'outputs' must be a list of instances of Output.")
 
 
-def _map_container_config_step_parameter(container_config):
+def _map_container_config_step_parameter(container_config: ContainerConfig):
     """
     Maps container config with :obj:`CREATION_PARAMETER_VALUE` enum values to final
     container config. `None` is considered to be equivalent to
@@ -57,12 +125,12 @@ def _map_container_config_step_parameter(container_config):
 
 
 def _prepare_create_step_data(
-    function_path,
-    function_name,
-    description,
-    container_config,
-    inputs,
-    outputs,
+    function_path: Union[str, None],
+    function_name: Union[str, None],
+    description: Union[str, None],
+    container_config: ContainerConfig,
+    inputs: Union[list[Input], None],
+    outputs: Union[list[Output], None],
     **step_or_pipeline_name,
 ):
     assert step_or_pipeline_name.keys() == {
@@ -87,7 +155,12 @@ def _prepare_create_step_data(
     return data
 
 
-def _prepare_create_step_files(sdk, container_config, data, function_path):
+def _prepare_create_step_files(
+    sdk: BaseCraftAiSdk,
+    container_config: ContainerConfig,
+    data: dict,
+    function_path: Union[str, None],
+):
     if "local_folder" not in container_config:
         return {}
 
@@ -136,15 +209,15 @@ def _remove_id_from_step(step):
 @log_func_result("Steps creation")
 def create_step(
     sdk: BaseCraftAiSdk,
-    step_name,
-    function_path=None,
-    function_name=None,
-    description=None,
-    container_config=None,
-    inputs=None,
-    outputs=None,
+    step_name: str,
+    function_path: Union[str, None] = None,
+    function_name: Union[str, None] = None,
+    description: Union[str, None] = None,
+    container_config: Union[ContainerConfig, None] = None,
+    inputs: Union[list[Input], None] = None,
+    outputs: Union[list[Output], None] = None,
     wait_for_completion=True,
-    timeout_s=None,
+    timeout_s: Union[int, None] = None,
 ):
     """Create pipeline step from a function located on a remote repository or locally.
 
@@ -277,7 +350,11 @@ def create_step(
             ``"git_repository"`` or ``"local"``.
     """
 
-    container_config = {} if container_config is None else container_config.copy()
+    container_config = (
+        cast(ContainerConfig, {})
+        if container_config is None
+        else cast(ContainerConfig, container_config.copy())
+    )
     _validate_create_step_parameters(inputs, outputs, timeout_s=timeout_s)
 
     url = f"{sdk.base_environment_api_url}/steps"
@@ -302,7 +379,12 @@ def create_step(
     return get_step(sdk, step_name, wait_for_completion, timeout_s)
 
 
-def get_step(sdk: BaseCraftAiSdk, step_name, wait_for_completion=False, timeout_s=None):
+def get_step(
+    sdk: BaseCraftAiSdk,
+    step_name: str,
+    wait_for_completion=False,
+    timeout_s: Union[int, None] = None,
+) -> Step:
     """Get a single step if it exists.
 
     Args:
@@ -391,7 +473,7 @@ def get_step(sdk: BaseCraftAiSdk, step_name, wait_for_completion=False, timeout_
     return _remove_id_from_step(step)
 
 
-def list_steps(sdk: BaseCraftAiSdk):
+def list_steps(sdk: BaseCraftAiSdk) -> list[StepListItem]:
     """Get the list of all steps.
 
     Returns:
@@ -408,7 +490,9 @@ def list_steps(sdk: BaseCraftAiSdk):
 
 
 @log_func_result("Step deletion")
-def delete_step(sdk: BaseCraftAiSdk, step_name, force_dependents_deletion=False):
+def delete_step(
+    sdk: BaseCraftAiSdk, step_name: str, force_dependents_deletion=False
+) -> StepDeleted:
     """Delete one step.
 
     Args:
@@ -431,14 +515,14 @@ def delete_step(sdk: BaseCraftAiSdk, step_name, force_dependents_deletion=False)
     return sdk._delete(url, params=params)
 
 
-def _get_download_presigned_url(sdk: BaseCraftAiSdk, step_name):
+def _get_download_presigned_url(sdk: BaseCraftAiSdk, step_name: str) -> str:
     url = f"{sdk.base_environment_api_url}/steps/{step_name}/download"
     presigned_url = sdk._get(url)["signed_url"]
     return presigned_url
 
 
 @log_func_result("Step download")
-def download_step_local_folder(sdk: BaseCraftAiSdk, step_name, folder):
+def download_step_local_folder(sdk: BaseCraftAiSdk, step_name: str, folder: str):
     """Download a step's local folder as a `.tgz` archive.
 
     Only available if the step's ``origin`` is ``"local_folder"``. This archive
@@ -468,11 +552,11 @@ def download_step_local_folder(sdk: BaseCraftAiSdk, step_name, folder):
 @log_func_result("Step logs")
 def get_step_logs(
     sdk: BaseCraftAiSdk,
-    step_name,
-    from_datetime=None,
-    to_datetime=None,
-    limit=None,
-):
+    step_name: str,
+    from_datetime: Union[datetime, None] = None,
+    to_datetime: Union[datetime, None] = None,
+    limit: Union[int, None] = None,
+) -> list[Log]:
     """Get the logs of a step.
 
     Args:

{craft_ai_sdk-0.63.1rc1 → craft_ai_sdk-0.64.0rc1}/craft_ai_sdk/core/users.py

@@ -1,7 +1,14 @@
+from typing import TypedDict
 from ..sdk import BaseCraftAiSdk
 
 
-def get_user(sdk: BaseCraftAiSdk, user_id):
+class User(TypedDict):
+    id: str
+    name: str
+    email: str
+
+
+def get_user(sdk: BaseCraftAiSdk, user_id: str) -> User:
     """Get information about a user.
 
     Args:

{craft_ai_sdk-0.63.1rc1 → craft_ai_sdk-0.64.0rc1}/craft_ai_sdk/core/vector_database.py

@@ -1,10 +1,16 @@
+from typing import TypedDict
 from craft_ai_sdk.shared.environments import get_environment_id
 
 from ..sdk import BaseCraftAiSdk
 from ..shared.logger import log_action, log_func_result
 
 
-def get_vector_database_credentials(sdk: BaseCraftAiSdk):
+class VectorDatabaseCredentials(TypedDict):
+    vector_database_url: str
+    vector_database_token: str
+
+
+def get_vector_database_credentials(sdk: BaseCraftAiSdk) -> VectorDatabaseCredentials:
     """Get the credentials of the vector database.
 
     Returns:
@@ -62,6 +68,9 @@ def get_weaviate_client(sdk: BaseCraftAiSdk):
         auth_credentials=None,
     )
 
-    log_action(sdk, "Connected to Weaviate, using version", weaviate.__version__)
+    log_action(
+        sdk,
+        f"Connected to Weaviate, using version {weaviate.__version__}",
+    )
 
     return weaviate_client

{craft_ai_sdk-0.63.1rc1 → craft_ai_sdk-0.64.0rc1}/craft_ai_sdk/io.py

@@ -1,3 +1,5 @@
+from typing import Any, TypedDict, cast
+from typing_extensions import NotRequired
 import warnings
 
 from strenum import LowercaseStrEnum
@@ -82,6 +84,18 @@ class Output:
         return remove_none_values(output)
 
 
+class InputSourceDict(TypedDict):
+    pipeline_input_name: str
+    description: str
+    constant_value: NotRequired[Any]
+    environment_variable_name: NotRequired[str]
+    endpoint_input_name: NotRequired[str]
+    is_null: NotRequired[bool]
+    datastore_path: NotRequired[str]
+    is_required: NotRequired[bool]
+    default_value: NotRequired[Any]
+
+
 class InputSource:
     """Class to specify to which source a step input should be mapped when creating
     a deployment (cf. :meth:`.CraftAiSdk.create_deployment`). The different sources can
@@ -167,7 +181,7 @@ class InputSource:
         self.datastore_path = datastore_path
         self.step_input_name = step_input_name
 
-    def to_dict(self):
+    def to_dict(self) -> InputSourceDict:
         input_mapping_dict = {
             "pipeline_input_name": self.pipeline_input_name,
             "endpoint_input_name": self.endpoint_input_name,
@@ -180,7 +194,15 @@ class InputSource:
             "step_input_name": self.step_input_name,
         }
 
-        return remove_none_values(input_mapping_dict)
+        return cast(InputSourceDict, remove_none_values(input_mapping_dict))
+
+
+class OutputDestinationDict(TypedDict):
+    pipeline_output_name: str
+    endpoint_output_name: NotRequired[str]
+    is_null: NotRequired[bool]
+    datastore_path: NotRequired[str]
+    step_output_name: NotRequired[str]
 
 
 class OutputDestination:
@@ -252,7 +274,7 @@ class OutputDestination:
         self.datastore_path = datastore_path
         self.step_output_name = step_output_name
 
-    def to_dict(self):
+    def to_dict(self) -> OutputDestinationDict:
         output_mapping_dict = {
             "pipeline_output_name": self.pipeline_output_name,
             "endpoint_output_name": self.endpoint_output_name,
@@ -261,7 +283,7 @@ class OutputDestination:
             "step_output_name": self.step_output_name,
         }
 
-        return remove_none_values(output_mapping_dict)
+        return cast(OutputDestinationDict, remove_none_values(output_mapping_dict))
 
 
 def _format_execution_output(name, output):

{craft_ai_sdk-0.63.1rc1 → craft_ai_sdk-0.64.0rc1}/craft_ai_sdk/sdk.py

@@ -23,6 +23,7 @@ class BaseCraftAiSdk(ABC):
     _MULTIPART_THRESHOLD: int
     _MULTIPART_PART_SIZE: int
     _version: str
+    warn_on_metric_outside_of_step: bool
 
     @abstractmethod
     def _get(self, url: str, params=None, **kwargs) -> Any:
@@ -137,7 +138,7 @@ class CraftAiSdk(BaseCraftAiSdk):
         os.environ.get("CRAFT_AI__MULTIPART_PART_SIZE__B", str(38 * 256 * 1024))
     )
     _access_token_margin = timedelta(seconds=30)
-    _version = "0.63.1rc1"  # Would be better to share it somewhere
+    _version = "0.64.0rc1"  # Would be better to share it somewhere
 
     def __init__(
         self,

craft_ai_sdk-0.64.0rc1/craft_ai_sdk/shared/types.py

@@ -0,0 +1,6 @@
+from typing import TypedDict
+
+
+class Log(TypedDict):
+    timestamp: str
+    message: str

{craft_ai_sdk-0.63.1rc1 → craft_ai_sdk-0.64.0rc1}/craft_ai_sdk/utils/datetime_utils.py

@@ -2,13 +2,13 @@ import re
 from datetime import datetime
 
 
-def datetime_to_timestamp_in_ms(dt):
+def datetime_to_timestamp_in_ms(dt: datetime) -> int:
     if not isinstance(dt, datetime):
         raise ValueError("Parameter must be a datetime.datetime object.")
     return int(1_000 * dt.timestamp())
 
 
-def parse_isodate(date_string):
+def parse_isodate(date_string: str):
     """_summary_
 
     Args:

{craft_ai_sdk-0.63.1rc1 → craft_ai_sdk-0.64.0rc1}/craft_ai_sdk/utils/dict_utils.py

@@ -1,7 +1,7 @@
 from typing import Union
 
 
-def remove_none_values(obj):
+def remove_none_values(obj: dict):
     return {key: value for key, value in obj.items() if value is not None}
 
 

{craft_ai_sdk-0.63.1rc1 → craft_ai_sdk-0.64.0rc1}/craft_ai_sdk/utils/file_utils.py

@@ -1,5 +1,5 @@
 from io import BytesIO, IOBase, StringIO
-from typing import Iterable, Union
+from typing import Iterable, TypedDict, Union
 
 
 def merge_paths(prefix, path):
@@ -8,7 +8,13 @@ def merge_paths(prefix, path):
 
 
 # From https://stackoverflow.com/a/58767245/4839162
-def chunk_buffer(buffer: IOBase, size: int) -> Iterable[Union[BytesIO, StringIO]]:
+class ChunkedIO(TypedDict):
+    chunk: Union[BytesIO, StringIO]
+    len: int
+    lastChunk: bool
+
+
+def chunk_buffer(buffer: IOBase, size: int) -> Iterable[ChunkedIO]:
     size_int = int(size)
     b = buffer.read(size_int)
     next_data = None
@@ -22,7 +28,7 @@ def chunk_buffer(buffer: IOBase, size: int) -> Iterable[Union[BytesIO, StringIO]
 
         next_data = buffer.read(1)
 
-        data = {
+        data: ChunkedIO = {
             "chunk": chunk,
             "len": len(b) + (len(previous_data) if previous_data else 0),
             "lastChunk": len(next_data) == 0,
@@ -32,7 +38,7 @@ def chunk_buffer(buffer: IOBase, size: int) -> Iterable[Union[BytesIO, StringIO]
         b = buffer.read(size_int - 1)
 
 
-def convert_size(size_in_bytes):
+def convert_size(size_in_bytes: Union[int, float]):
     """
     Convert a size in bytes to a human readable string.
     """

{craft_ai_sdk-0.63.1rc1 → craft_ai_sdk-0.64.0rc1}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "craft-ai-sdk"
-version = "0.63.1rc1"
+version = "0.64.0rc1"
 description = "Craft AI MLOps platform SDK"
 license = "Apache-2.0"
 authors = ["Craft AI <contact@craft.ai>"]