waveium 0.2.0__py3-none-any.whl

waveium/resources/_stream_ops.py

@@ -0,0 +1,318 @@
+"""Reusable streaming operation mixins for resources."""
+
+import logging
+import time
+from typing import (
+    TYPE_CHECKING,
+    AsyncIterator,
+    Callable,
+    FrozenSet,
+    Iterator,
+    Optional,
+)
+
+logger = logging.getLogger("waveium")
+
+from .._sse import AsyncSSEConnection, SSEConnection
+from ..exceptions import StreamTimeoutError
+from ..models.events import (
+    ErrorEvent,
+    TERMINAL_STATUSES,
+    ServerEvent,
+    StatusEvent,
+    parse_server_event,
+)
+
+
+def _log_event(event: ServerEvent) -> None:
+    """Log an SSE event at the appropriate level."""
+    if isinstance(event, StatusEvent):
+        msg = f"[{event.status}]"
+        if event.progress and event.progress.overall_percent:
+            msg += f" {event.progress.overall_percent}"
+            if event.progress.overall_waypoints:
+                msg += f" (waypoints: {event.progress.overall_waypoints})"
+        if event.status in TERMINAL_STATUSES:
+            logger.info("%s %s", event.resource_id, msg)
+        else:
+            logger.info("%s %s", event.resource_id, msg)
+    elif isinstance(event, ErrorEvent):
+        logger.warning("SSE error: %s", event.message)
+
+
+if TYPE_CHECKING:
+    from .._http import AsyncHTTPClient, HTTPClient
+
+
+class EventStream:
+    """Synchronous iterator over typed SSE events.
+
+    Wraps an SSEConnection and parses raw events into typed
+    ServerEvent models. Supports context manager protocol.
+
+    Usage:
+        with stream_ops.stream(resource_id) as events:
+            for event in events:
+                if event.event_type == "status":
+                    print(event.status)
+    """
+
+    def __init__(self, connection: SSEConnection) -> None:
+        self._connection = connection
+
+    def __iter__(self) -> Iterator[ServerEvent]:
+        for raw_event in self._connection:
+            try:
+                event = parse_server_event(raw_event.event, raw_event.data)
+            except ValueError:
+                # Skip unparseable events rather than crashing
+                continue
+            _log_event(event)
+            yield event
+
+    def close(self) -> None:
+        """Close the underlying SSE connection."""
+        self._connection.close()
+
+    def __enter__(self) -> "EventStream":
+        return self
+
+    def __exit__(self, *args: object) -> None:
+        self.close()
+
+
+class AsyncEventStream:
+    """Asynchronous iterator over typed SSE events.
+
+    Async counterpart of EventStream.
+
+    Usage:
+        async with stream_ops.stream(resource_id) as events:
+            async for event in events:
+                ...
+    """
+
+    def __init__(self, connection: AsyncSSEConnection) -> None:
+        self._connection = connection
+
+    async def __aiter__(self) -> AsyncIterator[ServerEvent]:
+        async for raw_event in self._connection:
+            try:
+                event = parse_server_event(raw_event.event, raw_event.data)
+            except ValueError:
+                continue
+            _log_event(event)
+            yield event
+
+    async def close(self) -> None:
+        """Close the underlying SSE connection."""
+        await self._connection.close()
+
+    async def __aenter__(self) -> "AsyncEventStream":
+        return self
+
+    async def __aexit__(self, *args: object) -> None:
+        await self.close()
+
+
+class StreamOperations:
+    """Synchronous streaming helpers, composed into resource classes.
+
+    Follows the same pattern as FileOperations and UploadOperations.
+    Each resource class instantiates this with its SSE path template.
+    """
+
+    def __init__(
+        self,
+        http: "HTTPClient",
+        sse_path_template: str,
+    ) -> None:
+        """Initialize stream operations.
+
+        Args:
+            http: HTTP client instance.
+            sse_path_template: Path template with {resource_id}
+                placeholder (e.g. "/v0/events/environments/{resource_id}").
+        """
+        self._http = http
+        self._sse_path_template = sse_path_template
+
+    def stream(
+        self,
+        resource_id: str,
+        *,
+        max_total_seconds: float = 7200.0,
+    ) -> EventStream:
+        """Open a typed SSE event stream for a resource.
+
+        Args:
+            resource_id: The resource identifier.
+            max_total_seconds: Maximum total stream duration
+                including reconnections. Defaults to 2 hours.
+
+        Returns:
+            An EventStream (iterator + context manager).
+        """
+        from .._http import _build_headers
+
+        path = self._sse_path_template.format(resource_id=resource_id)
+        connection = SSEConnection(
+            client=self._http._get_client(),
+            url=path,
+            headers=_build_headers(self._http.config),
+            max_total_seconds=max_total_seconds,
+        )
+        return EventStream(connection)
+
+    def wait(
+        self,
+        resource_id: str,
+        *,
+        timeout: float = 3600.0,
+        on_event: Optional[Callable[[ServerEvent], None]] = None,
+        terminal_statuses: FrozenSet[str] = TERMINAL_STATUSES,
+    ) -> StatusEvent:
+        """Block until a terminal status is received via SSE.
+
+        Consumes the event stream internally, calling on_event
+        for each event if provided. Returns the final StatusEvent.
+
+        Args:
+            resource_id: The resource identifier.
+            timeout: Maximum wait time in seconds. Defaults to 1 hour.
+            on_event: Optional callback invoked for each event.
+            terminal_statuses: Set of statuses that end the wait.
+
+        Returns:
+            The final StatusEvent with terminal status.
+
+        Raises:
+            StreamTimeoutError: If timeout is exceeded.
+        """
+        start_time = time.monotonic()
+        last_status_event: Optional[StatusEvent] = None
+
+        with self.stream(
+            resource_id,
+            max_total_seconds=timeout,
+        ) as events:
+            for event in events:
+                elapsed = time.monotonic() - start_time
+                if elapsed >= timeout:
+                    raise StreamTimeoutError(
+                        f"Timed out after {timeout}s waiting "
+                        f"for resource {resource_id}"
+                    )
+
+                if on_event is not None:
+                    on_event(event)
+
+                if isinstance(event, StatusEvent):
+                    last_status_event = event
+                    if event.status in terminal_statuses:
+                        return event
+
+        # Stream ended without terminal status
+        if last_status_event is not None:
+            return last_status_event
+
+        raise StreamTimeoutError(
+            f"SSE stream ended without terminal status "
+            f"for resource {resource_id}"
+        )
+
+
+class AsyncStreamOperations:
+    """Asynchronous streaming helpers, composed into resource classes.
+
+    Async counterpart of StreamOperations.
+    """
+
+    def __init__(
+        self,
+        http: "AsyncHTTPClient",
+        sse_path_template: str,
+    ) -> None:
+        self._http = http
+        self._sse_path_template = sse_path_template
+
+    def stream(
+        self,
+        resource_id: str,
+        *,
+        max_total_seconds: float = 7200.0,
+    ) -> AsyncEventStream:
+        """Open a typed async SSE event stream for a resource.
+
+        Args:
+            resource_id: The resource identifier.
+            max_total_seconds: Maximum total stream duration.
+
+        Returns:
+            An AsyncEventStream (async iterator + context manager).
+        """
+        from .._http import _build_headers
+
+        path = self._sse_path_template.format(resource_id=resource_id)
+        connection = AsyncSSEConnection(
+            client=self._http._get_client(),
+            url=path,
+            headers=_build_headers(self._http.config),
+            max_total_seconds=max_total_seconds,
+        )
+        return AsyncEventStream(connection)
+
+    async def wait(
+        self,
+        resource_id: str,
+        *,
+        timeout: float = 3600.0,
+        on_event: Optional[Callable[[ServerEvent], None]] = None,
+        terminal_statuses: FrozenSet[str] = TERMINAL_STATUSES,
+    ) -> StatusEvent:
+        """Block until a terminal status is received via SSE.
+
+        Async counterpart of StreamOperations.wait().
+
+        Args:
+            resource_id: The resource identifier.
+            timeout: Maximum wait time in seconds.
+            on_event: Optional callback invoked for each event.
+            terminal_statuses: Set of statuses that end the wait.
+
+        Returns:
+            The final StatusEvent with terminal status.
+
+        Raises:
+            StreamTimeoutError: If timeout is exceeded.
+        """
+        start_time = time.monotonic()
+        last_status_event: Optional[StatusEvent] = None
+
+        async with self.stream(
+            resource_id,
+            max_total_seconds=timeout,
+        ) as events:
+            async for event in events:
+                elapsed = time.monotonic() - start_time
+                if elapsed >= timeout:
+                    raise StreamTimeoutError(
+                        f"Timed out after {timeout}s waiting "
+                        f"for resource {resource_id}"
+                    )
+
+                if on_event is not None:
+                    on_event(event)
+
+                if isinstance(event, StatusEvent):
+                    last_status_event = event
+                    if event.status in terminal_statuses:
+                        return event
+
+        if last_status_event is not None:
+            return last_status_event
+
+        raise StreamTimeoutError(
+            f"SSE stream ended without terminal status "
+            f"for resource {resource_id}"
+        )

