oagi-core 0.10.1__py3-none-any.whl

oagi/server/socketio_server.py

@@ -0,0 +1,405 @@
+# -----------------------------------------------------------------------------
+#  Copyright (c) OpenAGI Foundation
+#  All rights reserved.
+#
+#  This file is part of the official API project.
+#  Licensed under the MIT License.
+# -----------------------------------------------------------------------------
+
+import asyncio
+import logging
+from datetime import datetime
+from typing import Any
+
+from pydantic import ValidationError
+
+from ..agent import AsyncDefaultAgent, create_agent
+from ..client import AsyncClient
+from ..exceptions import check_optional_dependency
+from ..types.models.action import Action, ActionType
+from .agent_wrappers import SocketIOActionHandler, SocketIOImageProvider
+from .config import ServerConfig
+from .models import (
+    BaseActionEventData,
+    ClickEventData,
+    DragEventData,
+    ErrorEventData,
+    FinishEventData,
+    HotkeyEventData,
+    InitEventData,
+    ScrollEventData,
+    TypeEventData,
+    WaitEventData,
+)
+from .session_store import Session, session_store
+
+check_optional_dependency("socketio", "Server features", "server")
+import socketio  # noqa: E402
+
+logger = logging.getLogger(__name__)
+
+sio = socketio.AsyncServer(
+    async_mode="asgi",
+    cors_allowed_origins="*",
+    logger=False,
+    engineio_logger=False,
+)
+
+
+class SessionNamespace(socketio.AsyncNamespace):
+    def __init__(self, namespace: str, config: ServerConfig):
+        super().__init__(namespace)
+        self.config = config
+        self.background_tasks: dict[str, asyncio.Task] = {}
+
+    async def on_connect(self, sid: str, environ: dict, auth: dict | None) -> bool:
+        session_id = self.namespace.split("/")[-1]
+        logger.info(f"Client {sid} connected to session {session_id}")
+
+        session = session_store.get_session(session_id)
+        if session:
+            session.socket_id = sid
+            session.namespace = self.namespace
+            session_store.update_activity(session_id)
+
+            # Create OAGI client if not exists
+            if not session.oagi_client:
+                session.oagi_client = AsyncClient(
+                    base_url=self.config.oagi_base_url,
+                    api_key=self.config.oagi_api_key,
+                )
+        else:
+            logger.warning(f"Connection to non-existent session {session_id}")
+            # Create session on connect if it doesn't exist
+            session = Session(
+                session_id=session_id,
+                instruction="",
+                mode="actor",  # Default mode
+                model=self.config.default_model,
+                temperature=self.config.default_temperature,
+            )
+            session.socket_id = sid
+            session.namespace = self.namespace
+            session.oagi_client = AsyncClient(
+                base_url=self.config.oagi_base_url,
+                api_key=self.config.oagi_api_key,
+            )
+            session_store.sessions[session_id] = session
+
+        return True
+
+    async def on_disconnect(self, sid: str) -> None:
+        session_id = self.namespace.split("/")[-1]
+        logger.info(f"Client {sid} disconnected from session {session_id}")
+
+        # Cancel any background tasks
+        if sid in self.background_tasks:
+            self.background_tasks[sid].cancel()
+            del self.background_tasks[sid]
+
+        # Start cleanup task
+        asyncio.create_task(self._cleanup_after_timeout(session_id))
+
+    async def _cleanup_after_timeout(self, session_id: str) -> None:
+        await asyncio.sleep(self.config.session_timeout_seconds)
+
+        session = session_store.get_session(session_id)
+        if session:
+            current_time = datetime.now().timestamp()
+            if (
+                current_time - session.last_activity
+                >= self.config.session_timeout_seconds
+            ):
+                logger.info(f"Session {session_id} timed out, cleaning up")
+
+                # Close OAGI client
+                if session.oagi_client:
+                    await session.oagi_client.close()
+
+                session_store.delete_session(session_id)
+
+    async def on_init(self, sid: str, data: dict) -> None:
+        try:
+            session_id = self.namespace.split("/")[-1]
+            logger.info(f"Initializing session {session_id}")
+
+            # Validate input
+            event_data = InitEventData(**data)
+
+            # Get or create session
+            session = session_store.get_session(session_id)
+            if not session:
+                logger.error(f"Session {session_id} not found")
+                await self.emit(
+                    "error",
+                    ErrorEventData(
+                        message=f"Session {session_id} not found"
+                    ).model_dump(),
+                    room=sid,
+                )
+                return
+
+            # Update session with init data
+            session.instruction = event_data.instruction
+            if event_data.mode:
+                session.mode = event_data.mode
+            if event_data.model:
+                session.model = event_data.model
+            if event_data.temperature is not None:
+                session.temperature = event_data.temperature
+            session.status = "running"
+            session_store.update_activity(session_id)
+
+            logger.info(
+                f"Session {session_id} initialized with: {event_data.instruction} "
+                f"(mode={event_data.mode}, model={event_data.model})"
+            )
+
+            # Create agent and wrappers
+            agent = create_agent(
+                mode=session.mode,
+                api_key=self.config.oagi_api_key,
+                base_url=self.config.oagi_base_url,
+                max_steps=self.config.max_steps,
+                model=event_data.model,
+                temperature=event_data.temperature,
+            )
+
+            action_handler = SocketIOActionHandler(self, session)
+            image_provider = SocketIOImageProvider(self, session, session.oagi_client)
+
+            # Start execution in background using agent
+            task = asyncio.create_task(
+                self._run_agent_task(
+                    agent,
+                    session,
+                    action_handler,
+                    image_provider,
+                    event_data.instruction,
+                )
+            )
+            self.background_tasks[sid] = task
+
+        except ValidationError as e:
+            logger.error(f"Invalid init data: {e}")
+            await self.emit(
+                "error",
+                ErrorEventData(
+                    message="Invalid init data",
+                    details={"validation_errors": e.errors()},
+                ).model_dump(),
+                room=sid,
+            )
+        except Exception as e:
+            logger.error(f"Error in init: {e}", exc_info=True)
+            await self.emit(
+                "error",
+                ErrorEventData(message=str(e)).model_dump(),
+                room=sid,
+            )
+
+    async def _run_agent_task(
+        self,
+        agent: AsyncDefaultAgent,
+        session: Session,
+        action_handler: SocketIOActionHandler,
+        image_provider: SocketIOImageProvider,
+        instruction: str,
+    ) -> None:
+        try:
+            # Execute task using agent
+            success = await agent.execute(
+                instruction=instruction,
+                action_handler=action_handler,
+                image_provider=image_provider,
+            )
+
+            # Update session status
+            if success:
+                session.status = "completed"
+                logger.info(
+                    f"Task completed successfully for session {session.session_id}"
+                )
+
+                # Emit finish event
+                await self.call(
+                    "finish",
+                    FinishEventData(index=0, total=1).model_dump(),
+                    to=session.socket_id,
+                    timeout=self.config.socketio_timeout,
+                )
+            else:
+                session.status = "failed"
+                logger.warning(f"Task failed for session {session.session_id}")
+
+            session_store.update_activity(session.session_id)
+
+        except asyncio.CancelledError:
+            logger.info(f"Agent task cancelled for session {session.session_id}")
+            session.status = "cancelled"
+        except Exception as e:
+            logger.error(f"Error in agent task: {e}", exc_info=True)
+            session.status = "failed"
+            if session.socket_id:
+                await self.emit(
+                    "error",
+                    ErrorEventData(message=f"Execution failed: {str(e)}").model_dump(),
+                    room=session.socket_id,
+                )
+
+    async def _emit_actions(self, session: Session, actions: list[Action]) -> None:
+        total = len(actions)
+
+        for i, action in enumerate(actions):
+            try:
+                ack = await self._emit_single_action(session, action, i, total)
+                session.actions_executed += 1
+
+                if ack and not ack.get("success"):
+                    logger.warning(f"Action {i} failed: {ack.get('error')}")
+
+            except Exception as e:
+                logger.error(f"Error emitting action {i}: {e}", exc_info=True)
+
+    async def _emit_single_action(
+        self, session: Session, action: Action, index: int, total: int
+    ) -> dict | None:
+        arg = action.argument.strip("()")
+        common = BaseActionEventData(index=index, total=total).model_dump()
+
+        logger.info(f"Emitting action {index + 1}/{total}: {action.type.value} {arg}")
+        match action.type:
+            case (
+                ActionType.CLICK
+                | ActionType.LEFT_DOUBLE
+                | ActionType.LEFT_TRIPLE
+                | ActionType.RIGHT_SINGLE
+            ):
+                coords = arg.split(",")
+                if len(coords) >= 2:
+                    x, y = int(coords[0]), int(coords[1])
+                else:
+                    logger.warning(f"Invalid action coordinates: {arg}")
+                    return None
+
+                return await self.call(
+                    action.type.value,
+                    ClickEventData(**common, x=x, y=y).model_dump(),
+                    to=session.socket_id,
+                    timeout=self.config.socketio_timeout,
+                )
+
+            case ActionType.DRAG:
+                coords = arg.split(",")
+                if len(coords) >= 4:
+                    x1, y1, x2, y2 = (int(coords[i]) for i in range(4))
+                else:
+                    logger.warning(f"Invalid drag coordinates: {arg}")
+                    return None
+
+                return await self.call(
+                    "drag",
+                    DragEventData(**common, x1=x1, y1=y1, x2=x2, y2=y2).model_dump(),
+                    to=session.socket_id,
+                    timeout=self.config.socketio_timeout,
+                )
+
+            case ActionType.HOTKEY:
+                combo = arg.strip()
+                count = action.count or 1
+
+                return await self.call(
+                    "hotkey",
+                    HotkeyEventData(**common, combo=combo, count=count).model_dump(),
+                    to=session.socket_id,
+                    timeout=self.config.socketio_timeout,
+                )
+
+            case ActionType.TYPE:
+                text = arg.strip()
+
+                return await self.call(
+                    "type",
+                    TypeEventData(**common, text=text).model_dump(),
+                    to=session.socket_id,
+                    timeout=self.config.socketio_timeout,
+                )
+
+            case ActionType.SCROLL:
+                parts = arg.split(",")
+                if len(parts) >= 3:
+                    x, y = int(parts[0]), int(parts[1])
+                    direction = parts[2].strip().lower()
+                else:
+                    logger.warning(f"Invalid scroll coordinates: {arg}")
+                    return None
+
+                count = action.count or 1
+
+                return await self.call(
+                    "scroll",
+                    ScrollEventData(
+                        **common,
+                        x=x,
+                        y=y,
+                        direction=direction,
+                        count=count,  # type: ignore
+                    ).model_dump(),
+                    to=session.socket_id,
+                    timeout=self.config.socketio_timeout,
+                )
+
+            case ActionType.WAIT:
+                try:
+                    duration_ms = int(arg) if arg else 1000
+                except (ValueError, TypeError):
+                    duration_ms = 1000
+
+                return await self.call(
+                    "wait",
+                    WaitEventData(**common, duration_ms=duration_ms).model_dump(),
+                    to=session.socket_id,
+                    timeout=self.config.socketio_timeout,
+                )
+
+            case ActionType.FINISH:
+                return await self.call(
+                    "finish",
+                    FinishEventData(**common).model_dump(),
+                    to=session.socket_id,
+                    timeout=self.config.socketio_timeout,
+                )
+
+            case _:
+                logger.warning(f"Unknown action type: {action.type}")
+                return None
+
+
+# Dynamic namespace registration
+_registered_namespaces: dict[str, SessionNamespace] = {}
+
+
+def get_or_create_namespace(namespace: str, config: ServerConfig) -> SessionNamespace:
+    if namespace not in _registered_namespaces:
+        ns = SessionNamespace(namespace, config)
+        sio.register_namespace(ns)
+        _registered_namespaces[namespace] = ns
+        logger.info(f"Registered namespace: {namespace}")
+    return _registered_namespaces[namespace]
+
+
+# Patch connect handler for dynamic registration
+original_connect = sio._handle_connect
+
+
+async def _patched_handle_connect(eio_sid: str, namespace: str, data: Any) -> Any:
+    if namespace and namespace.startswith("/session/"):
+        config = ServerConfig()
+        get_or_create_namespace(namespace, config)
+    return await original_connect(eio_sid, namespace, data)
+
+
+sio._handle_connect = _patched_handle_connect
+
+# Create ASGI app
+socket_app = socketio.ASGIApp(sio, socketio_path="socket.io")

