unrealon 1.1.6__py3-none-any.whl → 2.0.5__py3-none-any.whl

unrealon_driver/decorators/timing.py

@@ -0,0 +1,132 @@
+"""
+Clean timing decorator for performance monitoring.
+"""
+
+import asyncio
+import logging
+import time
+from functools import wraps
+from typing import Callable, Optional, Dict, Any
+
+logger = logging.getLogger(__name__)
+
+
+def timing(
+    log_result: bool = True,
+    log_level: str = "INFO",
+    include_args: bool = False,
+    threshold: Optional[float] = None
+):
+    """
+    Clean timing decorator.
+    
+    Measures and logs function execution time.
+    
+    Args:
+        log_result: Whether to log the timing result
+        log_level: Log level for timing messages
+        include_args: Include function arguments in log
+        threshold: Only log if execution time exceeds threshold (seconds)
+    """
+    
+    def decorator(func: Callable) -> Callable:
+        @wraps(func)
+        async def async_wrapper(*args, **kwargs):
+            """Async timing wrapper."""
+            
+            start_time = asyncio.get_event_loop().time()
+            
+            try:
+                result = await func(*args, **kwargs)
+                success = True
+                error = None
+            except Exception as e:
+                result = None
+                success = False
+                error = str(e)
+                raise
+            finally:
+                end_time = asyncio.get_event_loop().time()
+                duration = end_time - start_time
+                
+                # Log timing if enabled and threshold met
+                if log_result and (threshold is None or duration >= threshold):
+                    _log_timing(func, duration, success, error, args, kwargs, include_args, log_level)
+            
+            return result
+        
+        @wraps(func)
+        def sync_wrapper(*args, **kwargs):
+            """Sync timing wrapper."""
+            
+            start_time = time.perf_counter()
+            
+            try:
+                result = func(*args, **kwargs)
+                success = True
+                error = None
+            except Exception as e:
+                result = None
+                success = False
+                error = str(e)
+                raise
+            finally:
+                end_time = time.perf_counter()
+                duration = end_time - start_time
+                
+                # Log timing if enabled and threshold met
+                if log_result and (threshold is None or duration >= threshold):
+                    _log_timing(func, duration, success, error, args, kwargs, include_args, log_level)
+            
+            return result
+        
+        # Return appropriate wrapper based on function type
+        if asyncio.iscoroutinefunction(func):
+            return async_wrapper
+        else:
+            return sync_wrapper
+    
+    return decorator
+
+
+def _log_timing(
+    func: Callable,
+    duration: float,
+    success: bool,
+    error: Optional[str],
+    args: tuple,
+    kwargs: dict,
+    include_args: bool,
+    log_level: str
+):
+    """Log timing information."""
+    
+    # Format duration
+    if duration < 0.001:
+        duration_str = f"{duration * 1000000:.0f}μs"
+    elif duration < 1.0:
+        duration_str = f"{duration * 1000:.1f}ms"
+    else:
+        duration_str = f"{duration:.2f}s"
+    
+    # Build log message
+    status = "✓" if success else "✗"
+    message = f"{status} {func.__name__} took {duration_str}"
+    
+    if error:
+        message += f" (failed: {error})"
+    
+    if include_args:
+        args_str = ", ".join([repr(arg) for arg in args])
+        kwargs_str = ", ".join([f"{k}={repr(v)}" for k, v in kwargs.items()])
+        
+        if args_str and kwargs_str:
+            message += f" | Args: ({args_str}, {kwargs_str})"
+        elif args_str:
+            message += f" | Args: ({args_str})"
+        elif kwargs_str:
+            message += f" | Args: ({kwargs_str})"
+    
+    # Log at specified level
+    log_func = getattr(logger, log_level.lower(), logger.info)
+    log_func(message)

unrealon_driver/driver/__init__.py

@@ -0,0 +1,20 @@
+"""
+UniversalDriver - Modular driver system.
+
+Clean, organized driver architecture with:
+- Core driver functionality
+- Lifecycle management
+- Communication layer
+- Factory pattern for managers
+- Health monitoring
+- Utilities for common operations
+"""
+
+from .core.driver import UniversalDriver
+from .core.config import DriverConfig, DriverMode
+
+__all__ = [
+    "UniversalDriver",
+    "DriverConfig", 
+    "DriverMode"
+]

unrealon_driver/driver/communication/__init__.py

