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: 1 addition & 1 deletion src/celeste/protocols/chatcompletions/streaming.py
Original file line number Diff line number Diff line change
Expand Up @@ -89,7 +89,7 @@ def _parse_chunk(self, event_data: dict[str, Any]) -> Any: # noqa: ANN401
if choices and isinstance(choices[0], dict):
delta = choices[0].get("delta", {})
if isinstance(delta, dict):
for tc_delta in delta.get("tool_calls", []):
for tc_delta in delta.get("tool_calls") or []:
idx = tc_delta.get("index", 0)
if idx not in self._tool_call_deltas:
self._tool_call_deltas[idx] = {
Expand Down
2 changes: 1 addition & 1 deletion src/celeste/protocols/openresponses/tools.py
Original file line number Diff line number Diff line change
Expand Up @@ -28,7 +28,7 @@ def map_tool(self, tool: Tool) -> dict[str, Any]:
result: dict[str, Any] = {"type": "web_search"}
if tool.allowed_domains is not None:
result.setdefault("filters", {})["allowed_domains"] = tool.allowed_domains
for field in tool.model_fields:
for field in type(tool).model_fields:
if field not in self._supported_fields and getattr(tool, field) is not None:
warnings.warn(
f"WebSearch.{field} is not supported by OpenResponses "
Expand Down
2 changes: 1 addition & 1 deletion src/celeste/providers/google/generate_content/tools.py
Original file line number Diff line number Diff line change
Expand Up @@ -18,7 +18,7 @@ def map_tool(self, tool: Tool) -> dict[str, Any]:
config: dict[str, Any] = {}
if tool.blocked_domains is not None:
config["exclude_domains"] = tool.blocked_domains
for field in tool.model_fields:
for field in type(tool).model_fields:
if field not in self._supported_fields and getattr(tool, field) is not None:
warnings.warn(
f"WebSearch.{field} is not supported by Google "
Expand Down
2 changes: 1 addition & 1 deletion src/celeste/providers/groq/chat/tools.py
Original file line number Diff line number Diff line change
Expand Up @@ -15,7 +15,7 @@ class WebSearchMapper(ToolMapper):

def map_tool(self, tool: Tool) -> dict[str, Any]:
assert isinstance(tool, WebSearch)
for field in tool.model_fields:
for field in type(tool).model_fields:
if field not in self._supported_fields and getattr(tool, field) is not None:
warnings.warn(
f"WebSearch.{field} is not supported by Groq "
Expand Down
2 changes: 1 addition & 1 deletion src/celeste/providers/moonshot/chat/tools.py
Original file line number Diff line number Diff line change
Expand Up @@ -15,7 +15,7 @@ class WebSearchMapper(ToolMapper):

def map_tool(self, tool: Tool) -> dict[str, Any]:
assert isinstance(tool, WebSearch)
for field in tool.model_fields:
for field in type(tool).model_fields:
if field not in self._supported_fields and getattr(tool, field) is not None:
warnings.warn(
f"WebSearch.{field} is not supported by Moonshot "
Expand Down
59 changes: 23 additions & 36 deletions src/celeste/streaming.py
Original file line number Diff line number Diff line change
@@ -1,13 +1,13 @@
"""Streaming support for Celeste."""

from abc import ABC, abstractmethod
from collections.abc import AsyncIterator, Callable
from contextlib import AbstractContextManager, suppress
from collections.abc import AsyncIterator, Callable, Iterator
from contextlib import suppress
from types import TracebackType
from typing import Any, ClassVar, Self, Unpack

import httpx
from anyio.from_thread import BlockingPortal, start_blocking_portal
from anyio.from_thread import start_blocking_portal

from celeste.exceptions import StreamEventError, StreamNotExhaustedError
from celeste.io import Chunk as ChunkBase
Expand Down Expand Up @@ -63,9 +63,8 @@ def __init__(
self._parameters = parameters
self._transform_output = transform_output
self._stream_metadata = stream_metadata or {}
# Sync iteration state
self._portal: BlockingPortal | None = None
self._portal_cm: AbstractContextManager[BlockingPortal] | None = None
# Sync iteration state (portal lifecycle managed by __iter__ generator)
self._sync_generator: Iterator[Chunk] | 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."""
Expand Down Expand Up @@ -267,38 +266,26 @@ async def __anext__(self) -> Chunk:
raise StopAsyncIteration

# Iterator protocol (sync)
def __iter__(self) -> Self:
"""Return self as sync iterator with dedicated event loop.
def __iter__(self) -> Iterator[Chunk]:
"""Sync iterator using a generator with try/finally for guaranteed cleanup.

Creates a blocking portal that maintains a persistent event loop
in a dedicated thread for consistent async context.
Creates a blocking portal in a dedicated thread. The generator's finally
block ensures the portal is cleaned up on exhaustion, break, exception,
or garbage collection — unlike __next__ which only cleans up on exhaustion.
"""
if self._portal is None:
self._portal_cm = start_blocking_portal()
self._portal = self._portal_cm.__enter__()
return self

def __next__(self) -> Chunk:
"""Yield next chunk via portal's persistent event loop."""
if self._portal is None:
self.__iter__()

portal_cm = start_blocking_portal()
portal = portal_cm.__enter__()
try:
return self._portal.call(self.__anext__) # type: ignore[union-attr,no-any-return]
except StopAsyncIteration:
self._cleanup_portal()
raise StopIteration from None

def _cleanup_portal(self) -> None:
"""Clean up the blocking portal and its thread."""
if self._portal_cm is not None:
# Close stream via portal before exiting (ensures _closed = True)
if self._portal is not None and not self._closed:
while True:
try:
yield portal.call(self.__anext__)
except StopAsyncIteration:
return
finally:
if not self._closed:
with suppress(RuntimeError):
self._portal.call(self.aclose)
self._portal_cm.__exit__(None, None, None)
self._portal = None
self._portal_cm = None
portal.call(self.aclose)
portal_cm.__exit__(None, None, None)

# AsyncContextManager protocol
async def __aenter__(self) -> Self:
Expand Down Expand Up @@ -326,8 +313,8 @@ def __exit__(
exc_val: BaseException | None,
exc_tb: TracebackType | None,
) -> None:
"""Exit sync context - ensure cleanup."""
self._cleanup_portal()
"""Exit sync context. Portal cleanup is handled by __iter__ generator's finally."""
return

@property
def output(self) -> Out:
Expand Down
Loading