oagi/task/__init__.py

@@ -0,0 +1,21 @@
+# -----------------------------------------------------------------------------
+#  Copyright (c) OpenAGI Foundation
+#  All rights reserved.
+#
+#  This file is part of the official API project.
+#  Licensed under the MIT License.
+# -----------------------------------------------------------------------------
+
+from .async_ import AsyncActor, AsyncTask
+from .async_short import AsyncShortTask
+from .short import ShortTask
+from .sync import Actor, Task
+
+__all__ = [
+    "Actor",
+    "AsyncActor",
+    "Task",  # Deprecated: Use Actor instead
+    "AsyncTask",  # Deprecated: Use AsyncActor instead
+    "ShortTask",  # Deprecated
+    "AsyncShortTask",  # Deprecated
+]

oagi/task/async_.py

@@ -0,0 +1,101 @@
+# -----------------------------------------------------------------------------
+#  Copyright (c) OpenAGI Foundation
+#  All rights reserved.
+#
+#  This file is part of the official API project.
+#  Licensed under the MIT License.
+# -----------------------------------------------------------------------------
+
+import warnings
+
+from ..client import AsyncClient
+from ..types import URL, Image, Step
+from .base import BaseActor
+
+
+class AsyncActor(BaseActor):
+    """Async base class for task automation with the OAGI API."""
+
+    def __init__(
+        self,
+        api_key: str | None = None,
+        base_url: str | None = None,
+        model: str = "lux-actor-1",
+        temperature: float | None = None,
+    ):
+        super().__init__(api_key, base_url, model, temperature)
+        self.client = AsyncClient(base_url=base_url, api_key=api_key)
+        self.api_key = self.client.api_key
+        self.base_url = self.client.base_url
+
+    async def init_task(
+        self,
+        task_desc: str,
+        max_steps: int = 20,
+    ):
+        """Initialize a new task with the given description.
+
+        Args:
+            task_desc: Task description
+            max_steps: Maximum number of steps allowed
+        """
+        self._prepare_init_task(task_desc, max_steps)
+
+    async def step(
+        self,
+        screenshot: Image | URL | bytes,
+        instruction: str | None = None,
+        temperature: float | None = None,
+    ) -> Step:
+        """Send screenshot to the server and get the next actions.
+
+        Args:
+            screenshot: Screenshot as Image object or raw bytes
+            instruction: Optional additional instruction for this step
+            temperature: Sampling temperature for this step (overrides task default if provided)
+
+        Returns:
+            Step: The actions and reasoning for this step
+        """
+        kwargs = self._prepare_step(
+            screenshot, instruction, temperature, prefix="async "
+        )
+
+        try:
+            response = await self.client.create_message(**kwargs)
+            return self._build_step_response(response, prefix="Async ")
+        except Exception as e:
+            self._handle_step_error(e, prefix="async ")
+
+    async def close(self):
+        """Close the underlying HTTP client to free resources."""
+        await self.client.close()
+
+    async def __aenter__(self):
+        return self
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb):
+        await self.close()
+
+
+class AsyncTask(AsyncActor):
+    """Deprecated: Use AsyncActor instead.
+
+    This class is deprecated and will be removed in a future version.
+    Please use AsyncActor instead.
+    """
+
+    def __init__(
+        self,
+        api_key: str | None = None,
+        base_url: str | None = None,
+        model: str = "lux-actor-1",
+        temperature: float | None = None,
+    ):
+        warnings.warn(
+            "AsyncTask is deprecated and will be removed in a future version. "
+            "Please use AsyncActor instead.",
+            DeprecationWarning,
+            stacklevel=2,
+        )
+        super().__init__(api_key, base_url, model, temperature)

