omniai-kit 1.0.0__py3-none-any.whl

omniai/resources/tasks.py

@@ -0,0 +1,171 @@
+"""Background tasks resource for OmniAI."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, Type, TypeVar, Union
+from pydantic import BaseModel
+
+from omniai.types.tasks import TaskJob, AsyncTaskJob, TaskStatus
+
+if TYPE_CHECKING:
+    from omniai._base import BaseClient, AsyncBaseClient
+
+T = TypeVar("T")
+
+
+class TasksResource:
+    """Synchronous background tasks resource."""
+
+    def __init__(self, client: BaseClient):
+        self._client = client
+
+    def submit(
+        self,
+        *,
+        model: str,
+        prompt: Optional[str] = None,
+        messages: Optional[List[Dict[str, str]]] = None,
+        response_model: Optional[Type[T]] = None,
+        tools: Optional[List[Any]] = None,
+        output_format: Optional[Union[str, Dict[str, Any]]] = None,
+        system_instruction: Optional[str] = None,
+        temperature: float = 0.1,
+        max_tokens: Optional[int] = None,
+        timeout: Optional[float] = None,
+    ) -> TaskJob[T]:
+        """Submits an asynchronous task to Celery queue and returns a TaskJob handle."""
+        payload: Dict[str, Any] = {
+            "model": model,
+            "temperature": temperature,
+        }
+        if prompt is not None:
+            payload["prompt"] = prompt
+        if messages is not None:
+            payload["messages"] = messages
+        if system_instruction is not None:
+            payload["system_instruction"] = system_instruction
+        if max_tokens is not None:
+            payload["max_tokens"] = max_tokens
+        if tools:
+            payload["tools"] = tools
+        if output_format:
+            payload["output_format"] = output_format
+        elif response_model and issubclass(response_model, BaseModel):
+            payload["output_format"] = response_model.model_json_schema()
+
+        response = self._client.client.post(
+            "/v1/chat/completions/async",
+            json=payload,
+            timeout=timeout or self._client.timeout,
+        )
+        if response.status_code not in (200, 202):
+            self._client._handle_error_response(response)
+
+        data = response.json()
+        task_id = data.get("task_id")
+
+        return TaskJob(
+            task_id=task_id,
+            status_checker=self.get_status,
+            canceller=self.cancel,
+            response_model=response_model,
+        )
+
+    def get_status(self, task_id: str, timeout: Optional[float] = None) -> TaskStatus:
+        """Checks status of background task."""
+        response = self._client.client.get(
+            f"/v1/chat/completions/tasks/{task_id}",
+            timeout=timeout or self._client.timeout,
+        )
+        if response.status_code != 200:
+            self._client._handle_error_response(response)
+        return TaskStatus.model_validate(response.json())
+
+    def cancel(self, task_id: str, timeout: Optional[float] = None) -> Dict[str, Any]:
+        """Cancels a background task."""
+        response = self._client.client.delete(
+            f"/v1/chat/completions/tasks/{task_id}",
+            timeout=timeout or self._client.timeout,
+        )
+        if response.status_code != 200:
+            self._client._handle_error_response(response)
+        return response.json()
+
+
+class AsyncTasksResource:
+    """Asynchronous background tasks resource."""
+
+    def __init__(self, client: AsyncBaseClient):
+        self._client = client
+
+    async def submit(
+        self,
+        *,
+        model: str,
+        prompt: Optional[str] = None,
+        messages: Optional[List[Dict[str, str]]] = None,
+        response_model: Optional[Type[T]] = None,
+        tools: Optional[List[Any]] = None,
+        output_format: Optional[Union[str, Dict[str, Any]]] = None,
+        system_instruction: Optional[str] = None,
+        temperature: float = 0.1,
+        max_tokens: Optional[int] = None,
+        timeout: Optional[float] = None,
+    ) -> AsyncTaskJob[T]:
+        """Submits an asynchronous task to Celery queue and returns an AsyncTaskJob handle."""
+        payload: Dict[str, Any] = {
+            "model": model,
+            "temperature": temperature,
+        }
+        if prompt is not None:
+            payload["prompt"] = prompt
+        if messages is not None:
+            payload["messages"] = messages
+        if system_instruction is not None:
+            payload["system_instruction"] = system_instruction
+        if max_tokens is not None:
+            payload["max_tokens"] = max_tokens
+        if tools:
+            payload["tools"] = tools
+        if output_format:
+            payload["output_format"] = output_format
+        elif response_model and issubclass(response_model, BaseModel):
+            payload["output_format"] = response_model.model_json_schema()
+
+        response = await self._client.client.post(
+            "/v1/chat/completions/async",
+            json=payload,
+            timeout=timeout or self._client.timeout,
+        )
+        if response.status_code not in (200, 202):
+            self._client._handle_error_response(response)
+
+        data = response.json()
+        task_id = data.get("task_id")
+
+        return AsyncTaskJob(
+            task_id=task_id,
+            status_checker=self.get_status,
+            canceller=self.cancel,
+            response_model=response_model,
+        )
+
+    async def get_status(self, task_id: str, timeout: Optional[float] = None) -> TaskStatus:
+        """Checks status of background task asynchronously."""
+        response = await self._client.client.get(
+            f"/v1/chat/completions/tasks/{task_id}",
+            timeout=timeout or self._client.timeout,
+        )
+        if response.status_code != 200:
+            self._client._handle_error_response(response)
+        return TaskStatus.model_validate(response.json())
+
+    async def cancel(self, task_id: str, timeout: Optional[float] = None) -> Dict[str, Any]:
+        """Cancels a background task asynchronously."""
+        response = await self._client.client.delete(
+            f"/v1/chat/completions/tasks/{task_id}",
+            timeout=timeout or self._client.timeout,
+        )
+        if response.status_code != 200:
+            self._client._handle_error_response(response)
+        return response.json()

omniai/types/__init__.py

@@ -0,0 +1,27 @@
+"""Unified type definitions for the OmniAI Python SDK."""
+
+from omniai.types.chat import (
+    ChatMessage,
+    ChatCompletionChoice,
+    ChatCompletionUsage,
+    ChatCompletionResponse,
+    ChatCompletionChunk,
+)
+from omniai.types.models import ModelInfo, ModelList
+from omniai.types.livekit import LiveSessionInfo
+from omniai.types.tasks import TaskStatus, TaskJob, AsyncTaskJob, clean_json_text
+
+__all__ = [
+    "ChatMessage",
+    "ChatCompletionChoice",
+    "ChatCompletionUsage",
+    "ChatCompletionResponse",
+    "ChatCompletionChunk",
+    "ModelInfo",
+    "ModelList",
+    "LiveSessionInfo",
+    "TaskStatus",
+    "TaskJob",
+    "AsyncTaskJob",
+    "clean_json_text",
+]

omniai/types/chat.py

@@ -0,0 +1,48 @@
+"""Type definitions for Chat Completions."""
+
+from __future__ import annotations
+
+from typing import Any, Dict, List, Optional, Union
+from pydantic import BaseModel, Field
+
+
+class ChatMessage(BaseModel):
+    """A single chat message in a conversation thread."""
+    role: str = Field(default="user", description="Role of the message author (e.g. user, system, assistant)")
+    content: str = Field(..., description="The content of the message")
+    name: Optional[str] = None
+
+
+class ChatCompletionChoice(BaseModel):
+    """A choice in a chat completion response."""
+    index: int = 0
+    message: ChatMessage
+    finish_reason: Optional[str] = None
+
+
+class ChatCompletionUsage(BaseModel):
+    """Usage statistics for a completion request."""
+    prompt_tokens: int = 0
+    completion_tokens: int = 0
+    total_tokens: int = 0
+
+
+class ChatCompletionResponse(BaseModel):
+    """Full chat completion response from the OmniAI gateway."""
+    id: Optional[str] = None
+    object: str = "chat.completion"
+    created: Optional[int] = None
+    model: str
+    choices: List[ChatCompletionChoice] = Field(default_factory=list)
+    usage: Optional[ChatCompletionUsage] = None
+    result: Optional[Union[str, Dict[str, Any]]] = None
+
+
+class ChatCompletionChunk(BaseModel):
+    """Streaming chunk of a chat completion."""
+    id: Optional[str] = None
+    object: str = "chat.completion.chunk"
+    created: Optional[int] = None
+    model: Optional[str] = None
+    delta: Optional[str] = None
+    finish_reason: Optional[str] = None

omniai/types/livekit.py

@@ -0,0 +1,17 @@
+"""Type definitions for LiveKit multimodal audio/video sessions."""
+
+from __future__ import annotations
+
+from typing import Optional
+from pydantic import BaseModel, Field
+
+
+class LiveSessionInfo(BaseModel):
+    """LiveKit WebRTC multimodal audio/video session credentials and connection info."""
+    server_url: str = Field(..., description="WebSocket/WebRTC URL of the LiveKit server")
+    room_name: str = Field(..., description="Unique room name for the multimodal session")
+    participant_token: str = Field(..., description="JWT token for authenticating to the room")
+    participant_identity: str = Field(..., description="Unique identity assigned to the participant")
+    model_used: str = Field(default="gemini:gemini-2.5-flash", description="Underlying model configured for voice/video agent")
+    voice: str = Field(default="Aoede", description="Voice personality for realtime synthesis")
+    video_input: bool = Field(default=False, description="Whether realtime video streaming input is active")

omniai/types/models.py

@@ -0,0 +1,30 @@
+"""Type definitions for OmniAI Model catalog and metadata."""
+
+from __future__ import annotations
+
+from typing import List, Optional
+from pydantic import BaseModel, Field
+
+
+class ModelInfo(BaseModel):
+    """Metadata describing an AI model supported by OmniAI."""
+    id: str
+    object: str = "model"
+    created: Optional[int] = None
+    created_at: Optional[str] = None
+    release_date: Optional[str] = None
+    owned_by: Optional[str] = None
+    provider: Optional[str] = None
+    name: Optional[str] = None
+    description: Optional[str] = None
+    context_window: Optional[int] = None
+    max_output_tokens: Optional[int] = None
+    supports_vision: bool = False
+    supports_function_calling: bool = False
+    supports_streaming: bool = True
+
+
+class ModelList(BaseModel):
+    """List of supported AI models."""
+    object: str = "list"
+    data: List[ModelInfo] = Field(default_factory=list)

omniai/types/tasks.py

@@ -0,0 +1,208 @@
+"""Type definitions for asynchronous Celery background tasks."""
+
+from __future__ import annotations
+
+import asyncio
+import logging
+import re
+import time
+from typing import Any, Callable, Generic, Optional, Type, TypeVar
+from pydantic import BaseModel
+
+from omniai.exceptions import (
+    OmniaiExecutionError,
+    OmniaiTimeoutError,
+    OmniaiValidationError,
+)
+
+logger = logging.getLogger("omniai.types.tasks")
+
+T = TypeVar("T")
+
+
+def clean_json_text(text: str) -> str:
+    """Strips markdown code blocks (e.g. ```json ... ```) from model output."""
+    stripped = text.strip()
+    match = re.search(r"```(?:json)?\s*([\s\S]*?)\s*```", stripped, re.IGNORECASE)
+    if match:
+        return match.group(1).strip()
+    return stripped
+
+
+class TaskStatus(BaseModel):
+    """Current status and progress of a background Celery task."""
+    task_id: str
+    status: str
+    ready: bool = False
+    result: Optional[Any] = None
+    error: Optional[str] = None
+
+
+class TaskJob(Generic[T]):
+    """Synchronous handle to a background task running on the OmniAI Celery queue."""
+
+    def __init__(
+        self,
+        task_id: str,
+        status_checker: Callable[[str], TaskStatus],
+        canceller: Callable[[str], dict],
+        response_model: Optional[Type[T]] = None,
+    ):
+        self.task_id = task_id
+        self._status_checker = status_checker
+        self._canceller = canceller
+        self._response_model = response_model
+        self._last_status: Optional[TaskStatus] = None
+
+    def get_status(self) -> TaskStatus:
+        self._last_status = self._status_checker(self.task_id)
+        return self._last_status
+
+    def wait(
+        self,
+        timeout: float = 120.0,
+        interval: float = 1.5,
+        max_consecutive_poll_errors: int = 3,
+    ) -> T:
+        start_time = time.time()
+        consecutive_errors = 0
+
+        while time.time() - start_time < timeout:
+            try:
+                status = self.get_status()
+                consecutive_errors = 0
+            except Exception as poll_err:
+                consecutive_errors += 1
+                logger.warning(
+                    f"Transient error polling task {self.task_id} ({consecutive_errors}/{max_consecutive_poll_errors}): {poll_err}"
+                )
+                if consecutive_errors >= max_consecutive_poll_errors:
+                    raise OmniaiExecutionError(
+                        f"Failed polling task {self.task_id} after {consecutive_errors} consecutive errors: {poll_err}"
+                    ) from poll_err
+                time.sleep(interval)
+                continue
+
+            if status.ready:
+                if status.status == "SUCCESS":
+                    raw_result = status.result
+                    if self._response_model and issubclass(self._response_model, BaseModel):
+                        if isinstance(raw_result, dict):
+                            try:
+                                return self._response_model.model_validate(raw_result)
+                            except Exception as val_err:
+                                raise OmniaiValidationError(
+                                    f"Task result validation failed: {val_err}",
+                                    raw_output=str(raw_result),
+                                    details={"error": str(val_err)},
+                                ) from val_err
+                        elif isinstance(raw_result, str):
+                            cleaned = clean_json_text(raw_result)
+                            try:
+                                return self._response_model.model_validate_json(cleaned)
+                            except Exception as val_err:
+                                raise OmniaiValidationError(
+                                    f"Task JSON string validation failed: {val_err}",
+                                    raw_output=raw_result,
+                                    details={"error": str(val_err)},
+                                ) from val_err
+                    return raw_result
+                else:
+                    raise OmniaiExecutionError(
+                        f"Task {self.task_id} failed with status {status.status}: {status.error or 'Unknown error'}",
+                        details={"status": status.status, "error": status.error},
+                    )
+
+            time.sleep(interval)
+
+        raise OmniaiTimeoutError(
+            f"Task {self.task_id} timed out after {timeout} seconds (last status: {self._last_status.status if self._last_status else 'UNKNOWN'})"
+        )
+
+    def cancel(self) -> dict:
+        return self._canceller(self.task_id)
+
+
+class AsyncTaskJob(Generic[T]):
+    """Asynchronous handle to a background task running on the OmniAI Celery queue."""
+
+    def __init__(
+        self,
+        task_id: str,
+        status_checker: Callable[[str], Any],
+        canceller: Callable[[str], Any],
+        response_model: Optional[Type[T]] = None,
+    ):
+        self.task_id = task_id
+        self._status_checker = status_checker
+        self._canceller = canceller
+        self._response_model = response_model
+        self._last_status: Optional[TaskStatus] = None
+
+    async def get_status(self) -> TaskStatus:
+        self._last_status = await self._status_checker(self.task_id)
+        return self._last_status
+
+    async def wait(
+        self,
+        timeout: float = 120.0,
+        interval: float = 1.5,
+        max_consecutive_poll_errors: int = 3,
+    ) -> T:
+        start_time = time.time()
+        consecutive_errors = 0
+
+        while time.time() - start_time < timeout:
+            try:
+                status = await self.get_status()
+                consecutive_errors = 0
+            except Exception as poll_err:
+                consecutive_errors += 1
+                logger.warning(
+                    f"Transient error polling async task {self.task_id} ({consecutive_errors}/{max_consecutive_poll_errors}): {poll_err}"
+                )
+                if consecutive_errors >= max_consecutive_poll_errors:
+                    raise OmniaiExecutionError(
+                        f"Failed polling async task {self.task_id} after {consecutive_errors} consecutive errors: {poll_err}"
+                    ) from poll_err
+                await asyncio.sleep(interval)
+                continue
+
+            if status.ready:
+                if status.status == "SUCCESS":
+                    raw_result = status.result
+                    if self._response_model and issubclass(self._response_model, BaseModel):
+                        if isinstance(raw_result, dict):
+                            try:
+                                return self._response_model.model_validate(raw_result)
+                            except Exception as val_err:
+                                raise OmniaiValidationError(
+                                    f"Async task result validation failed: {val_err}",
+                                    raw_output=str(raw_result),
+                                    details={"error": str(val_err)},
+                                ) from val_err
+                        elif isinstance(raw_result, str):
+                            cleaned = clean_json_text(raw_result)
+                            try:
+                                return self._response_model.model_validate_json(cleaned)
+                            except Exception as val_err:
+                                raise OmniaiValidationError(
+                                    f"Async task JSON string validation failed: {val_err}",
+                                    raw_output=raw_result,
+                                    details={"error": str(val_err)},
+                                ) from val_err
+                    return raw_result
+                else:
+                    raise OmniaiExecutionError(
+                        f"Task {self.task_id} failed with status {status.status}: {status.error or 'Unknown error'}",
+                        details={"status": status.status, "error": status.error},
+                    )
+
+            await asyncio.sleep(interval)
+
+        raise OmniaiTimeoutError(
+            f"Async task {self.task_id} timed out after {timeout} seconds (last status: {self._last_status.status if self._last_status else 'UNKNOWN'})"
+        )
+
+    async def cancel(self) -> dict:
+        return await self._canceller(self.task_id)