Source code for olo.platform

"""Platform catalog APIs — durable org artifacts via server Connect RPC."""

import os
from collections.abc import Callable
from typing import Optional

from olo_protos.platform.v1 import platform_pb2

from olo.core._payloads import COMPRESSED_IMAGE_TYPE, payload_to_proto
from olo.core.images import Image, compressed_image_to_message
from olo.errors import OloPlatformNotConfigured
from olo.platform._transport import ConnectUnaryTransport
from olo.platform.types import (
    BagInfo,
    CatalogLocation,
    ImageCatalogInfo,
    ImageDescription,
    ImageDetections,
    MapInfo,
    NavConfigInfo,
    NavConfigSource,
    NavigationDeploymentInfo,
    NavigationDesiredSpec,
    NavigationDesiredState,
    NavigationMode,
    NavigationObservedState,
    NavigationObservedStatus,
    Pose2D,
    VideoCatalogInfo,
    bag_from_proto,
    deployment_from_proto,
    image_catalog_from_proto,
    image_description_from_proto,
    image_detections_from_proto,
    map_from_proto,
    nav_config_from_proto,
    navigation_mode_to_proto,
    video_catalog_from_proto,
)

NamespaceResolver = Callable[[Optional[str]], str]

ENV_PLATFORM_URL = "OLO_SDK_PLATFORM_URL"
ENV_PLATFORM_TOKEN = "OLO_SDK_PLATFORM_TOKEN"
ENV_ROBOT_ID = "OLO_SDK_ROBOT_ID"

_DEFAULT_TIMEOUT_SEC = 30.0
_DESCRIBE_TIMEOUT_SEC = 120.0


def _resolve_config(
    url: Optional[str],
    token: Optional[str],
) -> tuple[str, str]:
    resolved_url = (url or os.environ.get(ENV_PLATFORM_URL) or "").strip()
    resolved_token = (token or os.environ.get(ENV_PLATFORM_TOKEN) or "").strip()
    if not resolved_url or not resolved_token:
        raise OloPlatformNotConfigured(
            "client.platform requires OLO_SDK_PLATFORM_URL and OLO_SDK_PLATFORM_TOKEN "
            "(or url=/token= kwargs). Create a personal access token at /developer."
        )
    return resolved_url, resolved_token


def resolve_robot_id(robot_id: Optional[str] = None) -> str:
    """Resolve the Platform robot UUID from an explicit value or ``OLO_SDK_ROBOT_ID``."""
    resolved = (robot_id or os.environ.get(ENV_ROBOT_ID) or "").strip()
    if not resolved:
        raise OloPlatformNotConfigured(
            "Platform robot id required. Pass robot_id=… or set OLO_SDK_ROBOT_ID "
            "(injected automatically for Appliance-managed runs)."
        )
    return resolved


def _normalize_rpc_base(url: str) -> str:
    """Ensure the transport base ends with /rpc."""
    trimmed = url.rstrip("/")
    if trimmed.endswith("/rpc"):
        return trimmed
    return f"{trimmed}/rpc"




