feat: Add XAI provider for text generation

[Kamilbenkirane]

Add XAI provider for text generation.

Changes

Added XAI provider with support for 6 Grok models:

• grok-4-1-fast-reasoning: Max tokens 30,000, reasoning capabilities, streaming support, output schema support
• grok-4-1-fast-non-reasoning: Max tokens 30,000, streaming support, output schema support
• grok-4-fast-reasoning: Max tokens 30,000, reasoning capabilities, streaming support, output schema support
• grok-4-fast-non-reasoning: Max tokens 30,000, streaming support, output schema support
• grok-4-0709: Max tokens 64,000, streaming support, output schema support
• grok-3-mini: Max tokens 16,000, thinking_level parameter (low/high), streaming support, output schema support

All models support temperature (0.0-2.0), max_tokens, and output_schema parameters. Integration tests added to test_generate.py and test_stream.py using grok-3-mini.

[claude]

Pull Request Review: Add XAI Provider for Text Generation

Overview

This PR successfully adds XAI (xAI) provider support for text generation with 6 Grok models. The implementation follows the established patterns in the codebase and maintains consistency with other providers like Mistral and Cohere. Overall, this is high-quality work with good attention to detail.

✅ Strengths

Code Quality

• Excellent consistency: The implementation closely follows established patterns from Mistral/Cohere providers
• Clean architecture: Proper separation of concerns across client.py, parameters.py, streaming.py, models.py, and config.py
• Comprehensive documentation: Thorough PR description with examples, model specifications, and parameter details
• Good type hints: Proper use of type annotations throughout

Implementation Highlights

• Robust streaming support: Well-implemented SSE parsing with proper chunk handling
• Usage tracking: Comprehensive token usage including reasoning tokens
• Structured output: Full support for JSON schema with output_schema parameter
• Parameter mapping: Clean mapper implementations for temperature, max_tokens, thinking_level, and output_schema

Testing

• Integration tests added: XAI properly added to both test_generate.py and test_stream.py
• Cost-effective testing: Smart use of grok-3-mini for tests

🔍 Issues & Concerns

1. JSON Parsing Error Handling (Medium Priority)

Location: parameters.py:49 - No error handling for malformed JSON responses. If the API returns invalid JSON, this will raise an unhandled json.JSONDecodeError.

Recommendation: Add try-except block with informative error message

2. Schema Transformation Edge Cases (Low-Medium Priority)

Location: parameters.py:78-117 (_transform_schema method)

Questions:

1. Has this been tested with deeply nested schemas with circular references?
2. Does XAI API require additionalProperties: false on all objects?
3. What happens if a schema has a $ref pointing to an external URI?

Recommendation: Add edge case tests for complex nested schemas and document the additionalProperties behavior

3. List Schema Handling (Low Priority)

Location: parameters.py:51-52 - The list schema wrapping/unwrapping logic assumes XAI wraps responses in {"items": [...]} format. This coupling is implicit.

Recommendation: Add validation and error handling for unexpected formats

4. Empty Content Handling (Low Priority)

Location: client.py:69 - Silent conversion of None to empty string may hide API issues

5. Missing Error Context (Low Priority)

Location: client.py:64-66 - Error messages could include more context from response_data for debugging

🔒 Security Review

✅ No security concerns identified:

• API key properly handled via SecretStr with .get_secret_value()
• Bearer token authentication correctly implemented
• No hardcoded credentials
• No injection vulnerabilities (parameters are passed as JSON)
• Proper use of httpx client for requests

🚀 Performance Considerations

✅ Generally good:

• Streaming properly implemented with async iterators
• No unnecessary data copying or transformations
• Schema transformation happens once per request

⚠️ Minor concern: The _transform_schema method recursively processes the entire schema on every request. For large schemas, consider caching.

📊 Test Coverage

Current state: ✅ Basic integration tests added

Gaps:

1. No unit tests for the complex OutputSchemaMapper._transform_schema logic
2. No tests for error conditions (malformed JSON, missing fields)
3. No tests for the thinking_level parameter with grok-3-mini
4. No tests for structured output with nested schemas

Recommendation: Add unit tests for parameter mappers with edge cases

🎯 Best Practices Compliance

✅ Follows codebase conventions:

• Consistent with other provider implementations
• Proper module structure and exports
• Correct use of abstract base classes
• Type hints throughout

📝 Documentation

✅ Excellent PR description: Comprehensive with model specs, parameters, and examples

⚠️ Missing: Inline docstrings could be more detailed for complex methods like _transform_schema

🎉 Verdict

Status: ✅ Approve with minor suggestions

This is solid work that maintains the high quality of the codebase. The implementation is clean, follows established patterns, and provides comprehensive functionality. The issues identified are mostly edge cases and opportunities for improvement rather than blocking problems.

Priority Action Items (Before Merge):

1. Add error handling for JSON parsing in OutputSchemaMapper.parse_output() (Medium priority)
2. Add inline documentation for the schema transformation logic (Low priority)

Follow-up Tasks (Can be separate PR):

1. Add unit tests for OutputSchemaMapper with complex schemas
2. Add test coverage for thinking_level parameter
3. Test error conditions and edge cases

Great work overall! 🚀

[claude]

Pull Request Review: XAI Provider Implementation

Summary

This PR adds XAI (Grok) provider support for text generation, implementing 6 Grok models with comprehensive features including streaming, structured output, reasoning capabilities, and parameter mappings. The implementation follows the established patterns in the codebase and integrates well with the existing architecture.

✅ Strengths

Code Quality

• Consistent architecture: The implementation follows the established provider pattern perfectly (client.py, parameters.py, models.py, streaming.py, config.py)
• Type safety: Proper use of type hints throughout, including Unpack[TextGenerationParameters] and generic types
• Documentation: Well-documented classes and methods with clear docstrings
• Error handling: Appropriate error handling in _parse_content and _parse_chunk methods

