From fef3e19561b2f9486ac46a265db6e30abf6618ce Mon Sep 17 00:00:00 2001 From: viambot <79611529+viambot@users.noreply.github.com> Date: Thu, 6 Nov 2025 02:20:08 +0000 Subject: [PATCH] [WORKFLOW] AI update based on proto changes from commit 26d3c076f0e269740a491b2f2143b415e427f6fc --- src/viam/components/arm/arm.py | 30 ++++++++++++++- src/viam/components/arm/client.py | 17 ++++++++- src/viam/components/arm/service.py | 11 ++++++ src/viam/utils.py | 27 +++++++++++++- tests/mocks/components.py | 11 +++++- tests/test_arm.py | 60 +++++++++++++++++++++++++----- 6 files changed, 139 insertions(+), 17 deletions(-) diff --git a/src/viam/components/arm/arm.py b/src/viam/components/arm/arm.py index 5569bfafb..c6d301ff0 100644 --- a/src/viam/components/arm/arm.py +++ b/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 `_. """ ... + + @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 `_. + """ + ... \ No newline at end of file diff --git a/src/viam/components/arm/client.py b/src/viam/components/arm/client.py index 37ad25e08..5e03a4818 100644 --- a/src/viam/components/arm/client.py +++ b/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) diff --git a/src/viam/components/arm/service.py b/src/viam/components/arm/service.py index 77661b8ce..8da7473c9 100644 --- a/src/viam/components/arm/service.py +++ b/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) diff --git a/src/viam/utils.py b/src/viam/utils.py index e8fa4e1b8..90a29deeb 100644 --- a/src/viam/utils.py +++ b/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. diff --git a/tests/mocks/components.py b/tests/mocks/components.py index 1e3f471ba..5bc5ab81a 100644 --- a/tests/mocks/components.py +++ b/tests/mocks/components.py @@ -35,7 +35,8 @@ 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 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 from viam.proto.component.encoder import PositionType @@ -50,6 +51,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): @@ -58,6 +61,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) @@ -115,6 +119,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} diff --git a/tests/test_arm.py b/tests/test_arm.py index 9a2bfd07c..c122a5df1 100644 --- a/tests/test_arm.py +++ b/tests/test_arm.py @@ -9,6 +9,9 @@ GetGeometriesResponse, GetKinematicsRequest, GetKinematicsResponse, + Get3DModelsRequest, + Get3DModelsResponse, + Mesh, Pose, ) from viam.proto.component.arm import ( @@ -28,7 +31,7 @@ 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: @@ -36,6 +39,7 @@ class TestArm: 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)