oagi-core 0.9.0__py3-none-any.whl

oagi/server/models.py

@@ -0,0 +1,98 @@
+# -----------------------------------------------------------------------------
+#  Copyright (c) OpenAGI Foundation
+#  All rights reserved.
+#
+#  This file is part of the official API project.
+#  Licensed under the MIT License.
+# -----------------------------------------------------------------------------
+
+from typing import Literal
+
+from pydantic import BaseModel, Field
+
+
+# Client-to-server events
+class InitEventData(BaseModel):
+    instruction: str = Field(...)
+    mode: str | None = Field(default="actor")
+    model: str | None = Field(default="lux-v1")
+    temperature: float | None = Field(default=0.1, ge=0.0, le=2.0)
+
+
+# Server-to-client events
+class BaseActionEventData(BaseModel):
+    index: int = Field(..., ge=0)
+    total: int = Field(..., ge=1)
+
+
+class ClickEventData(BaseActionEventData):
+    x: int = Field(..., ge=0, le=1000)
+    y: int = Field(..., ge=0, le=1000)
+
+
+class DragEventData(BaseActionEventData):
+    x1: int = Field(..., ge=0, le=1000)
+    y1: int = Field(..., ge=0, le=1000)
+    x2: int = Field(..., ge=0, le=1000)
+    y2: int = Field(..., ge=0, le=1000)
+
+
+class HotkeyEventData(BaseActionEventData):
+    combo: str = Field(...)
+    count: int = Field(default=1, ge=1)
+
+
+class TypeEventData(BaseActionEventData):
+    text: str = Field(...)
+
+
+class ScrollEventData(BaseActionEventData):
+    x: int = Field(..., ge=0, le=1000)
+    y: int = Field(..., ge=0, le=1000)
+    direction: Literal["up", "down"] = Field(...)
+    count: int = Field(default=1, ge=1)
+
+
+class WaitEventData(BaseActionEventData):
+    duration_ms: int = Field(default=1000, ge=0)
+
+
+class FinishEventData(BaseActionEventData):
+    pass
+
+
+# Screenshot request/response
+class ScreenshotRequestData(BaseModel):
+    presigned_url: str = Field(...)
+    uuid: str = Field(...)
+    expires_at: str = Field(...)
+
+
+class ScreenshotResponseData(BaseModel):
+    success: bool = Field(...)
+    error: str | None = Field(None)
+
+
+# Action acknowledgement
+class ActionAckData(BaseModel):
+    action_index: int = Field(...)
+    success: bool = Field(...)
+    error: str | None = Field(None)
+    execution_time_ms: int | None = Field(None)
+
+
+# Session status
+class SessionStatusData(BaseModel):
+    session_id: str = Field(...)
+    status: Literal["initialized", "running", "completed", "failed"] = Field(...)
+    instruction: str = Field(...)
+    created_at: str = Field(...)
+    actions_executed: int = Field(default=0)
+    last_activity: str = Field(...)
+
+
+# Error event
+class ErrorEventData(BaseModel):
+    message: str = Field(...)
+    code: str | None = Field(None)
+    details: dict | None = Field(None)

oagi/server/session_store.py

