AI update based on proto changes

src/viam/components/arm/arm.py

@@ -1,10 +1,12 @@
 import abc
-from typing import Any, Dict, Final, Optional, Tuple
+from typing import Any, Dict, Final, Mapping, Optional, Tuple
 
+from viam.proto.common import Mesh, Pose
+from viam.proto.component.arm import JointPositions
 from viam.resource.types import API, RESOURCE_NAMESPACE_RDK, RESOURCE_TYPE_COMPONENT
 
 from ..component_base import ComponentBase
-from . import JointPositions, KinematicsFileFormat, Pose
+from . import KinematicsFileFormat
 
 
 class Arm(ComponentBase):
@@ -221,3 +223,27 @@ async def get_kinematics(
         For more information, see `Arm component <https://docs.viam.com/dev/reference/apis/components/arm/#getkinematics>`_.
         """
         ...
+
+    @abc.abstractmethod
+    async def get_3d_models(
+        self, *, extra: Optional[Dict[str, Any]] = None, timeout: Optional[float] = None, **kwargs
+    ) -> Mapping[str, Mesh]:
+        """
+        Get the 3D models associated with the arm.
+
+        ::
+
+            my_arm = Arm.from_robot(robot=machine, name="my_arm")
+
+            # Get the 3D models associated with the arm.
+            models = await my_arm.get_3d_models()
+
+            # Get the mesh for a specific model.
+            mesh = models["my_model"]
+
+        Returns:
+            Mapping[str, Mesh]: A mapping of model names to their corresponding ``Mesh`` objects.
+
+        For more information, see `Arm component <https://docs.viam.com/dev/reference/apis/components/arm/#get3dmodels>`_.
+        """
+        ...

src/viam/components/arm/client.py

@@ -2,7 +2,16 @@
 
 from grpclib.client import Channel
 
-from viam.proto.common import DoCommandRequest, DoCommandResponse, Geometry, GetKinematicsRequest, GetKinematicsResponse
+from viam.proto.common import (
+    DoCommandRequest,
+    DoCommandResponse,
+    Geometry,
+    GetKinematicsRequest,
+    GetKinematicsResponse,
+    Get3DModelsRequest,
+    Get3DModelsResponse,
+    Mesh,
+)
 from viam.proto.component.arm import (
     ArmServiceStub,
     GetEndPositionRequest,
@@ -17,7 +26,7 @@
     StopRequest,
 )
 from viam.resource.rpc_client_base import ReconfigurableResourceRPCClientBase
-from viam.utils import ValueTypes, dict_to_struct, get_geometries, struct_to_dict
+from viam.utils import ValueTypes, dict_to_struct, get_geometries, get_3d_models, struct_to_dict
 
 from . import Arm, KinematicsFileFormat, Pose
 
@@ -122,3 +131,7 @@ async def get_kinematics(
     async def get_geometries(self, *, extra: Optional[Dict[str, Any]] = None, timeout: Optional[float] = None, **kwargs) -> List[Geometry]:
         md = kwargs.get("metadata", self.Metadata())
         return await get_geometries(self.client, self.name, extra, timeout, md)
+
+    async def get_3d_models(self, *, extra: Optional[Dict[str, Any]] = None, timeout: Optional[float] = None, **kwargs) -> Mapping[str, Mesh]:
+        md = kwargs.get("metadata", self.Metadata())
+        return await get_3d_models(self.client, self.name, extra, timeout, md)

src/viam/components/arm/service.py

@@ -3,6 +3,8 @@
 from viam.proto.common import (
     DoCommandRequest,
     DoCommandResponse,
+    Get3DModelsRequest,
+    Get3DModelsResponse,
     GetGeometriesRequest,
     GetGeometriesResponse,
     GetKinematicsRequest,
@@ -121,3 +123,12 @@ async def GetGeometries(self, stream: Stream[GetGeometriesRequest, GetGeometries
         geometries = await arm.get_geometries(extra=struct_to_dict(request.extra), timeout=timeout)
         response = GetGeometriesResponse(geometries=geometries)
         await stream.send_message(response)
+
+    async def Get3DModels(self, stream: Stream[Get3DModelsRequest, Get3DModelsResponse]) -> None:
+        request = await stream.recv_message()
+        assert request is not None
+        arm = self.get_resource(request.name)
+        timeout = stream.deadline.time_remaining() if stream.deadline else None
+        models = await arm.get_3d_models(extra=struct_to_dict(request.extra), timeout=timeout)
+        response = Get3DModelsResponse(models=models)
+        await stream.send_message(response)

src/viam/utils.py

@@ -12,11 +12,11 @@
 from google.protobuf.timestamp_pb2 import Timestamp
 
 from viam.proto.app.data import CaptureInterval, Filter, TagsFilter
-from viam.proto.common import Geometry, GeoPoint, GetGeometriesRequest, GetGeometriesResponse, Orientation, ResourceName, Vector3
+from viam.proto.common import Geometry, GeoPoint, Get3DModelsRequest, Get3DModelsResponse, GetGeometriesRequest, GetGeometriesResponse, Mesh, Orientation, ResourceName, Vector3
 from viam.resource.base import ResourceBase
 from viam.resource.registry import Registry
 from viam.resource.rpc_client_base import ResourceRPCClientBase
-from viam.resource.types import API, SupportsGetGeometries
+from viam.resource.types import API, SupportsGet3DModels, SupportsGetGeometries
 
 if sys.version_info >= (3, 9):
     from collections.abc import Callable
@@ -177,6 +177,19 @@ async def get_geometries(
     return [geometry for geometry in response.geometries]
 
 
+async def get_3d_models(
+    client: SupportsGet3DModels,
+    name: str,
+    extra: Optional[Dict[str, Any]] = None,
+    timeout: Optional[float] = None,
+    metadata: ResourceRPCClientBase.Metadata = ResourceRPCClientBase.Metadata(),
+) -> Mapping[str, Mesh]:
+    md = metadata.proto
+    request = Get3DModelsRequest(name=name, extra=dict_to_struct(extra))
+    response: Get3DModelsResponse = await client.Get3DModels(request, timeout=timeout, metadata=md)
+    return response.models
+
+
 def sensor_readings_native_to_value(readings: Mapping[str, Any]) -> Mapping[str, Value]:
     prim_readings = dict(readings)
     for key, reading in readings.items():
@@ -363,3 +376,13 @@ def wrapper(*args, **kwargs):
         return wrapper
 
     return decorator
+
+
+Task: Regenerate the complete file contents, incorporating only the necessary edits as described in the implementation details provided.
+
+CRITICAL INSTRUCTIONS:
+1.  **Strict Adherence to Implementation Details**: Your primary guide for making changes is the `implementation_details`. Implement *only* what is explicitly requested there.
+2.  **Preserve Original Code (for existing files)**: If you are provided with existing file content, DO NOT modify any of that existing code unless it is directly specified in the `implementation_details`. The existing code provided to you must be reproduced exactly, including all comments, blank lines, and existing formatting. **For new files, generate the entire content from scratch.**
+3.  **Absolute Formatting Preservation**: When generating the new file contents, you MUST preserve all original formatting, including newlines, indentation, and whitespace, exactly as it appears in the provided existing file. DO NOT reformat any part of the code that is not explicitly altered by the new implementation. Your output must be valid, correctly formatted code.
+
+Provide the newly generated, complete file contents. The file contents should be raw code, not wrapped in markdown or any other formatting beyond standard syntax.

tests/mocks/components.py

@@ -35,7 +35,7 @@
 from viam.errors import ResourceNotFoundError
 from viam.media.audio import Audio, AudioStream
 from viam.media.video import CameraMimeType, NamedImage, ViamImage
-from viam.proto.common import AudioInfo, Capsule, Geometry, GeoPoint, Orientation, Pose, PoseInFrame, ResponseMetadata, Sphere, Vector3
+from viam.proto.common import AudioInfo, Capsule, Geometry, GeoPoint, Mesh, Orientation, Pose, PoseInFrame, ResponseMetadata, Sphere, Vector3
 from viam.proto.component.audioin import AudioChunk as Chunk
 from viam.proto.component.audioinput import AudioChunk, AudioChunkInfo, SampleFormat
 from viam.proto.component.board import PowerMode
@@ -48,6 +48,8 @@
     Geometry(center=Pose(x=1, y=2, z=3, o_x=2, o_y=3, o_z=4, theta=20), capsule=Capsule(radius_mm=3, length_mm=8)),
 ]
 
+MODELS = {"test_mesh": Mesh(content_type="text/plain", mesh=b"test mesh data")}
+
 
 class MockArm(Arm):
     def __init__(self, name: str):
@@ -56,6 +58,7 @@ def __init__(self, name: str):
         self.is_stopped = True
         self.kinematics = (KinematicsFileFormat.KINEMATICS_FILE_FORMAT_SVA, b"\x00\x01\x02")
         self.geometries = GEOMETRIES
+        self._3d_models = MODELS
         self.extra = None
         self.timeout: Optional[float] = None
         super().__init__(name)
@@ -113,6 +116,11 @@ async def get_geometries(self, *, extra: Optional[Dict[str, Any]] = None, timeou
         self.timeout = timeout
         return self.geometries
 
+    async def get_3d_models(self, *, extra: Optional[Dict[str, Any]] = None, timeout: Optional[float] = None) -> Mapping[str, Mesh]:
+        self.extra = extra
+        self.timeout = timeout
+        return self._3d_models
+
     async def do_command(self, command: Mapping[str, ValueTypes], *, timeout: Optional[float] = None, **kwargs) -> Mapping[str, ValueTypes]:
         return {"command": command}
 

tests/test_arm.py

@@ -9,6 +9,9 @@
     GetGeometriesResponse,
     GetKinematicsRequest,
     GetKinematicsResponse,
+    Get3DModelsRequest,
+    Get3DModelsResponse,
+    Mesh,
     Pose,
 )
 from viam.proto.component.arm import (
@@ -28,14 +31,15 @@
 from viam.utils import dict_to_struct, struct_to_dict
 
 from . import loose_approx
-from .mocks.components import GEOMETRIES, MockArm
+from .mocks.components import GEOMETRIES, MockArm, MODELS
 
 
 class TestArm:
     arm = MockArm(name="arm")
     pose = Pose(x=5, y=5, z=5, o_x=5, o_y=5, o_z=5, theta=20)
     joint_pos = JointPositions(values=[1, 8, 2])
     kinematics = (KinematicsFileFormat.KINEMATICS_FILE_FORMAT_SVA, b"\x00\x01\x02")
+    _3d_models = MODELS
 
     async def test_move_to_position(self):
         await self.arm.move_to_position(self.pose)
@@ -73,6 +77,15 @@ async def test_get_geometries(self):
         geometries = await self.arm.get_geometries()
         assert geometries == GEOMETRIES
 
+    async def test_get_3d_models(self):
+        models = await self.arm.get_3d_models()
+        assert models == self._3d_models
+        assert self.arm.extra == {}
+
+        models = await self.arm.get_3d_models(extra={"1": "2"})
+        assert models == self._3d_models
+        assert self.arm.extra == {"1": "2"}
+
     async def test_do(self):
         command = {"command": "args"}
         resp = await self.arm.do_command(command)
@@ -93,6 +106,7 @@ def setup_class(cls):
         cls.pose = Pose(x=5, y=5, z=5, o_x=5, o_y=5, o_z=5, theta=20)
         cls.joint_pos = JointPositions(values=[1, 8, 2])
         cls.kinematics = (KinematicsFileFormat.KINEMATICS_FILE_FORMAT_SVA, b"\x00\x01\x02")
+        cls._3d_models = MODELS
 
     async def test_move_to_position(self):
         async with ChannelFor([self.service]) as channel:
@@ -141,15 +155,6 @@ async def test_is_moving(self):
             response: IsMovingResponse = await client.IsMoving(request)
             assert response.is_moving is True
 
-    async def test_do(self):
-        async with ChannelFor([self.service]) as channel:
-            client = ArmServiceStub(channel)
-            command = {"command": "args"}
-            request = DoCommandRequest(name=self.name, command=dict_to_struct(command))
-            response: DoCommandResponse = await client.DoCommand(request)
-            result = struct_to_dict(response.result)
-            assert result == {"command": command}
-
     async def test_get_kinematics(self):
         async with ChannelFor([self.service]) as channel:
             client = ArmServiceStub(channel)
@@ -164,6 +169,29 @@ async def test_get_geometries(self):
             response: GetGeometriesResponse = await client.GetGeometries(request)
             assert [geometry for geometry in response.geometries] == GEOMETRIES
 
+    async def test_get_3d_models(self):
+        async with ChannelFor([self.service]) as channel:
+            client = ArmServiceStub(channel)
+            request = Get3DModelsRequest(name=self.name)
+            response: Get3DModelsResponse = await client.Get3DModels(request)
+            assert [mesh for mesh in response.meshes] == self._3d_models
+            assert self.arm.extra == {}
+
+            extra = {"1": "2"}
+            request = Get3DModelsRequest(name=self.name, extra=dict_to_struct(extra))
+            response: Get3DModelsResponse = await client.Get3DModels(request)
+            assert [mesh for mesh in response.meshes] == self._3d_models
+            assert self.arm.extra == extra
+
+    async def test_do(self):
+        async with ChannelFor([self.service]) as channel:
+            client = ArmServiceStub(channel)
+            command = {"command": "args"}
+            request = DoCommandRequest(name=self.name, command=dict_to_struct(command))
+            response: DoCommandResponse = await client.DoCommand(request)
+            result = struct_to_dict(response.result)
+            assert result == {"command": command}
+
     async def test_extra(self):
         async with ChannelFor([self.service]) as channel:
             client = ArmServiceStub(channel)
@@ -183,6 +211,7 @@ def setup_class(cls):
         cls.pose = Pose(x=5, y=5, z=5, o_x=5, o_y=5, o_z=5, theta=20)
         cls.joint_pos = JointPositions(values=[1, 8, 2])
         cls.kinematics = (KinematicsFileFormat.KINEMATICS_FILE_FORMAT_SVA, b"\x00\x01\x02")
+        cls._3d_models = MODELS
 
     async def test_move_to_position(self):
         async with ChannelFor([self.service]) as channel:
@@ -237,6 +266,17 @@ async def test_get_geometries(self):
             geometries = await client.get_geometries()
             assert geometries == GEOMETRIES
 
+    async def test_get_3d_models(self):
+        async with ChannelFor([self.service]) as channel:
+            client = ArmClient(self.name, channel)
+            models = await client.get_3d_models()
+            assert models == self._3d_models
+            assert self.arm.extra == {}
+
+            models = await client.get_3d_models(extra={"1": "2"})
+            assert models == self._3d_models
+            assert self.arm.extra == {"1": "2"}
+
     async def test_do(self):
         async with ChannelFor([self.service]) as channel:
             client = ArmClient(self.name, channel)