Merge fix-system-prompt-exposure: system prompt fix and image regeneration

Author: hunterjam

content-gen/src/app/frontend/src/App.tsx

@@ -292,6 +292,115 @@ function App() {
           timestamp: new Date().toISOString(),
         };
         setMessages(prev => [...prev, assistantMessage]);
+      } else if (generatedContent && confirmedBrief) {
+        // Content has been generated - check if user wants to modify the image
+        const imageModificationKeywords = [
+          'change', 'modify', 'update', 'replace', 'show', 'display', 'use', 
+          'instead', 'different', 'another', 'make it', 'make the', 
+          'kitchen', 'dining', 'living', 'bedroom', 'bathroom', 'outdoor', 'office',
+          'room', 'scene', 'setting', 'background', 'style', 'color', 'lighting'
+        ];
+        const isImageModification = imageModificationKeywords.some(kw => content.toLowerCase().includes(kw));
+        
+        if (isImageModification) {
+          // User wants to modify the image - use regeneration endpoint
+          const { streamRegenerateImage } = await import('./api');
+          
+          setGenerationStatus('Regenerating image with your changes...');
+          
+          let responseData: GeneratedContent | null = null;
+          let messageContent = '';
+          
+          // Get previous prompt from image_content if available
+          const previousPrompt = generatedContent.image_content?.prompt_used;
+          
+          for await (const response of streamRegenerateImage(
+            content,
+            confirmedBrief,
+            selectedProducts,
+            previousPrompt,
+            conversationId,
+            userId,
+            signal
+          )) {
+            if (response.type === 'heartbeat') {
+              setGenerationStatus(response.message || 'Regenerating image...');
+            } else if (response.type === 'agent_response' && response.is_final) {
+              try {
+                const parsedContent = JSON.parse(response.content);
+                
+                // Update generatedContent with new image
+                if (parsedContent.image_url || parsedContent.image_base64) {
+                  responseData = {
+                    ...generatedContent,
+                    image_content: {
+                      ...generatedContent.image_content,
+                      image_url: parsedContent.image_url || generatedContent.image_content?.image_url,
+                      image_base64: parsedContent.image_base64,
+                      prompt_used: parsedContent.image_prompt || generatedContent.image_content?.prompt_used,
+                    },
+                  };
+                  setGeneratedContent(responseData);
+                  messageContent = parsedContent.message || 'Image regenerated with your requested changes.';
+                } else if (parsedContent.error) {
+                  messageContent = parsedContent.error;
+                } else {
+                  messageContent = parsedContent.message || 'I processed your request.';
+                }
+              } catch {
+                messageContent = response.content || 'Image regenerated.';
+              }
+            } else if (response.type === 'error') {
+              messageContent = response.content || 'An error occurred while regenerating the image.';
+            }
+          }
+          
+          setGenerationStatus('');
+          
+          const assistantMessage: ChatMessage = {
+            id: uuidv4(),
+            role: 'assistant',
+            content: messageContent,
+            agent: 'ImageAgent',
+            timestamp: new Date().toISOString(),
+          };
+          setMessages(prev => [...prev, assistantMessage]);
+        } else {
+          // General question after content generation - use regular chat
+          let fullContent = '';
+          let currentAgent = '';
+          let messageAdded = false;
+          
+          setGenerationStatus('Processing your request...');
+          for await (const response of streamChat(content, conversationId, userId, signal)) {
+            if (response.type === 'agent_response') {
+              fullContent = response.content;
+              currentAgent = response.agent || '';
+              
+              if ((response.is_final || response.requires_user_input) && !messageAdded) {
+                const assistantMessage: ChatMessage = {
+                  id: uuidv4(),
+                  role: 'assistant',
+                  content: fullContent,
+                  agent: currentAgent,
+                  timestamp: new Date().toISOString(),
+                };
+                setMessages(prev => [...prev, assistantMessage]);
+                messageAdded = true;
+              }
+            } else if (response.type === 'error') {
+              const errorMessage: ChatMessage = {
+                id: uuidv4(),
+                role: 'assistant',
+                content: response.content || 'An error occurred.',
+                timestamp: new Date().toISOString(),
+              };
+              setMessages(prev => [...prev, errorMessage]);
+              messageAdded = true;
+            }
+          }
+          setGenerationStatus('');
+        }
       } else {
         // Check if this looks like a creative brief
         const briefKeywords = ['campaign', 'marketing', 'target audience', 'objective', 'deliverable'];

content-gen/src/app/frontend/src/api/index.ts

@@ -287,3 +287,65 @@ export async function* streamGenerateContent(
 
   throw new Error('Generation timed out after 10 minutes');
 }
+/**
+ * Regenerate image with a modification request
+ * Used when user wants to change the generated image after initial content generation
+ */
+export async function* streamRegenerateImage(
+  modificationRequest: string,
+  brief: CreativeBrief,
+  products?: Product[],
+  previousImagePrompt?: string,
+  conversationId?: string,
+  userId?: string,
+  signal?: AbortSignal
+): AsyncGenerator<AgentResponse> {
+  const response = await fetch(`${API_BASE}/regenerate`, {
+    signal,
+    method: 'POST',
+    headers: { 'Content-Type': 'application/json' },
+    body: JSON.stringify({
+      modification_request: modificationRequest,
+      brief,
+      products: products || [],
+      previous_image_prompt: previousImagePrompt,
+      conversation_id: conversationId,
+      user_id: userId || 'anonymous',
+    }),
+  });
+
+  if (!response.ok) {
+    throw new Error(`Regeneration request failed: ${response.statusText}`);
+  }
+
+  const reader = response.body?.getReader();
+  if (!reader) {
+    throw new Error('No response body');
+  }
+
+  const decoder = new TextDecoder();
+  let buffer = '';
+
+  while (true) {
+    const { done, value } = await reader.read();
+    if (done) break;
+
+    buffer += decoder.decode(value, { stream: true });
+    const lines = buffer.split('\n\n');
+    buffer = lines.pop() || '';
+
+    for (const line of lines) {
+      if (line.startsWith('data: ')) {
+        const data = line.slice(6);
+        if (data === '[DONE]') {
+          return;
+        }
+        try {
+          yield JSON.parse(data) as AgentResponse;
+        } catch {
+          console.error('Failed to parse SSE data:', data);
+        }
+      }
+    }
+  }
+}

content-gen/src/backend/app.py

@@ -857,6 +857,170 @@ async def generate():
     )
 
 