class PlatformNavigation:
    """Durable navigation artifact catalogs."""



    def __init__(
        self,
        transport: ConnectUnaryTransport,
        namespace_resolver: Optional[NamespaceResolver] = None,
    ):
        self._transport = transport
        self._namespace_resolver = namespace_resolver


    def _resolve_namespace(self, namespace: Optional[str]) -> str:
        """Resolve like ``client.robot()``; standalone Platform falls back to ``\"\"``."""
        if self._namespace_resolver is None:
            return namespace if namespace is not None else ""
        return self._namespace_resolver(namespace)



    def get_configs(
        self,
        mode: NavigationMode,
        *,
        timeout: float = _DEFAULT_TIMEOUT_SEC,
    ) -> list[NavConfigInfo]:
        """List org nav config summaries for a navigation mode (metadata only, never YAML)."""
        response = self._transport.unary(
            "ListNavConfigs",
            platform_pb2.ListNavConfigsRequest(mode=navigation_mode_to_proto(mode)),
            platform_pb2.ListNavConfigsResponse,
            timeout=timeout,
        )
        return [nav_config_from_proto(item) for item in response.configs]




    def list_maps(self, *, timeout: float = _DEFAULT_TIMEOUT_SEC) -> list[MapInfo]:
        """List org map summaries (metadata only, never map blobs)."""
        response = self._transport.unary(
            "ListMaps",
            platform_pb2.ListMapsRequest(),
            platform_pb2.ListMapsResponse,
            timeout=timeout,
        )
        return [map_from_proto(item) for item in response.maps]




    def start(
        self,
        robot_id: Optional[str] = None,
        *,
        mode: NavigationMode,
        config_id: str,
        map_id: str = "",
        initial_pose: Optional[Pose2D] = None,
        lidar3d: bool = False,
        has_laser_scan: bool = False,
        namespace: Optional[str] = None,
        timeout: float = _DEFAULT_TIMEOUT_SEC,
    ) -> NavigationDeploymentInfo:
        """Start a navigation deployment.

        ``robot_id`` defaults to ``OLO_SDK_ROBOT_ID`` (Appliance-managed runs
        inject this automatically). ``namespace`` defaults to the unambiguous
        appliance robot namespace (same as ``client.robot()``); pass ``\"\"`` for
        a global launch or an explicit value when multiple robots share the
        appliance.
        """
        resolved_robot_id = resolve_robot_id(robot_id)
        pose_pb = None
        if initial_pose is not None:
            pose_pb = platform_pb2.Pose2D(
                x=initial_pose.x, y=initial_pose.y, theta=initial_pose.yaw
            )
        response = self._transport.unary(
            "StartNavigationDeployment",
            platform_pb2.StartNavigationDeploymentRequest(
                robot_id=resolved_robot_id,
                namespace=self._resolve_namespace(namespace),
                mode=navigation_mode_to_proto(mode),
                config_id=config_id,
                map_id=map_id,
                initial_pose=pose_pb,
                lidar3d=lidar3d,
                has_laser_scan=has_laser_scan,
            ),
            platform_pb2.NavigationDeploymentResponse,
            timeout=timeout,
        )
        return deployment_from_proto(response.deployment)




    def stop(
        self,
        robot_id: Optional[str] = None,
        *,
        namespace: Optional[str] = None,
        timeout: float = _DEFAULT_TIMEOUT_SEC,
    ) -> NavigationDeploymentInfo:
        """Stop a navigation deployment.

        ``robot_id`` defaults to ``OLO_SDK_ROBOT_ID``. ``namespace`` resolves
        like :meth:`start`.
        """
        response = self._transport.unary(
            "StopNavigationDeployment",
            platform_pb2.StopNavigationDeploymentRequest(
                robot_id=resolve_robot_id(robot_id),
                namespace=self._resolve_namespace(namespace),
            ),
            platform_pb2.NavigationDeploymentResponse,
            timeout=timeout,
        )
        return deployment_from_proto(response.deployment)




    def deployment(
        self,
        robot_id: Optional[str] = None,
        *,
        namespace: Optional[str] = None,
        timeout: float = _DEFAULT_TIMEOUT_SEC,
    ) -> NavigationDeploymentInfo:
        """Get current navigation deployment state for a robot.

        ``robot_id`` defaults to ``OLO_SDK_ROBOT_ID``. ``namespace`` resolves
        like :meth:`start`.
        """
        response = self._transport.unary(
            "GetNavigationDeployment",
            platform_pb2.GetNavigationDeploymentRequest(
                robot_id=resolve_robot_id(robot_id),
                namespace=self._resolve_namespace(namespace),
            ),
            platform_pb2.NavigationDeploymentResponse,
            timeout=timeout,
        )
        return deployment_from_proto(response.deployment)