Implementation Quality

• Parameter mapping: Comprehensive mappers for temperature, max_tokens, thinking_level, and output_schema
• Schema transformation: Sophisticated _transform_schema implementation that properly handles $defs and nested structures
• Streaming support: Complete streaming implementation with proper chunk parsing and usage aggregation
• Usage tracking: Detailed usage metrics including cached_tokens and reasoning_tokens

🔍 Issues & Concerns

1. Critical: Missing XAI Parameter Validation ⚠️

Location: packages/text-generation/src/celeste_text_generation/providers/xai/parameters.py:30

The OutputSchemaMapper duplicates the OpenAI implementation almost exactly, but XAI uses a different API format:

• XAI: response_format.json_schema
• OpenAI: text.format

Issue: The XAI implementation correctly uses response_format (good!), but the method names reference OpenAI:

schema = self._convert_to_openai_schema(validated_value)  # Line 30
json_schema = self._transform_schema_for_openai(json_schema)  # Line 74

Recommendation: Rename methods to be provider-agnostic:

• _convert_to_openai_schema → _convert_to_json_schema
• _transform_schema_for_openai → _transform_schema

This is more of a naming consistency issue than a functional bug, but it could cause confusion.

2. Code Duplication - Schema Transformation Logic

Locations:

• packages/text-generation/src/celeste_text_generation/providers/xai/parameters.py:76-142
• packages/text-generation/src/celeste_text_generation/providers/openai/parameters.py:76-142

The schema transformation logic (_transform_schema, _collect_all_defs, _get_schema_name) is ~140 lines of identical code duplicated between OpenAI and XAI providers.

Impact:

• Maintenance burden - bugs/improvements need fixing in multiple places
• Code bloat - increases package size unnecessarily

Recommendation: Extract shared schema transformation logic to a base class or utility module:

# celeste_text_generation/parameters/schema_utils.py
class JSONSchemaTransformer:
    @staticmethod
    def transform_schema(schema, defs=None):
        # Shared logic
        
    @staticmethod
    def collect_all_defs(schema):
        # Shared logic
        
    @staticmethod
    def get_schema_name(output_schema):
        # Shared logic

3. Test Coverage - Only Basic Integration Tests

Locations:

• packages/text-generation/tests/integration_tests/test_text_generation/test_generate.py:16
• packages/text-generation/tests/integration_tests/test_text_generation/test_stream.py:16

Only basic integration tests added (1 line each). Missing tests for:

• ✗ Output schema validation with complex nested types
• ✗ Reasoning capabilities (grok-4-1-fast-reasoning models)
• ✗ Thinking level parameter (grok-3-mini specific)
• ✗ Error handling (invalid parameters, API errors)
• ✗ Usage metadata parsing (cached_tokens, reasoning_tokens)
• ✗ List[BaseModel] schema transformations

Recommendation: Add comprehensive unit tests for parameter mappers and schema transformations. These don't require API calls and catch edge cases.

4. Minor: Inconsistent Comment Style

Location: packages/text-generation/src/celeste_text_generation/providers/mistral/config.py:6

The PR changes:

-STREAM_ENDPOINT = ENDPOINT  # Same endpoint
+STREAM_ENDPOINT = ENDPOINT

This is in the Mistral provider file - appears to be an unrelated formatting change. While not harmful, it's best to keep PRs focused on the feature being added.

5. Documentation: Missing Provider-Specific Details

The PR description is excellent, but the code lacks comments explaining XAI-specific quirks:

• Why does grok-3-mini have thinking_level instead of thinking_budget?
• Are there any XAI-specific API behaviors to be aware of?
• What's the difference between reasoning and non-reasoning models?

Recommendation: Add comments in models.py or a provider-level docstring explaining XAI's unique characteristics.

🔒 Security Review

✅ No security concerns identified

• API key properly handled via SecretStr and get_secret_value()
• No credential leakage in error messages or logs
• Input validation through Pydantic models
• No injection vulnerabilities in request construction

⚡ Performance Considerations

✅ Good performance practices

• Efficient streaming with SSE iterator
• Minimal object allocations in hot paths
• No unnecessary data copies
• Proper use of or "" for null content handling

Minor optimization opportunity: In streaming.py:117-118, filtering empty chunks could be done during iteration instead of list comprehension:

# Current (creates intermediate list)
content_chunks = [chunk for chunk in chunks if chunk.content]
content = "".join(chunk.content for chunk in content_chunks)

# Alternative (single pass)
content = "".join(chunk.content for chunk in chunks if chunk.content)

📋 Recommendations Summary

Must Fix (before merge):

None - the implementation is functional

Should Fix (high priority):

1. ⚠️ Rename OpenAI-specific method names to be provider-agnostic
2. 📦 Extract duplicated schema transformation logic to shared utility
3. 🧪 Add unit tests for parameter mappers and edge cases

Nice to Have (low priority):

4. 📝 Add provider-specific documentation/comments
5. 🧹 Remove unrelated Mistral config change
6. ⚡ Minor streaming performance optimization

🎯 Verdict

Recommendation: Approve with minor suggestions

This is a high-quality implementation that follows established patterns and adds valuable functionality. The code is clean, well-structured, and production-ready. The issues identified are mostly about maintainability and testing rather than correctness.

The XAI provider will work correctly as-is. The recommendations above would improve long-term maintainability, especially the code deduplication suggestion which benefits both XAI and OpenAI providers.

Great work on the implementation! 🚀

---

Reviewed files: 15 changed files (+598, -7 lines)
Review focus: Architecture, security, performance, test coverage, code quality