oagi/task/async_short.py

@@ -0,0 +1,76 @@
+# -----------------------------------------------------------------------------
+#  Copyright (c) OpenAGI Foundation
+#  All rights reserved.
+#
+#  This file is part of the official API project.
+#  Licensed under the MIT License.
+# -----------------------------------------------------------------------------
+
+import warnings
+
+from ..logging import get_logger
+from ..types import AsyncActionHandler, AsyncImageProvider
+from .async_ import AsyncActor
+from .base import BaseAutoMode
+
+logger = get_logger("async_short_task")
+
+
+class AsyncShortTask(AsyncActor, BaseAutoMode):
+    """Deprecated: This class is deprecated and will be removed in a future version.
+
+    Async task implementation with automatic mode for short-duration tasks.
+    Please use AsyncActor directly with custom automation logic instead.
+    """
+
+    def __init__(
+        self,
+        api_key: str | None = None,
+        base_url: str | None = None,
+        model: str = "lux-actor-1",
+        temperature: float | None = None,
+    ):
+        warnings.warn(
+            "AsyncShortTask is deprecated and will be removed in a future version. "
+            "Please use AsyncActor with custom automation logic instead.",
+            DeprecationWarning,
+            stacklevel=2,
+        )
+        super().__init__(
+            api_key=api_key, base_url=base_url, model=model, temperature=temperature
+        )
+
+    async def auto_mode(
+        self,
+        task_desc: str,
+        max_steps: int = 20,
+        executor: AsyncActionHandler = None,
+        image_provider: AsyncImageProvider = None,
+        temperature: float | None = None,
+    ) -> bool:
+        """Run the task in automatic mode with the provided executor and image provider.
+
+        Args:
+            task_desc: Task description
+            max_steps: Maximum number of steps
+            executor: Async handler to execute actions
+            image_provider: Async provider for screenshots
+            temperature: Sampling temperature for all steps (overrides task default if provided)
+        """
+        self._log_auto_mode_start(task_desc, max_steps, prefix="async ")
+
+        await self.init_task(task_desc, max_steps=max_steps)
+
+        for i in range(max_steps):
+            self._log_auto_mode_step(i + 1, max_steps, prefix="async ")
+            image = await image_provider()
+            step = await self.step(image, temperature=temperature)
+            if executor:
+                self._log_auto_mode_actions(len(step.actions), prefix="async ")
+                await executor(step.actions)
+            if step.stop:
+                self._log_auto_mode_completion(i + 1, prefix="async ")
+                return True
+
+        self._log_auto_mode_max_steps(max_steps, prefix="async ")
+        return False