+@app.route("/api/regenerate", methods=["POST"])
+async def regenerate_content():
+    """
+    Regenerate image based on user modification request.
+    
+    This endpoint is called when the user wants to modify the generated image
+    after initial content generation (e.g., "show a kitchen instead of dining room").
+    
+    Request body:
+    {
+        "modification_request": "User's modification request",
+        "brief": { ... CreativeBrief fields ... },
+        "products": [ ... Product list ... ],
+        "previous_image_prompt": "Previous image prompt (optional)",
+        "conversation_id": "uuid"
+    }
+    
+    Returns regenerated image with the modification applied.
+    """
+    import asyncio
+    
+    data = await request.get_json()
+    
+    modification_request = data.get("modification_request", "")
+    brief_data = data.get("brief", {})
+    products_data = data.get("products", [])
+    previous_image_prompt = data.get("previous_image_prompt")
+    conversation_id = data.get("conversation_id") or str(uuid.uuid4())
+    user_id = data.get("user_id", "anonymous")
+    
+    if not modification_request:
+        return jsonify({"error": "modification_request is required"}), 400
+    
+    try:
+        brief = CreativeBrief(**brief_data)
+    except Exception as e:
+        return jsonify({"error": f"Invalid brief format: {str(e)}"}), 400
+    
+    # Save user request for regeneration
+    try:
+        cosmos_service = await get_cosmos_service()
+        await cosmos_service.add_message_to_conversation(
+            conversation_id=conversation_id,
+            user_id=user_id,
+            message={
+                "role": "user",
+                "content": f"Modify image: {modification_request}",
+                "timestamp": datetime.now(timezone.utc).isoformat()
+            }
+        )
+    except Exception as e:
+        logger.warning(f"Failed to save regeneration request to CosmosDB: {e}")
+    
+    orchestrator = get_orchestrator()
+    
+    async def generate():
+        """Stream regeneration responses with keepalive heartbeats."""
+        logger.info(f"Starting image regeneration for conversation {conversation_id}")
+        regeneration_task = None
+        
+        try:
+            # Create a task for the regeneration
+            regeneration_task = asyncio.create_task(
+                orchestrator.regenerate_image(
+                    modification_request=modification_request,
+                    brief=brief,
+                    products=products_data,
+                    previous_image_prompt=previous_image_prompt
+                )
+            )
+            
+            # Send keepalive heartbeats while regeneration is running
+            heartbeat_count = 0
+            while not regeneration_task.done():
+                for _ in range(30):  # 15 seconds
+                    if regeneration_task.done():
+                        break
+                    await asyncio.sleep(0.5)
+                
+                if not regeneration_task.done():
+                    heartbeat_count += 1
+                    yield f"data: {json.dumps({'type': 'heartbeat', 'count': heartbeat_count, 'message': 'Regenerating image...'})}\n\n"
+            
+        except asyncio.CancelledError:
+            logger.warning(f"Regeneration cancelled for conversation {conversation_id}")
+            if regeneration_task and not regeneration_task.done():
+                regeneration_task.cancel()
+            raise
+        except GeneratorExit:
+            logger.warning(f"Regeneration closed by client for conversation {conversation_id}")
+            if regeneration_task and not regeneration_task.done():
+                regeneration_task.cancel()
+            return
+        
+        # Get the result
+        try:
+            response = regeneration_task.result()
+            logger.info(f"Regeneration complete. Response keys: {list(response.keys()) if response else 'None'}")
+            
+            # Check for RAI block
+            if response.get("rai_blocked"):
+                yield f"data: {json.dumps({'type': 'error', 'content': response.get('error', 'Request blocked by content safety'), 'rai_blocked': True, 'is_final': True})}\n\n"
+                yield "data: [DONE]\n\n"
+                return
+            
+            # Handle image URL from orchestrator's blob save
+            if response.get("image_blob_url"):
+                blob_url = response["image_blob_url"]
+                parts = blob_url.split("/")
+                filename = parts[-1]
+                conv_folder = parts[-2]
+                response["image_url"] = f"/api/images/{conv_folder}/{filename}"
+                del response["image_blob_url"]
+            elif response.get("image_base64"):
+                # Save to blob storage
+                try:
+                    blob_service = await get_blob_service()
+                    blob_url = await blob_service.save_generated_image(
+                        conversation_id=conversation_id,
+                        image_base64=response["image_base64"]
+                    )
+                    if blob_url:
+                        parts = blob_url.split("/")
+                        filename = parts[-1]
+                        response["image_url"] = f"/api/images/{conversation_id}/{filename}"
+                        del response["image_base64"]
+                except Exception as e:
+                    logger.warning(f"Failed to save regenerated image to blob: {e}")
+            
+            # Save assistant response
+            try:
+                cosmos_service = await get_cosmos_service()
+                await cosmos_service.add_message_to_conversation(
+                    conversation_id=conversation_id,
+                    user_id=user_id,
+                    message={
+                        "role": "assistant",
+                        "content": response.get("message", "Image regenerated based on your request."),
+                        "agent": "ImageAgent",
+                        "timestamp": datetime.now(timezone.utc).isoformat()
+                    }
+                )
+            except Exception as e:
+                logger.warning(f"Failed to save regeneration response to CosmosDB: {e}")
+            
+            yield f"data: {json.dumps({'type': 'agent_response', 'content': json.dumps(response), 'is_final': True})}\n\n"
+        except Exception as e:
+            logger.exception(f"Error in regeneration: {e}")
+            yield f"data: {json.dumps({'type': 'error', 'content': str(e), 'is_final': True})}\n\n"
+        
+        yield "data: [DONE]\n\n"
+    
+    return Response(
+        generate(),
+        mimetype="text/event-stream",
+        headers={
+            "Cache-Control": "no-cache, no-store, must-revalidate",
+            "X-Accel-Buffering": "no",
+            "Connection": "keep-alive",
+            "Content-Type": "text/event-stream; charset=utf-8",
+        }
+    )
+
+
 # ==================== Image Proxy Endpoints ====================
 
 @app.route("/api/images/<conversation_id>/<filename>", methods=["GET"])