diff --git a/src/celeste/client.py b/src/celeste/client.py index d4c220b8..213fc5db 100644 --- a/src/celeste/client.py +++ b/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) diff --git a/src/celeste/modalities/text/providers/anthropic/client.py b/src/celeste/modalities/text/providers/anthropic/client.py index d9f46554..e28b289d 100644 --- a/src/celeste/modalities/text/providers/anthropic/client.py +++ b/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.""" diff --git a/src/celeste/modalities/text/providers/cohere/client.py b/src/celeste/modalities/text/providers/cohere/client.py index 08f0d459..cdfb793f 100644 --- a/src/celeste/modalities/text/providers/cohere/client.py +++ b/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.""" diff --git a/src/celeste/modalities/text/providers/deepseek/client.py b/src/celeste/modalities/text/providers/deepseek/client.py index 38c25fac..40e6010e 100644 --- a/src/celeste/modalities/text/providers/deepseek/client.py +++ b/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.""" diff --git a/src/celeste/modalities/text/providers/google/client.py b/src/celeste/modalities/text/providers/google/client.py index 2ab31afa..0d3f0236 100644 --- a/src/celeste/modalities/text/providers/google/client.py +++ b/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.""" diff --git a/src/celeste/modalities/text/providers/groq/client.py b/src/celeste/modalities/text/providers/groq/client.py index 031ba819..b48ea373 100644 --- a/src/celeste/modalities/text/providers/groq/client.py +++ b/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.""" diff --git a/src/celeste/modalities/text/providers/huggingface/client.py b/src/celeste/modalities/text/providers/huggingface/client.py index 8c47a959..c403c07d 100644 --- a/src/celeste/modalities/text/providers/huggingface/client.py +++ b/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.""" diff --git a/src/celeste/modalities/text/providers/mistral/client.py b/src/celeste/modalities/text/providers/mistral/client.py index ea03a552..2a30fc88 100644 --- a/src/celeste/modalities/text/providers/mistral/client.py +++ b/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.""" diff --git a/src/celeste/modalities/text/providers/moonshot/client.py b/src/celeste/modalities/text/providers/moonshot/client.py index f9073c5f..165d72e3 100644 --- a/src/celeste/modalities/text/providers/moonshot/client.py +++ b/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.""" diff --git a/src/celeste/modalities/text/providers/openai/client.py b/src/celeste/modalities/text/providers/openai/client.py index 5be04f1c..51156916 100644 --- a/src/celeste/modalities/text/providers/openai/client.py +++ b/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.""" diff --git a/src/celeste/modalities/text/providers/openresponses/client.py b/src/celeste/modalities/text/providers/openresponses/client.py index 90b8714d..0185c1fc 100644 --- a/src/celeste/modalities/text/providers/openresponses/client.py +++ b/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.""" diff --git a/src/celeste/modalities/text/providers/xai/client.py b/src/celeste/modalities/text/providers/xai/client.py index 3a1c3eba..91f0ad07 100644 --- a/src/celeste/modalities/text/providers/xai/client.py +++ b/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.""" diff --git a/src/celeste/parameters.py b/src/celeste/parameters.py index 3be12812..fa5fcac9 100644 --- a/src/celeste/parameters.py +++ b/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 diff --git a/src/celeste/providers/anthropic/messages/client.py b/src/celeste/providers/anthropic/messages/client.py index 634ab01e..bf7298d0 100644 --- a/src/celeste/providers/anthropic/messages/client.py +++ b/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 "" """ diff --git a/src/celeste/providers/cohere/chat/client.py b/src/celeste/providers/cohere/chat/client.py index 9f3576b5..7f8d08ae 100644 --- a/src/celeste/providers/cohere/chat/client.py +++ b/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"} diff --git a/src/celeste/providers/google/generate_content/client.py b/src/celeste/providers/google/generate_content/client.py index d537763d..734c373e 100644 --- a/src/celeste/providers/google/generate_content/client.py +++ b/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"} diff --git a/src/celeste/providers/google/interactions/client.py b/src/celeste/providers/google/interactions/client.py index 9b095e90..d490ef6b 100644 --- a/src/celeste/providers/google/interactions/client.py +++ b/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"} diff --git a/tests/unit_tests/test_client.py b/tests/unit_tests/test_client.py index 7364b939..87e4712a 100644 --- a/tests/unit_tests/test_client.py +++ b/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."""