oagi 0.2.1__py3-none-any.whl → 0.4.0__py3-none-any.whl

oagi/__init__.py

@@ -6,6 +6,12 @@
 #  Licensed under the MIT License.
 # -----------------------------------------------------------------------------
 
+from oagi.async_client import AsyncClient
+from oagi.async_pyautogui_action_handler import AsyncPyautoguiActionHandler
+from oagi.async_screenshot_maker import AsyncScreenshotMaker
+from oagi.async_short_task import AsyncShortTask
+from oagi.async_single_step import async_single_step
+from oagi.async_task import AsyncTask
 from oagi.exceptions import (
     APIError,
     AuthenticationError,
@@ -18,23 +24,45 @@ from oagi.exceptions import (
     ServerError,
     ValidationError,
 )
-from oagi.pyautogui_action_handler import PyautoguiActionHandler
+from oagi.pil_image import PILImage
+from oagi.pyautogui_action_handler import PyautoguiActionHandler, PyautoguiConfig
 from oagi.screenshot_maker import ScreenshotMaker
 from oagi.short_task import ShortTask
 from oagi.single_step import single_step
 from oagi.sync_client import ErrorDetail, ErrorResponse, LLMResponse, SyncClient
 from oagi.task import Task
+from oagi.types import (
+    AsyncActionHandler,
+    AsyncImageProvider,
+    ImageConfig,
+)
 
 __all__ = [
-    # Core classes
+    # Core sync classes
     "Task",
     "ShortTask",
     "SyncClient",
+    # Core async classes
+    "AsyncTask",
+    "AsyncShortTask",
+    "AsyncClient",
     # Functions
     "single_step",
+    "async_single_step",
+    # Image classes
+    "PILImage",
     # Handler classes
     "PyautoguiActionHandler",
+    "PyautoguiConfig",
     "ScreenshotMaker",
+    # Async handler classes
+    "AsyncPyautoguiActionHandler",
+    "AsyncScreenshotMaker",
+    # Async protocols
+    "AsyncActionHandler",
+    "AsyncImageProvider",
+    # Configuration
+    "ImageConfig",
     # Response models
     "LLMResponse",
     "ErrorResponse",

oagi/async_client.py

@@ -0,0 +1,239 @@
+# -----------------------------------------------------------------------------
+#  Copyright (c) OpenAGI Foundation
+#  All rights reserved.
+#
+#  This file is part of the official API project.
+#  Licensed under the MIT License.
+# -----------------------------------------------------------------------------
+
+import os
+from functools import wraps
+
+import httpx
+
+from .exceptions import (
+    APIError,
+    AuthenticationError,
+    ConfigurationError,
+    NetworkError,
+    NotFoundError,
+    RateLimitError,
+    RequestTimeoutError,
+    ServerError,
+    ValidationError,
+)
+from .logging import get_logger
+from .sync_client import ErrorResponse, LLMResponse
+
+logger = get_logger("async_client")
+
+
+def async_log_trace_on_failure(func):
+    """Async decorator that logs trace ID when a method fails."""
+
+    @wraps(func)
+    async def wrapper(*args, **kwargs):
+        try:
+            return await func(*args, **kwargs)
+        except Exception as e:
+            # Try to get response from the exception if it has one
+            if (response := getattr(e, "response", None)) is not None:
+                logger.error(f"Request Id: {response.headers.get('x-request-id', '')}")
+                logger.error(f"Trace Id: {response.headers.get('x-trace-id', '')}")
+            raise
+
+    return wrapper
+
+
+class AsyncClient:
+    """Async HTTP client for the OAGI API."""
+
+    def __init__(self, base_url: str | None = None, api_key: str | None = None):
+        # Get from environment if not provided
+        self.base_url = base_url or os.getenv("OAGI_BASE_URL")
+        self.api_key = api_key or os.getenv("OAGI_API_KEY")
+
+        # Validate required configuration
+        if not self.base_url:
+            raise ConfigurationError(
+                "OAGI base URL must be provided either as 'base_url' parameter or "
+                "OAGI_BASE_URL environment variable"
+            )
+
+        if not self.api_key:
+            raise ConfigurationError(
+                "OAGI API key must be provided either as 'api_key' parameter or "
+                "OAGI_API_KEY environment variable"
+            )
+
+        self.base_url = self.base_url.rstrip("/")
+        self.client = httpx.AsyncClient(base_url=self.base_url)
+        self.timeout = 60
+
+        logger.info(f"AsyncClient initialized with base_url: {self.base_url}")
+
+    async def __aenter__(self):
+        return self
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb):
+        await self.client.aclose()
+
+    async def close(self):
+        """Close the underlying httpx async client"""
+        await self.client.aclose()
+
+    @async_log_trace_on_failure
+    async def create_message(
+        self,
+        model: str,
+        screenshot: str,  # base64 encoded
+        task_description: str | None = None,
+        task_id: str | None = None,
+        instruction: str | None = None,
+        max_actions: int | None = 5,
+        api_version: str | None = None,
+    ) -> LLMResponse:
+        """
+        Call the /v1/message endpoint to analyze task and screenshot
+
+        Args:
+            model: The model to use for task analysis
+            screenshot: Base64-encoded screenshot image
+            task_description: Description of the task (required for new sessions)
+            task_id: Task ID for continuing existing task
+            instruction: Additional instruction when continuing a session (only works with task_id)
+            max_actions: Maximum number of actions to return (1-20)
+            api_version: API version header
+
+        Returns:
+            LLMResponse: The response from the API
+
+        Raises:
+            httpx.HTTPStatusError: For HTTP error responses
+        """
+        headers = {}
+        if api_version:
+            headers["x-api-version"] = api_version
+        if self.api_key:
+            headers["x-api-key"] = self.api_key
+
+        payload = {"model": model, "screenshot": screenshot}
+
+        if task_description is not None:
+            payload["task_description"] = task_description
+        if task_id is not None:
+            payload["task_id"] = task_id
+        if instruction is not None:
+            payload["instruction"] = instruction
+        if max_actions is not None:
+            payload["max_actions"] = max_actions
+
+        logger.info(f"Making async API request to /v1/message with model: {model}")
+        logger.debug(
+            f"Request includes task_description: {task_description is not None}, task_id: {task_id is not None}"
+        )
+
+        try:
+            response = await self.client.post(
+                "/v1/message", json=payload, headers=headers, timeout=self.timeout
+            )
+        except httpx.TimeoutException as e:
+            logger.error(f"Request timed out after {self.timeout} seconds")
+            raise RequestTimeoutError(
+                f"Request timed out after {self.timeout} seconds", e
+            )
+        except httpx.NetworkError as e:
+            logger.error(f"Network error: {e}")
+            raise NetworkError(f"Network error: {e}", e)
+
+        try:
+            response_data = response.json()
+        except ValueError:
+            # If response is not JSON, raise API error
+            logger.error(f"Non-JSON API response: {response.status_code}")
+            raise APIError(
+                f"Invalid response format (status {response.status_code})",
+                status_code=response.status_code,
+                response=response,
+            )
+
+        # Check if it's an error response (non-200 status or has error field)
+        if response.status_code != 200:
+            error_resp = ErrorResponse(**response_data)
+            if error_resp.error:
+                error_code = error_resp.error.code
+                error_msg = error_resp.error.message
+                logger.error(f"API Error [{error_code}]: {error_msg}")
+
+                # Map to specific exception types based on status code
+                exception_class = self._get_exception_class(response.status_code)
+                raise exception_class(
+                    error_msg,
+                    code=error_code,
+                    status_code=response.status_code,
+                    response=response,
+                )
+            else:
+                # Error response without error details
+                logger.error(
+                    f"API error response without details: {response.status_code}"
+                )
+                exception_class = self._get_exception_class(response.status_code)
+                raise exception_class(
+                    f"API error (status {response.status_code})",
+                    status_code=response.status_code,
+                    response=response,
+                )
+
+        # Parse successful response
+        result = LLMResponse(**response_data)
+
+        # Check if the response contains an error (even with 200 status)
+        if result.error:
+            logger.error(
+                f"API Error in response: [{result.error.code}]: {result.error.message}"
+            )
+            raise APIError(
+                result.error.message,
+                code=result.error.code,
+                status_code=200,
+                response=response,
+            )
+
+        logger.info(
+            f"Async API request successful - task_id: {result.task_id}, step: {result.current_step}, complete: {result.is_complete}"
+        )
+        logger.debug(f"Response included {len(result.actions)} actions")
+        return result
+
+    def _get_exception_class(self, status_code: int) -> type[APIError]:
+        """Get the appropriate exception class based on status code."""
+        status_map = {
+            401: AuthenticationError,
+            404: NotFoundError,
+            422: ValidationError,
+            429: RateLimitError,
+        }
+
+        if status_code >= 500:
+            return ServerError
+
+        return status_map.get(status_code, APIError)
+
+    async def health_check(self) -> dict:
+        """
+        Call the /health endpoint for health check
+
+        Returns:
+            dict: Health check response
+        """
+        logger.debug("Making async health check request")
+        try:
+            response = await self.client.get("/health")
+            response.raise_for_status()
+            result = response.json()
+            logger.debug("Async health check successful")
+            return result
+        except httpx.HTTPStatusError as e:
+            logger.warning(f"Async health check failed: {e}")
+            raise

oagi/async_pyautogui_action_handler.py

@@ -0,0 +1,44 @@
+# -----------------------------------------------------------------------------
+#  Copyright (c) OpenAGI Foundation
+#  All rights reserved.
+#
+#  This file is part of the official API project.
+#  Licensed under the MIT License.
+# -----------------------------------------------------------------------------
+
+import asyncio
+
+from .pyautogui_action_handler import PyautoguiActionHandler, PyautoguiConfig
+from .types import Action
+
+
+class AsyncPyautoguiActionHandler:
+    """
+    Async wrapper for PyautoguiActionHandler that runs actions in a thread pool.
+
+    This allows PyAutoGUI operations to be non-blocking in async contexts,
+    enabling concurrent execution of other async tasks while GUI actions are performed.
+    """
+
+    def __init__(self, config: PyautoguiConfig | None = None):
+        """Initialize with optional configuration.
+
+        Args:
+            config: PyautoguiConfig instance for customizing behavior
+        """
+        self.sync_handler = PyautoguiActionHandler(config=config)
+        self.config = config or PyautoguiConfig()
+
+    async def __call__(self, actions: list[Action]) -> None:
+        """
+        Execute actions asynchronously using a thread pool executor.
+
+        This prevents PyAutoGUI operations from blocking the async event loop,
+        allowing other coroutines to run while GUI actions are being performed.
+
+        Args:
+            actions: List of actions to execute
+        """
+        loop = asyncio.get_event_loop()
+        # Run the synchronous handler in a thread pool to avoid blocking
+        await loop.run_in_executor(None, self.sync_handler, actions)

oagi/async_screenshot_maker.py

@@ -0,0 +1,47 @@
+# -----------------------------------------------------------------------------
+#  Copyright (c) OpenAGI Foundation
+#  All rights reserved.
+#
+#  This file is part of the official API project.
+#  Licensed under the MIT License.
+# -----------------------------------------------------------------------------
+
+import asyncio
+
+from .screenshot_maker import ScreenshotMaker
+from .types import Image, ImageConfig
+
+
+class AsyncScreenshotMaker:
+    """
+    Async wrapper for ScreenshotMaker that captures screenshots in a thread pool.
+
+    This allows screenshot capture to be non-blocking in async contexts,
+    enabling concurrent execution of other async tasks while screenshots are taken.
+    """
+
+    def __init__(self, config: ImageConfig | None = None):
+        """Initialize with optional image configuration.
+
+        Args:
+            config: ImageConfig instance for customizing screenshot format and quality
+        """
+        self.sync_screenshot_maker = ScreenshotMaker(config=config)
+        self.config = config
+
+    async def __call__(self) -> Image:
+        """
+        Capture a screenshot asynchronously using a thread pool executor.
+
+        This prevents screenshot capture from blocking the async event loop,
+        allowing other coroutines to run while the screenshot is being taken.
+
+        Returns:
+            Image: The captured screenshot as a PILImage
+        """
+        loop = asyncio.get_event_loop()
+        # Run the synchronous screenshot capture in a thread pool to avoid blocking
+        return await loop.run_in_executor(None, self.sync_screenshot_maker)
+
+    async def last_image(self) -> Image:
+        return self.sync_screenshot_maker.last_image()

oagi/async_short_task.py

@@ -0,0 +1,56 @@
+# -----------------------------------------------------------------------------
+#  Copyright (c) OpenAGI Foundation
+#  All rights reserved.
+#
+#  This file is part of the official API project.
+#  Licensed under the MIT License.
+# -----------------------------------------------------------------------------
+
+from .async_task import AsyncTask
+from .logging import get_logger
+from .types import AsyncActionHandler, AsyncImageProvider
+
+logger = get_logger("async_short_task")
+
+
+class AsyncShortTask(AsyncTask):
+    """Async task implementation with automatic mode for short-duration tasks."""
+
+    def __init__(
+        self,
+        api_key: str | None = None,
+        base_url: str | None = None,
+        model: str = "vision-model-v1",
+    ):
+        super().__init__(api_key=api_key, base_url=base_url, model=model)
+
+    async def auto_mode(
+        self,
+        task_desc: str,
+        max_steps: int = 5,
+        executor: AsyncActionHandler = None,
+        image_provider: AsyncImageProvider = None,
+    ) -> bool:
+        """Run the task in automatic mode with the provided executor and image provider."""
+        logger.info(
+            f"Starting async auto mode for task: '{task_desc}' (max_steps: {max_steps})"
+        )
+        await self.init_task(task_desc, max_steps=max_steps)
+
+        for i in range(max_steps):
+            logger.debug(f"Async auto mode step {i + 1}/{max_steps}")
+            image = await image_provider()
+            step = await self.step(image)
+            if step.stop:
+                logger.info(
+                    f"Async auto mode completed successfully after {i + 1} steps"
+                )
+                return True
+            if executor:
+                logger.debug(f"Executing {len(step.actions)} actions asynchronously")
+                await executor(step.actions)
+
+        logger.warning(
+            f"Async auto mode reached max steps ({max_steps}) without completion"
+        )
+        return False

oagi/async_single_step.py

@@ -0,0 +1,83 @@
+# -----------------------------------------------------------------------------
+#  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 .async_task import AsyncTask
+from .pil_image import PILImage
+from .types import Image, Step
+
+
+async def async_single_step(
+    task_description: str,
+    screenshot: str | bytes | Path | Image,
+    instruction: str | None = None,
+    api_key: str | None = None,
+    base_url: str | None = None,
+) -> Step:
+    """
+    Perform a single-step inference asynchronously 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)
+
+    Returns:
+        Step: Object containing reasoning, actions, and completion status
+
+    Example:
+        >>> # Using with bytes
+        >>> import asyncio
+        >>> async def main():
+        ...     with open("screenshot.png", "rb") as f:
+        ...         screenshot_bytes = f.read()
+        ...     step = await async_single_step(
+        ...         "Click the submit button",
+        ...         screenshot=screenshot_bytes
+        ...     )
+        ...     print(f"Actions: {step.actions}")
+        >>> asyncio.run(main())
+
+        >>> # Using with file path
+        >>> step = await async_single_step(
+        ...     "Find the search box",
+        ...     screenshot="screenshot.png"
+        ... )
+
+        >>> # Using with PILImage
+        >>> image = PILImage.from_file("screenshot.png")
+        >>> step = await async_single_step(
+        ...     "Click next page",
+        ...     screenshot=image
+        ... )
+    """
+    # Handle different screenshot input types
+    if isinstance(screenshot, (str, Path)):
+        # Convert file path to PILImage
+        screenshot = PILImage.from_file(str(screenshot))
+    elif isinstance(screenshot, bytes):
+        # Convert bytes to PILImage
+        screenshot = PILImage.from_bytes(screenshot)
+
+    # Create a temporary task instance
+    task = AsyncTask(api_key=api_key, base_url=base_url)
+
+    try:
+        # Initialize task and perform single step
+        await task.init_task(task_description)
+        result = await task.step(screenshot, instruction=instruction)
+        return result
+    finally:
+        # Clean up resources
+        await task.close()

oagi/async_task.py

@@ -0,0 +1,117 @@
+# -----------------------------------------------------------------------------
+#  Copyright (c) OpenAGI Foundation
+#  All rights reserved.
+#
+#  This file is part of the official API project.
+#  Licensed under the MIT License.
+# -----------------------------------------------------------------------------
+
+from .async_client import AsyncClient
+from .logging import get_logger
+from .sync_client import encode_screenshot_from_bytes
+from .types import Image, Step
+
+logger = get_logger("async_task")
+
+
+class AsyncTask:
+    """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 = "vision-model-v1",
+    ):
+        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
+        self.task_id: str | None = None
+        self.task_description: str | None = None
+        self.model = model
+
+    async def init_task(self, task_desc: str, max_steps: int = 5):
+        """Initialize a new task with the given description."""
+        self.task_description = task_desc
+        response = await self.client.create_message(
+            model=self.model,
+            screenshot="",
+            task_description=self.task_description,
+            task_id=None,
+        )
+        self.task_id = response.task_id  # Reset task_id for new task
+        logger.info(f"Async task initialized: '{task_desc}' (max_steps: {max_steps})")
+
+    async def step(
+        self, screenshot: Image | bytes, instruction: str | 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 (only works with existing task_id)
+
+        Returns:
+            Step: The actions and reasoning for this step
+        """
+        if not self.task_description:
+            raise ValueError("Task description must be set. Call init_task() first.")
+
+        logger.debug(f"Executing async step for task: '{self.task_description}'")
+
+        try:
+            # Convert Image to bytes using the protocol
+            if isinstance(screenshot, Image):
+                screenshot_bytes = screenshot.read()
+            else:
+                screenshot_bytes = screenshot
+            screenshot_b64 = encode_screenshot_from_bytes(screenshot_bytes)
+
+            # Call API
+            response = await self.client.create_message(
+                model=self.model,
+                screenshot=screenshot_b64,
+                task_description=self.task_description,
+                task_id=self.task_id,
+                instruction=instruction,
+            )
+
+            # Update task_id from response
+            if self.task_id != response.task_id:
+                if self.task_id is None:
+                    logger.debug(f"Task ID assigned: {response.task_id}")
+                else:
+                    logger.debug(
+                        f"Task ID changed: {self.task_id} -> {response.task_id}"
+                    )
+                self.task_id = response.task_id
+
+            # Convert API response to Step
+            result = Step(
+                reason=response.reason,
+                actions=response.actions,
+                stop=response.is_complete,
+            )
+
+            if response.is_complete:
+                logger.info(f"Async task completed after {response.current_step} steps")
+            else:
+                logger.debug(
+                    f"Async step {response.current_step} completed with {len(response.actions)} actions"
+                )
+
+            return result
+
+        except Exception as e:
+            logger.error(f"Error during async step execution: {e}")
+            raise
+
+    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()