oagi-core 0.11.0__py3-none-any.whl → 0.12.1__py3-none-any.whl

oagi/client/async_.py

@@ -9,17 +9,19 @@
 from functools import wraps
 
 import httpx
+from httpx import AsyncHTTPTransport
+from openai import AsyncOpenAI
 
 from ..constants import (
-    API_HEALTH_ENDPOINT,
     API_V1_FILE_UPLOAD_ENDPOINT,
     API_V1_GENERATE_ENDPOINT,
-    API_V2_MESSAGE_ENDPOINT,
+    DEFAULT_MAX_RETRIES,
     HTTP_CLIENT_TIMEOUT,
 )
 from ..logging import get_logger
 from ..types import Image
-from ..types.models import GenerateResponse, LLMResponse, UploadFileResponse
+from ..types.models import GenerateResponse, UploadFileResponse, Usage
+from ..types.models.step import Step
 from .base import BaseClient
 
 logger = get_logger("async_client")
@@ -35,8 +37,7 @@ def async_log_trace_on_failure(func):
         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', '')}")
+                BaseClient._log_trace_id(response)
             raise
 
     return wrapper
@@ -45,115 +46,72 @@ def async_log_trace_on_failure(func):
 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)
-        self.upload_client = httpx.AsyncClient(timeout=HTTP_CLIENT_TIMEOUT)
+    def __init__(
+        self,
+        base_url: str | None = None,
+        api_key: str | None = None,
+        max_retries: int = DEFAULT_MAX_RETRIES,
+    ):
+        super().__init__(base_url, api_key, max_retries)
+
+        # OpenAI client for chat completions (with retries)
+        self.openai_client = AsyncOpenAI(
+            api_key=self.api_key,
+            base_url=f"{self.base_url}/v1",
+            max_retries=self.max_retries,
+        )
+
+        # httpx clients for S3 uploads and other endpoints (with retries)
+        transport = AsyncHTTPTransport(retries=self.max_retries)
+        self.http_client = httpx.AsyncClient(
+            transport=transport, base_url=self.base_url
+        )
+        self.upload_client = httpx.AsyncClient(
+            transport=transport, timeout=HTTP_CLIENT_TIMEOUT
+        )
+
         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()
-        await self.upload_client.aclose()
+        await self.close()
 
     async def close(self):
-        """Close the underlying httpx async clients."""
-        await self.client.aclose()
+        """Close the underlying async clients."""
+        await self.openai_client.close()
+        await self.http_client.aclose()
         await self.upload_client.aclose()
 
-    @async_log_trace_on_failure
-    async def create_message(
+    async def chat_completion(
         self,
         model: str,
-        screenshot: bytes | None = None,
-        screenshot_url: str | None = None,
-        task_description: str | None = None,
-        task_id: str | None = None,
-        instruction: str | None = None,
-        messages_history: list | None = None,
+        messages: list,
         temperature: float | None = None,
-        api_version: str | None = None,
-    ) -> "LLMResponse":
+        task_id: str | None = None,
+    ) -> tuple[Step, str, Usage | None]:
         """
-        Call the /v2/message endpoint to analyze task and screenshot
+        Call OpenAI-compatible /v1/chat/completions endpoint.
 
         Args:
-            model: The model to use for task analysis
-            screenshot: Screenshot image bytes (mutually exclusive with screenshot_url)
-            screenshot_url: Direct URL to screenshot (mutually exclusive with screenshot)
-            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
-            messages_history: OpenAI-compatible chat message history
-            temperature: Sampling temperature (0.0-2.0) for LLM inference
-            api_version: API version header
+            model: Model to use for inference
+            messages: Full message history (OpenAI-compatible format)
+            temperature: Sampling temperature (0.0-2.0)
+            task_id: Optional task ID for multi-turn conversations
 
         Returns:
-            LLMResponse: The response from the API
-
-        Raises:
-            ValueError: If both or neither screenshot and screenshot_url are provided
-            httpx.HTTPStatusError: For HTTP error responses
+            Tuple of (Step, raw_output, Usage)
+            - Step: Parsed actions and reasoning
+            - raw_output: Raw model output string (for message history)
+            - Usage: Token usage statistics (or None if not available)
         """