class PlatformArchive:
    """Cloud archive catalog (bags, videos, and images)."""



    def __init__(self, transport: ConnectUnaryTransport):
        self._transport = transport


    def list_bags(self, *, timeout: float = _DEFAULT_TIMEOUT_SEC) -> list[BagInfo]:
        response = self._transport.unary(
            "ListBags",
            platform_pb2.ListBagsRequest(),
            platform_pb2.ListBagsResponse,
            timeout=timeout,
        )
        return [bag_from_proto(item) for item in response.bags]

    def list_videos(self, *, timeout: float = _DEFAULT_TIMEOUT_SEC) -> list[VideoCatalogInfo]:
        response = self._transport.unary(
            "ListVideos",
            platform_pb2.ListVideosRequest(),
            platform_pb2.ListVideosResponse,
            timeout=timeout,
        )
        return [video_catalog_from_proto(item) for item in response.videos]

    def list_images(self, *, timeout: float = _DEFAULT_TIMEOUT_SEC) -> list[ImageCatalogInfo]:
        response = self._transport.unary(
            "ListImages",
            platform_pb2.ListImagesRequest(),
            platform_pb2.ListImagesResponse,
            timeout=timeout,
        )
        return [image_catalog_from_proto(item) for item in response.images]





class PlatformAi:
    """Cloud model inference, metered against the organization's AI credit."""



    def __init__(self, transport: ConnectUnaryTransport):
        self._transport = transport




    def describe_image(
        self, image: Image, *, prompt: str = "", timeout: float = _DESCRIBE_TIMEOUT_SEC
    ) -> ImageDescription:
        """Describe an 8-bit image, e.g. from ``robot.core.get_image(topic)``.

        Without ``prompt`` the result describes the scene; with one it answers
        that request. Check ``image_usable`` before trusting the description.
        """
        response = self._request(image, platform_pb2.DESCRIBE_TASK_DESCRIBE, prompt, timeout)
        return image_description_from_proto(response)




    def detect_objects(
        self, image: Image, *, prompt: str = "", timeout: float = _DESCRIBE_TIMEOUT_SEC
    ) -> ImageDetections:
        """Locate objects in an 8-bit image, e.g. from ``robot.core.get_image(topic)``.

        ``prompt`` names what to find ("fire extinguishers"); without it every
        notable object is returned. Boxes are in the image's pixels and feed
        :func:`olo.vision.deproject_region` with an aligned depth image.
        """
        response = self._request(image, platform_pb2.DESCRIBE_TASK_DETECT, prompt, timeout)
        return image_detections_from_proto(response)


    def _request(
        self, image: Image, task: platform_pb2.DescribeTask, prompt: str, timeout: float
    ) -> platform_pb2.DescribeImageResponse:
        if image.array.dtype.itemsize != 1:
            raise ValueError("platform.ai needs an 8-bit image; convert depth images first")
        frame = Image(array=image.array, format="jpeg", header=image.header)
        return self._transport.unary(
            "DescribeImage",
            platform_pb2.DescribeImageRequest(
                image=payload_to_proto(COMPRESSED_IMAGE_TYPE, compressed_image_to_message(frame)),
                task=task,
                prompt=prompt,
            ),
            platform_pb2.DescribeImageResponse,
            timeout=timeout,
        )





class Platform:
    """Server-direct platform surface (`client.platform`)."""



    def __init__(
        self,
        url: Optional[str] = None,
        token: Optional[str] = None,
        *,
        namespace_resolver: Optional[NamespaceResolver] = None,
    ):
        resolved_url, resolved_token = _resolve_config(url, token)
        transport = ConnectUnaryTransport(_normalize_rpc_base(resolved_url), resolved_token)
        self.navigation = PlatformNavigation(transport, namespace_resolver=namespace_resolver)
        self.archive = PlatformArchive(transport)
        self.ai = PlatformAi(transport)


    @property
    def robot_id(self) -> str:
        """Platform robot UUID for this runtime (from ``OLO_SDK_ROBOT_ID``)."""
        return resolve_robot_id()



__all__ = [
    "ENV_PLATFORM_TOKEN",
    "ENV_PLATFORM_URL",
    "ENV_ROBOT_ID",
    "BagInfo",
    "ImageCatalogInfo",
    "MapInfo",
    "NavConfigInfo",
    "NavConfigSource",
    "NavigationDeploymentInfo",
    "NavigationDesiredSpec",
    "NavigationDesiredState",
    "NavigationMode",
    "NavigationObservedState",
    "NavigationObservedStatus",
    "Platform",
    "PlatformArchive",
    "PlatformNavigation",
    "Pose2D",
    "CatalogLocation",
    "ImageDescription",
    "ImageDetections",
    "PlatformAi",
    "VideoCatalogInfo",
    "resolve_robot_id",
]