Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions src/celeste/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -24,6 +24,7 @@
MissingCredentialsError,
ModelNotFoundError,
StreamEmptyError,
StreamEventError,
StreamingNotSupportedError,
StreamNotExhaustedError,
UnsupportedCapabilityError,
Expand Down Expand Up @@ -264,6 +265,7 @@ def create_client(
"RefResolvingJsonSchemaGenerator",
"Role",
"StreamEmptyError",
"StreamEventError",
"StreamNotExhaustedError",
"StreamingNotSupportedError",
"StrictJsonSchemaGenerator",
Expand Down
24 changes: 24 additions & 0 deletions src/celeste/exceptions.py
Original file line number Diff line number Diff line change
@@ -1,5 +1,7 @@
"""Custom exceptions for Celeste."""

from typing import Any


class Error(Exception):
"""Base exception for all Celeste errors."""
Expand Down Expand Up @@ -170,6 +172,27 @@ def __init__(self) -> None:
super().__init__("Stream completed but no chunks were produced")


class StreamEventError(StreamingError):
"""Raised when the provider sends an error event during streaming."""

def __init__(
self,
message: str,
*,
error_type: str | None = None,
event_data: dict[str, Any] | None = None,
provider: str | None = None,
) -> None:
"""Initialize with error details from the stream event."""
self.error_type = error_type
self.event_data = event_data or {}
self.provider = provider

prefix = f"{provider} stream error" if provider else "Stream error"
suffix = f" [{error_type}]" if error_type else ""
super().__init__(f"{prefix}{suffix}: {message}")


class MissingDependencyError(Error):
"""Raised when a required optional dependency is not installed."""

Expand Down Expand Up @@ -231,6 +254,7 @@ def __init__(self, parameter: str, model_id: str) -> None:
"ModalityNotFoundError",
"ModelNotFoundError",
"StreamEmptyError",
"StreamEventError",
"StreamNotExhaustedError",
"StreamingNotSupportedError",
"UnsupportedCapabilityError",
Expand Down
1 change: 1 addition & 0 deletions src/celeste/http.py
Original file line number Diff line number Diff line change
Expand Up @@ -185,6 +185,7 @@ async def stream_post(
headers=headers,
timeout=timeout,
) as event_source:
event_source.response.raise_for_status()
async for sse in event_source.aiter_sse():
try:
yield json.loads(sse.data)
Expand Down
10 changes: 9 additions & 1 deletion src/celeste/protocols/openresponses/streaming.py
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
"""OpenResponses protocol SSE parsing for streaming."""

from typing import Any
from typing import Any, ClassVar

from celeste.io import FinishReason

Expand All @@ -22,6 +22,14 @@ class OpenResponsesStream:
Modality streams call super() methods which resolve to this via MRO.
"""

_error_type_fields: ClassVar[tuple[str, ...]] = ("code",)

def _parse_stream_error(self, event_data: dict[str, Any]) -> dict[str, Any] | None:
"""Detect Responses API error events (flat shape: code/message at root level)."""
if event_data.get("type") == "error":
return self._build_error_from_value(event_data) # type: ignore[attr-defined, no-any-return]
return None

def _parse_chunk_content(self, event_data: dict[str, Any]) -> str | None:
"""Extract content from SSE event."""
event_type = event_data.get("type")
Expand Down
1 change: 0 additions & 1 deletion src/celeste/providers/anthropic/messages/config.py
Original file line number Diff line number Diff line change
Expand Up @@ -48,4 +48,3 @@ class VertexAnthropicEndpoint(StrEnum):
SSE_EVENT_CONTENT_BLOCK_STOP = "content_block_stop"
SSE_EVENT_MESSAGE_DELTA = "message_delta"
SSE_EVENT_MESSAGE_STOP = "message_stop"
SSE_EVENT_ERROR = "error"
4 changes: 3 additions & 1 deletion src/celeste/providers/google/generate_content/streaming.py
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
"""Google GenerateContent SSE parsing for streaming."""

from typing import Any
from typing import Any, ClassVar

from celeste.io import FinishReason

Expand All @@ -18,6 +18,8 @@ class GoogleGenerateContentStream:
Modality streams call super() methods which resolve to this via MRO.
"""

_error_type_fields: ClassVar[tuple[str, ...]] = ("status", "code")

def _parse_chunk_content(self, event_data: dict[str, Any]) -> str | None:
"""Extract content from SSE event."""
candidates = event_data.get("candidates", [])
Expand Down
11 changes: 10 additions & 1 deletion src/celeste/providers/google/interactions/streaming.py
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
"""Google Interactions SSE parsing for streaming."""

from typing import Any
from typing import Any, ClassVar

from celeste.io import FinishReason

Expand All @@ -22,6 +22,15 @@ class GoogleInteractionsStream:
Modality streams call super() methods which resolve to this via MRO.
"""

_error_type_fields: ClassVar[tuple[str, ...]] = ("status", "code")

def _parse_stream_error(self, event_data: dict[str, Any]) -> dict[str, Any] | None:
"""Detect Interactions error events (dual event_type/type field)."""
event_type = event_data.get("event_type") or event_data.get("type")
if event_type == "error":
return self._build_error_from_value(event_data.get("error", {}))
return None

def _parse_chunk_content(self, event_data: dict[str, Any]) -> str | None:
"""Extract content from SSE event.

Expand Down
50 changes: 49 additions & 1 deletion src/celeste/streaming.py
Original file line number Diff line number Diff line change
Expand Up @@ -8,7 +8,7 @@

from anyio.from_thread import BlockingPortal, start_blocking_portal

from celeste.exceptions import StreamNotExhaustedError
from celeste.exceptions import StreamEventError, StreamNotExhaustedError
from celeste.io import Chunk as ChunkBase
from celeste.io import FinishReason, Output, Usage
from celeste.parameters import Parameters
Expand All @@ -31,6 +31,7 @@ class Stream[Out: Output, Params: Parameters, Chunk: ChunkBase](ABC):
_chunk_class: ClassVar[type[ChunkBase]]
_output_class: ClassVar[type[Output]]
_empty_content: ClassVar[Any]
_error_type_fields: ClassVar[tuple[str, ...]] = ("type", "code")

def __init__(
self,
Expand All @@ -51,6 +52,45 @@ def __init__(
self._portal: BlockingPortal | None = None
self._portal_cm: AbstractContextManager[BlockingPortal] | None = None

def _build_error_from_value(self, error: Any) -> dict[str, Any]: # noqa: ANN401
"""Extract {type, message} from an error value using _error_type_fields."""
if isinstance(error, dict):
error_type = None
for field in self._error_type_fields:
val = error.get(field)
if val is not None:
error_type = str(val)
break
return {
"type": error_type,
"message": error.get("message", "Unknown error"),
}
# Non-dict error value (e.g., plain string from Cohere)
return {"message": str(error) if error else "Unknown error"}

def _parse_stream_error(self, event_data: dict[str, Any]) -> dict[str, Any] | None:
"""Detect error events in the SSE stream.

Handles two generic SSE error patterns:
1. Type-based: {"type": "error", "error": {"type": "...", "message": "..."}}
2. Field-based: {"error": {"message": "...", "type": "..."}}

Override in provider mixin for non-standard error shapes.
"""
error = None

# Pattern 1: Type-based — event has "type": "error" with nested error
if event_data.get("type") == "error":
error = event_data.get("error", {})
# Pattern 2: Field-based — event has top-level "error" dict
elif isinstance(event_data.get("error"), dict):
error = event_data["error"]

if error is None:
return None

return self._build_error_from_value(error)

def _parse_chunk_content(self, event_data: dict[str, Any]) -> Any | None: # noqa: ANN401
"""Parse content from chunk event. Override in provider mixin."""
return None
Expand All @@ -61,6 +101,14 @@ def _wrap_chunk_content(self, raw_content: Any) -> Any: # noqa: ANN401

def _parse_chunk(self, event: dict[str, Any]) -> Chunk | None:
"""Parse SSE event into Chunk (returns None to filter lifecycle events)."""
error = self._parse_stream_error(event)
if error is not None:
raise StreamEventError(
message=error.get("message", "Unknown stream error"),
error_type=error.get("type"),
event_data=event,
provider=self._stream_metadata.get("provider"),
)
content = self._parse_chunk_content(event)
usage = self._get_chunk_usage(event)
finish_reason = self._get_chunk_finish_reason(event)
Expand Down
160 changes: 158 additions & 2 deletions tests/unit_tests/test_streaming.py
Original file line number Diff line number Diff line change
@@ -1,13 +1,13 @@
"""High-value tests for Stream - focusing on lifecycle, resource cleanup, and state management."""

from collections.abc import AsyncIterator
from typing import Any, Unpack
from typing import Any, ClassVar, Unpack
from unittest.mock import AsyncMock

import pytest
from pydantic import Field

from celeste.exceptions import StreamNotExhaustedError
from celeste.exceptions import StreamEventError, StreamNotExhaustedError
from celeste.io import Chunk, FinishReason, Output, Usage
from celeste.parameters import Parameters
from celeste.streaming import Stream
Expand Down Expand Up @@ -720,3 +720,159 @@ def _parse_output( # type: ignore[override]
assert output.finish_reason is not None
assert isinstance(output.finish_reason, TypedFinishReason)
assert output.finish_reason.reason == "stop"


class PipelineStream(Stream[ConcreteOutput, Parameters, Chunk]):
"""Stream that uses the base _parse_chunk pipeline (for testing error detection).

Unlike ConcreteStream which overrides _parse_chunk entirely, this class
only overrides _parse_chunk_content and _aggregate_content, so the base
_parse_stream_error → StreamEventError pipeline is exercised.
"""

_chunk_class: ClassVar[type[Chunk]] = Chunk
_output_class: ClassVar[type[Output]] = ConcreteOutput
_empty_content: ClassVar[str] = ""

def _aggregate_content(self, chunks: list[Chunk]) -> str:
"""Aggregate content from chunks."""
return "".join(str(chunk.content) for chunk in chunks)

def _parse_chunk_content(self, event_data: dict[str, Any]) -> str | None:
"""Extract content from delta field."""
return event_data.get("delta") or None


class TestStreamErrorDetection:
"""Test Stream error detection via base _parse_stream_error pipeline."""

async def test_type_based_error_raises_stream_event_error(self) -> None:
"""Type-based error pattern (Anthropic) must raise StreamEventError."""
events = [
{
"type": "error",
"error": {"type": "overloaded_error", "message": "Server overloaded"},
},
]
stream = PipelineStream(
_async_iter(events),
stream_metadata={"provider": "anthropic"},
)
with pytest.raises(StreamEventError, match="Server overloaded") as exc_info:
async for _ in stream:
pass
assert exc_info.value.error_type == "overloaded_error"
assert exc_info.value.provider == "anthropic"
assert exc_info.value.event_data == events[0]

async def test_field_based_error_raises_stream_event_error(self) -> None:
"""Field-based error pattern (ChatCompletions) must raise StreamEventError."""
events = [
{"error": {"type": "invalid_request", "message": "Bad request"}},
]
stream = PipelineStream(
_async_iter(events),
stream_metadata={"provider": "openai"},
)
with pytest.raises(StreamEventError, match="Bad request") as exc_info:
async for _ in stream:
pass
assert exc_info.value.error_type == "invalid_request"
assert exc_info.value.provider == "openai"

async def test_field_based_error_falls_back_to_code_field(self) -> None:
"""Field-based error without 'type' must fall back to 'code' field."""
events = [
{"error": {"code": "rate_limit_exceeded", "message": "Rate limited"}},
]
stream = PipelineStream(_async_iter(events))
with pytest.raises(StreamEventError) as exc_info:
async for _ in stream:
pass
assert exc_info.value.error_type == "rate_limit_exceeded"

async def test_type_based_error_with_string_error_value(self) -> None:
"""Type-based error with non-dict error value must use string fallback."""
events = [
{"type": "error", "error": "Something went wrong"},
]
stream = PipelineStream(_async_iter(events))
with pytest.raises(StreamEventError, match="Something went wrong") as exc_info:
async for _ in stream:
pass
assert exc_info.value.error_type is None

async def test_error_type_fields_classvar_override(self) -> None:
"""ClassVar override of _error_type_fields must change field lookup order."""

class GoogleLikeStream(PipelineStream):
_error_type_fields: ClassVar[tuple[str, ...]] = ("status", "code")

events = [
{
"error": {
"status": "PERMISSION_DENIED",
"code": 403,
"message": "Forbidden",
},
},
]
stream = GoogleLikeStream(_async_iter(events))
with pytest.raises(StreamEventError) as exc_info:
async for _ in stream:
pass
assert exc_info.value.error_type == "PERMISSION_DENIED"

async def test_non_error_events_pass_through(self) -> None:
"""Normal events must not trigger error detection."""
events = [
{"delta": "Hello"},
{"delta": " world"},
]
stream = PipelineStream(_async_iter(events))
chunks = [chunk async for chunk in stream]
assert len(chunks) == 2
assert stream.output.content == "Hello world"

async def test_error_after_successful_chunks(self) -> None:
"""Error mid-stream (after successful chunks) must raise StreamEventError."""
events = [
{"delta": "Hello"},
{
"type": "error",
"error": {"type": "server_error", "message": "Internal error"},
},
]
stream = PipelineStream(
_async_iter(events),
stream_metadata={"provider": "test"},
)
chunks: list[Chunk] = []
with pytest.raises(StreamEventError, match="Internal error"):
async for chunk in stream:
chunks.append(chunk)
assert len(chunks) == 1
assert chunks[0].content == "Hello"

async def test_error_with_no_message_uses_default(self) -> None:
"""Error event without message field must use 'Unknown error' default."""
events = [
{"error": {"type": "mystery_error"}},
]
stream = PipelineStream(_async_iter(events))
with pytest.raises(StreamEventError, match="Unknown error"):
async for _ in stream:
pass

async def test_error_provides_full_event_data(self) -> None:
"""StreamEventError must include the full original event data."""
event = {
"type": "error",
"error": {"type": "api_error", "message": "Fail"},
"extra": "data",
}
stream = PipelineStream(_async_iter([event]))
with pytest.raises(StreamEventError) as exc_info:
async for _ in stream:
pass
assert exc_info.value.event_data == event
Loading