luminarycloud 0.22.0__py3-none-any.whl → 0.22.2__py3-none-any.whl

luminarycloud/_client/authentication_plugin.py

@@ -38,3 +38,52 @@ class AuthenticationPlugin(grpc.AuthMetadataPlugin):
             callback(metadata, None)
         except Exception as err:
             callback(None, err)
+
+
+class AuthInterceptor(
+    grpc.UnaryUnaryClientInterceptor,
+    grpc.UnaryStreamClientInterceptor,
+    grpc.StreamUnaryClientInterceptor,
+    grpc.StreamStreamClientInterceptor,
+):
+    """
+    I need this as a workaround for container-to-host connections because I need to create a channel
+    that uses CallCredentials but doesn't use any ChannelCredentials. I.e. I need to authenticate
+    the requests, but I need the connection to be unencrypted. This is because the grpc server on
+    the native host isn't using SSL, so I can't use grpc.ssl_channel_credentials(), but it's also
+    not reachable on a loopback interface, so I can't use grpc.local_channel_credentials() either.
+    So I need to use a grpc.insecure_channel(), but you can't use CallCredentials with an insecure
+    channel. So the workaround is to use an interceptor instead of CallCredentials.
+
+    Also, I don't care about auth0, so I'm only supporting an API key.
+    """
+
+    def __init__(self, api_key: str):
+        self._api_key = api_key
+
+    def _augment(self, metadata):
+        return metadata + [("x-api-key", self._api_key)]
+
+    def intercept_unary_unary(self, continuation, client_call_details, request):
+        new_details = client_call_details._replace(
+            metadata=self._augment(client_call_details.metadata or [])
+        )
+        return continuation(new_details, request)
+
+    def intercept_unary_stream(self, continuation, client_call_details, request):
+        new_details = client_call_details._replace(
+            metadata=self._augment(client_call_details.metadata or [])
+        )
+        return continuation(new_details, request)
+
+    def intercept_stream_unary(self, continuation, client_call_details, request_iter):
+        new_details = client_call_details._replace(
+            metadata=self._augment(client_call_details.metadata or [])
+        )
+        return continuation(new_details, request_iter)
+
+    def intercept_stream_stream(self, continuation, client_call_details, request_iter):
+        new_details = client_call_details._replace(
+            metadata=self._augment(client_call_details.metadata or [])
+        )
+        return continuation(new_details, request_iter)

luminarycloud/_client/client.py

@@ -34,7 +34,9 @@ from .._proto.api.v0.luminarycloud.named_variable_set.named_variable_set_pb2_grp
 from .._proto.api.v0.luminarycloud.physics_ai.physics_ai_pb2_grpc import (
     PhysicsAiServiceStub,
 )
-from .._proto.api.v0.luminarycloud.inference.inference_pb2_grpc import InferenceServiceStub
+from .._proto.api.v0.luminarycloud.physicsaiinference.physicsaiinference_pb2_grpc import (
+    PhysicsAiInferenceServiceStub,
+)
 from .._proto.api.v0.luminarycloud.thirdpartyintegration.onshape.onshape_pb2_grpc import (
     OnshapeServiceStub,
 )
@@ -47,7 +49,7 @@ from .._proto.api.v0.luminarycloud.vis.vis_pb2_grpc import VisAPIServiceStub
 from .._proto.api.v0.luminarycloud.feature_flag.feature_flag_pb2_grpc import (
     FeatureFlagServiceStub,
 )
-from .authentication_plugin import AuthenticationPlugin
+from .authentication_plugin import AuthenticationPlugin, AuthInterceptor
 from .config import LC_DOMAIN, LC_API_KEY
 from .logging_interceptor import LoggingInterceptor
 from .retry_interceptor import RetryInterceptor
@@ -71,8 +73,8 @@ class Client(
     StoppingConditionServiceStub,
     NamedVariableSetServiceStub,
     PhysicsAiServiceStub,
-    InferenceServiceStub,
     OnshapeServiceStub,
+    PhysicsAiInferenceServiceStub,
     ProjectUIStateServiceStub,
     FeatureFlagServiceStub,
 ):
