From b22c21fcd49f1f3dfecf697a7426ab4fc773c2b9 Mon Sep 17 00:00:00 2001 From: kamilbenkirane Date: Mon, 15 Dec 2025 14:02:22 +0100 Subject: [PATCH 1/4] refactor(text-generation): use StructuredOutput and improve API - Update _parse_content to return StructuredOutput instead of str | BaseModel - Make _parse_finish_reason return non-optional TextGenerationFinishReason - Refactor _create_inputs to extract prompt from parameters dict - Consolidate anthropic parameter imports - Move model parameter to client initialization in tests --- .../src/celeste_text_generation/client.py | 12 +++++------- .../providers/anthropic/parameters.py | 6 ------ .../test_text_generation/__init__.py | 2 +- .../test_text_generation/test_generate.py | 2 +- .../test_text_generation/test_stream.py | 2 +- 5 files changed, 8 insertions(+), 16 deletions(-) diff --git a/packages/capabilities/text-generation/src/celeste_text_generation/client.py b/packages/capabilities/text-generation/src/celeste_text_generation/client.py index 4b08856c..03cc2db6 100644 --- a/packages/capabilities/text-generation/src/celeste_text_generation/client.py +++ b/packages/capabilities/text-generation/src/celeste_text_generation/client.py @@ -4,10 +4,10 @@ from typing import Any, Unpack import httpx -from pydantic import BaseModel from celeste.client import Client from celeste.exceptions import ValidationError +from celeste.types import StructuredOutput from celeste_text_generation.io import ( TextGenerationFinishReason, TextGenerationInput, @@ -35,24 +35,22 @@ def _parse_content( self, response_data: dict[str, Any], **parameters: Unpack[TextGenerationParameters], - ) -> str | BaseModel: + ) -> StructuredOutput: """Parse content from provider response.""" @abstractmethod def _parse_finish_reason( self, response_data: dict[str, Any] - ) -> TextGenerationFinishReason | None: + ) -> TextGenerationFinishReason: """Parse finish reason from provider response.""" def _create_inputs( - self, - *args: str, - prompt: str | None = None, - **parameters: Unpack[TextGenerationParameters], + self, *args: str, **parameters: Unpack[TextGenerationParameters] ) -> TextGenerationInput: """Map positional arguments to Input type.""" if args: return TextGenerationInput(prompt=args[0]) + prompt: str | None = parameters.get("prompt") if prompt is None: msg = ( "prompt is required (either as positional argument or keyword argument)" diff --git a/packages/capabilities/text-generation/src/celeste_text_generation/providers/anthropic/parameters.py b/packages/capabilities/text-generation/src/celeste_text_generation/providers/anthropic/parameters.py index 4cd0ee9c..6d51e15e 100644 --- a/packages/capabilities/text-generation/src/celeste_text_generation/providers/anthropic/parameters.py +++ b/packages/capabilities/text-generation/src/celeste_text_generation/providers/anthropic/parameters.py @@ -4,14 +4,8 @@ from celeste_anthropic.messages.parameters import ( MaxTokensMapper as _MaxTokensMapper, -) -from celeste_anthropic.messages.parameters import ( OutputSchemaMapper as _OutputSchemaMapper, -) -from celeste_anthropic.messages.parameters import ( TemperatureMapper as _TemperatureMapper, -) -from celeste_anthropic.messages.parameters import ( ThinkingMapper as _ThinkingMapper, ) diff --git a/packages/capabilities/text-generation/tests/integration_tests/test_text_generation/__init__.py b/packages/capabilities/text-generation/tests/integration_tests/test_text_generation/__init__.py index f264f567..432e7b47 100644 --- a/packages/capabilities/text-generation/tests/integration_tests/test_text_generation/__init__.py +++ b/packages/capabilities/text-generation/tests/integration_tests/test_text_generation/__init__.py @@ -1 +1 @@ -"""Integration tests for text generation capability.""" +"""Text generation integration test module.""" diff --git a/packages/capabilities/text-generation/tests/integration_tests/test_text_generation/test_generate.py b/packages/capabilities/text-generation/tests/integration_tests/test_text_generation/test_generate.py index d3f57e15..2a06f34d 100644 --- a/packages/capabilities/text-generation/tests/integration_tests/test_text_generation/test_generate.py +++ b/packages/capabilities/text-generation/tests/integration_tests/test_text_generation/test_generate.py @@ -31,6 +31,7 @@ async def test_generate(provider: Provider, model: str, parameters: dict) -> Non client = create_client( capability=Capability.TEXT_GENERATION, provider=provider, + model=model, ) prompt = "Hi" max_tokens = 30 @@ -38,7 +39,6 @@ async def test_generate(provider: Provider, model: str, parameters: dict) -> Non # Act response = await client.generate( prompt=prompt, - model=model, max_tokens=max_tokens, **parameters, ) diff --git a/packages/capabilities/text-generation/tests/integration_tests/test_text_generation/test_stream.py b/packages/capabilities/text-generation/tests/integration_tests/test_text_generation/test_stream.py index b67b1ad6..f0a6e760 100644 --- a/packages/capabilities/text-generation/tests/integration_tests/test_text_generation/test_stream.py +++ b/packages/capabilities/text-generation/tests/integration_tests/test_text_generation/test_stream.py @@ -37,6 +37,7 @@ async def test_stream(provider: Provider, model: str, parameters: dict) -> None: client = create_client( capability=Capability.TEXT_GENERATION, provider=provider, + model=model, ) prompt = "Hi" max_tokens = 30 @@ -44,7 +45,6 @@ async def test_stream(provider: Provider, model: str, parameters: dict) -> None: # Act - Create stream stream = client.stream( prompt=prompt, - model=model, max_tokens=max_tokens, **parameters, ) From 1f880c8ed26094bfd87ef412533ccfad9e838167 Mon Sep 17 00:00:00 2001 From: kamilbenkirane Date: Mon, 15 Dec 2025 14:04:05 +0100 Subject: [PATCH 2/4] fix(text-generation): sort imports in anthropic parameters --- .../providers/anthropic/parameters.py | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/packages/capabilities/text-generation/src/celeste_text_generation/providers/anthropic/parameters.py b/packages/capabilities/text-generation/src/celeste_text_generation/providers/anthropic/parameters.py index 6d51e15e..105aabd1 100644 --- a/packages/capabilities/text-generation/src/celeste_text_generation/providers/anthropic/parameters.py +++ b/packages/capabilities/text-generation/src/celeste_text_generation/providers/anthropic/parameters.py @@ -2,16 +2,15 @@ from typing import Any +from celeste.core import Parameter +from celeste.models import Model +from celeste.parameters import ParameterMapper from celeste_anthropic.messages.parameters import ( MaxTokensMapper as _MaxTokensMapper, OutputSchemaMapper as _OutputSchemaMapper, TemperatureMapper as _TemperatureMapper, ThinkingMapper as _ThinkingMapper, ) - -from celeste.core import Parameter -from celeste.models import Model -from celeste.parameters import ParameterMapper from celeste_text_generation.parameters import TextGenerationParameter From 33fccedd1789e1383d029e8c5041d8c443b20a55 Mon Sep 17 00:00:00 2001 From: kamilbenkirane Date: Mon, 15 Dec 2025 14:05:07 +0100 Subject: [PATCH 3/4] fix(text-generation): reorder imports alphabetically for anthropic parameters --- .../providers/anthropic/parameters.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/packages/capabilities/text-generation/src/celeste_text_generation/providers/anthropic/parameters.py b/packages/capabilities/text-generation/src/celeste_text_generation/providers/anthropic/parameters.py index 105aabd1..780e31c4 100644 --- a/packages/capabilities/text-generation/src/celeste_text_generation/providers/anthropic/parameters.py +++ b/packages/capabilities/text-generation/src/celeste_text_generation/providers/anthropic/parameters.py @@ -2,15 +2,15 @@ from typing import Any -from celeste.core import Parameter -from celeste.models import Model -from celeste.parameters import ParameterMapper from celeste_anthropic.messages.parameters import ( MaxTokensMapper as _MaxTokensMapper, OutputSchemaMapper as _OutputSchemaMapper, TemperatureMapper as _TemperatureMapper, ThinkingMapper as _ThinkingMapper, ) +from celeste.core import Parameter +from celeste.models import Model +from celeste.parameters import ParameterMapper from celeste_text_generation.parameters import TextGenerationParameter From 38b41d30216b41bcbbaf8dd536dee05d2b2c0b58 Mon Sep 17 00:00:00 2001 From: kamilbenkirane Date: Mon, 15 Dec 2025 14:06:12 +0100 Subject: [PATCH 4/4] fix(text-generation): consolidate anthropic parameter imports --- .../providers/anthropic/parameters.py | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/packages/capabilities/text-generation/src/celeste_text_generation/providers/anthropic/parameters.py b/packages/capabilities/text-generation/src/celeste_text_generation/providers/anthropic/parameters.py index 780e31c4..4cd0ee9c 100644 --- a/packages/capabilities/text-generation/src/celeste_text_generation/providers/anthropic/parameters.py +++ b/packages/capabilities/text-generation/src/celeste_text_generation/providers/anthropic/parameters.py @@ -4,10 +4,17 @@ from celeste_anthropic.messages.parameters import ( MaxTokensMapper as _MaxTokensMapper, +) +from celeste_anthropic.messages.parameters import ( OutputSchemaMapper as _OutputSchemaMapper, +) +from celeste_anthropic.messages.parameters import ( TemperatureMapper as _TemperatureMapper, +) +from celeste_anthropic.messages.parameters import ( ThinkingMapper as _ThinkingMapper, ) + from celeste.core import Parameter from celeste.models import Model from celeste.parameters import ParameterMapper