Enhance image generation functionality: Add console tracking to monitor progress and errors during image processing, and include image queue in response metadata for better integration.
This commit is contained in:
@@ -84,7 +84,8 @@ class BaseAIFunction(ABC):
|
|||||||
original_data: Any,
|
original_data: Any,
|
||||||
account=None,
|
account=None,
|
||||||
progress_tracker=None,
|
progress_tracker=None,
|
||||||
step_tracker=None
|
step_tracker=None,
|
||||||
|
console_tracker=None
|
||||||
) -> Dict:
|
) -> Dict:
|
||||||
"""
|
"""
|
||||||
Save parsed results to database.
|
Save parsed results to database.
|
||||||
|
|||||||
@@ -348,8 +348,8 @@ class AIEngine:
|
|||||||
return self._handle_error(error_msg, fn)
|
return self._handle_error(error_msg, fn)
|
||||||
|
|
||||||
# Phase 5: SAVE - Database Operations (85-98%)
|
# Phase 5: SAVE - Database Operations (85-98%)
|
||||||
# Pass step_tracker to save_output so it can add validation steps
|
# Pass step_tracker and console_tracker to save_output so it can add validation steps and log
|
||||||
save_result = fn.save_output(parsed, data, self.account, self.tracker, step_tracker=self.step_tracker)
|
save_result = fn.save_output(parsed, data, self.account, self.tracker, step_tracker=self.step_tracker, console_tracker=self.console_tracker)
|
||||||
clusters_created = save_result.get('clusters_created', 0)
|
clusters_created = save_result.get('clusters_created', 0)
|
||||||
keywords_updated = save_result.get('keywords_updated', 0)
|
keywords_updated = save_result.get('keywords_updated', 0)
|
||||||
count = save_result.get('count', 0)
|
count = save_result.get('count', 0)
|
||||||
|
|||||||
@@ -146,15 +146,24 @@ class GenerateImagesFromPromptsFunction(BaseAIFunction):
|
|||||||
original_data: Dict,
|
original_data: Dict,
|
||||||
account=None,
|
account=None,
|
||||||
progress_tracker=None,
|
progress_tracker=None,
|
||||||
step_tracker=None
|
step_tracker=None,
|
||||||
|
console_tracker=None
|
||||||
) -> Dict:
|
) -> Dict:
|
||||||
"""
|
"""
|
||||||
Process all images sequentially and generate them.
|
Process all images sequentially and generate them.
|
||||||
This method handles the loop and makes AI calls directly.
|
This method handles the loop and makes AI calls directly.
|
||||||
"""
|
"""
|
||||||
|
function_name = self.get_name()
|
||||||
|
|
||||||
|
if console_tracker:
|
||||||
|
console_tracker.save(f"[{function_name}] Starting image generation queue")
|
||||||
|
|
||||||
images = original_data.get('images', [])
|
images = original_data.get('images', [])
|
||||||
if not images:
|
if not images:
|
||||||
raise ValueError("No images to process")
|
error_msg = "[{function_name}] No images to process"
|
||||||
|
if console_tracker:
|
||||||
|
console_tracker.error('ValidationError', error_msg)
|
||||||
|
raise ValueError(error_msg)
|
||||||
|
|
||||||
provider = original_data.get('provider', 'openai')
|
provider = original_data.get('provider', 'openai')
|
||||||
model = original_data.get('model', 'dall-e-3')
|
model = original_data.get('model', 'dall-e-3')
|
||||||
@@ -169,6 +178,9 @@ class GenerateImagesFromPromptsFunction(BaseAIFunction):
|
|||||||
images_failed = 0
|
images_failed = 0
|
||||||
errors = []
|
errors = []
|
||||||
|
|
||||||
|
if console_tracker:
|
||||||
|
console_tracker.prep(f"[{function_name}] Preparing {total_images} image{'s' if total_images != 1 else ''} for generation")
|
||||||
|
|
||||||
# Initialize image queue in meta for frontend
|
# Initialize image queue in meta for frontend
|
||||||
image_queue = []
|
image_queue = []
|
||||||
for idx, img in enumerate(images, 1):
|
for idx, img in enumerate(images, 1):
|
||||||
@@ -198,6 +210,9 @@ class GenerateImagesFromPromptsFunction(BaseAIFunction):
|
|||||||
initial_meta['image_queue'] = image_queue
|
initial_meta['image_queue'] = image_queue
|
||||||
progress_tracker.update("PREP", 10, f"Preparing to generate {total_images} image{'s' if total_images != 1 else ''}", meta=initial_meta)
|
progress_tracker.update("PREP", 10, f"Preparing to generate {total_images} image{'s' if total_images != 1 else ''}", meta=initial_meta)
|
||||||
|
|
||||||
|
if console_tracker:
|
||||||
|
console_tracker.prep(f"[{function_name}] Image queue initialized with {total_images} image{'s' if total_images != 1 else ''}")
|
||||||
|
|
||||||
# Process each image sequentially
|
# Process each image sequentially
|
||||||
for index, image in enumerate(images, 1):
|
for index, image in enumerate(images, 1):
|
||||||
queue_item = image_queue[index - 1]
|
queue_item = image_queue[index - 1]
|
||||||
@@ -216,6 +231,9 @@ class GenerateImagesFromPromptsFunction(BaseAIFunction):
|
|||||||
else:
|
else:
|
||||||
content_title = content.title or content.meta_title or "Content"
|
content_title = content.title or content.meta_title or "Content"
|
||||||
|
|
||||||
|
if console_tracker:
|
||||||
|
console_tracker.prep(f"[{function_name}] Processing image {index}/{total_images}: {image.image_type} for '{content_title}'")
|
||||||
|
|
||||||
# Format prompt using template
|
# Format prompt using template
|
||||||
if image_prompt_template:
|
if image_prompt_template:
|
||||||
try:
|
try:
|
||||||
@@ -224,12 +242,18 @@ class GenerateImagesFromPromptsFunction(BaseAIFunction):
|
|||||||
image_prompt=image.prompt,
|
image_prompt=image.prompt,
|
||||||
image_type=image_type
|
image_type=image_type
|
||||||
)
|
)
|
||||||
|
if console_tracker:
|
||||||
|
console_tracker.prep(f"[{function_name}] Formatted prompt using template (length: {len(formatted_prompt)})")
|
||||||
except KeyError as e:
|
except KeyError as e:
|
||||||
logger.warning(f"Template formatting error: {e}, using simple format")
|
logger.warning(f"Template formatting error: {e}, using simple format")
|
||||||
formatted_prompt = f"Create a high-quality {image_type} image: {image.prompt}"
|
formatted_prompt = f"Create a high-quality {image_type} image: {image.prompt}"
|
||||||
|
if console_tracker:
|
||||||
|
console_tracker.prep(f"[{function_name}] Template formatting error, using fallback prompt")
|
||||||
else:
|
else:
|
||||||
# Fallback template
|
# Fallback template
|
||||||
formatted_prompt = f"Create a high-quality {image_type} image: {image.prompt}"
|
formatted_prompt = f"Create a high-quality {image_type} image: {image.prompt}"
|
||||||
|
if console_tracker:
|
||||||
|
console_tracker.prep(f"[{function_name}] Using fallback prompt template")
|
||||||
|
|
||||||
# Update progress: PREP phase for this image
|
# Update progress: PREP phase for this image
|
||||||
if progress_tracker and step_tracker:
|
if progress_tracker and step_tracker:
|
||||||
@@ -261,6 +285,9 @@ class GenerateImagesFromPromptsFunction(BaseAIFunction):
|
|||||||
|
|
||||||
# Generate image (this is the actual API call)
|
# Generate image (this is the actual API call)
|
||||||
# Frontend will simulate smooth progress from 50% to 95% while waiting
|
# Frontend will simulate smooth progress from 50% to 95% while waiting
|
||||||
|
if console_tracker:
|
||||||
|
console_tracker.ai_call(f"[{function_name}] Calling {provider}/{model} API for image {index}/{total_images}")
|
||||||
|
|
||||||
result = ai_core.generate_image(
|
result = ai_core.generate_image(
|
||||||
prompt=formatted_prompt,
|
prompt=formatted_prompt,
|
||||||
provider=provider,
|
provider=provider,
|
||||||
@@ -286,11 +313,14 @@ class GenerateImagesFromPromptsFunction(BaseAIFunction):
|
|||||||
image.status = 'failed'
|
image.status = 'failed'
|
||||||
image.save(update_fields=['status', 'updated_at'])
|
image.save(update_fields=['status', 'updated_at'])
|
||||||
|
|
||||||
error_msg = f"Image {index} failed: {result['error']}"
|
error_msg = f"[{function_name}] Image {index}/{total_images} failed: {result['error']}"
|
||||||
errors.append(error_msg)
|
errors.append(error_msg)
|
||||||
images_failed += 1
|
images_failed += 1
|
||||||
logger.error(f"Image generation failed for image {image.id}: {result['error']}")
|
logger.error(f"Image generation failed for image {image.id}: {result['error']}")
|
||||||
|
|
||||||
|
if console_tracker:
|
||||||
|
console_tracker.error('ImageGenerationError', error_msg)
|
||||||
|
|
||||||
if progress_tracker and step_tracker:
|
if progress_tracker and step_tracker:
|
||||||
parse_msg = f"Image {index} failed: {result['error']}"
|
parse_msg = f"Image {index} failed: {result['error']}"
|
||||||
step_tracker.add_response_step("PARSE", "error", parse_msg)
|
step_tracker.add_response_step("PARSE", "error", parse_msg)
|
||||||
@@ -311,11 +341,14 @@ class GenerateImagesFromPromptsFunction(BaseAIFunction):
|
|||||||
image.status = 'failed'
|
image.status = 'failed'
|
||||||
image.save(update_fields=['status', 'updated_at'])
|
image.save(update_fields=['status', 'updated_at'])
|
||||||
|
|
||||||
error_msg = f"Image {index} failed: No URL returned"
|
error_msg = f"[{function_name}] Image {index}/{total_images} failed: No URL returned"
|
||||||
errors.append(error_msg)
|
errors.append(error_msg)
|
||||||
images_failed += 1
|
images_failed += 1
|
||||||
logger.error(f"No image URL returned for image {image.id}")
|
logger.error(f"No image URL returned for image {image.id}")
|
||||||
|
|
||||||
|
if console_tracker:
|
||||||
|
console_tracker.error('ImageGenerationError', error_msg)
|
||||||
|
|
||||||
if progress_tracker and step_tracker:
|
if progress_tracker and step_tracker:
|
||||||
parse_msg = f"Image {index} failed: No URL returned"
|
parse_msg = f"Image {index} failed: No URL returned"
|
||||||
step_tracker.add_response_step("PARSE", "error", parse_msg)
|
step_tracker.add_response_step("PARSE", "error", parse_msg)
|
||||||
@@ -336,6 +369,9 @@ class GenerateImagesFromPromptsFunction(BaseAIFunction):
|
|||||||
progress_pct = 70 + int((index - 1) / total_images * 15) # 70-85% for PARSE
|
progress_pct = 70 + int((index - 1) / total_images * 15) # 70-85% for PARSE
|
||||||
progress_tracker.update("PARSE", progress_pct, parse_msg, meta=meta)
|
progress_tracker.update("PARSE", progress_pct, parse_msg, meta=meta)
|
||||||
|
|
||||||
|
if console_tracker:
|
||||||
|
console_tracker.parse(f"[{function_name}] Image {index}/{total_images} generated successfully: {image_url[:50]}...")
|
||||||
|
|
||||||
# Update image record
|
# Update image record
|
||||||
with transaction.atomic():
|
with transaction.atomic():
|
||||||
image.image_url = image_url
|
image.image_url = image_url
|
||||||
@@ -349,6 +385,9 @@ class GenerateImagesFromPromptsFunction(BaseAIFunction):
|
|||||||
images_generated += 1
|
images_generated += 1
|
||||||
logger.info(f"Image {image.id} ({image.image_type}) generated successfully: {image_url}")
|
logger.info(f"Image {image.id} ({image.image_type}) generated successfully: {image_url}")
|
||||||
|
|
||||||
|
if console_tracker:
|
||||||
|
console_tracker.save(f"[{function_name}] Saved image {index}/{total_images} to database (ID: {image.id})")
|
||||||
|
|
||||||
# Update progress: SAVE phase
|
# Update progress: SAVE phase
|
||||||
if progress_tracker and step_tracker:
|
if progress_tracker and step_tracker:
|
||||||
save_msg = f"Saved image {index} of {total_images}"
|
save_msg = f"Saved image {index} of {total_images}"
|
||||||
@@ -360,21 +399,40 @@ class GenerateImagesFromPromptsFunction(BaseAIFunction):
|
|||||||
|
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
# Mark as failed
|
# Mark as failed
|
||||||
|
queue_item['status'] = 'failed'
|
||||||
|
queue_item['progress'] = 100
|
||||||
|
queue_item['error'] = str(e)
|
||||||
with transaction.atomic():
|
with transaction.atomic():
|
||||||
image.status = 'failed'
|
image.status = 'failed'
|
||||||
image.save(update_fields=['status', 'updated_at'])
|
image.save(update_fields=['status', 'updated_at'])
|
||||||
|
|
||||||
error_msg = f"Image {index} failed: {str(e)}"
|
error_msg = f"[{function_name}] Image {index}/{total_images} exception: {str(e)}"
|
||||||
errors.append(error_msg)
|
errors.append(error_msg)
|
||||||
images_failed += 1
|
images_failed += 1
|
||||||
logger.error(f"Exception generating image {image.id}: {str(e)}", exc_info=True)
|
logger.error(f"Exception generating image {image.id}: {str(e)}", exc_info=True)
|
||||||
|
|
||||||
|
if console_tracker:
|
||||||
|
console_tracker.error('Exception', error_msg)
|
||||||
|
|
||||||
continue
|
continue
|
||||||
|
|
||||||
# Final progress update
|
# Final progress update
|
||||||
if progress_tracker and step_tracker:
|
if progress_tracker and step_tracker:
|
||||||
final_msg = f"Generated {images_generated} of {total_images} images"
|
final_msg = f"Generated {images_generated} of {total_images} images"
|
||||||
step_tracker.add_request_step("SAVE", "success", final_msg)
|
step_tracker.add_request_step("SAVE", "success", final_msg)
|
||||||
progress_tracker.update("SAVE", 98, final_msg, meta=step_tracker.get_meta())
|
meta = step_tracker.get_meta()
|
||||||
|
meta['image_queue'] = image_queue
|
||||||
|
progress_tracker.update("SAVE", 98, final_msg, meta=meta)
|
||||||
|
|
||||||
|
if console_tracker:
|
||||||
|
if images_generated > 0:
|
||||||
|
console_tracker.save(f"[{function_name}] SUCCESS: Generated {images_generated}/{total_images} image{'s' if images_generated != 1 else ''} successfully")
|
||||||
|
if images_failed > 0:
|
||||||
|
console_tracker.error('ImageGenerationError', f"[{function_name}] FAILED: {images_failed}/{total_images} image{'s' if images_failed != 1 else ''} failed")
|
||||||
|
if images_generated == total_images:
|
||||||
|
console_tracker.done(f"[{function_name}] All {total_images} image{'s' if total_images != 1 else ''} generated successfully")
|
||||||
|
else:
|
||||||
|
console_tracker.done(f"[{function_name}] Completed: {images_generated} succeeded, {images_failed} failed out of {total_images} total")
|
||||||
|
|
||||||
return {
|
return {
|
||||||
'count': images_generated,
|
'count': images_generated,
|
||||||
|
|||||||
@@ -1001,6 +1001,9 @@ class IntegrationSettingsViewSet(viewsets.ViewSet):
|
|||||||
response_meta['request_steps'] = meta['request_steps']
|
response_meta['request_steps'] = meta['request_steps']
|
||||||
if 'response_steps' in meta:
|
if 'response_steps' in meta:
|
||||||
response_meta['response_steps'] = meta['response_steps']
|
response_meta['response_steps'] = meta['response_steps']
|
||||||
|
# Include image_queue if available (for image generation)
|
||||||
|
if 'image_queue' in meta:
|
||||||
|
response_meta['image_queue'] = meta['image_queue']
|
||||||
return Response({
|
return Response({
|
||||||
'state': task_state,
|
'state': task_state,
|
||||||
'meta': response_meta
|
'meta': response_meta
|
||||||
|
|||||||
Reference in New Issue
Block a user