refactor(text-generation): use StructuredOutput and improve API

[Kamilbenkirane]

Summary

Refactors text generation API to use StructuredOutput return types and improves code organization.

Changes

• 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

Testing

• All pre-commit hooks passed
• Type checking passed
• Tests passed

[claude]

Code Review

Thank you for this refactoring PR! I've reviewed the changes and have the following feedback:

✅ Positive Changes

1. Type Safety Improvement: Changing _parse_content return type from str | BaseModel to StructuredOutput is excellent. This aligns with the framework's type system (StructuredOutput = str | JsonValue | BaseModel | list[BaseModel]) and provides better type coverage.

2. Code Organization: Consolidating the anthropic parameter imports is a good cleanup - the original multi-line imports were unnecessarily verbose.

3. Test Improvements: Moving model parameter to client initialization is the right approach and follows better separation of concerns.

⚠️ Critical Issue: Breaking Change

Problem: The changes to _parse_finish_reason and _parse_content introduce a breaking change that affects multiple provider implementations.

In client.py:44, you changed:

def _parse_finish_reason(self, response_data: dict[str, Any]) -> TextGenerationFinishReason:

However, several provider implementations still return TextGenerationFinishReason | None:

1. OpenAI provider (providers/openai/client.py:85-98):

   def _parse_finish_reason(self, response_data: dict[str, Any]) -> TextGenerationFinishReason | None:
       # Returns None in multiple code paths
       return None

2. Cohere provider (providers/cohere/client.py:81-88):

   def _parse_finish_reason(self, response_data: dict[str, Any]) -> TextGenerationFinishReason | None:
       return TextGenerationFinishReason(reason=finish_reason_str) if finish_reason_str else None

Impact: This creates a type mismatch that will:

• Cause mypy/type checker failures
• Create runtime issues when these providers return None
• Break the contract defined in the base class

Also note: The base Client class at src/celeste/client.py:129-137 has _parse_finish_reason returning FinishReason | None, which suggests that None is a valid return value for cases where finish reason isn't available.

📝 Recommendations

Option 1 (Recommended): Keep _parse_finish_reason return type as optional

def _parse_finish_reason(
    self, response_data: dict[str, Any]
) -> TextGenerationFinishReason | None:

Option 2: Update ALL provider implementations to never return None (requires changes to OpenAI, Cohere, and potentially other providers).

🧪 Testing Concern

The PR description mentions "Tests passed," but the changes affect multiple providers. Please confirm:

1. Did integration tests run for all providers (OpenAI, Cohere, Google, Mistral, xAI, Anthropic)?
2. Did type checking pass with the new signature?

🔍 Minor: Code Organization Question

In client.py:53:

prompt: str | None = parameters.get("prompt")

This extracts prompt from the parameters dict. While this works, it assumes prompt exists in the TypedDict. Consider whether this should use .pop() to remove it from parameters after extraction, or if the current approach is intentional for the parameter flow.

Summary

The refactoring direction is good, but the _parse_finish_reason change needs to be reverted or applied consistently across all providers to avoid type safety issues and potential runtime errors.

Please address the breaking change concern before merging. Happy to review again once updated!

[claude]

Code Review for PR #76

I've reviewed the refactoring changes to the text generation API. Overall, this is a solid refactoring that improves type safety and API consistency. Here are my findings:

✅ Strengths

1. Improved Type Safety: The change from str | BaseModel to StructuredOutput in _parse_content (packages/capabilities/text-generation/src/celeste_text_generation/client.py:38) is more explicit and aligns with the actual type definition in celeste.types.

2. Better API Design: Moving the model parameter to client initialization (test_generate.py:34, test_stream.py:37) is a good improvement - it makes sense for the model to be a client property rather than a per-request parameter.

3. Cleaner Parameter Handling: The refactored _create_inputs (client.py:47-59) that extracts prompt from the parameters dict is cleaner and removes redundancy.