@@ -0,0 +1,10 @@
+"""
+Communication layer.
+
+WebSocket client and RPC session management.
+"""
+
+from .websocket_client import WebSocketClient
+from .session import DriverSession
+
+__all__ = ["WebSocketClient", "DriverSession"]

unrealon_driver/driver/communication/session.py

@@ -0,0 +1,203 @@
+"""
+Clean driver session management.
+"""
+
+import asyncio
+import logging
+from datetime import datetime
+from typing import Dict, Any, Optional, List, Callable
+from enum import Enum
+
+from pydantic import BaseModel, Field
+
+from unrealon_core.models.websocket import (
+    TaskAssignmentData,
+    TaskResultData,
+    WebSocketMessage
+)
+from unrealon_core.enums import DriverStatus, TaskStatus, MessageType
+from unrealon_driver.utils.time import utc_now
+
+logger = logging.getLogger(__name__)
+
+
+class SessionStatus(str, Enum):
+    """Clean session status."""
+    DISCONNECTED = "disconnected"
+    CONNECTING = "connecting" 
+    CONNECTED = "connected"
+    REGISTERED = "registered"
+    ACTIVE = "active"
+    ERROR = "error"
+
+
+class SessionStats(BaseModel):
+    """Session statistics."""
+    connected_at: Optional[datetime] = None
+    registered_at: Optional[datetime] = None
+    tasks_completed: int = 0
+    tasks_failed: int = 0
+    last_activity: Optional[datetime] = None
+
+
+class DriverSession:
+    """
+    Clean session manager.
+    
+    Handles connection, registration, and task processing
+    with simple, clear interface.
+    """
+    
+    def __init__(self, driver_id: str, websocket_client=None):
+        self.driver_id = driver_id
+        self.websocket_client = websocket_client
+        self.status = SessionStatus.DISCONNECTED
+        self.stats = SessionStats()
+        
+        # Task handling
+        self.task_handlers: Dict[str, Callable] = {}
+        
+        # Setup WebSocket message handler
+        if self.websocket_client:
+            self.websocket_client.on_message = self._handle_websocket_message
+    
+    async def start_session(self, capabilities: List[str] = None) -> bool:
+        """Start session with registration."""
+        if not self.websocket_client:
+            logger.error("No WebSocket client configured")
+            return False
+        
+        try:
+            # Connect WebSocket
+            self.status = SessionStatus.CONNECTING
+            if not await self.websocket_client.connect():
+                self.status = SessionStatus.ERROR
+                return False
+            
+            self.status = SessionStatus.CONNECTED
+            self.stats.connected_at = utc_now()
+            
+            # Register driver
+            if await self.register(capabilities or []):
+                self.status = SessionStatus.REGISTERED
+                self.stats.registered_at = utc_now()
+                logger.info(f"Session started for driver: {self.driver_id}")
+                return True
+            else:
+                self.status = SessionStatus.ERROR
+                return False
+                
+        except Exception as e:
+            logger.error(f"Session start failed: {e}")
+            self.status = SessionStatus.ERROR
+            return False
+    
+    async def stop_session(self):
+        """Stop session cleanly."""
+        try:
+            if self.websocket_client:
+                await self.websocket_client.disconnect()
+            
+            self.status = SessionStatus.DISCONNECTED
+            logger.info(f"Session stopped for driver: {self.driver_id}")
+            
+        except Exception as e:
+            logger.error(f"Session stop error: {e}")
+    
+    async def register(self, capabilities: List[str]) -> bool:
+        """Register driver with capabilities."""
+        if not self.websocket_client:
+            return False
+        
+        try:
+            success = await self.websocket_client.register_driver(capabilities)
+            if success:
+                logger.info(f"Driver registered: {self.driver_id} with capabilities: {capabilities}")
+            return success
+        except Exception as e:
+            logger.error(f"Registration failed: {e}")
+            return False
+    
+    def register_task_handler(self, task_type: str, handler: Callable):
+        """Register handler for task type."""
+        self.task_handlers[task_type] = handler
+        logger.debug(f"Registered handler for task type: {task_type}")
+    
+    def _handle_websocket_message(self, message: WebSocketMessage):
+        """Handle incoming WebSocket messages."""
+        try:
+            self.stats.last_activity = utc_now()
+            
+            if message.type == MessageType.TASK_ASSIGN:
+                asyncio.create_task(self._handle_task_assignment(message))
+            else:
+                logger.debug(f"Received message type: {message.type}")
+                
+        except Exception as e:
+            logger.error(f"Error handling WebSocket message: {e}")
+    
+    async def _handle_task_assignment(self, message: WebSocketMessage):
+        """Handle task assignment."""
+        try:
+            # Validate task data
+            task_data = TaskAssignmentData.model_validate(message.data)
+            
+            # Find handler
+            handler = self.task_handlers.get(task_data.task_type)
+            if not handler:
+                logger.warning(f"No handler for task type: {task_data.task_type}")
+                await self._send_task_result(
+                    task_data.task_id,
+                    TaskStatus.FAILED,
+                    error="No handler for task type"
+                )
+                return
+            
+            # Execute task
+            try:
+                result = await handler(task_data)
+                await self._send_task_result(
+                    task_data.task_id,
+                    TaskStatus.COMPLETED,
+                    result=result
+                )
+                self.stats.tasks_completed += 1
+                
+            except Exception as e:
+                logger.error(f"Task execution failed: {e}")
+                await self._send_task_result(
+                    task_data.task_id,
+                    TaskStatus.FAILED,
+                    error=str(e)
+                )
+                self.stats.tasks_failed += 1
+                
+        except Exception as e:
+            logger.error(f"Task assignment handling failed: {e}")
+    
+    async def _send_task_result(self, task_id: str, status: TaskStatus, result: Any = None, error: str = None):
+        """Send task result back."""
+        try:
+            result_data = TaskResultData(
+                task_id=task_id,
+                driver_id=self.driver_id,
+                status=status,
+                result=result,
+                error_message=error,
+                completed_at=utc_now()
+            )
+            
+            if self.websocket_client:
+                self.websocket_client.send(result_data)
+                
+        except Exception as e:
+            logger.error(f"Failed to send task result: {e}")
+    
+    def get_status(self) -> Dict[str, Any]:
+        """Get session status."""
+        return {
+            "driver_id": self.driver_id,
+            "status": self.status.value,
+            "stats": self.stats.model_dump(),
+            "handlers": list(self.task_handlers.keys())
+        }