waveium/resources/_upload_ops.py

@@ -0,0 +1,286 @@
+"""Reusable upload operation mixins for resources."""
+
+import mimetypes
+from pathlib import Path
+from typing import TYPE_CHECKING, Any, Dict, Optional, Union
+
+import anyio
+import httpx
+
+from ..exceptions import FileUploadError
+
+if TYPE_CHECKING:
+    from .._http import AsyncHTTPClient, HTTPClient
+
+
+class UploadOperations:
+    """Sync upload operations for a resource."""
+
+    def __init__(self, http: "HTTPClient", base_path: str) -> None:
+        self._http = http
+        self._base_path = base_path
+
+    def initiate_upload(
+        self,
+        resource_id: str,
+        filename: str,
+        content_type: Optional[str] = None,
+        response_model: type = None,  # type: ignore[assignment]
+    ) -> Any:
+        """Initiate a file upload session.
+
+        Args:
+            resource_id: The resource identifier.
+            filename: Name of the file to upload.
+            content_type: Optional MIME type override.
+            response_model: Pydantic model for response.
+
+        Returns:
+            Upload session response with signed URL.
+        """
+        body: Dict[str, Any] = {"filename": filename}
+        if content_type:
+            body["content_type"] = content_type
+        return self._http.request(
+            "POST",
+            f"{self._base_path}/{resource_id}/uploads",
+            json_data=body,
+            response_model=response_model,
+        )
+
+    def complete_upload(
+        self,
+        resource_id: str,
+        upload_id: str,
+        purpose: Optional[str] = None,
+        response_model: type = None,  # type: ignore[assignment]
+    ) -> Any:
+        """Mark an upload session as complete.
+
+        Args:
+            resource_id: The resource identifier.
+            upload_id: The upload session identifier.
+            purpose: Optional purpose tag for the file.
+            response_model: Pydantic model for response.
+
+        Returns:
+            Completed upload response.
+        """
+        params = {}
+        if purpose:
+            params["purpose"] = purpose
+        return self._http.request(
+            "POST",
+            (
+                f"{self._base_path}/{resource_id}"
+                f"/uploads/{upload_id}/complete"
+            ),
+            params=params if params else None,
+            response_model=response_model,
+        )
+
+    def upload(
+        self,
+        resource_id: str,
+        file_path: Union[str, Path],
+        content_type: Optional[str] = None,
+        purpose: Optional[str] = None,
+        initiate_model: type = None,  # type: ignore[assignment]
+        complete_model: type = None,  # type: ignore[assignment]
+    ) -> Any:
+        """Upload a file end-to-end.
+
+        Initiates an upload session, transfers the file
+        to the signed URL, then completes the session.
+
+        Args:
+            resource_id: The resource identifier.
+            file_path: Local path to the file.
+            content_type: Optional MIME type override.
+            purpose: Optional purpose tag for the file.
+            initiate_model: Model for initiate response.
+            complete_model: Model for complete response.
+
+        Returns:
+            Completed upload response.
+
+        Raises:
+            FileNotFoundError: If file does not exist.
+            FileUploadError: If upload fails.
+        """
+        file_path = Path(file_path)
+        if not file_path.exists():
+            raise FileNotFoundError(f"File not found: {file_path}")
+        filename = file_path.name
+        if content_type is None:
+            content_type, _ = mimetypes.guess_type(str(file_path))
+            if content_type is None:
+                content_type = "application/octet-stream"
+
+        upload_resp = self.initiate_upload(
+            resource_id=resource_id,
+            filename=filename,
+            content_type=content_type,
+            response_model=initiate_model,
+        )
+
+        try:
+            with (
+                open(file_path, "rb") as f,
+                httpx.Client() as client,
+            ):
+                r = client.request(
+                    method=upload_resp.method,
+                    url=upload_resp.upload_url,
+                    content=f,
+                    headers={"Content-Type": content_type},
+                )
+                r.raise_for_status()
+
+            return self.complete_upload(
+                resource_id=resource_id,
+                upload_id=upload_resp.upload_id,
+                purpose=purpose,
+                response_model=complete_model,
+            )
+        except FileUploadError:
+            raise
+        except Exception as e:
+            raise FileUploadError(f"Failed to upload file: {e}")
+
+
+class AsyncUploadOperations:
+    """Async upload operations for a resource."""
+
+    def __init__(self, http: "AsyncHTTPClient", base_path: str) -> None:
+        self._http = http
+        self._base_path = base_path
+
+    async def initiate_upload(
+        self,
+        resource_id: str,
+        filename: str,
+        content_type: Optional[str] = None,
+        response_model: type = None,  # type: ignore[assignment]
+    ) -> Any:
+        """Initiate a file upload session.
+
+        Args:
+            resource_id: The resource identifier.
+            filename: Name of the file to upload.
+            content_type: Optional MIME type override.
+            response_model: Pydantic model for response.
+
+        Returns:
+            Upload session response with signed URL.
+        """
+        body: Dict[str, Any] = {"filename": filename}
+        if content_type:
+            body["content_type"] = content_type
+        return await self._http.request(
+            "POST",
+            f"{self._base_path}/{resource_id}/uploads",
+            json_data=body,
+            response_model=response_model,
+        )
+
+    async def complete_upload(
+        self,
+        resource_id: str,
+        upload_id: str,
+        purpose: Optional[str] = None,
+        response_model: type = None,  # type: ignore[assignment]
+    ) -> Any:
+        """Mark an upload session as complete.
+
+        Args:
+            resource_id: The resource identifier.
+            upload_id: The upload session identifier.
+            purpose: Optional purpose tag for the file.
+            response_model: Pydantic model for response.
+
+        Returns:
+            Completed upload response.
+        """
+        params = {}
+        if purpose:
+            params["purpose"] = purpose
+        return await self._http.request(
+            "POST",
+            (
+                f"{self._base_path}/{resource_id}"
+                f"/uploads/{upload_id}/complete"
+            ),
+            params=params if params else None,
+            response_model=response_model,
+        )
+
+    async def upload(
+        self,
+        resource_id: str,
+        file_path: Union[str, Path],
+        content_type: Optional[str] = None,
+        purpose: Optional[str] = None,
+        initiate_model: type = None,  # type: ignore[assignment]
+        complete_model: type = None,  # type: ignore[assignment]
+    ) -> Any:
+        """Upload a file end-to-end.
+
+        Initiates an upload session, transfers the file
+        to the signed URL, then completes the session.
+
+        Args:
+            resource_id: The resource identifier.
+            file_path: Local path to the file.
+            content_type: Optional MIME type override.
+            purpose: Optional purpose tag for the file.
+            initiate_model: Model for initiate response.
+            complete_model: Model for complete response.
+
+        Returns:
+            Completed upload response.
+
+        Raises:
+            FileNotFoundError: If file does not exist.
+            FileUploadError: If upload fails.
+        """
+        file_path = Path(file_path)
+        if not file_path.exists():
+            raise FileNotFoundError(f"File not found: {file_path}")
+        filename = file_path.name
+        if content_type is None:
+            content_type, _ = mimetypes.guess_type(str(file_path))
+            if content_type is None:
+                content_type = "application/octet-stream"
+
+        upload_resp = await self.initiate_upload(
+            resource_id=resource_id,
+            filename=filename,
+            content_type=content_type,
+            response_model=initiate_model,
+        )
+
+        try:
+            async with (
+                await anyio.open_file(file_path, "rb") as f,
+                httpx.AsyncClient() as client,
+            ):
+                file_content = await f.read()
+                r = await client.request(
+                    method=upload_resp.method,
+                    url=upload_resp.upload_url,
+                    content=file_content,
+                    headers={"Content-Type": content_type},
+                )
+                r.raise_for_status()
+
+            return await self.complete_upload(
+                resource_id=resource_id,
+                upload_id=upload_resp.upload_id,
+                purpose=purpose,
+                response_model=complete_model,
+            )
+        except FileUploadError:
+            raise
+        except Exception as e:
+            raise FileUploadError(f"Failed to upload file: {e}")