Add Image Generation from Prompts: Implement new functionality to generate images from prompts, including backend processing, API integration, and frontend handling with progress modal. Update settings and registry for new AI function.

2025-11-11 20:49:11 +00:00
parent 5f11da03e4
commit 5638ea78df
11 changed files with 1511 additions and 129 deletions
--- a/frontend/src/components/common/ProgressModal.tsx
+++ b/frontend/src/components/common/ProgressModal.tsx
@@ -73,7 +73,8 @@ const getSuccessMessage = (functionId?: string, title?: string, stepLogs?: any[]
    }
    return 'Article drafted successfully.';
  }
-  if (funcName.includes('image')) {
+  if (funcName.includes('image') && (funcName.includes('prompt') || funcName.includes('extract'))) {
+    // Image prompt generation
    // Try to extract from SAVE step message first (most reliable)
    const saveStepLog = stepLogs?.find(log => log.stepName === 'SAVE');
    if (saveStepLog?.message) {
@@ -119,6 +120,13 @@ const getSuccessMessage = (functionId?: string, title?: string, stepLogs?: any[]
    
    // Default message
    return 'Featured Image and X In‑article Image Prompts ready for image generation';
+  } else if (funcName.includes('image') && funcName.includes('from')) {
+    // Image generation from prompts
+    const imageCount = extractCount(/(\d+)\s+image/i, stepLogs || []);
+    if (imageCount) {
+      return `${imageCount} image${imageCount !== '1' ? 's' : ''} generated successfully`;
+    }
+    return 'Images generated successfully';
  }
  return 'Task completed successfully.';
 };
@@ -157,7 +165,8 @@ const getStepsForFunction = (functionId?: string, title?: string): Array<{phase:
    ];
  }
  
-  if (funcName.includes('image')) {
+  if (funcName.includes('image') && (funcName.includes('prompt') || funcName.includes('extract'))) {
+    // Image prompt generation
    return [
      { phase: 'INIT', label: 'Checking content and image slots' },
      { phase: 'PREP', label: 'Mapping Content for X Image Prompts' },
@@ -165,6 +174,15 @@ const getStepsForFunction = (functionId?: string, title?: string): Array<{phase:
      { phase: 'PARSE', label: 'Writing X In‑article Image Prompts' },
      { phase: 'SAVE', label: 'Assigning Prompts to Dedicated Slots' },
    ];
+  } else if (funcName.includes('image') && funcName.includes('from')) {
+    // Image generation from prompts
+    return [
+      { phase: 'INIT', label: 'Validating image prompts' },
+      { phase: 'PREP', label: 'Preparing image generation queue' },
+      { phase: 'AI_CALL', label: 'Generating images with AI' },
+      { phase: 'PARSE', label: 'Processing image URLs' },
+      { phase: 'SAVE', label: 'Saving image URLs' },
+    ];
  }
  
  // Default fallback
--- a/frontend/src/config/pages/images.config.tsx
+++ b/frontend/src/config/pages/images.config.tsx
@@ -49,6 +49,7 @@ export const createImagesPageConfig = (
    setStatusFilter: (value: string) => void;
    setCurrentPage: (page: number) => void;
    maxInArticleImages?: number; // Optional: max in-article images to display
+    onGenerateImages?: (contentId: number) => void; // Handler for generate images button
  }
 ): ImagesPageConfig => {
  const maxImages = handlers.maxInArticleImages || 5; // Default to 5 in-article images
@@ -99,13 +100,13 @@ export const createImagesPageConfig = (
    });
  }

-  // Add overall status column
+  // Add overall status column with Generate Images button
  columns.push({
    key: 'overall_status',
    label: 'Status',
    sortable: false,
-    width: '120px',
-    render: (value: string) => {
+    width: '180px',
+    render: (value: string, row: ContentImagesGroup) => {
      const statusColors: Record<string, 'success' | 'warning' | 'error' | 'info'> = {
        'complete': 'success',
        'partial': 'info',
@@ -118,13 +119,34 @@ export const createImagesPageConfig = (
        'pending': 'Pending',
        'failed': 'Failed',
      };
+      
+      // Check if there are any pending images with prompts
+      const hasPendingImages = 
+        (row.featured_image?.status === 'pending' && row.featured_image?.prompt) ||
+        row.in_article_images.some(img => img.status === 'pending' && img.prompt);
+      
      return (
-        <Badge
-          color={statusColors[value] || 'warning'}
-          size="sm"
-        >
-          {labels[value] || value}
-        </Badge>
+        <div className="flex items-center gap-2">
+          <Badge
+            color={statusColors[value] || 'warning'}
+            size="sm"
+          >
+            {labels[value] || value}
+          </Badge>
+          {hasPendingImages && handlers.onGenerateImages && (
+            <button
+              onClick={(e) => {
+                e.stopPropagation();
+                handlers.onGenerateImages!(row.content_id);
+              }}
+              className="inline-flex items-center gap-1 px-2 py-1 text-xs font-medium text-white bg-brand-500 hover:bg-brand-600 rounded transition-colors"
+              title="Generate Images"
+            >
+              <BoltIcon className="w-3 h-3" />
+              Generate
+            </button>
+          )}
+        </div>
      );
    },
  });
--- a/frontend/src/pages/Writer/Images.tsx
+++ b/frontend/src/pages/Writer/Images.tsx
@@ -3,19 +3,26 @@
 * Shows content images grouped by content - one row per content
 */

-import { useState, useEffect, useMemo, useCallback } from 'react';
+import { useState, useEffect, useMemo, useCallback, useRef } from 'react';
 import TablePageTemplate from '../../templates/TablePageTemplate';
 import {
  fetchContentImages,
  ContentImagesGroup,
  ContentImagesResponse,
+  generateImages,
 } from '../../services/api';
 import { useToast } from '../../components/ui/toast/ToastContainer';
-import { FileIcon, DownloadIcon } from '../../icons';
+import { FileIcon, DownloadIcon, BoltIcon } from '../../icons';
 import { createImagesPageConfig } from '../../config/pages/images.config';
+import ProgressModal from '../../components/common/ProgressModal';
+import { useProgressModal } from '../../hooks/useProgressModal';

 export default function Images() {
  const toast = useToast();
+  
+  // Progress modal for AI functions
+  const progressModal = useProgressModal();
+  const hasReloadedRef = useRef(false);

  // Data state
  const [images, setImages] = useState<ContentImagesGroup[]>([]);
@@ -136,6 +143,62 @@ export default function Images() {
    toast.info(`Bulk action "${action}" for ${ids.length} items`);
  }, [toast]);

+  // Generate images handler
+  const handleGenerateImages = useCallback(async (contentId: number) => {
+    try {
+      // Get all pending images for this content
+      const contentImages = images.find(g => g.content_id === contentId);
+      if (!contentImages) {
+        toast.error('Content not found');
+        return;
+      }
+      
+      // Collect all image IDs with prompts and pending status
+      const imageIds: number[] = [];
+      if (contentImages.featured_image?.id && 
+          contentImages.featured_image.status === 'pending' && 
+          contentImages.featured_image.prompt) {
+        imageIds.push(contentImages.featured_image.id);
+      }
+      contentImages.in_article_images.forEach(img => {
+        if (img.id && img.status === 'pending' && img.prompt) {
+          imageIds.push(img.id);
+        }
+      });
+      
+      if (imageIds.length === 0) {
+        toast.info('No pending images with prompts found for this content');
+        return;
+      }
+      
+      const result = await generateImages(imageIds);
+      if (result.success) {
+        if (result.task_id) {
+          // Open progress modal for async task
+          progressModal.openModal(
+            result.task_id,
+            'Generate Images',
+            'ai-generate-images-from-prompts-01-desktop'
+          );
+        } else {
+          // Synchronous completion
+          const generated = result.images_generated || 0;
+          const failed = result.images_failed || 0;
+          if (generated > 0) {
+            toast.success(`Images generated: ${generated} image${generated !== 1 ? 's' : ''} created${failed > 0 ? `, ${failed} failed` : ''}`);
+          } else {
+            toast.error(`Image generation failed: ${failed} image${failed !== 1 ? 's' : ''} failed`);
+          }
+          loadImages(); // Reload to show new images
+        }
+      } else {
+        toast.error(result.error || 'Failed to generate images');
+      }
+    } catch (error: any) {
+      toast.error(`Failed to generate images: ${error.message}`);
+    }
+  }, [toast, progressModal, loadImages, images]);
+
  // Get max in-article images from the data (to determine column count)
  const maxInArticleImages = useMemo(() => {
    if (images.length === 0) return 5; // Default
@@ -152,8 +215,9 @@ export default function Images() {
      setStatusFilter,
      setCurrentPage,
      maxInArticleImages,
+      onGenerateImages: handleGenerateImages,
    });
-  }, [searchTerm, statusFilter, maxInArticleImages]);
+  }, [searchTerm, statusFilter, maxInArticleImages, handleGenerateImages]);

  // Calculate header metrics
  const headerMetrics = useMemo(() => {
@@ -166,6 +230,7 @@ export default function Images() {
  }, [pageConfig?.headerMetrics, images, totalCount]);

  return (
+    <>
    <TablePageTemplate
      title="Content Images"
      titleIcon={<FileIcon className="text-purple-500 size-5" />}
@@ -218,5 +283,30 @@ export default function Images() {
        setCurrentPage(1);
      }}
    />
+    
+    {/* Progress Modal for AI Functions */}
+    <ProgressModal
+      isOpen={progressModal.isOpen}
+      title={progressModal.title}
+      percentage={progressModal.progress.percentage}
+      status={progressModal.progress.status}
+      message={progressModal.progress.message}
+      details={progressModal.progress.details}
+      taskId={progressModal.taskId || undefined}
+      functionId={progressModal.functionId}
+      onClose={() => {
+        const wasCompleted = progressModal.progress.status === 'completed';
+        progressModal.closeModal();
+        // Reload data after modal closes (if completed)
+        if (wasCompleted && !hasReloadedRef.current) {
+          hasReloadedRef.current = true;
+          loadImages();
+          setTimeout(() => {
+            hasReloadedRef.current = false;
+          }, 1000);
+        }
+      }}
+    />
+    </>
  );
 }
--- a/frontend/src/services/api.ts
+++ b/frontend/src/services/api.ts
@@ -1032,6 +1032,13 @@ export async function fetchContentImages(): Promise<ContentImagesResponse> {
  return fetchAPI('/v1/writer/images/content_images/');
 }

+export async function generateImages(imageIds: number[]): Promise<any> {
+  return fetchAPI('/v1/writer/images/generate_images/', {
+    method: 'POST',
+    body: JSON.stringify({ ids: imageIds }),
+  });
+}
+
 export async function deleteTaskImage(id: number): Promise<void> {
  return fetchAPI(`/v1/writer/images/${id}/`, {
    method: 'DELETE',