refactor: add Chunk type parameter to Stream class for type-safe streaming (#42)

Make Stream generic over Chunk type to enable capability-specific chunk
types, improving type safety and IDE support across all streaming
implementations.

Changes:
- Add Chunk type parameter to Stream[Out, Params, Chunk] base class
- Update all provider streaming implementations to use specific chunk types
  (TextGenerationChunk, ImageGenerationChunk) instead of generic Chunk
- Update Client.stream() and _stream_class() return types to include Chunk
- Remove generic Chunk imports from provider implementations
- Update all tests to use new Stream signature
- Remove streaming modules from mypy ignore lists (now properly typed)
- Use ApplicationMimeType.JSON constant in HTTP tests

Benefits:
- Type-safe chunk parsing per capability
- Better IDE autocomplete and type checking
- Cleaner provider code without generic Chunk imports
- Improved mypy compliance

Author: Kamilbenkirane

packages/image-generation/src/celeste_image_generation/providers/bytedance/streaming.py

@@ -5,7 +5,6 @@
 from typing import Any
 
 from celeste.artifacts import ImageArtifact
-from celeste.io import Chunk
 from celeste.mime_types import ImageMimeType
 from celeste_image_generation.io import ImageGenerationChunk, ImageGenerationUsage
 from celeste_image_generation.streaming import ImageGenerationStream
@@ -21,7 +20,7 @@ def __init__(self, sse_iterator: AsyncIterator[dict[str, Any]]) -> None:
         super().__init__(sse_iterator)
         self._completed_usage: ImageGenerationUsage | None = None
 
-    def _parse_chunk(self, chunk_data: dict[str, Any]) -> Chunk | None:
+    def _parse_chunk(self, chunk_data: dict[str, Any]) -> ImageGenerationChunk | None:
         """Parse chunk from SSE event."""
         event_type = chunk_data.get("type")
 

packages/image-generation/src/celeste_image_generation/providers/openai/streaming.py

@@ -5,7 +5,6 @@
 from typing import Any
 
 from celeste.artifacts import ImageArtifact
-from celeste.io import Chunk
 from celeste_image_generation.io import ImageGenerationChunk, ImageGenerationUsage
 from celeste_image_generation.streaming import ImageGenerationStream
 
@@ -15,7 +14,7 @@
 class OpenAIImageGenerationStream(ImageGenerationStream):
     """OpenAI streaming for image generation."""
 
-    def _parse_chunk(self, chunk_data: dict[str, Any]) -> Chunk | None:
+    def _parse_chunk(self, chunk_data: dict[str, Any]) -> ImageGenerationChunk | None:
         """Parse chunk from SSE event.
 
         OpenAI returns two event types:

packages/image-generation/src/celeste_image_generation/streaming.py

@@ -12,7 +12,9 @@
 from celeste_image_generation.parameters import ImageGenerationParameters
 
 
-class ImageGenerationStream(Stream[ImageGenerationOutput, ImageGenerationParameters]):
+class ImageGenerationStream(
+    Stream[ImageGenerationOutput, ImageGenerationParameters, ImageGenerationChunk]
+):
     """Streaming for image generation."""
 
     def _parse_output(

packages/text-generation/src/celeste_text_generation/providers/anthropic/streaming.py

@@ -3,7 +3,6 @@
 from collections.abc import Callable
 from typing import Any, Unpack
 
-from celeste.io import Chunk
 from celeste_text_generation.io import (
     TextGenerationChunk,
     TextGenerationFinishReason,
@@ -34,7 +33,7 @@ def __init__(
         self._transform_output = transform_output
         self._last_finish_reason: TextGenerationFinishReason | None = None
 
-    def _parse_chunk(self, event: dict[str, Any]) -> Chunk | None:
+    def _parse_chunk(self, event: dict[str, Any]) -> TextGenerationChunk | None:
         """Parse SSE event into Chunk."""
         event_type = event.get("type")
         if not event_type:

packages/text-generation/src/celeste_text_generation/providers/cohere/streaming.py

@@ -4,7 +4,6 @@
 from collections.abc import Callable
 from typing import Any, Unpack
 
-from celeste.io import Chunk
 from celeste_text_generation.io import (
     TextGenerationChunk,
     TextGenerationFinishReason,
@@ -36,7 +35,7 @@ def __init__(
         super().__init__(sse_iterator, **parameters)
         self._transform_output = transform_output
 
-    def _parse_chunk(self, event: dict[str, Any]) -> Chunk | None:
+    def _parse_chunk(self, event: dict[str, Any]) -> TextGenerationChunk | None:
         """Parse SSE event into Chunk, extracting text deltas and metadata."""
         event_type = event.get("type")
 

packages/text-generation/src/celeste_text_generation/providers/google/streaming.py

@@ -3,7 +3,6 @@
 from collections.abc import Callable
 from typing import Any, Unpack
 
-from celeste.io import Chunk
 from celeste_text_generation.io import (
     TextGenerationChunk,
     TextGenerationFinishReason,
@@ -33,7 +32,7 @@ def __init__(
         super().__init__(sse_iterator, **parameters)
         self._transform_output = transform_output
 
-    def _parse_chunk(self, event: dict[str, Any]) -> Chunk | None:
+    def _parse_chunk(self, event: dict[str, Any]) -> TextGenerationChunk | None:
         """Parse SSE event into Chunk.
 
         Extract text delta from candidates[0].content.parts[0].text.

packages/text-generation/src/celeste_text_generation/providers/mistral/streaming.py

@@ -4,7 +4,6 @@
 from collections.abc import Callable
 from typing import Any, Unpack
 
-from celeste.io import Chunk
 from celeste_text_generation.io import (
     TextGenerationChunk,
     TextGenerationFinishReason,
@@ -36,7 +35,7 @@ def __init__(
         super().__init__(sse_iterator, **parameters)
         self._transform_output = transform_output
 
-    def _parse_chunk(self, event: dict[str, Any]) -> Chunk | None:
+    def _parse_chunk(self, event: dict[str, Any]) -> TextGenerationChunk | None:
         """Parse chunk from SSE event.
 
         Extract from choices[0].delta.content (content delta events).

packages/text-generation/src/celeste_text_generation/providers/openai/streaming.py

@@ -3,7 +3,6 @@
 from collections.abc import Callable
 from typing import Any, Unpack
 
-from celeste.io import Chunk
 from celeste_text_generation.io import (
     TextGenerationChunk,
     TextGenerationFinishReason,
@@ -27,7 +26,7 @@ def __init__(
         super().__init__(sse_iterator, **parameters)
         self._transform_output = transform_output
 
-    def _parse_chunk(self, event: dict[str, Any]) -> Chunk | None:
+    def _parse_chunk(self, event: dict[str, Any]) -> TextGenerationChunk | None:
         """Parse SSE event into Chunk."""
         event_type = event.get("type")
         if not event_type:

packages/text-generation/src/celeste_text_generation/streaming.py

@@ -3,7 +3,6 @@
 from abc import abstractmethod
 from typing import Any, Unpack
 
-from celeste.io import Chunk
 from celeste.streaming import Stream
 from celeste_text_generation.io import (
     TextGenerationChunk,
@@ -13,11 +12,13 @@
 from celeste_text_generation.parameters import TextGenerationParameters
 
 
-class TextGenerationStream(Stream[TextGenerationOutput, TextGenerationParameters]):
+class TextGenerationStream(
+    Stream[TextGenerationOutput, TextGenerationParameters, TextGenerationChunk]
+):
     """Streaming for text generation."""
 
     @abstractmethod
-    def _parse_chunk(self, event: dict[str, Any]) -> Chunk | None:
+    def _parse_chunk(self, event: dict[str, Any]) -> TextGenerationChunk | None:
         """Parse SSE event into Chunk (provider-specific)."""
 
     def _parse_output(

pyproject.toml

@@ -165,7 +165,6 @@ ignore_missing_imports = true
 module = [
     "celeste_text_generation.*",
     "celeste_text_generation.client",
-    "celeste_text_generation.streaming",
     "celeste_text_generation.providers.*",
 ]
 disable_error_code = ["override", "return-value", "arg-type", "call-arg", "assignment", "no-any-return"]
@@ -174,7 +173,6 @@ disable_error_code = ["override", "return-value", "arg-type", "call-arg", "assig
 module = [
     "celeste_image_generation.*",
     "celeste_image_generation.client",
-    "celeste_image_generation.streaming",
     "celeste_image_generation.providers.*",
 ]
 disable_error_code = ["override", "return-value", "arg-type", "call-arg", "assignment", "no-any-return"]