unrealon_driver/driver/communication/websocket_client.py

@@ -0,0 +1,205 @@
+"""
+Clean WebSocket client with auto-reconnection and message queuing.
+"""
+
+import asyncio
+import json
+import logging
+from typing import Optional, Callable, List
+from collections import deque
+
+import websockets
+try:
+    from websockets.asyncio.client import ClientConnection
+except ImportError:
+    # Fallback for older websockets versions
+    try:
+        from websockets.client import WebSocketClientProtocol as ClientConnection
+    except ImportError:
+        # Ultimate fallback
+        ClientConnection = None
+
+from unrealon_core.models.websocket import (
+    DriverRegistrationMessage,
+    DriverRegistrationData,
+    WebSocketMessage
+)
+from unrealon_core.enums.types import MessageType
+
+logger = logging.getLogger(__name__)
+
+
+class WebSocketClient:
+    """
+    Clean WebSocket client with robust connection management.
+    
+    Features:
+    - Auto-reconnection with exponential backoff
+    - Message queuing during disconnections
+    - Background message processing
+    - Clean error handling
+    """
+    
+    def __init__(self, websocket_url: str, driver_id: str):
+        self.websocket_url = websocket_url
+        self.driver_id = driver_id
+        self.websocket: Optional[ClientConnection] = None
+        self.connected = False
+        self.running = False
+        
+        # Message handling
+        self.on_message: Optional[Callable[[WebSocketMessage], None]] = None
+        self.message_queue: deque = deque()
+        
+        # Background tasks
+        self._sender_task: Optional[asyncio.Task] = None
+        self._receiver_task: Optional[asyncio.Task] = None
+        self._monitor_task: Optional[asyncio.Task] = None
+        
+        # Reconnection settings
+        self.reconnect_delay = 1.0
+        self.max_reconnect_delay = 60.0
+        self.reconnect_attempts = 0
+    
+    async def connect(self) -> bool:
+        """Connect and start background tasks."""
+        if self.running:
+            return True
+            
+        self.running = True
+        
+        # Start background tasks
+        self._sender_task = asyncio.create_task(self._message_sender())
+        self._receiver_task = asyncio.create_task(self._message_receiver()) 
+        self._monitor_task = asyncio.create_task(self._connection_monitor())
+        
+        return await self._establish_connection()
+    
+    async def disconnect(self):
+        """Clean disconnect and stop all tasks."""
+        self.running = False
+        
+        # Cancel background tasks
+        for task in [self._sender_task, self._receiver_task, self._monitor_task]:
+            if task and not task.done():
+                task.cancel()
+                try:
+                    await task
+                except asyncio.CancelledError:
+                    pass
+        
+        # Close WebSocket
+        if self.websocket:
+            await self.websocket.close()
+            self.websocket = None
+        
+        self.connected = False
+        logger.info("WebSocket client stopped")
+    
+    async def _establish_connection(self) -> bool:
+        """Establish WebSocket connection."""
+        try:
+            logger.info(f"Connecting to WebSocket: {self.websocket_url}")
+            self.websocket = await websockets.connect(self.websocket_url)
+            self.connected = True
+            self.reconnect_attempts = 0
+            self.reconnect_delay = 1.0
+            logger.info("WebSocket connected successfully")
+            return True
+        except Exception as e:
+            logger.error(f"WebSocket connection failed: {e}")
+            self.connected = False
+            return False
+    
+    async def _connection_monitor(self):
+        """Monitor connection and handle reconnection."""
+        while self.running:
+            if not self.connected and self.running:
+                logger.info(f"Attempting reconnection (attempt {self.reconnect_attempts + 1})")
+                
+                if await self._establish_connection():
+                    logger.info("Reconnection successful")
+                else:
+                    self.reconnect_attempts += 1
+                    # Exponential backoff
+                    self.reconnect_delay = min(
+                        self.reconnect_delay * 2, 
+                        self.max_reconnect_delay
+                    )
+                    await asyncio.sleep(self.reconnect_delay)
+            else:
+                await asyncio.sleep(5.0)  # Check every 5 seconds
+    
+    async def _message_sender(self):
+        """Background task to send queued messages."""
+        while self.running:
+            if self.connected and self.websocket and self.message_queue:
+                try:
+                    message = self.message_queue.popleft()
+                    await self.websocket.send(message)
+                except (websockets.exceptions.ConnectionClosed, ConnectionResetError):
+                    self.connected = False
+                    # Put message back in queue
+                    self.message_queue.appendleft(message)
+                except Exception as e:
+                    logger.error(f"Error sending message: {e}")
+            else:
+                await asyncio.sleep(0.1)
+    
+    async def _message_receiver(self):
+        """Background task to receive messages."""
+        while self.running:
+            if self.connected and self.websocket:
+                try:
+                    message_str = await self.websocket.recv()
+                    
+                    # Parse and validate message
+                    message_data = json.loads(message_str)
+                    message = WebSocketMessage.model_validate(message_data)
+                    
+                    # Handle message
+                    if self.on_message:
+                        self.on_message(message)
+                        
+                except (websockets.exceptions.ConnectionClosed, ConnectionResetError):
+                    self.connected = False
+                    logger.warning("WebSocket connection lost")
+                except Exception as e:
+                    logger.error(f"Error receiving message: {e}")
+            else:
+                await asyncio.sleep(0.1)
+    
+    def send(self, message_data) -> None:
+        """Queue message for sending."""
+        if hasattr(message_data, 'model_dump_json'):
+            # Pydantic model
+            message_json = message_data.model_dump_json()
+        else:
+            # Raw data
+            message_json = json.dumps(message_data)
+        
+        self.message_queue.append(message_json)
+    
+    async def register_driver(self, capabilities: List[str]) -> bool:
+        """Register driver with capabilities."""
+        try:
+            # Create registration data
+            registration_data = DriverRegistrationData(
+                driver_id=self.driver_id,
+                driver_name=self.driver_id,
+                driver_type="universal",
+                capabilities=capabilities
+            )
+            
+            # Create registration message with correct type
+            registration_message = DriverRegistrationMessage(data=registration_data)
+            
+            # Queue for sending
+            self.send(registration_message)
+            
+            logger.info(f"Driver registration queued: {self.driver_id}")
+            return True
+            
+        except Exception as e:
+            logger.error(f"Driver registration failed: {e}")
+            return False

unrealon_driver/driver/core/__init__.py

@@ -0,0 +1,10 @@
+"""
+Core driver components.
+
+Contains the main UniversalDriver class and configuration.
+"""
+
+from .driver import UniversalDriver
+from .config import DriverConfig, DriverMode
+
+__all__ = ["UniversalDriver", "DriverConfig", "DriverMode"]