4. Consistent Import Style: Consolidating the anthropic parameter imports (parameters.py:5-10) improves readability.

⚠️ Issues Found

CRITICAL: Type Signature Inconsistency

There's a breaking change in the _parse_finish_reason return type that creates inconsistency across provider implementations:

In client.py:42-44, the abstract method now requires non-optional TextGenerationFinishReason:

def _parse_finish_reason(
    self, response_data: dict[str, Any]
) -> TextGenerationFinishReason:

However, other provider implementations still return TextGenerationFinishReason | None:

• OpenAI client (openai/client.py:85): -> TextGenerationFinishReason | None
• Google client (google/client.py:79): -> TextGenerationFinishReason | None
• Mistral client (mistral/client.py:79): -> TextGenerationFinishReason | None
• Cohere client (cohere/client.py:81): -> TextGenerationFinishReason | None
• xAI client (xai/client.py:74): -> TextGenerationFinishReason | None

Only the Anthropic client was updated to match the new signature (anthropic/client.py:55).

Impact: This creates a type safety issue where:

1. The abstract base class declares a non-optional return type
2. Multiple concrete implementations return optional types
3. Type checkers would flag this as a Liskov Substitution Principle violation

Recommendation: Either:

• Revert the abstract method to -> TextGenerationFinishReason | None to match existing implementations
• Update ALL provider implementations to ensure finish_reason is always returned (may require default values when providers don't supply one)
• Verify this aligns with the base Client._parse_finish_reason in src/celeste/client.py:129-137 which returns FinishReason | None

Minor: Incomplete Refactoring

The OpenAI client (openai/client.py:53) still has the old return type:

def _parse_content(...) -> str | BaseModel:

This should be updated to -> StructuredOutput for consistency.

🔍 Questions/Clarifications

1. Docstring Update: The module docstring change in test_text_generation/__init__.py from "Integration tests for text generation capability." to "Text generation integration test module." seems unnecessary. Was this intentional?

2. Type Safety Verification: Did the pre-commit type checking pass with the _parse_finish_reason inconsistency? If so, you might want to enable stricter mypy settings (e.g., --strict or --disallow-untyped-overrides).

📋 Test Coverage

The test changes look appropriate:

• Moving model to client initialization is covered by existing tests
• The tests still verify the same behavior

However, consider adding tests that verify:

• The new StructuredOutput type is properly handled
• Finish reason is always present (if that's the new contract)

🔒 Security & Performance

No security concerns identified. The changes are purely structural and don't affect data validation or error handling.

Performance is neutral - no algorithmic changes.

✨ Recommendations

1. Fix the type inconsistency with _parse_finish_reason across all providers
2. Update OpenAI client to use StructuredOutput return type
3. Add migration notes if the finish_reason change from optional to required is intentional API behavior
4. Consider stricter type checking to catch these inconsistencies earlier

Let me know if you'd like me to help identify all the files that need updating for consistency!

[claude]

PR Review: refactor(text-generation): use StructuredOutput and improve API

Overview

This PR refactors the text generation API to use StructuredOutput return types and improves code organization. Overall, the changes are well-structured and improve type consistency across the codebase.

---

✅ Strengths

1. Improved Type Safety: The change from str | BaseModel to StructuredOutput in _parse_content provides better type consistency across the codebase, since StructuredOutput is a defined type alias that includes these possibilities plus more (lines 38 in client.py).

2. Better API Design: Moving the model parameter to client initialization in tests is a good architectural improvement - it aligns with the principle that model selection is a client configuration concern, not a per-request parameter.

3. Code Organization: Consolidating the anthropic parameter imports is cleaner and more maintainable (parameters.py:8-13).

4. Consistency: The refactoring makes the return types more consistent with the _transform_output method which already returns object type (client.py:206-216).

---

⚠️ Issues & Concerns

1. Type Inconsistency Across Provider Implementations (High Priority)

Issue: While the abstract method in TextGenerationClient._parse_content now returns StructuredOutput (client.py:38), not all provider implementations have been updated:

• ✅ AnthropicTextGenerationClient._parse_content - Returns StructuredOutput
• ❌ OpenAITextGenerationClient._parse_content - Still returns str | BaseModel (line 53)
• ❌ CohereTextGenerationClient._parse_content - Still returns str | BaseModel (line 66)
• ❌ GoogleTextGenerationClient._parse_content - Still returns str | BaseModel (line 57)

Impact: This creates a type mismatch that may not be caught by type checkers depending on configuration, and violates the Liskov Substitution Principle.

Recommendation: Update all provider implementations to return StructuredOutput. Example:

# In openai/client.py, cohere/client.py, google/client.py, etc.
from celeste.types import StructuredOutput

def _parse_content(
    self,
    response_data: dict[str, Any],
    **parameters: Unpack[TextGenerationParameters],
) -> StructuredOutput:  # Changed from str | BaseModel
    """Parse content from response."""
    # ... existing implementation

---

2. Non-Optional Return Type May Break Existing Behavior (Medium Priority)

Issue: The _parse_finish_reason method signature changed from returning TextGenerationFinishReason | None to just TextGenerationFinishReason (client.py:42-44).

However:

• The base Client class default implementation returns None (client.py:130-138)
• Some provider implementations still return None:
  • OpenAITextGenerationClient._parse_finish_reason returns TextGenerationFinishReason | None (line 85-98)
  • CohereTextGenerationClient._parse_finish_reason returns TextGenerationFinishReason | None (line 81-88)
  • GoogleTextGenerationClient._parse_finish_reason returns TextGenerationFinishReason | None (line 79-91)

Impact: Type mismatch between abstract method and implementations. The base class implementation (client.py:68-73) expects the return value to be used directly, so returning None could cause issues.

Recommendation: Either:

1. Revert to TextGenerationFinishReason | None in the abstract method, OR
2. Update all implementations to guarantee non-None returns (e.g., return a default/unknown finish reason instead of None)

---

3. Subtle Breaking Change in _create_inputs (Low Priority)

Issue: The refactored _create_inputs method (client.py:47-59) changes how prompt is extracted:

Before:

def _create_inputs(
    self,
    *args: str,
    prompt: str | None = None,  # explicit parameter
    **parameters: Unpack[TextGenerationParameters],
) -> TextGenerationInput:

After:

def _create_inputs(
    self, *args: str, **parameters: Unpack[TextGenerationParameters]
) -> TextGenerationInput:
    if args:
        return TextGenerationInput(prompt=args[0])
    prompt: str | None = parameters.get("prompt")  # extracted from dict

Impact: This is functionally equivalent but changes the method signature. If any code was relying on the explicit prompt parameter in the signature (e.g., for inspection or documentation generation), this could be a breaking change.

Note: The change is cleaner since prompt is part of TextGenerationParameters, but it's worth documenting.

---

🧪 Test Coverage

Positive: The integration tests have been updated to move the model parameter to client initialization, which validates that the new API works correctly.

Concern: No unit tests appear to directly test the _parse_content and _parse_finish_reason methods. Consider adding unit tests that:

1. Verify the return type of _parse_content matches StructuredOutput
2. Test edge cases for _parse_finish_reason (especially when providers return different or missing finish reasons)
3. Validate that _create_inputs handles both positional and keyword prompt arguments correctly

---

🔒 Security

No security concerns identified. The changes are refactoring-focused and don't introduce new attack vectors.

---

⚡ Performance

No performance impact expected. The changes are primarily type-level refactoring with equivalent runtime behavior.

---

📝 Code Quality

1. Import Organization: The import consolidation in parameters.py is excellent and follows Python best practices.

2. Docstring Update: The docstring for test_generate.py module changed from "Integration tests for text generation capability." to "Text generation integration test module." - this is fine but seems unnecessary.

---

🎯 Recommendations Summary

Must Fix (Before Merge):

1. Update all provider _parse_content implementations to return StructuredOutput
2. Resolve the _parse_finish_reason return type inconsistency

Should Fix:
3. Add unit tests for the refactored methods
4. Consider documenting the signature change in _create_inputs

Nice to Have:
5. Run mypy or similar type checker in strict mode to catch these inconsistencies automatically in CI

---

✨ Overall Assessment

This is a solid refactoring that improves type consistency and API design. The core changes are well-thought-out, but the PR is incomplete as it doesn't update all provider implementations to match the new abstract method signatures. Once the type inconsistencies across providers are resolved, this will be a valuable improvement to the codebase.

Recommendation: Request changes to address the type inconsistencies before merging.

[claude]

Pull Request Review: refactor(text-generation): use StructuredOutput and improve API

Overview

This PR makes several improvements to the text-generation API, including type safety enhancements and test organization. Overall, the changes are well-structured and improve code quality.

✅ Strengths

1. Type Safety Improvement: The _parse_content return type change from str | BaseModel to StructuredOutput is excellent. This leverages the centralized type alias from celeste.types and makes the type system more maintainable.

2. Test Organization: Moving the model parameter to client initialization (lines test_generate.py:34, test_stream.py:40) is the correct approach. This aligns with the create_client API design where model selection happens at client creation time.

3. Code Cleanup: Removing the now-unused pydantic.BaseModel import (client.py:7) keeps imports clean.

4. Parameter Extraction: The refactoring of _create_inputs (client.py:47-53) to extract prompt from the parameters dict is cleaner and more consistent.

⚠️ Critical Issue Found

_parse_finish_reason Return Type Change

Location: packages/capabilities/text-generation/src/celeste_text_generation/client.py:42-44

The change from TextGenerationFinishReason | None to TextGenerationFinishReason (non-optional) introduces a contract violation with existing provider implementations.

Problem: 5 out of 6 text-generation providers currently return None in certain cases:

• OpenAI: Returns None when status != "completed" (client.py:88, 98)
• Google: Returns None when no candidates or no finishReason (client.py:83, 89)
• Mistral: Returns None in certain conditions (client.py:83, 90)
• Cohere: Returns None when finish_reason is missing (client.py:84, 87)
• XAI: Returns None in certain conditions (client.py:78, 84)

Only Anthropic consistently returns a non-None TextGenerationFinishReason.

Impact:

• The abstract method contract now requires a non-None return value
• Provider implementations still return None in their type signatures
• This creates a type inconsistency that mypy should catch, but may not be enforced at runtime
• The base Client.generate() method (celeste/client.py:70) passes this value directly to the Output constructor, which expects FinishReason | None

Recommendation:
You have two options:

1. Keep the optional type (revert this change):

   def _parse_finish_reason(
       self, response_data: dict[str, Any]
   ) -> TextGenerationFinishReason | None:

2. Update all providers to always return a non-None value (e.g., with a default reason like "unknown" when the provider doesn't supply one)

I recommend option 1 unless there's a strong reason to guarantee non-None finish reasons, as some providers genuinely don't always provide this information.

Minor Observations

1. Docstring changes: The test module docstring change (test_text_generation/init.py:1) is purely cosmetic. Consider whether this adds value.

Testing

✅ The PR author reports:

• All pre-commit hooks passed
• Type checking passed
• Tests passed

However, the type checking passing is surprising given the finish reason issue. Verify that mypy is properly checking the provider implementations against the abstract base class.

Security & Performance

✅ No security concerns identified
✅ No performance regressions expected
✅ Changes are refactoring-only with no behavioral modifications

Verdict

Request Changes: The _parse_finish_reason return type change needs to be addressed before merging. All other changes look good!

---

📍 Key files reviewed:

• packages/capabilities/text-generation/src/celeste_text_generation/client.py:34-44
• packages/capabilities/text-generation/tests/integration_tests/test_text_generation/test_generate.py:31-43
• packages/capabilities/text-generation/tests/integration_tests/test_text_generation/test_stream.py:37-45