diff --git a/packages/text-generation/src/celeste_text_generation/__init__.py b/packages/text-generation/src/celeste_text_generation/__init__.py index 70cdb9e5..e49ad72e 100644 --- a/packages/text-generation/src/celeste_text_generation/__init__.py +++ b/packages/text-generation/src/celeste_text_generation/__init__.py @@ -18,6 +18,8 @@ def register_package() -> None: # Import after register_package is defined to avoid circular imports from celeste_text_generation.io import ( # noqa: E402 + TextGenerationChunk, + TextGenerationFinishReason, TextGenerationInput, TextGenerationOutput, TextGenerationUsage, @@ -25,6 +27,8 @@ def register_package() -> None: from celeste_text_generation.streaming import TextGenerationStream # noqa: E402 __all__ = [ + "TextGenerationChunk", + "TextGenerationFinishReason", "TextGenerationInput", "TextGenerationOutput", "TextGenerationStream", diff --git a/packages/text-generation/src/celeste_text_generation/providers/anthropic/streaming.py b/packages/text-generation/src/celeste_text_generation/providers/anthropic/streaming.py index b4f912b3..6422961a 100644 --- a/packages/text-generation/src/celeste_text_generation/providers/anthropic/streaming.py +++ b/packages/text-generation/src/celeste_text_generation/providers/anthropic/streaming.py @@ -38,6 +38,7 @@ def __init__( self._tool_use_blocks: list[dict[str, Any]] = [] self._current_tool_use: dict[str, Any] | None = None self._current_tool_use_partial_json: str = "" + self._last_finish_reason: TextGenerationFinishReason | None = None def _parse_chunk(self, event: dict[str, Any]) -> Chunk | None: """Parse SSE event into Chunk.""" @@ -208,6 +209,7 @@ def _parse_chunk(self, event: dict[str, Any]) -> Chunk | None: finish_reason: TextGenerationFinishReason | None = None if stop_reason is not None: finish_reason = TextGenerationFinishReason(reason=stop_reason) + self._last_finish_reason = finish_reason usage = self._parse_usage_from_event(event) @@ -223,7 +225,7 @@ def _parse_chunk(self, event: dict[str, Any]) -> Chunk | None: return TextGenerationChunk( content="", - finish_reason=None, + finish_reason=self._last_finish_reason, usage=usage, )