@@ -93,11 +95,20 @@ class Client(
         The URL of the HTTP REST server. If not provided, it will default to the `target`.
     localhost : bool
         True if the API server is running locally.
+    insecure_grpc_channel : bool
+        True to use an unencrypted gRPC channel, even though requests are authenticated. There's no
+        legitimate reason to do this outside of a local development situation where the SDK is
+        running from a container and connecting to an API server that is running on the host.
     grpc_channel_options : Optional[Iterable[tuple[str, str]]]
         A list of gRPC channel args. The full list is available here:
         https://github.com/grpc/grpc/blob/v1.46.x/include/grpc/impl/codegen/grpc_types.h
     api_key : Optional[str]
         The API key to use for authentication.
+    log_retries : bool
+        True to log each retriable error response. There are some errors the API server may return
+        that are known to be transient, and the client will always retry requests when it gets one
+        of them. By default, the client retries silently. Set this to True to log the error
+        responses (at INFO level) for the retriable errors.
     **kwargs : dict, optional
         Additional arguments are passed to Auth0Client. See _auth/auth.py.
 
@@ -116,9 +127,11 @@ class Client(
         target: str = LC_DOMAIN,
         http_target: str | None = None,
         localhost: bool = False,
+        insecure_grpc_channel: bool = False,
         grpc_channel_options: Optional[Iterable[tuple[str, Union[str, int]]]] = None,
         channel_credentials: Optional[grpc.ChannelCredentials] = None,
         api_key: Optional[str] = LC_API_KEY,
+        log_retries: bool = False,
         **kwargs: Any,
     ):
         self._target = target
@@ -139,7 +152,12 @@ class Client(
         if grpc_channel_options:
             grpc_channel_options_with_keep_alive.extend(grpc_channel_options)
         self._channel = self._create_channel(
-            localhost, grpc_channel_options_with_keep_alive, channel_credentials, api_key
+            localhost,
+            insecure_grpc_channel,
+            grpc_channel_options_with_keep_alive,
+            channel_credentials,
+            api_key,
+            log_retries,
         )
         self._context_tokens: list[Token] = []
         self.__register_rpcs()
@@ -193,9 +211,11 @@ class Client(
     def _create_channel(
         self,
         localhost: bool = False,
+        insecure: bool = False,
         grpc_channel_options: Optional[Iterable[tuple[str, Union[str, int]]]] = None,
         channel_credentials: Optional[grpc.ChannelCredentials] = None,
         api_key: Optional[str] = None,
+        log_retries: bool = False,
     ) -> grpc.Channel:
         if channel_credentials is None:
             if localhost:
@@ -217,15 +237,22 @@ class Client(
         call_creds = grpc.metadata_call_credentials(auth_plugin)
         composite_creds = grpc.composite_channel_credentials(channel_credentials, call_creds)
         options = grpc_channel_options and list(grpc_channel_options)
-        channel = grpc.secure_channel(
-            self._target,
-            composite_creds,
-            options=options,
-        )
+        if insecure:
+            channel = grpc.insecure_channel(self._target, options=options)
+            channel = grpc.intercept_channel(
+                channel,
+                AuthInterceptor(api_key),
+            )
+        else:
+            channel = grpc.secure_channel(
+                self._target,
+                composite_creds,
+                options=options,
+            )
         intercepted_channel = grpc.intercept_channel(
             channel,
             LoggingInterceptor(),
-            RetryInterceptor(),
+            RetryInterceptor(log_retries),
         )
         return add_instrumentation(
             intercepted_channel,
@@ -248,7 +275,7 @@ class Client(
         OutputDefinitionServiceStub.__init__(self, self._channel)
         StoppingConditionServiceStub.__init__(self, self._channel)
         PhysicsAiServiceStub.__init__(self, self._channel)
-        InferenceServiceStub.__init__(self, self._channel)
+        PhysicsAiInferenceServiceStub.__init__(self, self._channel)
         NamedVariableSetServiceStub.__init__(self, self._channel)
         OnshapeServiceStub.__init__(self, self._channel)
         ProjectUIStateServiceStub.__init__(self, self._channel)

luminarycloud/_client/http_client.py

@@ -32,7 +32,7 @@ class HttpClient:
         api_key: str | None = None,
         auth0_client: Auth0Client | None = None,
         *,
-        timeout: int = 10,
+        timeout: int | None = 300,
         retries: int = 3,
         backoff_factor: float = 0.3,
         retriable_status_codes: tuple = (500, 502, 503, 504, 429),

luminarycloud/_client/retry_interceptor.py

@@ -2,19 +2,68 @@
 from collections.abc import Callable
 from time import sleep
 from typing import Any
+import logging
 
 import grpc
 from grpc import (
     ClientCallDetails,
     UnaryUnaryClientInterceptor,
 )
+from grpc_status.rpc_status import GRPC_DETAILS_METADATA_KEY
 
 from luminarycloud.exceptions import AuthenticationError
+from luminarycloud._proto.base import base_pb2
+
+
+def _is_rate_limited(call: grpc.Call) -> bool:
+    """
+    Check if a gRPC call failed due to rate limiting.
+
+    Rate limit errors are identified with the SUBCODE_RATE_LIMITED subcode and UNAVAILABLE status.
+
+    Args:
+        call: The gRPC call to check
+
+    Returns:
+        True if the error is a rate limit error, False otherwise
+    """
+    if call.code() != grpc.StatusCode.UNAVAILABLE:
+        return False
+
+    try:
+        # Get the trailing metadata which contains error details
+        # Metadata is a sequence of tuples (key, value)
+        for key, value in call.trailing_metadata() or []:
+            if key != GRPC_DETAILS_METADATA_KEY or not isinstance(value, bytes):
+                continue
+
+            status = base_pb2.Status()
+            status.ParseFromString(value)
+            for any_detail in status.details:
+                if any_detail.Is(base_pb2.StatusPayload.DESCRIPTOR):
+                    payload = base_pb2.StatusPayload()
+                    any_detail.Unpack(payload)
+                    if payload.subcode == base_pb2.SUBCODE_RATE_LIMITED:
+                        return True
+    except Exception:
+        pass
+    return False
+
+
+logger = logging.getLogger(__name__)
 
 
 class RetryInterceptor(UnaryUnaryClientInterceptor):
+    def __init__(self, log_retries: bool = False):
+        self.log_retries = log_retries
+        super().__init__()
+
     """
-    A retry interceptor that retries on status code RESOURCE_EXHAUSTED (i.e. rate-limited).
+    A retry interceptor that retries on rate limit errors and other retryable errors.
+
+    This interceptor handles:
+    1. Rate limit errors (UNAVAILABLE with SUBCODE_RATE_LIMITED) - always retried
+    2. [grpc.StatusCode.RESOURCE_EXHAUSTED, grpc.StatusCode.UNAVAILABLE] - retried
 
     This is required because, while the retry policy for the gRPC client is configurable via
     https://github.com/grpc/grpc-proto/blob/master/grpc/service_config/service_config.proto,
@@ -43,7 +92,10 @@ class RetryInterceptor(UnaryUnaryClientInterceptor):
         n_max_retries = 20
         max_retry_seconds = 20
         backoffs = [min(i * 2, max_retry_seconds) for i in range(1, n_max_retries)]
-        for backoff in backoffs:  # in seconds
+        backoff_index = 0
+        while True:
+            if backoff_index >= len(backoffs):
+                break
             call = continuation(client_call_details, request)
             if call.code() not in retryable_codes:
                 break
@@ -54,7 +106,17 @@ class RetryInterceptor(UnaryUnaryClientInterceptor):
                 details = call.details() or ""
                 if "InteractiveAuthException" in details:
                     break
+            backoff = backoffs[backoff_index]
+            if self.log_retries:
+                logger.info(
+                    f"Retrying {client_call_details.method} in {backoff} seconds (last response: {call.code()}, {call.details()})"
+                )
             sleep(backoff)
+            # Keep retrying rate-limited calls while increasing the backoff up to the max.
+            backoff_index += 1
+            if _is_rate_limited(call):
+                backoff_index = max(min(backoff_index, len(backoffs) - 2), 0)
+
         try:
             call.result()
         except grpc.RpcError as e:

luminarycloud/_helpers/__init__.py

@@ -2,6 +2,12 @@
 from ._create_geometry import (
     create_geometry as create_geometry,
 )
+from ._inference_jobs import (
+    create_inference_job as create_inference_job,
+    get_inference_job as get_inference_job,
+    list_inference_jobs as list_inference_jobs,
+    SurfaceForInference as SurfaceForInference,
+)
 from ._create_simulation import (
     create_simulation as create_simulation,
 )
@@ -24,6 +30,9 @@ from ._simulation_params_from_json import (
 from ._timestamp_to_datetime import (
     timestamp_to_datetime as timestamp_to_datetime,
 )
+from ._parse_iso_datetime import (
+    parse_iso_datetime as parse_iso_datetime,
+)
 from .upload import (
     upload_file as upload_file,
 )

luminarycloud/_helpers/_inference_jobs.py

@@ -0,0 +1,227 @@
+# Copyright 2025 Luminary Cloud, Inc. All Rights Reserved.
+from typing import Any, Dict, Optional
+from json import dumps as json_dumps
+import os
+import uuid
+from .._client import get_default_client
+from .._proto.api.v0.luminarycloud.physicsaiinference import (
+    physicsaiinference_pb2 as physicsaiinferencepb,
+)
+from .upload import upload_file, uploadpb
+from ..types import PhysicsAiModelVersionID, PhysicsAiInferenceJobID
+from ..physics_ai.inference import VisualizationExport, InferenceJob, SurfaceForInference
+
+
+# Helper function to upload an STL file if it is not a GCS URL
+def _upload_if_file(project_id: str, fname: str) -> str:
+    if not fname.split(".")[-1].lower() == "stl":
+        raise RuntimeError("Unsupported file for inference")
+    if fname.startswith("gs://"):
+        return fname
+    if os.path.exists(fname) and os.path.isfile(fname):
+        params = uploadpb.ResourceParams()
+        client = get_default_client()
+        result = upload_file(client, project_id, params, fname)
+        return result[1].url
+    raise RuntimeError("Unsupported file for inference")
+
+
+def _build_inference_request(
+    project_id: str,
+    geometry: str,
+    model_version_id: PhysicsAiModelVersionID,
+    conditions: Optional[Dict[str, Any]] = None,
+    settings: Optional[Dict[str, Any]] = None,
+    surfaces: Optional[list[SurfaceForInference]] = None,
+    inference_fields: Optional[list[str]] = None,
+    per_surface_visualizations: Optional[list[VisualizationExport]] = None,
+    merged_visualizations: Optional[list[VisualizationExport]] = None,
+) -> physicsaiinferencepb.CreateInferenceServiceJobRequest:
+    """Helper function to build an inference service job request.
+
+    Parameters
+    ----------
+    project_id : str
+        Reference to a project.
+    geometry : str
+        Path to STL file or GCS URL (gs://) of the geometry to run inference on.
+        If a local file path is provided, it will be uploaded to the project.
+    model_version_id : PhysicsAiModelVersionID
+        The ID of the trained model version to use for inference.
+    conditions : Dict[str, Any], optional
+        Dictionary of conditions to be passed to the inference service (e.g., alpha, beta, etc.).
+    settings : Dict[str, Any], optional
+        Dictionary of settings to be passed to inference service (e.g., stencil_size)
+    surfaces : list[SurfaceForInference], optional
+        List of surfaces for inference, each with 'name' and 'url' keys.
+    inference_fields : list[str], optional
+        Specific fields within the trained model to return inference results for.
+    per_surface_visualizations : list[VisualizationExport], optional
+        Types of visualization to write for each surface (e.g., LUMINARY, VTK).
+    merged_visualizations : list[VisualizationExport], optional
+        Types of merged visualization to write across all surfaces.
+
+    Returns
+    -------
+    CreateInferenceServiceJobRequest
+        The constructed protobuf request object.
+    """
+
+    geometry_url = _upload_if_file(project_id, geometry)
+
+    # Embed settings and store as bytes
+    settings_bytes = b""
+    if settings is not None:
+        settings_bytes = json_dumps(settings).encode("utf-8")
+
+    # Convert parameters dict to bytes if provided
+    conditions_bytes = b""
+    if conditions is not None:
+        conditions_bytes = json_dumps(conditions).encode("utf-8")
+
+    # Generate a unique request_id for deduplication and to satisfy the database constraint
+    # The backend uses request_id as the Name field, which must be non-empty
+    request_id = str(uuid.uuid4())
+
+    # Build request with base parameters
+    req_params = {
+        "request_id": request_id,
+        "geometry": geometry_url,
+        "model_version_id": str(model_version_id),
+        "conditions": conditions_bytes,
+        "settings": settings_bytes,
+        "project_id": project_id,
+    }
+
+    # Add optional inference fields
+    if inference_fields is not None:
+        req_params["inference_fields"] = inference_fields
+
+    # Add optional per-surface visualizations
+    if per_surface_visualizations is not None:
+        req_params["per_surface_visualizations"] = per_surface_visualizations
+
+    # Add optional merged visualizations
+    if merged_visualizations is not None:
+        req_params["merged_visualizations"] = merged_visualizations
+
+    # Add optional surfaces
+    if surfaces is not None:
+        surfaces_proto: list[physicsaiinferencepb.SurfaceForInference] = []
+        for surface in surfaces:
+            surfaces_proto.append(
+                physicsaiinferencepb.SurfaceForInference(
+                    name=surface["name"], url=_upload_if_file(project_id, surface["url"])
+                )
+            )
+        req_params["surfaces"] = surfaces_proto
+
+    return physicsaiinferencepb.CreateInferenceServiceJobRequest(**req_params)
+
+
+def create_inference_job(
+    project_id: str,
+    geometry: str,
+    model_version_id: PhysicsAiModelVersionID,
+    synchronous: bool = False,
+    conditions: Optional[Dict[str, Any]] = None,
+    settings: Optional[Dict[str, Any]] = None,
+    surfaces: Optional[list[SurfaceForInference]] = None,
+    inference_fields: Optional[list[str]] = None,
+    per_surface_visualizations: Optional[list[VisualizationExport]] = None,
+    merged_visualizations: Optional[list[VisualizationExport]] = None,
+) -> InferenceJob:
+    """Creates a synchronous inference service job.
+
+    Parameters
+    ----------
+    project_id : str
+        Reference to a project.
+    geometry : str
+        Path to STL file or GCS URL (gs://) of the geometry to run inference on.
+        If a local file path is provided, it will be uploaded to the project.
+    model_version_id : PhysicsAiModelVersionID
+        The ID of the trained model version to use for inference.
+    synchronous: bool = False,
+        Whether to wait for the job to complete before returning the result.
+    conditions : Dict[str, Any], optional
+        Dictionary of conditions to be passed to the inference service (e.g., alpha, beta, etc.).
+    settings : Dict[str, Any], optional
+        Dictionary of settings to be passed to inference service (e.g., stencil_size)
+    surfaces : list[SurfaceForInference], optional
+        List of surfaces for inference, each with 'name' and 'url' keys.
+    inference_fields : list[str], optional
+        Specific fields within the trained model to return inference results for.
+    per_surface_visualizations : list[VisualizationOutput], optional
+        Types of visualization to write for each surface (e.g., LUMINARY, VTK).
+    merged_visualizations : list[VisualizationOutput], optional
+        Types of merged visualization to write across all surfaces.
+
+    Returns
+    -------
+    dict[str, Any]
+        Response from the server containing results, with keys mapping to:
+        - Numeric results: float/vector values
+        - Surface/volume results: URLs to data files
+        - Visualization results: URLs to visualization files
+
+    warning:: This feature is experimental and may change or be removed without notice.
+    """
+    req = _build_inference_request(
+        project_id,
+        geometry,
+        model_version_id,
+        conditions,
+        settings,
+        surfaces,
+        inference_fields,
+        per_surface_visualizations,
+        merged_visualizations,
+    )
+    if synchronous:
+        res: physicsaiinferencepb.GetInferenceServiceJobResponse = (
+            get_default_client().CreateInferenceServiceJob(req)
+        )
+    else:  # Asynchronous inference
+        res: physicsaiinferencepb.CreateInferenceServiceJobAsyncResponse = (
+            get_default_client().CreateInferenceServiceJobAsync(req)
+        )
+    return InferenceJob(res.job)
+
+
+def get_inference_job(job_id: str) -> InferenceJob:
+    """Retrieves an inference service job by its ID.
+
+    Parameters
+    ----------
+    job_id : str
+        The ID of the inference job to retrieve.
+
+    Returns
+    -------
+    dict[str, Any]
+        The inference job details including results and status.
+
+    warning:: This feature is experimental and may change or be removed without notice.
+    """
+
+    req = physicsaiinferencepb.GetInferenceServiceJobRequest(job_id=job_id)
+    res: physicsaiinferencepb.GetInferenceServiceJobResponse = (
+        get_default_client().GetInferenceServiceJob(req)
+    )
+    return InferenceJob(res.job)
+
+
+def list_inference_jobs(project_id: str) -> list[InferenceJob]:
+    """Lists all inference service jobs for a project.
+
+    Parameters
+    ----------
+    project_id : str
+        The project to list inference jobs for.
+    """
+    req = physicsaiinferencepb.ListInferenceServiceJobsRequest(project_id=project_id)
+    res: physicsaiinferencepb.ListInferenceServiceJobsResponse = (
+        get_default_client().ListInferenceServiceJobs(req)
+    )
+    return [InferenceJob(job) for job in res.jobs]

luminarycloud/_helpers/_parse_iso_datetime.py

@@ -0,0 +1,54 @@
+# Copyright 2023-2025 Luminary Cloud, Inc. All Rights Reserved.
+from datetime import datetime
+
+
+def parse_iso_datetime(iso_str: str) -> datetime:
+    """
+    Parse an ISO format datetime string, handling 'Z' timezone indicator.
+
+    This function acts as a compatibility shim for Python < 3.11, which doesn't
+    support 'Z' in fromisoformat. It normalizes 'Z' (or 'z') to '+00:00' before parsing,
+    making it compatible with all Python versions (3.7+).
+
+    Parameters
+    ----------
+    iso_str : str
+        ISO format datetime string, optionally ending with 'Z' or 'z' for UTC.
+
+    Returns
+    -------
+    datetime
+        Parsed datetime object.
+
+    Raises
+    ------
+    ValueError
+        If the string is not a valid ISO format datetime string.
+    TypeError
+        If iso_str is not a string.
+
+    Examples
+    --------
+    >>> parse_iso_datetime("2023-07-31T13:54:12Z")
+    datetime.datetime(2023, 7, 31, 13, 54, 12, tzinfo=datetime.timezone.utc)
+    >>> parse_iso_datetime("2023-07-31T13:54:12+00:00")
+    datetime.datetime(2023, 7, 31, 13, 54, 12, tzinfo=datetime.timezone.utc)
+    """
+    if not isinstance(iso_str, str):
+        raise TypeError(f"parse_iso_datetime expects a string, got {type(iso_str).__name__}")
+
+    if not iso_str:
+        raise ValueError("parse_iso_datetime: empty string is not a valid ISO format datetime")
+
+    # Normalize 'Z' or 'z' at the end to '+00:00'
+    # Strip whitespace first to handle cases like "2023-01-01T00:00:00Z "
+    iso_str = iso_str.strip()
+
+    if iso_str.endswith(("Z", "z")):
+        # Only normalize if the string is longer than just 'Z'/'z'
+        if len(iso_str) > 1:
+            iso_str = iso_str[:-1] + "+00:00"
+        else:
+            raise ValueError("parse_iso_datetime: 'Z' alone is not a valid ISO format datetime")
+
+    return datetime.fromisoformat(iso_str)

luminarycloud/_helpers/download.py

@@ -17,6 +17,7 @@ from .._proto.api.v0.luminarycloud.solution.solution_pb2 import (
 )
 from .._proto.api.v0.luminarycloud.physics_ai.physics_ai_pb2 import (
     GetSolutionDataPhysicsAIRequest,
+    SurfaceGroup,
 )
 from .._client import Client
 from ..enum.quantity_type import QuantityType
@@ -140,6 +141,7 @@ def download_solution_physics_ai(
     process_volume: bool = False,
     single_precision: bool = False,
     internal_options: Optional[Dict[str, str]] = None,
+    export_surface_groups: Optional[Dict[str, List[str]]] = None,
 ) -> Optional[FileChunkStream]:
     """
     Returns the download as a file-like object, or None if destination_url is provided.
@@ -165,6 +167,9 @@ def download_solution_physics_ai(
         Whether to process volume meshes during physics AI processing.
     single_precision: bool
         If True, the solution will be downloaded in single precision.
+    export_surface_groups: Optional[Dict[str, List[str]]]
+        Dictionary mapping group names to lists of surface names.
+        Each group will be exported as an individual STL file.
 
     Examples
     --------
@@ -173,6 +178,11 @@ def download_solution_physics_ai(
     ...         fp.write(dl.read())
     """
 
+    surface_groups_pb = []
+    if export_surface_groups:
+        for group_name, surfaces in export_surface_groups.items():
+            surface_groups_pb.append(SurfaceGroup(name=group_name, surfaces=surfaces))
+
     request = GetSolutionDataPhysicsAIRequest(
         solution_id=solution_id,
         exclude_surfaces=exclude_surfaces or [],
@@ -186,6 +196,7 @@ def download_solution_physics_ai(
         process_volume=process_volume,
         single_precision=single_precision,
         internal_options=internal_options or {},
+        export_surface_groups=surface_groups_pb,
     )
     response = client.GetSolutionDataPhysicsAI(request)