Skip to content

feat(providers): add Google GenerateContent and Imagen API packages#79

Merged
Kamilbenkirane merged 5 commits into
mainfrom
api/google_generate_content_and_imagen
Dec 15, 2025
Merged

feat(providers): add Google GenerateContent and Imagen API packages#79
Kamilbenkirane merged 5 commits into
mainfrom
api/google_generate_content_and_imagen

Conversation

@Kamilbenkirane

Copy link
Copy Markdown
Member

Summary

Introduces the celeste-google provider package with shared API implementations for Google's GenerateContent and Imagen APIs. This follows the same mixin pattern established with celeste-anthropic, enabling capability-agnostic code reuse across text generation, image generation, and future capabilities.

Architecture

packages/providers/google/          # NEW: API layer (celeste-google)
├── generate_content/               # GenerateContent API mixin
│   ├── client.py                   # GoogleGenerateContentClient
│   ├── parameters.py               # Temperature, TopP, Tools, etc.
│   ├── streaming.py                # GoogleGenerateContentStream
│   └── config.py                   # Endpoints, API version
└── imagen/                         # Imagen API mixin
    ├── client.py                   # GoogleImagenClient
    ├── parameters.py               # SampleCount, AspectRatio, ImageSize
    └── config.py                   # Endpoints

packages/capabilities/
├── text-generation/.../google/     # Uses GoogleGenerateContentClient mixin
└── image-generation/.../google/    # Uses both mixins via strategy pattern

Changes

1. Google GenerateContent API (celeste-google.generate_content)

Client Mixin:

  • HTTP POST/streaming to /v1beta/models/{model}:generateContent
  • Usage parsing: promptTokenCount, candidatesTokenCount, totalTokenCount
  • Content extraction from candidates[0].content.parts
  • API key auth via x-goog-api-key header

Parameters:

  • TemperatureMapper, TopPMapper, TopKMapper, MaxTokensMapper
  • StopSequencesMapper, SystemInstructionMapper
  • ResponseFormatMapper (JSON schema via responseSchema)
  • ToolsMapper, ToolChoiceMapper

Streaming:

  • SSE parsing for streamGenerateContent endpoint
  • Text delta extraction, finish reason tracking

2. Google Imagen API (celeste-google.imagen)

