fix: move _transform_output to base _predict pipeline

src/celeste/client.py

@@ -17,7 +17,7 @@
 from celeste.models import Model
 from celeste.parameters import ParameterMapper, Parameters
 from celeste.streaming import Stream
-from celeste.types import RawUsage, TextContent
+from celeste.types import RawUsage
 
 
 class APIMixin(ABC):
@@ -178,8 +178,10 @@ async def _predict(
         response_data = await self._make_request(
             request_body, endpoint=endpoint, **parameters
         )
+        content = self._parse_content(response_data, **parameters)
+        content = self._transform_output(content, **parameters)
         return self._output_class()(
-            content=self._parse_content(response_data, **parameters),
+            content=content,
             usage=self._get_usage(response_data),
             finish_reason=self._get_finish_reason(response_data),
             metadata=self._build_metadata(response_data),
@@ -345,9 +347,9 @@ def _handle_error_response(self, response: httpx.Response) -> None:
 
     def _transform_output(
         self,
-        content: TextContent,
+        content: Content,
         **parameters: Unpack[Params],  # type: ignore[misc]
-    ) -> TextContent:
+    ) -> Content:
         """Transform content using parameter mapper output transformations."""
         for mapper in self.parameter_mappers():
             value = parameters.get(mapper.name)

src/celeste/modalities/text/providers/anthropic/client.py

@@ -164,7 +164,7 @@ def _parse_content(
                 text_content = content_block.get("text") or ""
                 break
 
-        return self._transform_output(text_content, **parameters)
+        return text_content
 
     def _stream_class(self) -> type[TextStream]:
         """Return the Stream class for this provider."""

src/celeste/modalities/text/providers/cohere/client.py

@@ -86,7 +86,7 @@ def _parse_content(
         content_array = super()._parse_content(response_data)
         first_content = content_array[0]
         text = first_content.get("text") or ""
-        return self._transform_output(text, **parameters)
+        return text
 
     def _stream_class(self) -> type[TextStream]:
         """Return the Stream class for this provider."""

src/celeste/modalities/text/providers/deepseek/client.py

@@ -66,7 +66,7 @@ def _parse_content(
         choices = super()._parse_content(response_data)
         message = choices[0].get("message", {})
         content = message.get("content") or ""
-        return self._transform_output(content, **parameters)
+        return content
 
     def _stream_class(self) -> type[TextStream]:
         """Return the Stream class for this provider."""

src/celeste/modalities/text/providers/google/client.py

@@ -178,7 +178,7 @@ def _parse_content(
         candidates = super()._parse_content(response_data)
         parts = candidates[0].get("content", {}).get("parts", [])
         text = parts[0].get("text") if parts else ""
-        return self._transform_output(text or "", **parameters)
+        return text or ""
 
     def _stream_class(self) -> type[TextStream]:
         """Return the Stream class for this provider."""

src/celeste/modalities/text/providers/groq/client.py

@@ -84,7 +84,7 @@ def _parse_content(
         choices = super()._parse_content(response_data)
         message = choices[0].get("message", {})
         content = message.get("content") or ""
-        return self._transform_output(content, **parameters)
+        return content
 
     def _stream_class(self) -> type[TextStream]:
         """Return the Stream class for this provider."""

src/celeste/modalities/text/providers/huggingface/client.py

@@ -86,7 +86,7 @@ def _parse_content(
         choices = super()._parse_content(response_data)
         message = choices[0].get("message", {})
         content = message.get("content") or ""
-        return self._transform_output(content, **parameters)
+        return content
 
     def _stream_class(self) -> type[TextStream]:
         """Return the Stream class for this provider."""

src/celeste/modalities/text/providers/mistral/client.py

@@ -93,7 +93,7 @@ def _parse_content(
                     text_parts.append(block.get("text", ""))
             content = "".join(text_parts)
 
-        return self._transform_output(content, **parameters)
+        return content
 
     def _stream_class(self) -> type[TextStream]:
         """Return the Stream class for this provider."""

src/celeste/modalities/text/providers/moonshot/client.py

@@ -83,7 +83,7 @@ def _parse_content(
         choices = super()._parse_content(response_data)
         message = choices[0].get("message", {})
         content = message.get("content") or ""
-        return self._transform_output(content, **parameters)
+        return content
 
     def _stream_class(self) -> type[TextStream]:
         """Return the Stream class for this provider."""

src/celeste/modalities/text/providers/openai/client.py

@@ -89,9 +89,9 @@ def _parse_content(
                 for part in item.get("content", []):
                     if part.get("type") == "output_text":
                         text = part.get("text") or ""
-                        return self._transform_output(text, **parameters)
+                        return text
 
-        return self._transform_output("", **parameters)
+        return ""
 
     def _stream_class(self) -> type[TextStream]:
         """Return the Stream class for this provider."""

src/celeste/modalities/text/providers/openresponses/client.py

@@ -115,9 +115,9 @@ def _parse_content(
                 for part in item.get("content", []):
                     if part.get("type") == "output_text":
                         text = part.get("text") or ""
-                        return self._transform_output(text, **parameters)
+                        return text
 
-        return self._transform_output("", **parameters)
+        return ""
 
     def _stream_class(self) -> type[TextStream]:
         """Return the Stream class for this provider."""

src/celeste/modalities/text/providers/xai/client.py

@@ -87,9 +87,9 @@ def _parse_content(
                 for part in item.get("content", []):
                     if part.get("type") == "output_text":
                         text = part.get("text") or ""
-                        return self._transform_output(text, **parameters)
+                        return text
 
-        return self._transform_output("", **parameters)
+        return ""
 
     def _stream_class(self) -> type[TextStream]:
         """Return the Stream class for this provider."""

src/celeste/parameters.py

@@ -5,7 +5,6 @@
 from typing import Any, ClassVar, TypedDict
 
 from celeste.models import Model
-from celeste.types import TextContent
 
 
 class Parameters(TypedDict, total=False):
@@ -32,7 +31,7 @@ def map(self, request: dict[str, Any], value: Any, model: Model) -> dict[str, An
         """
         ...
 
-    def parse_output(self, content: TextContent, value: object | None) -> TextContent:
+    def parse_output(self, content: Any, value: object | None) -> Any:  # noqa: ANN401
         """Optionally transform parsed content based on parameter value (default: return unchanged)."""
         return content
 

src/celeste/providers/anthropic/messages/client.py

@@ -32,7 +32,7 @@ def _parse_content(self, response_data, **parameters):
                 content = super()._parse_content(response_data)  # Raw content array
                 for block in content:
                     if block.get("type") == "text":
-                        return self._transform_output(block.get("text") or "", **parameters)
+                        return block.get("text") or ""
                 return ""
     """
 

src/celeste/providers/cohere/chat/client.py

@@ -26,7 +26,7 @@ class CohereTextGenerationClient(CohereChatClient, TextGenerationClient):
             def _parse_content(self, response_data, **parameters):
                 content_array = super()._parse_content(response_data)
                 text = content_array[0].get("text") or ""
-                return self._transform_output(text, **parameters)
+                return text
     """
 
     _content_fields: ClassVar[set[str]] = {"message"}

src/celeste/providers/google/generate_content/client.py

@@ -33,7 +33,7 @@ class GoogleTextGenerationClient(GoogleGenerateContentClient, TextGenerationClie
             def _parse_content(self, response_data, **parameters):
                 parts = super()._parse_content(response_data)
                 text = parts[0].get("text") or ""
-                return self._transform_output(text, **parameters)
+                return text
     """
 
     _content_fields: ClassVar[set[str]] = {"candidates"}

src/celeste/providers/google/interactions/client.py

@@ -33,7 +33,7 @@ class GoogleInteractionsTextGenerationClient(
             def _parse_content(self, response_data, **parameters):
                 outputs = super()._parse_content(response_data)
                 text = "".join(o.get("text", "") for o in outputs if o.get("type") == "text")
-                return self._transform_output(text, **parameters)
+                return text
     """
 
     _content_fields: ClassVar[set[str]] = {"outputs"}

tests/unit_tests/test_client.py

@@ -264,6 +264,56 @@ def parameter_mappers(cls) -> list[ParameterMapper]:
         assert transformed == expected_output
 
 
+class TestPredictTransformOutput:
+    """Test that _predict applies _transform_output after _parse_content."""
+
+    @pytest.mark.asyncio
+    async def test_predict_applies_transform_output(
+        self, text_model: Model, api_key: str
+    ) -> None:
+        """_predict applies _transform_output to content from _parse_content."""
+
+        class ClientWithTransformMapper(ConcreteModalityClient):
+            @classmethod
+            def parameter_mappers(cls) -> list[ParameterMapper]:
+                return [_create_transform_mapper(ParamEnum.TEST_PARAM)]
+
+        client = ClientWithTransformMapper(
+            modality=Modality.TEXT,
+            model=text_model,
+            provider=text_model.provider,
+            auth=APIKey(secret=SecretStr(api_key)),
+        )
+
+        output = await client._predict(
+            _TestInput(prompt="test"), test_param="test_value"
+        )
+
+        assert output.content == "test content_transformed_with_test_value"
+
+    @pytest.mark.asyncio
+    async def test_predict_no_transform_without_param(
+        self, text_model: Model, api_key: str
+    ) -> None:
+        """_predict returns raw content when no transform param is provided."""
+
+        class ClientWithTransformMapper(ConcreteModalityClient):
+            @classmethod
+            def parameter_mappers(cls) -> list[ParameterMapper]:
+                return [_create_transform_mapper(ParamEnum.TEST_PARAM)]
+
+        client = ClientWithTransformMapper(
+            modality=Modality.TEXT,
+            model=text_model,
+            provider=text_model.provider,
+            auth=APIKey(secret=SecretStr(api_key)),
+        )
+
+        output = await client._predict(_TestInput(prompt="test"))
+
+        assert output.content == "test content"
+
+
 class TestHandleErrorResponse:
     """Test ModalityClient._handle_error_response edge cases."""