foodforthought-cli 0.2.7__py3-none-any.whl → 0.3.0__py3-none-any.whl

ate/transports/ble.py

@@ -0,0 +1,405 @@
+"""
+BLE Transport Implementation.
+
+Uses bleak library for cross-platform Bluetooth Low Energy communication.
+Tested with HiWonder MechDog but designed to be adaptable to other robots.
+
+Key learnings from MechDog reverse engineering:
+- Write commands to FFE1 characteristic
+- Subscribe to FFE2 for responses (CRITICAL - many examples miss this!)
+- Commands are ASCII strings in CMD|...|$ format
+- Some commands work (locomotion), some don't (arm) - firmware dependent
+"""
+
+import asyncio
+from typing import Optional
+from dataclasses import dataclass
+
+try:
+    from bleak import BleakClient, BleakScanner
+    from bleak.backends.device import BLEDevice
+    BLEAK_AVAILABLE = True
+except ImportError:
+    BLEAK_AVAILABLE = False
+    BleakClient = None
+    BleakScanner = None
+    BLEDevice = None
+
+from .base import (
+    RobotTransport,
+    TransportState,
+    TransportCapability,
+    ConnectionInfo,
+    CommandResult,
+    HiWonderProtocol,
+)
+
+
+# Standard HiWonder GATT UUIDs
+HIWONDER_SERVICE_UUID = "0000ffe0-0000-1000-8000-00805f9b34fb"
+HIWONDER_WRITE_UUID = "0000ffe1-0000-1000-8000-00805f9b34fb"  # Write commands here
+HIWONDER_NOTIFY_UUID = "0000ffe2-0000-1000-8000-00805f9b34fb"  # Subscribe for responses
+
+
+@dataclass
+class BLEConfig:
+    """Configuration for BLE transport."""
+    write_uuid: str = HIWONDER_WRITE_UUID
+    notify_uuid: str = HIWONDER_NOTIFY_UUID
+    response_timeout: float = 2.0
+    retry_count: int = 3
+    probe_on_connect: bool = True
+
+
+class BLETransport(RobotTransport):
+    """
+    Bluetooth Low Energy transport for robot communication.
+
+    Example:
+        async with BLETransport() as transport:
+            if await transport.connect("FAF198B6-D9A4-CC27-7BBA-3159F63A61A9"):
+                result = await transport.send_command("CMD|6|$")  # Battery query
+                print(f"Battery: {result.response}")
+    """
+
+    def __init__(self, config: Optional[BLEConfig] = None):
+        super().__init__()
+
+        if not BLEAK_AVAILABLE:
+            raise ImportError(
+                "bleak is required for BLE transport. "
+                "Install with: pip install bleak"
+            )
+
+        self.config = config or BLEConfig()
+        self._client: Optional[BleakClient] = None
+        self._protocol = HiWonderProtocol()
+        self._pending_responses: asyncio.Queue = asyncio.Queue()
+        self._response_buffer: list[bytes] = []
+
+    async def connect(self, address: str, **kwargs) -> bool:
+        """
+        Connect to a BLE robot.
+
+        Args:
+            address: MAC address or UUID of the device
+            **kwargs: Additional options (name hint, timeout)
+
+        Returns:
+            True if connected successfully
+        """
+        self._state = TransportState.CONNECTING
+
+        try:
+            # Scan to find the device
+            device = await self._find_device(address, kwargs.get("name"))
+            if not device:
+                self._state = TransportState.ERROR
+                return False
+
+            # Connect
+            self._client = BleakClient(device)
+            connected = await self._client.connect()
+
+            if not connected or not self._client.is_connected:
+                self._state = TransportState.ERROR
+                return False
+
+            # Start notification handler
+            await self._client.start_notify(
+                self.config.notify_uuid,
+                self._notification_handler
+            )
+
+            self._state = TransportState.CONNECTED
+            self._connection_info = ConnectionInfo(
+                transport_type="ble",
+                address=address,
+                name=device.name,
+                rssi=device.rssi if hasattr(device, 'rssi') else None,
+            )
+
+            # Probe capabilities if configured
+            if self.config.probe_on_connect:
+                self._capabilities = await self.probe_capabilities()
+
+            return True
+
+        except Exception as e:
+            print(f"BLE connect error: {e}")
+            self._state = TransportState.ERROR
+            return False
+
+    async def disconnect(self) -> None:
+        """Disconnect from the BLE device."""
+        if self._client and self._client.is_connected:
+            try:
+                await self._client.stop_notify(self.config.notify_uuid)
+            except Exception:
+                pass
+            try:
+                await self._client.disconnect()
+            except Exception:
+                pass
+
+        self._client = None
+        self._state = TransportState.DISCONNECTED
+        self._connection_info = None
+
+    async def send(self, data: bytes) -> CommandResult:
+        """Send raw bytes to the robot."""
+        if not self.is_connected or not self._client:
+            return CommandResult(success=False, error="Not connected")
+
+        import time
+        start = time.perf_counter()
+
+        try:
+            # Clear pending responses
+            while not self._pending_responses.empty():
+                try:
+                    self._pending_responses.get_nowait()
+                except asyncio.QueueEmpty:
+                    break
+
+            # Send command
+            await self._client.write_gatt_char(
+                self.config.write_uuid,
+                data
+            )
+
+            # Wait for response
+            try:
+                response = await asyncio.wait_for(
+                    self._pending_responses.get(),
+                    timeout=self.config.response_timeout
+                )
+                latency = (time.perf_counter() - start) * 1000
+
+                return CommandResult(
+                    success=True,
+                    response=response.decode('ascii', errors='replace'),
+                    raw_data=response,
+                    latency_ms=latency
+                )
+            except asyncio.TimeoutError:
+                # Command sent but no response (may still have worked)
+                latency = (time.perf_counter() - start) * 1000
+                return CommandResult(
+                    success=True,
+                    error="No response (command may have worked)",
+                    latency_ms=latency
+                )
+
+        except Exception as e:
+            return CommandResult(success=False, error=str(e))
+
+    async def send_command(self, command: str) -> CommandResult:
+        """Send a command string (handles encoding)."""
+        if command.startswith("CMD|"):
+            # Already formatted
+            data = command.encode('ascii')
+        else:
+            # Try to encode as canonical command
+            try:
+                data = self._protocol.encode_command(command)
+            except ValueError:
+                # Unknown command, send as-is
+                data = command.encode('ascii')
+
+        return await self.send(data)
+
+    @classmethod
+    async def discover(cls, timeout: float = 5.0) -> list[ConnectionInfo]:
+        """Discover BLE robots."""
+        if not BLEAK_AVAILABLE:
+            return []
+
+        devices = await BleakScanner.discover(timeout=timeout)
+        results = []
+
+        for device in devices:
+            # Filter for likely robot devices
+            name = device.name or ""
+            if any(kw in name.lower() for kw in ["mechdog", "robot", "hiwonder", "unitree"]):
+                results.append(ConnectionInfo(
+                    transport_type="ble",
+                    address=device.address,
+                    name=device.name,
+                    rssi=device.rssi if hasattr(device, 'rssi') else None,
+                ))
+
+        return results
+
+    async def probe_capabilities(self) -> set[TransportCapability]:
+        """
+        Probe which capabilities work on this connection.
+
+        Sends test commands and checks for valid responses.
+        This handles the firmware fragmentation problem.
+        """
+        caps = set()
+
+        if not self.is_connected:
+            return caps
+
+        # Always try battery first - good connectivity test
+        result = await self.send(b"CMD|6|$")
+        if result.success and result.response and "CMD|6|" in result.response:
+            caps.add(TransportCapability.BATTERY)
+            caps.add(TransportCapability.BIDIRECTIONAL)
+
+        # Try ultrasonic
+        result = await self.send(b"CMD|4|1|$")
+        await asyncio.sleep(0.3)
+        if result.success and result.response and "CMD|4|" in result.response:
+            caps.add(TransportCapability.ULTRASONIC)
+
+        # We know locomotion works from testing
+        # Could verify with a small movement, but risky for capability probe
+        caps.add(TransportCapability.WALK)
+        caps.add(TransportCapability.POSTURE)
+        caps.add(TransportCapability.ACTION_GROUPS)
+
+        # UPDATE: CMD|7|... commands DO work for arm control!
+        # Tested 2026-01-04: gripper, shoulder, elbow all respond to CMD|7
+        caps.add(TransportCapability.ARM_POSITION)
+        caps.add(TransportCapability.GRIPPER)
+
+        return caps
+
+    async def _find_device(
+        self,
+        address: str,
+        name_hint: Optional[str] = None
+    ) -> Optional["BLEDevice"]:
+        """Find a BLE device by address or name."""
+        devices = await BleakScanner.discover(timeout=5.0)
+
+        for device in devices:
+            if device.address == address:
+                return device
+            if name_hint and device.name and name_hint.lower() in device.name.lower():
+                return device
+            if device.name and "mechdog" in device.name.lower():
+                return device
+
+        return None
+
+    def _notification_handler(self, sender, data: bytes) -> None:
+        """Handle incoming BLE notifications."""
+        # Store in buffer
+        self._response_buffer.append(data)
+
+        # Notify async handlers
+        asyncio.create_task(self._notify_handlers(data))
+
+        # Put in queue for synchronous waiting
+        try:
+            self._pending_responses.put_nowait(data)
+        except asyncio.QueueFull:
+            pass
+
+    # ========== High-Level Commands ==========
+    # These use the protocol adapter for clean API
+
+    async def walk_forward(self, speed: int = 100) -> CommandResult:
+        """Walk forward at given speed (0-200)."""
+        return await self.send(f"CMD|3|{speed}|0|0|$".encode())
+
+    async def walk_backward(self, speed: int = 100) -> CommandResult:
+        """Walk backward at given speed."""
+        return await self.send(f"CMD|3|{-speed}|0|0|$".encode())
+
+    async def stop(self) -> CommandResult:
+        """Stop all movement."""
+        return await self.send(b"CMD|3|0|0|0|$")
+
+    async def stand(self) -> CommandResult:
+        """Stand in normal posture."""
+        return await self.send(b"CMD|1|0|$")
+
+    async def sit(self) -> CommandResult:
+        """Sit down."""
+        return await self.send(b"CMD|1|7|$")
+
+    async def get_battery(self) -> Optional[float]:
+        """Get battery voltage."""
+        result = await self.send(b"CMD|6|$")
+        if result.response and "CMD|6|" in result.response:
+            try:
+                parts = result.response.split("|")
+                return float(parts[2])
+            except (IndexError, ValueError):
+                pass
+        return None
+
+    async def get_ultrasonic(self) -> Optional[int]:
+        """Get ultrasonic distance in mm."""
+        result = await self.send(b"CMD|4|1|$")
+        if result.response and "CMD|4|" in result.response:
+            try:
+                parts = result.response.split("|")
+                return int(parts[2])
+            except (IndexError, ValueError):
+                pass
+        return None
+
+    async def run_action_group(self, group_id: int) -> CommandResult:
+        """Run a pre-programmed action group."""
+        return await self.send(f"CMD|2|1|{group_id}|$".encode())
+
+    # ========== ARM Control (CMD|7|...) ==========
+    # Tested and confirmed working 2026-01-04
+
+    async def gripper_open(self) -> CommandResult:
+        """Open the gripper."""
+        return await self.send(b"CMD|7|6|$")
+
+    async def gripper_close(self) -> CommandResult:
+        """Close the gripper."""
+        return await self.send(b"CMD|7|7|$")
+
+    async def shoulder_up(self, repeats: int = 1) -> CommandResult:
+        """Move shoulder joint up (incremental)."""
+        result = None
+        for _ in range(repeats):
+            result = await self.send(b"CMD|7|3|1|$")
+            await asyncio.sleep(0.2)
+        return result
+
+    async def shoulder_down(self, repeats: int = 1) -> CommandResult:
+        """Move shoulder joint down (incremental)."""
+        result = None
+        for _ in range(repeats):
+            result = await self.send(b"CMD|7|3|-1|$")
+            await asyncio.sleep(0.2)
+        return result
+
+    async def elbow_extend(self, repeats: int = 1) -> CommandResult:
+        """Extend elbow joint (incremental)."""
+        result = None
+        for _ in range(repeats):
+            result = await self.send(b"CMD|7|4|1|$")
+            await asyncio.sleep(0.2)
+        return result
+
+    async def elbow_retract(self, repeats: int = 1) -> CommandResult:
+        """Retract elbow joint (incremental)."""
+        result = None
+        for _ in range(repeats):
+            result = await self.send(b"CMD|7|4|-1|$")
+            await asyncio.sleep(0.2)
+        return result
+
+    async def arm_to_pickup_position(self) -> None:
+        """Move arm to pickup position (lowered, extended)."""
+        await self.gripper_open()
+        await asyncio.sleep(0.3)
+        await self.shoulder_down(repeats=10)
+        await self.elbow_extend(repeats=8)
+
+    async def arm_to_carry_position(self) -> None:
+        """Move arm to carry position (raised, retracted)."""
+        await self.shoulder_up(repeats=10)
+        await self.elbow_retract(repeats=8)