Client Mixin:

  • HTTP POST to /v1beta/models/{model}:predict
  • Predictions[] array parsing
  • No usage metadata (Imagen doesn't provide it)

Parameters:

  • SampleCountMapper, AspectRatioMapper, ImageSizeMapper

3. Text Generation Migration

  • GoogleTextGenerationClient now extends GoogleGenerateContentClient mixin
  • Delegates HTTP, usage parsing, metadata to API layer
  • Capability layer only handles TextGenerationInput → parts conversion
  • ~75% code reduction in capability client

4. Image Generation Migration

  • Strategy pattern: GoogleImageGenerationClient dispatches to model-specific strategies
  • GeminiImageGenerationClient uses GoogleGenerateContentClient for multimodal
  • ImagenImageGenerationClient uses GoogleImagenClient for Imagen API
  • Automatic model ID → strategy mapping

Benefits

  • DRY: API-specific code lives in one place, reused across capabilities
  • Consistency: Same pattern as Anthropic provider
  • Extensibility: New capabilities just extend the mixins
  • Maintainability: API changes only need updates in provider package

Test Plan

  • All unit tests pass
  • Type checking passes (mypy)
  • Linting passes (ruff)
  • Security scan passes (bandit)
  • Integration tests with real API calls

Add standalone provider package for Google Generate Content API with mixin
pattern for capability-agnostic reuse.

## Client (GoogleGenerateContentClient mixin)
- HTTP POST/streaming to /v1beta/models/{model}:generateContent endpoint
- Usage parsing: promptTokenCount, candidatesTokenCount, totalTokenCount
- Content extraction from candidates[0].content.parts
- Finish reason mapping (STOP, MAX_TOKENS, SAFETY, etc.)
- API key authentication via x-goog-api-key header

## Parameters
- TemperatureMapper: temperature float [0.0-2.0]
- TopPMapper: nucleus sampling top_p
- TopKMapper: top-k sampling
- MaxTokensMapper: maxOutputTokens
- StopSequencesMapper: custom stop sequences
- SystemInstructionMapper: system_instruction content block
- ResponseFormatMapper: JSON schema via responseSchema
  - Supports single BaseModel and list[BaseModel]
  - Uses StrictRefResolvingJsonSchemaGenerator for schema generation
- ToolsMapper: function declarations for tool use
- ToolChoiceMapper: tool_config with function_calling_config

## Streaming (GoogleGenerateContentStream mixin)
- SSE event parsing for streamGenerateContent endpoint
- Text delta extraction from candidates[0].content.parts
- Finish reason and usage tracking in final events

## Config
- API base URL: https://generativelanguage.googleapis.com
- API version: v1beta
- Endpoints: generateContent, streamGenerateContent
… API mixin

- Update GoogleTextGenerationClient to extend GoogleGenerateContentClient mixin
- Simplify client by delegating HTTP, usage parsing, and metadata to API layer
- Update parameter mappers to extend base API mappers
- Update streaming to extend base API stream class
- Add celeste-google as workspace dependency
## Client (GoogleImagenClient mixin)
- HTTP POST to /v1beta/models/{model_id}:predict endpoint
- Parse predictions[] array from response
- No usage metadata (Imagen API doesn't provide it)
- Metadata filtering for safety attributes

## Parameters
- SampleCountMapper: parameters.sampleCount for num_images
- AspectRatioMapper: parameters.aspectRatio
- ImageSizeMapper: parameters.imageSize

## Config
- API base URL: https://generativelanguage.googleapis.com
- Endpoint: /v1beta/models/{model_id}:predict
## Strategy Pattern
- GoogleImageGenerationClient delegates to model-specific strategies
- GeminiImageGenerationClient uses GoogleGenerateContentClient mixin
- ImagenImageGenerationClient uses GoogleImagenClient mixin
- Model ID → Strategy mapping for automatic dispatch

## Gemini Strategy
- Extends GoogleGenerateContentClient for multimodal image generation
- Uses responseModalities: ['Image'] for image output
- Parses inlineData from candidates[].content.parts[]

## Imagen Strategy
- Extends GoogleImagenClient for Imagen API
- Parses predictions[] array from response
- Maps to ImageArtifact with base64 data

## Parameters
- Split into gemini_parameters.py and imagen_parameters.py
- Extend base API mappers from celeste-google package

## Other Changes
- Add celeste-google as workspace dependency
- Remove old adapter pattern (gemini_api.py, imagen_api.py, config.py)
- Update tests for new strategy pattern
Without num_images parameter, Imagen may return multiple images (list)
instead of a single ImageArtifact, causing test assertion to fail.
@claude

claude Bot commented Dec 15, 2025

Copy link
Copy Markdown

PR Review: Google GenerateContent and Imagen API Packages

Summary

This PR introduces a well-architected mixin-based provider package for Google's APIs. The refactoring successfully reduces code duplication and improves maintainability. Overall, this is high-quality work with strong architectural decisions.

✅ Strengths

Excellent Architecture:

  • Mixin Pattern: Clean solution for sharing API-specific logic across capabilities
  • Strategy Pattern: Elegant dispatch between Gemini and Imagen APIs
  • Separation of Concerns: Clear boundaries between API layer and capability layer
  • Consistency: Follows the pattern established with celeste-anthropic

Code Quality:

  • Reduces ~75% code duplication by extracting shared API logic
  • Proper use of type hints throughout
  • Clear docstrings explaining mixin usage patterns
  • Maintainability: API changes only need updates in one place

Extensibility:

  • GOOGLE_MODEL_MAP makes adding new model types trivial
  • Well-structured for future Google API additions

⚠️ Issues & Recommendations

1. Potential Bug: Missing Error Handling
Location: packages/capabilities/image-generation/src/celeste_image_generation/providers/google/client.py:38

GOOGLE_MODEL_MAP[self.model.id] will raise KeyError if model ID isn't in map.
Recommendation: Add explicit error handling with ModelNotFoundError

2. Type Safety: Optional Strategy Attribute
Location: packages/capabilities/image-generation/src/celeste_image_generation/providers/google/client.py:32

All delegate methods use type: ignore[union-attr] because _strategy is typed as potentially None.
Recommendation: Use property with assertion or assert in init

3. Security: File Operations Without Validation
Location: packages/providers/google/src/celeste_google/generate_content/parameters.py:152-154

Opens files without path validation, size limits, or error handling.
Recommendation: Add path validation, size limits, and proper error handling

4. Inconsistent Return Type Handling
Location: packages/capabilities/image-generation/src/celeste_image_generation/providers/google/imagen.py:66-78

Complex return type logic makes API behavior less predictable.
Recommendation: Simplify to always return same type based on num_images parameter

5. Missing Parameter Validation
Parameter mappers don't validate value ranges (e.g., sampleCount should be >= 1)

6. Code Duplication in Parameter Mappers
Near-identical mappers that just wrap base classes. Consider factory function to reduce boilerplate.

🚀 Performance

Good: No vendor SDKs, mixin pattern has zero overhead, strategy selection happens once

Optimization: Consider caching compiled JSON schemas (OutputSchemaMapper rebuilds on every request)

🧪 Test Coverage

Observations:

  • Only 2 test files visible for image-generation Google provider
  • No visible tests for the new celeste-google provider package itself

Recommendations:

  1. Add unit tests for provider package mixins
  2. Add tests for edge cases (empty candidates, unexpected response structure, file errors)
  3. Consider integration tests with mocked API responses

🎯 Breaking Changes

  1. ImageGenerationOutput.content now returns ImageArtifact | list[ImageArtifact]
  2. ImageGenerationUsage.generated_images renamed to num_images
  3. _parse_finish_reason signature changed

Ensure these are documented and versioned properly.

🎉 Final Verdict

Excellent work that significantly improves the codebase architecture!

Blocking Issues: None - all issues are minor

Recommended Before Merge:

  1. Fix KeyError potential in GOOGLE_MODEL_MAP access
  2. Add file validation in MediaContentMapper
  3. Add unit tests for celeste-google package

Nice to Have:

  • Schema caching optimization
  • Simplify return type logic
  • Reduce parameter mapper boilerplate

Recommendation: APPROVE with minor fixes

The architecture is sound and sets a strong foundation for future Google API integrations.

@claude

claude Bot commented Dec 15, 2025

Copy link
Copy Markdown

Pull Request Review

Summary

This PR introduces the celeste-google provider package with shared API implementations for Google's GenerateContent and Imagen APIs. The architecture follows the established mixin pattern from celeste-anthropic, enabling code reuse across capabilities. Overall, this is a well-structured refactoring that significantly improves code organization and maintainability.


✅ Strengths

Architecture & Design

  • Excellent mixin pattern: The GoogleGenerateContentClient and GoogleImagenClient mixins provide clean separation between API layer and capability layer
  • Strategy pattern implementation: The GoogleImageGenerationClient properly delegates to model-specific strategies (GeminiImageGenerationClient vs ImagenImageGenerationClient)
  • DRY principle: Eliminates significant code duplication (~75% reduction in capability clients)
  • Extensibility: New capabilities can easily extend the mixins without reimplementing HTTP/parsing logic
  • Consistency: Mirrors the celeste-anthropic provider pattern, creating a consistent codebase structure

Code Quality

  • Clean delegation: The wrapper client properly delegates all methods to the strategy with appropriate type ignores
  • Type safety: Good use of type hints throughout, with type: ignore comments where mixin attributes are accessed
  • Clear separation of concerns: API-specific code lives in provider package, capability-specific transformations in capability packages

⚠️ Issues & Recommendations

1. Potential KeyError in Strategy Selection (High Priority)

Location: packages/capabilities/image-generation/src/celeste_image_generation/providers/google/client.py:38

StrategyClass = GOOGLE_MODEL_MAP[self.model.id]

Issue: Direct dictionary access will raise KeyError if an unknown model ID is passed.

Recommendation:

StrategyClass = GOOGLE_MODEL_MAP.get(self.model.id)
if StrategyClass is None:
    raise ModelNotFoundError(model_id=self.model.id, provider=Provider.GOOGLE)

2. Inconsistent Error Handling in Content Parsing (Medium Priority)

Locations:

  • packages/providers/google/src/celeste_google/generate_content/client.py:99-100
  • packages/providers/google/src/celeste_google/imagen/client.py:62-63

Issue: Both mixins raise ValueError when no candidates/predictions are found, but this may not be the most appropriate exception type. The downstream capability clients handle empty responses differently (returning empty ImageArtifact() in Gemini vs raising in the mixin).

Recommendation: Consider whether these should return empty data structures instead of raising, or use a more specific exception type like ProviderAPIError.

3. Complex Return Type Logic (Medium Priority)

Location: packages/capabilities/image-generation/src/celeste_image_generation/providers/google/imagen.py:66-78

The logic for deciding whether to return ImageArtifact or list[ImageArtifact] is complex:

num_images_requested = parameters.get("num_images")
if num_images_requested == 1:
    return images[0] if images else ImageArtifact()
if num_images_requested is not None and num_images_requested > 1:
    return images if images else []
# Not specified: return based on what provider actually returned
if len(images) == 1:
    return images[0]
return images if images else ImageArtifact()

Issue: This creates inconsistent return types based on runtime values, which can cause type checking issues.

Recommendation:

  • Consider always returning list[ImageArtifact] for consistency and simplifying downstream code
  • OR document this behavior very clearly in the docstring with examples
  • The empty case logic is also inconsistent (returns ImageArtifact() vs [])

4. Type Ignore Comments (Low Priority)

Location: Multiple files in the strategy classes

Multiple # type: ignore[union-attr] comments in delegation methods:

return self._strategy._init_request(inputs)  # type: ignore[union-attr]

Issue: While this is acceptable given the mixin pattern, it indicates that _strategy could be typed more precisely.

Recommendation: Consider using a Protocol or making _strategy non-optional (initialized to a sentinel or with __init__ instead of model_post_init).

5. Missing Return Type Logic for Index Out of Bounds (Medium Priority)

Location: packages/capabilities/image-generation/src/celeste_image_generation/providers/google/imagen.py:72

if num_images_requested == 1:
    return images[0] if images else ImageArtifact()

Issue: If num_images_requested == 1 but the API returns 0 images, this returns ImageArtifact(). However, the code doesn't validate whether images[0] exists when there are images.

Recommendation: More defensive:

if num_images_requested == 1:
    return images[0] if len(images) > 0 else ImageArtifact()

6. Base64 Decoding Without Error Handling (Low-Medium Priority)

Location:

  • packages/capabilities/image-generation/src/celeste_image_generation/providers/google/gemini.py:62
  • packages/capabilities/image-generation/src/celeste_image_generation/providers/google/imagen.py:63
image_bytes = base64.b64decode(base64_data)

Issue: base64.b64decode() can raise binascii.Error for invalid data. If the API returns malformed base64, this will crash.

Recommendation: Wrap in try-except to handle malformed data gracefully.

7. Empty Parameter Mappers (Low Priority)

Location: packages/capabilities/image-generation/src/celeste_image_generation/providers/google/gemini.py:25-26

@classmethod
def parameter_mappers(cls) -> list[ParameterMapper]:
    """Parameter mappers for Gemini image generation."""
    return []  # Parameter mapping handled by GoogleImageGenerationClient wrapper

Issue: The comment indicates that parameter mapping is handled by the wrapper, but this creates a confusing indirection.

Recommendation: Consider whether the strategy classes should directly return the appropriate mappers, or if the comment should be expanded to explain why this delegation pattern was chosen.


🔒 Security

✅ No Critical Issues Found

  • API key handling is delegated to self.auth.get_headers() (good practice)
  • No hardcoded credentials
  • Base64 decoding of images is safe (though error handling could be improved as noted above)
  • No injection vulnerabilities detected

🧪 Test Coverage

⚠️ Concerns

The PR description mentions:

  • ✅ All unit tests pass
  • ✅ Type checking passes (mypy)
  • ✅ Linting passes (ruff)
  • ✅ Security scan passes (bandit)
  • ❌ Integration tests with real API calls (unchecked)

Recommendation:

  • The new provider package (packages/providers/google) doesn't appear to have its own test directory based on my search
  • Consider adding unit tests for the mixin classes, especially:
    • GoogleGenerateContentClient._parse_usage()
    • GoogleGenerateContentClient._parse_content() with various response structures
    • GoogleImagenClient._parse_content() edge cases
    • Error handling paths (empty candidates, missing fields)

⚡ Performance

✅ Good Practices

  • No unnecessary deep copies or transformations
  • Efficient dictionary access patterns
  • Lazy evaluation where appropriate
  • Minimal overhead from delegation pattern

💡 Minor Optimization Opportunity

Location: packages/capabilities/image-generation/src/celeste_image_generation/providers/google/client.py:23-26

The GOOGLE_MODEL_MAP dictionary comprehension could be computed once at module load time (which it already is), but the comment suggests it's "extensible" - consider documenting how to extend it properly.


📝 Documentation

Areas for Improvement

  1. Mixin docstrings: The GoogleGenerateContentClient and GoogleImagenClient have good docstrings explaining usage, but could benefit from examples of how to extend them
  2. Return type variance: The complex return logic in imagen.py:66-78 needs clearer documentation
  3. Strategy pattern: Document how to add new model types to GOOGLE_MODEL_MAP

🎯 Overall Assessment

Recommendation: Approve with minor changes

This is a solid refactoring that improves code organization and follows established patterns in the codebase. The issues identified are mostly minor edge cases and potential improvements rather than blocking problems.

Before Merging:

  1. Must fix: Add KeyError handling for unknown model IDs
  2. Should fix: Clarify/simplify return type logic in ImagenImageGenerationClient._parse_content()
  3. Should fix: Add error handling for base64 decoding
  4. 💡 Nice to have: Add unit tests for the new provider package
  5. 💡 Nice to have: Document the parameter mapper delegation pattern

Final Notes

The mixin pattern and strategy implementation are excellent architectural choices that align with the project's goals. The code quality is high, and the reduction in duplication is significant. Great work! 🚀

@Kamilbenkirane Kamilbenkirane merged commit b851b10 into main Dec 15, 2025
18 of 19 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant