oagi 0.5.0__py3-none-any.whl → 0.6.1__py3-none-any.whl

oagi/__init__.py

@@ -5,13 +5,10 @@
 #  This file is part of the official API project.
 #  Licensed under the MIT License.
 # -----------------------------------------------------------------------------
+import importlib
 
-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.client import AsyncClient, SyncClient
 from oagi.exceptions import (
     APIError,
     AuthenticationError,
@@ -24,18 +21,35 @@ from oagi.exceptions import (
     ServerError,
     ValidationError,
 )
-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.task import AsyncShortTask, AsyncTask, ShortTask, Task
 from oagi.types import (
     AsyncActionHandler,
     AsyncImageProvider,
     ImageConfig,
 )
+from oagi.types.models import ErrorDetail, ErrorResponse, LLMResponse
+
+# Lazy imports for pyautogui-dependent modules
+# These will only be imported when actually accessed
+_LAZY_IMPORTS = {
+    "AsyncPyautoguiActionHandler": "oagi.async_pyautogui_action_handler",
+    "AsyncScreenshotMaker": "oagi.async_screenshot_maker",
+    "PILImage": "oagi.pil_image",
+    "PyautoguiActionHandler": "oagi.pyautogui_action_handler",
+    "PyautoguiConfig": "oagi.pyautogui_action_handler",
+    "ScreenshotMaker": "oagi.screenshot_maker",
+}
+
+
+def __getattr__(name: str):
+    """Lazy import for pyautogui-dependent modules."""
+    if name in _LAZY_IMPORTS:
+        module_name = _LAZY_IMPORTS[name]
+        module = importlib.import_module(module_name)
+        return getattr(module, name)
+    raise AttributeError(f"module '{__name__}' has no attribute '{name}'")
+
 
 __all__ = [
     # Core sync classes
@@ -49,15 +63,6 @@ __all__ = [
     # Functions
     "single_step",
     "async_single_step",
-    # Image classes
-    "PILImage",
-    # Handler classes
-    "PyautoguiActionHandler",
-    "PyautoguiConfig",
-    "ScreenshotMaker",
-    # Async handler classes
-    "AsyncPyautoguiActionHandler",
-    "AsyncScreenshotMaker",
     # Async protocols
     "AsyncActionHandler",
     "AsyncImageProvider",
@@ -78,4 +83,14 @@ __all__ = [
     "ServerError",
     "RequestTimeoutError",
     "ValidationError",
+    # Lazy imports
+    # Image classes
+    "PILImage",
+    # Handler classes
+    "PyautoguiActionHandler",
+    "PyautoguiConfig",
+    "ScreenshotMaker",
+    # Async handler classes
+    "AsyncPyautoguiActionHandler",
+    "AsyncScreenshotMaker",
 ]

oagi/async_single_step.py

@@ -8,8 +8,7 @@
 
 from pathlib import Path
 
-from .async_task import AsyncTask
-from .pil_image import PILImage
+from .task import AsyncTask
 from .types import Image, Step
 
 
@@ -19,6 +18,7 @@ async def async_single_step(
     instruction: str | None = None,
     api_key: str | None = None,
     base_url: str | None = None,
+    temperature: float | None = None,
 ) -> Step:
     """
     Perform a single-step inference asynchronously without maintaining task state.
@@ -32,6 +32,7 @@ async def async_single_step(
         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
@@ -62,16 +63,17 @@ async def async_single_step(
         ...     screenshot=image
         ... )
     """
+    # Lazy import PILImage only when needed
+    from .pil_image import PILImage  # noqa: PLC0415
+
     # 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)
+    task = AsyncTask(api_key=api_key, base_url=base_url, temperature=temperature)
 
     try:
         # Initialize task and perform single step

oagi/client/__init__.py

@@ -0,0 +1,12 @@
+# -----------------------------------------------------------------------------
+#  Copyright (c) OpenAGI Foundation
+#  All rights reserved.
+#
+#  This file is part of the official API project.
+#  Licensed under the MIT License.
+# -----------------------------------------------------------------------------
+
+from .async_ import AsyncClient
+from .sync import SyncClient
+
+__all__ = ["SyncClient", "AsyncClient"]

oagi/client/async_.py

@@ -0,0 +1,137 @@
+# -----------------------------------------------------------------------------
+#  Copyright (c) OpenAGI Foundation
+#  All rights reserved.
+#
+#  This file is part of the official API project.
+#  Licensed under the MIT License.
+# -----------------------------------------------------------------------------
+
+from functools import wraps
+
+import httpx
+
+from ..exceptions import NetworkError, RequestTimeoutError
+from ..logging import get_logger
+from ..types.models import LLMResponse
+from .base import BaseClient
+
+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(BaseClient[httpx.AsyncClient]):
+    """Asynchronous HTTP client for the OAGI API."""
+
+    def __init__(self, base_url: str | None = None, api_key: str | None = None):
+        super().__init__(base_url, api_key)
+        self.client = httpx.AsyncClient(base_url=self.base_url)
+        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,
+        last_task_id: str | None = None,
+        history_steps: int | None = None,
+        temperature: float | None = None,
+        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)
+            last_task_id: Previous task ID to retrieve history from (only works with task_id)
+            history_steps: Number of historical steps to include from last_task_id (default: 1, max: 10)
+            temperature: Sampling temperature (0.0-2.0) for LLM inference
+            api_version: API version header
+
+        Returns:
+            LLMResponse: The response from the API
+
+        Raises:
+            httpx.HTTPStatusError: For HTTP error responses
+        """
+        headers = self._build_headers(api_version)
+        payload = self._build_payload(
+            model=model,
+            screenshot=screenshot,
+            task_description=task_description,
+            task_id=task_id,
+            instruction=instruction,
+            max_actions=max_actions,
+            last_task_id=last_task_id,
+            history_steps=history_steps,
+            temperature=temperature,
+        )
+
+        self._log_request_info(model, task_description, task_id)
+
+        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)
+
+        return self._process_response(response)
+
+    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/client/base.py

@@ -0,0 +1,183 @@
+# -----------------------------------------------------------------------------
+#  Copyright (c) OpenAGI Foundation
+#  All rights reserved.
+#
+#  This file is part of the official API project.
+#  Licensed under the MIT License.
+# -----------------------------------------------------------------------------
+
+import os
+from typing import Any, Generic, TypeVar
+
+import httpx
+
+from ..exceptions import (
+    APIError,
+    AuthenticationError,
+    ConfigurationError,
+    NotFoundError,
+    RateLimitError,
+    ServerError,
+    ValidationError,
+)
+from ..logging import get_logger
+from ..types.models import ErrorResponse, LLMResponse
+
+logger = get_logger("client.base")
+
+# TypeVar for HTTP client type (httpx.Client or httpx.AsyncClient)
+HttpClientT = TypeVar("HttpClientT")
+
+
+class BaseClient(Generic[HttpClientT]):
+    """Base class with shared business logic for sync/async clients."""
+
+    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.timeout = 60
+        self.client: HttpClientT  # Will be set by subclasses
+
+        logger.info(f"Client initialized with base_url: {self.base_url}")
+
+    def _build_headers(self, api_version: str | None = None) -> dict[str, str]:
+        headers: dict[str, str] = {}
+        if api_version:
+            headers["x-api-version"] = api_version
+        if self.api_key:
+            headers["x-api-key"] = self.api_key
+        return headers
+
+    def _build_payload(
+        self,
+        model: str,
+        screenshot: str,
+        task_description: str | None = None,
+        task_id: str | None = None,
+        instruction: str | None = None,
+        max_actions: int | None = None,
+        last_task_id: str | None = None,
+        history_steps: int | None = None,
+        temperature: float | None = None,
+    ) -> dict[str, Any]:
+        payload: dict[str, Any] = {"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
+        if last_task_id is not None:
+            payload["last_task_id"] = last_task_id
+        if history_steps is not None:
+            payload["history_steps"] = history_steps
+        if temperature is not None:
+            payload["sampling_params"] = {"temperature": temperature}
+
+        return payload
+
+    def _handle_response_error(
+        self, response: httpx.Response, response_data: dict
+    ) -> None:
+        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,
+            )
+
+    def _get_exception_class(self, status_code: int) -> type[APIError]:
+        status_map = {
+            401: AuthenticationError,
+            404: NotFoundError,
+            422: ValidationError,
+            429: RateLimitError,
+        }
+
+        if status_code >= 500:
+            return ServerError
+
+        return status_map.get(status_code, APIError)
+
+    def _log_request_info(self, model: str, task_description: Any, task_id: Any):
+        logger.info(f"Making API request to /v1/message with model: {model}")
+        logger.debug(
+            f"Request includes task_description: {task_description is not None}, "
+            f"task_id: {task_id is not None}"
+        )
+
+    def _parse_response_json(self, response: httpx.Response) -> dict[str, Any]:
+        try:
+            return response.json()
+        except ValueError:
+            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,
+            )
+
+    def _process_response(self, response: httpx.Response) -> Any:
+        response_data = self._parse_response_json(response)
+
+        # Check if it's an error response (non-200 status)
+        if response.status_code != 200:
+            self._handle_response_error(response, response_data)
+
+        # 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"API request successful - task_id: {result.task_id}, "
+            f"step: {result.current_step}, complete: {result.is_complete}"
+        )
+        logger.debug(f"Response included {len(result.actions)} actions")
+        return result

oagi/client/sync.py

@@ -0,0 +1,142 @@
+# -----------------------------------------------------------------------------
+#  Copyright (c) OpenAGI Foundation
+#  All rights reserved.
+#
+#  This file is part of the official API project.
+#  Licensed under the MIT License.
+# -----------------------------------------------------------------------------
+
+from functools import wraps
+
+import httpx
+from httpx import Response
+
+from ..exceptions import NetworkError, RequestTimeoutError
+from ..logging import get_logger
+from ..types.models import LLMResponse
+from .base import BaseClient
+
+logger = get_logger("sync_client")
+
+
+def _log_trace_id(response: Response):
+    logger.error(f"Request Id: {response.headers.get('x-request-id', '')}")
+    logger.error(f"Trace Id: {response.headers.get('x-trace-id', '')}")
+
+
+def log_trace_on_failure(func):
+    """Decorator that logs trace ID when a method fails."""
+
+    @wraps(func)
+    def wrapper(*args, **kwargs):
+        try:
+            return 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:
+                _log_trace_id(response)
+            raise
+
+    return wrapper
+
+
+class SyncClient(BaseClient[httpx.Client]):
+    """Synchronous HTTP client for the OAGI API."""
+
+    def __init__(self, base_url: str | None = None, api_key: str | None = None):
+        super().__init__(base_url, api_key)
+        self.client = httpx.Client(base_url=self.base_url)
+        logger.info(f"SyncClient initialized with base_url: {self.base_url}")
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        self.client.close()
+
+    def close(self):
+        """Close the underlying httpx client."""
+        self.client.close()
+
+    @log_trace_on_failure
+    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,
+        last_task_id: str | None = None,
+        history_steps: int | None = None,
+        temperature: float | None = None,
+        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)
+            last_task_id: Previous task ID to retrieve history from (only works with task_id)
+            history_steps: Number of historical steps to include from last_task_id (default: 1, max: 10)
+            temperature: Sampling temperature (0.0-2.0) for LLM inference
+            api_version: API version header
+
+        Returns:
+            LLMResponse: The response from the API
+
+        Raises:
+            httpx.HTTPStatusError: For HTTP error responses
+        """
+        headers = self._build_headers(api_version)
+        payload = self._build_payload(
+            model=model,
+            screenshot=screenshot,
+            task_description=task_description,
+            task_id=task_id,
+            instruction=instruction,
+            max_actions=max_actions,
+            last_task_id=last_task_id,
+            history_steps=history_steps,
+            temperature=temperature,
+        )
+
+        self._log_request_info(model, task_description, task_id)
+
+        try:
+            response = 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)
+
+        return self._process_response(response)
+
+    def health_check(self) -> dict:
+        """
+        Call the /health endpoint for health check
+
+        Returns:
+            dict: Health check response
+        """
+        logger.debug("Making health check request")
+        try:
+            response = self.client.get("/health")
+            response.raise_for_status()
+            result = response.json()
+            logger.debug("Health check successful")
+            return result
+        except httpx.HTTPStatusError as e:
+            logger.warning(f"Health check failed: {e}")
+            raise

oagi/single_step.py

@@ -8,7 +8,6 @@
 
 from pathlib import Path
 
-from .pil_image import PILImage
 from .task import Task
 from .types import Image, Step
 
@@ -19,6 +18,7 @@ def single_step(
     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.
@@ -32,6 +32,7 @@ def single_step(
         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
@@ -60,6 +61,9 @@ def single_step(
         ...     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
@@ -78,6 +82,6 @@ def single_step(
         )
 
     # Use Task to perform single step
-    with Task(api_key=api_key, base_url=base_url) as task:
+    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)

oagi/task/__init__.py

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