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:
IGNY8 VPS (Salman)
2025-11-11 22:19:13 +00:00
parent 253d2e989d
commit ce9663438b
4 changed files with 71 additions and 9 deletions

View File

@@ -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.

View File

@@ -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)

View File

@@ -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,

View File

@@ -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