@@ -0,0 +1,116 @@
+# -----------------------------------------------------------------------------
+#  Copyright (c) OpenAGI Foundation
+#  All rights reserved.
+#
+#  This file is part of the official API project.
+#  Licensed under the MIT License.
+# -----------------------------------------------------------------------------
+
+import secrets
+from datetime import datetime
+from typing import Any
+from uuid import uuid4
+
+
+class Session:
+    def __init__(
+        self,
+        session_id: str,
+        instruction: str,
+        mode: str = "actor",
+        model: str = "lux-v1",
+        temperature: float = 0.0,
+    ):
+        self.session_id: str = session_id
+        self.instruction: str = instruction
+        self.mode: str = mode
+        self.model: str = model
+        self.temperature: float = temperature
+
+        # OAGI task state
+        self.task_id: str = uuid4().hex
+        self.message_history: list[dict[str, Any]] = []
+        self.current_screenshot_url: str | None = None
+
+        # Socket state
+        self.socket_id: str | None = None
+        self.namespace: str | None = None
+        self.last_activity: float = datetime.now().timestamp()
+
+        # Status tracking
+        self.status: str = "initialized"
+        self.created_at: str = datetime.now().isoformat()
+        self.actions_executed: int = 0
+
+        # OAGI client reference
+        self.oagi_client: Any | None = None
+
+
+class SessionStore:
+    def __init__(self):
+        self.sessions: dict[str, Session] = {}
+
+    def create_session(
+        self,
+        instruction: str,
+        mode: str = "actor",
+        model: str = "lux-v1",
+        temperature: float = 0.0,
+        session_id: str | None = None,
+    ) -> str:
+        if session_id is None:
+            session_id = f"ses_{secrets.token_urlsafe(16)}"
+
+        session = Session(session_id, instruction, mode, model, temperature)
+        self.sessions[session_id] = session
+        return session_id
+
+    def get_session(self, session_id: str) -> Session | None:
+        return self.sessions.get(session_id)
+
+    def get_session_by_socket_id(self, socket_id: str) -> Session | None:
+        for session in self.sessions.values():
+            if session.socket_id == socket_id:
+                return session
+        return None
+
+    def delete_session(self, session_id: str) -> bool:
+        if session_id in self.sessions:
+            self.sessions.pop(session_id)
+            return True
+        return False
+
+    def update_activity(self, session_id: str) -> None:
+        session = self.sessions.get(session_id)
+        if session:
+            session.last_activity = datetime.now().timestamp()
+
+    def list_sessions(self) -> list[dict[str, Any]]:
+        return [
+            {
+                "session_id": session.session_id,
+                "status": session.status,
+                "instruction": session.instruction,
+                "created_at": session.created_at,
+                "actions_executed": session.actions_executed,
+                "connected": session.socket_id is not None,
+            }
+            for session in self.sessions.values()
+        ]
+
+    def cleanup_inactive_sessions(self, timeout_seconds: float) -> int:
+        current_time = datetime.now().timestamp()
+        sessions_to_delete = []
+
+        for session_id, session in self.sessions.items():
+            if current_time - session.last_activity > timeout_seconds:
+                sessions_to_delete.append(session_id)
+
+        for session_id in sessions_to_delete:
+            self.delete_session(session_id)
+
+        return len(sessions_to_delete)
+
+
+# Global instance
+session_store = SessionStore()

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(action_index=0, total_actions=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/single_step.py

@@ -0,0 +1,87 @@
+# -----------------------------------------------------------------------------
+#  Copyright (c) OpenAGI Foundation
+#  All rights reserved.
+#
+#  This file is part of the official API project.
+#  Licensed under the MIT License.
+# -----------------------------------------------------------------------------
+
+from pathlib import Path
+
+from .task import Task
+from .types import Image, Step
+
+
+def single_step(
+    task_description: str,
+    screenshot: str | bytes | Path | Image,
+    instruction: str | None = None,
+    api_key: str | None = None,
+    base_url: str | None = None,
+    temperature: float | None = None,
+) -> Step:
+    """
+    Perform a single-step inference without maintaining task state.
+
+    This is useful for one-off analyses where you don't need to maintain
+    a conversation or task context across multiple steps.
+
+    Args:
+        task_description: Description of the task to perform
+        screenshot: Screenshot as Image, bytes, or file path
+        instruction: Optional additional instruction for the task
+        api_key: OAGI API key (uses environment variable if not provided)
+        base_url: OAGI base URL (uses environment variable if not provided)
+        temperature: Sampling temperature (0.0-2.0) for LLM inference
+
+    Returns:
+        Step: Object containing reasoning, actions, and completion status
+
+    Example:
+        >>> # Using with bytes
+        >>> with open("screenshot.png", "rb") as f:
+        ...     image_bytes = f.read()
+        >>> step = single_step(
+        ...     task_description="Click the submit button",
+        ...     screenshot=image_bytes
+        ... )
+
+        >>> # Using with file path
+        >>> step = single_step(
+        ...     task_description="Fill in the form",
+        ...     screenshot=Path("screenshot.png"),
+        ...     instruction="Use test@example.com for email"
+        ... )
+
+        >>> # Using with Image object
+        >>> from oagi.types import Image
+        >>> image = Image(...)
+        >>> step = single_step(
+        ...     task_description="Navigate to settings",
+        ...     screenshot=image
+        ... )
+    """
+    # Lazy import PILImage only when needed
+    from .pil_image import PILImage  # noqa: PLC0415
+
+    # Convert file paths to bytes using PILImage
+    if isinstance(screenshot, (str, Path)):
+        path = Path(screenshot) if isinstance(screenshot, str) else screenshot
+        if path.exists():
+            pil_image = PILImage.from_file(str(path))
+            screenshot_bytes = pil_image.read()
+        else:
+            raise FileNotFoundError(f"Screenshot file not found: {path}")
+    elif isinstance(screenshot, bytes):
+        screenshot_bytes = screenshot
+    elif isinstance(screenshot, Image):
+        screenshot_bytes = screenshot.read()
+    else:
+        raise ValueError(
+            f"screenshot must be Image, bytes, str, or Path, got {type(screenshot)}"
+        )
+
+    # Use Task to perform single step
+    with Task(api_key=api_key, base_url=base_url, temperature=temperature) as task:
+        task.init_task(task_description)
+        return task.step(screenshot_bytes, instruction=instruction)