-        # Validate that exactly one is provided
-        if (screenshot is None) == (screenshot_url is None):
-            raise ValueError(
-                "Exactly one of 'screenshot' or 'screenshot_url' must be provided"
-            )
-
-        self._log_request_info(model, task_description, task_id)
-
-        # Upload screenshot to S3 if bytes provided, otherwise use URL directly
-        upload_file_response = None
-        if screenshot is not None:
-            upload_file_response = await self.put_s3_presigned_url(
-                screenshot, api_version
-            )
-
-        # Prepare message payload
-        headers, payload = self._prepare_message_payload(
-            model=model,
-            upload_file_response=upload_file_response,
-            task_description=task_description,
-            task_id=task_id,
-            instruction=instruction,
-            messages_history=messages_history,
-            temperature=temperature,
-            api_version=api_version,
-            screenshot_url=screenshot_url,
+        logger.info(f"Making async chat completion request with model: {model}")
+        kwargs = self._build_chat_completion_kwargs(
+            model, messages, temperature, task_id
         )
-
-        # Make request
-        try:
-            response = await self.client.post(
-                API_V2_MESSAGE_ENDPOINT,
-                json=payload,
-                headers=headers,
-                timeout=self.timeout,
-            )
-            return self._process_response(response)
-        except (httpx.TimeoutException, httpx.NetworkError) as e:
-            self._handle_upload_http_errors(e)
-
-    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(API_HEALTH_ENDPOINT)
-            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
+        response = await self.openai_client.chat.completions.create(**kwargs)
+        return self._parse_chat_completion_response(response)
 
     async def get_s3_presigned_url(
         self,
@@ -172,7 +130,7 @@ class AsyncClient(BaseClient[httpx.AsyncClient]):
 
         try:
             headers = self._build_headers(api_version)
-            response = await self.client.get(
+            response = await self.http_client.get(
                 API_V1_FILE_UPLOAD_ENDPOINT, headers=headers, timeout=self.timeout
             )
             return self._process_upload_response(response)
@@ -292,7 +250,7 @@ class AsyncClient(BaseClient[httpx.AsyncClient]):
 
         # Make request
         try:
-            response = await self.client.post(
+            response = await self.http_client.post(
                 API_V1_GENERATE_ENDPOINT,
                 json=payload,
                 headers=headers,

oagi/client/base.py

@@ -11,7 +11,12 @@ from typing import Any, Generic, TypeVar
 
 import httpx
 
-from ..constants import API_KEY_HELP_URL, DEFAULT_BASE_URL, HTTP_CLIENT_TIMEOUT
+from ..constants import (
+    API_KEY_HELP_URL,
+    DEFAULT_BASE_URL,
+    DEFAULT_MAX_RETRIES,
+    HTTP_CLIENT_TIMEOUT,
+)
 from ..exceptions import (
     APIError,
     AuthenticationError,
@@ -27,9 +32,11 @@ from ..logging import get_logger
 from ..types.models import (
     ErrorResponse,
     GenerateResponse,
-    LLMResponse,
     UploadFileResponse,
+    Usage,
 )
+from ..types.models.step import Step
+from ..utils.output_parser import parse_raw_output
 
 logger = get_logger("client.base")
 
@@ -40,7 +47,12 @@ 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):
+    def __init__(
+        self,
+        base_url: str | None = None,
+        api_key: str | None = None,
+        max_retries: int = DEFAULT_MAX_RETRIES,
+    ):
         # Get from environment if not provided
         self.base_url = base_url or os.getenv("OAGI_BASE_URL") or DEFAULT_BASE_URL
         self.api_key = api_key or os.getenv("OAGI_API_KEY")
@@ -55,6 +67,7 @@ class BaseClient(Generic[HttpClientT]):
 
         self.base_url = self.base_url.rstrip("/")
         self.timeout = HTTP_CLIENT_TIMEOUT
+        self.max_retries = max_retries
         self.client: HttpClientT  # Will be set by subclasses
 
         logger.info(f"Client initialized with base_url: {self.base_url}")
@@ -67,39 +80,77 @@ class BaseClient(Generic[HttpClientT]):
             headers["x-api-key"] = self.api_key
         return headers
 
-    def _build_payload(
+    @staticmethod
+    def _log_trace_id(response) -> None:
+        """Log trace IDs from response headers for debugging."""
+        logger.error(f"Request Id: {response.headers.get('x-request-id', '')}")
+        logger.error(f"Trace Id: {response.headers.get('x-trace-id', '')}")
+
+    def _build_chat_completion_kwargs(
         self,
         model: str,
-        messages_history: list,
-        task_description: str | None = None,
-        task_id: str | None = None,
+        messages: list,
         temperature: float | None = None,
-    ) -> dict[str, Any]:
-        """Build OpenAI-compatible request payload.
+        task_id: str | None = None,
+    ) -> dict:
+        """Build kwargs dict for OpenAI chat completion call.
 
         Args:
-            model: Model to use
-            messages_history: OpenAI-compatible message history
-            task_description: Task description
-            task_id: Task ID for continuing session
-            temperature: Sampling temperature
+            model: Model to use for inference
+            messages: Full message history (OpenAI-compatible format)
+            temperature: Sampling temperature (0.0-2.0)
+            task_id: Optional task ID for multi-turn conversations
 
         Returns:
-            OpenAI-compatible request payload
+            Dict of kwargs for chat.completions.create()
         """
-        payload: dict[str, Any] = {
-            "model": model,
-            "messages": messages_history,
-        }
-
-        if task_description is not None:
-            payload["task_description"] = task_description
-        if task_id is not None:
-            payload["task_id"] = task_id
+        kwargs: dict = {"model": model, "messages": messages}
         if temperature is not None:
-            payload["temperature"] = temperature
+            kwargs["temperature"] = temperature
+        if task_id is not None:
+            kwargs["extra_body"] = {"task_id": task_id}
+        return kwargs
+
+    def _parse_chat_completion_response(
+        self, response
+    ) -> tuple[Step, str, Usage | None]:
+        """Extract and parse OpenAI chat completion response, and log success.
+
+        This is sync/async agnostic as it only processes the response object.
+
+        Args:
+            response: OpenAI ChatCompletion response object
 
-        return payload
+        Returns:
+            Tuple of (Step, raw_output, Usage)
+        """
+        raw_output = response.choices[0].message.content or ""
+        step = parse_raw_output(raw_output)
+
+        # Extract task_id from response (custom field from OAGI API)
+        task_id = getattr(response, "task_id", None)
+
+        usage = None
+        if response.usage:
+            usage = Usage(
+                prompt_tokens=response.usage.prompt_tokens,
+                completion_tokens=response.usage.completion_tokens,
+                total_tokens=response.usage.total_tokens,
+            )
+
+        # Log success with task_id and usage
+        usage_str = (
+            f", tokens: {usage.prompt_tokens}+{usage.completion_tokens}"
+            if usage
+            else ""
+        )
+        task_str = f"task_id: {task_id}, " if task_id else ""
+        logger.info(
+            f"Chat completion successful - {task_str}actions: {len(step.actions)}, "
+            f"stop: {step.stop}{usage_str}"
+        )
+
+        return step, raw_output, usage
 
     def _handle_response_error(
         self, response: httpx.Response, response_data: dict
@@ -141,84 +192,6 @@ class BaseClient(Generic[HttpClientT]):
 
         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 /v2/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 _build_user_message(
-        self, screenshot_url: str, instruction: str | None
-    ) -> dict[str, Any]:
-        """Build OpenAI-compatible user message with screenshot and optional instruction.
-
-        Args:
-            screenshot_url: URL of uploaded screenshot
-            instruction: Optional text instruction
-
-        Returns:
-            User message dict
-        """
-        content = [{"type": "image_url", "image_url": {"url": screenshot_url}}]
-        if instruction:
-            content.append({"type": "text", "text": instruction})
-        return {"role": "user", "content": content}
-
-    def _prepare_message_payload(
-        self,
-        model: str,
-        upload_file_response: UploadFileResponse | None,
-        task_description: str | None,
-        task_id: str | None,
-        instruction: str | None,
-        messages_history: list | None,
-        temperature: float | None,
-        api_version: str | None,
-        screenshot_url: str | None = None,
-    ) -> tuple[dict[str, str], dict[str, Any]]:
-        """Prepare headers and payload for /v2/message request.
-
-        Args:
-            model: Model to use
-            upload_file_response: Response from S3 upload (if screenshot was uploaded)
-            task_description: Task description
-            task_id: Task ID
-            instruction: Optional instruction
-            messages_history: Message history
-            temperature: Sampling temperature
-            api_version: API version
-            screenshot_url: Direct screenshot URL (alternative to upload_file_response)
-
-        Returns:
-            Tuple of (headers, payload)
-        """
-        # Use provided screenshot_url or get from upload_file_response
-        if screenshot_url is None:
-            if upload_file_response is None:
-                raise ValueError(
-                    "Either screenshot_url or upload_file_response must be provided"
-                )
-            screenshot_url = upload_file_response.download_url
-
-        # Build user message and append to history
-        if messages_history is None:
-            messages_history = []
-        user_message = self._build_user_message(screenshot_url, instruction)
-        messages_history.append(user_message)
-
-        # Build payload and headers
-        headers = self._build_headers(api_version)
-        payload = self._build_payload(
-            model=model,
-            messages_history=messages_history,
-            task_description=task_description,
-            task_id=task_id,
-            temperature=temperature,
-        )
-
-        return headers, payload
-
     def _parse_response_json(self, response: httpx.Response) -> dict[str, Any]:
         try:
             return response.json()
@@ -230,35 +203,6 @@ class BaseClient(Generic[HttpClientT]):
                 response=response,
             )
 
-    def _process_response(self, response: httpx.Response) -> "LLMResponse":
-        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"complete: {result.is_complete}"
-        )
-        logger.debug(f"Response included {len(result.actions)} actions")
-        return result
-
     def _process_upload_response(self, response: httpx.Response) -> UploadFileResponse:
         """Process response from /v1/file/upload endpoint.
 
@@ -449,7 +393,11 @@ class BaseClient(Generic[HttpClientT]):
         # Parse successful response
         result = GenerateResponse(**response_data)
 
+        # Capture request_id from response header
+        result.request_id = response.headers.get("X-Request-ID")
+
         logger.info(
             f"Generate request successful - tokens: {result.prompt_tokens}+{result.completion_tokens}, "
+            f"request_id: {result.request_id}"
         )
         return result