feat(image-generation): add BFL (Black Forest Labs) provider#54
feat(image-generation): add BFL (Black Forest Labs) provider#54Kamilbenkirane merged 3 commits intomainfrom
Conversation
Add support for BFL's FLUX.2 image generation models: - flux-2-pro: Professional model with automatic prompt upsampling - flux-2-flex: Flexible model with configurable steps and guidance Changes: - Add BFL provider with polling-based async generation - Add BFL to Provider enum and credential system - Add new ImageGenerationParameter values (guidance, steps, seed, etc.) - Add billed_units field to ImageGenerationUsage - Add BFL to integration tests and CI workflow - Support BFL_API_KEY environment variable
PR Review: BFL Provider ImplementationOverviewComprehensive BFL FLUX.2 implementation. Well-structured code following existing patterns. Critical Issues1. Response Construction Bug (client.py:131-135)httpx.Response does not accept json parameter. Current code will fail at runtime. Fix: Use content parameter with JSON-encoded bytes import json
content = json.dumps(final_data).encode('utf-8')
return httpx.Response(status_code=200, content=content, headers={'content-type': 'application/json'}, request=httpx.Request('GET', polling_url))2. Missing Input Validation (parameters.py:26-32)AspectRatioMapper lacks validation before indexing split result. Fix: Add length check before indexing parts array Should Fix3. API Key Pattern (client.py:80, 105)Use consistent pattern with AUTH_HEADER_PREFIX even when empty for maintainability 4. Timeout Configuration (config.py:12)120s polling timeout may conflict with httpx default 60s timeout. Needs documentation or fix. 5. Test CoverageAdd unit tests for parameter mappers, polling timeout, error handling Security ✅
Performance ✅
SummarySolid implementation with good architecture. Fix critical bugs (#1, #2) before merge. Address validation and consistency issues (#3, #4, #5) for production quality. Great work overall! |
Additional Review DetailsCode Quality ObservationsStrengths:
Detailed Issues: Issue #1: Response Construction (CRITICAL)File: packages/image-generation/src/celeste_image_generation/providers/bfl/client.py:131-135, 137-141 The httpx.Response() constructor signature is: Your code passes Working fix: final_data = {**poll_data, "_submit_metadata": submit_data}
content = json.dumps(final_data).encode('utf-8')
return httpx.Response(
status_code=200,
content=content,
headers={"content-type": "application/json"},
request=httpx.Request("GET", polling_url),
)Issue #2: Aspect Ratio Validation (HIGH)File: packages/image-generation/src/celeste_image_generation/providers/bfl/parameters.py:26-28 Current code assumes split always returns 2 parts: parts = validated_value.split("x")
width = int(parts[0]) # IndexError if only 1 part
height = int(parts[1]) # IndexError if only 1 partAdd validation: parts = validated_value.split("x")
if len(parts) != 2:
raise ValueError(f"Invalid aspect_ratio: {validated_value}. Expected WIDTHxHEIGHT")
try:
width = int(parts[0])
height = int(parts[1])
except ValueError as e:
raise ValueError(f"Invalid dimensions in aspect_ratio: {validated_value}") from eIssue #3: Polling Timeout vs HTTP TimeoutFile: packages/image-generation/src/celeste_image_generation/providers/bfl/config.py:12 POLLING_TIMEOUT = 120.0 but http_client default timeout is 60.0 seconds. Individual GET requests in polling loop may timeout before the 120s polling timeout. Consider passing timeout parameter to http_client.get() calls or document this behavior. Test Coverage Gaps
Minor: API Key Header PatternWhile functionally correct, using consistent pattern improves maintainability: # Current (works but inconsistent)
config.AUTH_HEADER_NAME: self.api_key.get_secret_value()
# Recommended (consistent with other providers)
config.AUTH_HEADER_NAME: f"{config.AUTH_HEADER_PREFIX}{self.api_key.get_secret_value()}"All other providers use this pattern even when prefix is empty. |
Positive Highlights & Final RecommendationsWhat Works Really Well ✅
Integration QualityThe provider integrates seamlessly:
Recommended Action PlanBefore Merge (Required):
Before Merge (Strongly Recommended): Post-Merge (Nice to Have): Estimated Fix Time
Final VerdictApprove with changes required. The implementation is architecturally sound and shows strong understanding of the codebase patterns. The critical issues are minor bugs that are easy to fix. Once addressed, this will be a high-quality addition to Celeste. Great work @Kamilbenkirane! 🎉 |
- celeste-ai: 0.2.11 → 0.2.12 - celeste-image-generation: 0.2.9 → 0.2.10
Code Review: BFL (Black Forest Labs) Provider ImplementationThank you for this well-structured contribution! This is a solid implementation that follows the existing codebase patterns. Below is my detailed review: ✅ StrengthsCode Quality
Architecture
🔍 Issues & RecommendationsCritical: httpx.Response Construction (client.py:129-134)Issue: The manual construction of return httpx.Response(
status_code=200,
json=final_data, # ❌ 'json' is not a valid constructor parameter
request=httpx.Request("GET", polling_url),
)Problem:
Recommended fix: import json as json_module
# For successful response
return httpx.Response(
status_code=200,
content=json_module.dumps(final_data).encode('utf-8'),
headers={"Content-Type": "application/json"},
request=httpx.Request("GET", polling_url),
)
# For error response
return httpx.Response(
status_code=400,
content=json_module.dumps(poll_data).encode('utf-8'),
headers={"Content-Type": "application/json"},
request=httpx.Request("GET", polling_url),
)References:
High: Error Handling ConsistencyIssue: Mixed exception types not aligned with codebase conventions (client.py:57, 101, 113) raise ValueError(msg) # Line 57, 101
raise TimeoutError(msg) # Line 113Recommendation: Use exceptions from
Example from ByteDance provider ( from celeste.exceptions import ValidationError
if not sample_url:
msg = f"No image URL in {self.provider} response"
raise ValidationError(msg)Medium: Missing Error Response HandlingIssue: Non-200 responses from submit endpoint return immediately without parsing (client.py:92-93) if submit_response.status_code != 200:
return submit_response # ❌ Error not handled by callerProblem: The base client's Recommendation: Either:
Medium: Missing NULL Safety CheckIssue: submit_data = submit_response.json() # Could raise JSONDecodeError
polling_url = submit_data.get("polling_url")
if not polling_url:
msg = f"No polling_url in {self.provider} response"
raise ValueError(msg)Recommendation: Wrap in try-except or check if response is valid JSON first: try:
submit_data = submit_response.json()
except json.JSONDecodeError as e:
msg = f"{self.provider} returned invalid JSON: {submit_response.text}"
raise ValidationError(msg) from eLow: Magic NumbersIssue: Hardcoded configuration values in models.py could be constants min_pixels=64 * 64, # 4,096
max_pixels=2048 * 2048, # 4,194,304 (4MP)
min_aspect_ratio=9 / 21, # ~0.429
max_aspect_ratio=21 / 9, # ~2.333Recommendation: Extract to config.py for maintainability: # config.py
MIN_PIXELS = 64 * 64 # 4,096
MAX_PIXELS = 2048 * 2048 # 4,194,304 (4MP)
MIN_ASPECT_RATIO = 9 / 21 # ~0.429
MAX_ASPECT_RATIO = 21 / 9 # ~2.333Low: Potential Race ConditionIssue: Polling loop doesn't handle connection errors (client.py:114-120) poll_response = await self.http_client.get(
polling_url,
headers=poll_headers,
)Recommendation: Consider adding retry logic for transient network errors: try:
poll_response = await self.http_client.get(...)
except httpx.RequestError as e:
# Log and retry or fail based on elapsed time
if elapsed >= config.POLLING_TIMEOUT - 5:
raise
await asyncio.sleep(config.POLLING_INTERVAL)
continue🔒 Security Review✅ Secure Practices Confirmed
|
httpx.Response doesn't accept json parameter - use content with JSON-encoded bytes instead.
PR Review: BFL (Black Forest Labs) Provider ImplementationOverviewThis PR adds support for BFL's FLUX.2 image generation models with a polling-based async implementation. The code is well-structured and follows the existing patterns in the codebase. ✅ StrengthsCode Quality
Implementation Details
Testing & Integration
🔍 Potential Issues & Suggestions1. Polling Loop - Potential Resource Issue (Medium Priority)Location: The polling loop creates a new HTTP request every 0.5 seconds without any backoff strategy or rate limiting consideration. Issues:
Suggestion: # Consider exponential backoff or jittered polling
import random
poll_interval = config.POLLING_INTERVAL
while True:
# ... existing timeout check ...
poll_response = await self.http_client.get(polling_url, headers=poll_headers)
# ... existing status checks ...
# Add jitter to avoid thundering herd
await asyncio.sleep(poll_interval + random.uniform(0, 0.1))
# Or exponential backoff: poll_interval = min(poll_interval * 1.5, 5.0)2. Missing Error Handling for Non-200 Poll Responses (Medium Priority)Location: When Suggestion: if poll_response.status_code in (429, 503): # Rate limit or service unavailable
await asyncio.sleep(poll_interval * 2) # Backoff
continue
elif poll_response.status_code != 200:
return poll_response # Permanent failure3. Aspect Ratio Rounding Edge Case (Low Priority)Location: The rounding logic Current behavior:
Suggestion: 4. Missing Unit Tests (Medium Priority)The PR only includes integration tests. Consider adding unit tests for:
Example: # tests/unit_tests/providers/bfl/test_aspect_ratio.py
def test_aspect_ratio_rounds_to_multiple_of_16():
mapper = AspectRatioMapper()
result = mapper.map({}, "1023x1023", model)
assert result["width"] == 1024
assert result["height"] == 10245. Hardcoded Status Strings (Low Priority)Location: Status strings like # config.py
STATUS_READY = "Ready"
STATUS_ERROR = "Error"
STATUS_FAILED = "Failed"This improves maintainability if BFL changes their API. 🔒 Security Review✅ Good Practices
|
1 participant
Summary
Add support for BFL (Black Forest Labs) FLUX.2 image generation models:
Changes
Providerenum and credential systemImageGenerationParametervalues:guidance,steps,seed,prompt_upsampling,safety_tolerance,output_formatbilled_unitsfield toImageGenerationUsagefor BFL cost trackingflux-2-prowith1024x1024for minimum cost)BFL_API_KEYto CI workflow for publish integration testsFiles Changed
src/celeste/core.py,src/celeste/credentials.pyproviders/bfl/(client, config, models, parameters)io.py,parameters.py,models.py,providers/__init__.pytest_generate.py.github/workflows/publish.yml,.env.exampleTest Plan
BFL_API_KEYsecret to GitHub repository for CI