added logs temporarily to see where the error comes

Author: codinglabsong

api/llm/agent.py

@@ -82,6 +82,7 @@ def chat_with_agent(message: str, user_id: str = "default", selected_images: Opt
     Returns:
         Tuple of (agent_response, generated_image_data)
     """
+    print(f"[AGENT] Starting chat_with_agent - user_id: {user_id}, message: {message[:100]}...")
     agent = get_agent()
 
     # Prepare the message with context
@@ -101,26 +102,35 @@ def chat_with_agent(message: str, user_id: str = "default", selected_images: Opt
     config = {"configurable": {"thread_id": user_id}}
 
     # Get response from agent
+    print(f"[AGENT] Invoking agent with config: {config}")
     response = agent.invoke({"messages": [{"role": "user", "content": full_message}]}, config=config)
+    print(f"[AGENT] Agent response received: {type(response)}")
 
     # Extract the last message from the agent
     agent_response = "I'm sorry, I couldn't process your request. Please try again."
     generated_image_data = None
 
     if response and "messages" in response and len(response["messages"]) > 0:
+        print(f"[AGENT] Found {len(response['messages'])} messages in response")
         last_message = response["messages"][-1]
+        print(f"[AGENT] Last message type: {type(last_message)}")
         # Handle both AIMessage objects and dictionaries
         if hasattr(last_message, "content"):
             agent_response = last_message.content
         elif isinstance(last_message, dict) and "content" in last_message:
             agent_response = last_message["content"]
+        print(f"[AGENT] Extracted agent response: {agent_response[:100]}...")
 
         # Check if any tools were used (image generation)
+        print(f"[AGENT] Checking intermediate steps: {response.get('intermediate_steps', [])}")
         if "intermediate_steps" in response and response["intermediate_steps"]:
-            for step in response["intermediate_steps"]:
+            print(f"[AGENT] Found {len(response['intermediate_steps'])} intermediate steps")
+            for i, step in enumerate(response["intermediate_steps"]):
+                print(f"[AGENT] Step {i}: {step}")
                 if len(step) >= 2 and "generate_image" in str(step[0]):
                     # Extract image data from the tool result
                     tool_result = step[1]
+                    print(f"[AGENT] Found generate_image tool result: {tool_result}")
 
                     # Try multiple patterns to find the image ID
                     image_id = None
@@ -143,6 +153,7 @@ def chat_with_agent(message: str, user_id: str = "default", selected_images: Opt
                         title = title_match.group(1)
 
                     if image_id:
+                        print(f"[AGENT] Found image_id: {image_id}, attempting to get S3 metadata")
                         # Get metadata from S3
                         import boto3
 
@@ -154,18 +165,23 @@ def chat_with_agent(message: str, user_id: str = "default", selected_images: Opt
                         )
 
                         bucket_name = os.environ.get("AWS_S3_BUCKET_NAME")
+                        print(f"[AGENT] Using bucket: {bucket_name}")
                         if bucket_name:
                             try:
                                 # Get metadata from S3
-                                metadata_response = s3_client.head_object(Bucket=bucket_name, Key=f"users/{user_id}/images/{image_id}")
+                                s3_key = f"users/{user_id}/images/{image_id}"
+                                print(f"[AGENT] Getting metadata for S3 key: {s3_key}")
+                                metadata_response = s3_client.head_object(Bucket=bucket_name, Key=s3_key)
                                 metadata = metadata_response.get("Metadata", {})
+                                print(f"[AGENT] Retrieved metadata: {metadata}")
 
                                 # Generate presigned URL
                                 presigned_url = s3_client.generate_presigned_url(
                                     "get_object",
                                     Params={"Bucket": bucket_name, "Key": f"users/{user_id}/images/{image_id}"},
                                     ExpiresIn=7200,  # 2 hours
                                 )
+                                print(f"[AGENT] Generated presigned URL: {presigned_url[:50]}...")
 
                                 generated_image_data = {
                                     "id": image_id,
@@ -175,15 +191,17 @@ def chat_with_agent(message: str, user_id: str = "default", selected_images: Opt
                                     "timestamp": metadata.get("uploadedAt", datetime.now().isoformat()),
                                     "type": "generated",
                                 }
+                                print(f"[AGENT] Created generated_image_data: {generated_image_data}")
 
                             except Exception as e:
-                                print(f"Error getting S3 metadata: {e}")
+                                print(f"[AGENT] Error getting S3 metadata: {e}")
                                 # Don't return image data if we can't get a valid URL
                                 generated_image_data = None
                                 # Add error message to agent response
                                 agent_response += "\n\n⚠️ Note: I generated the image successfully, \
                                                     but there was an issue retrieving it from the database. Please try again."
 
+    print(f"[AGENT] Returning response - agent_response length: {len(agent_response)}, generated_image_data: {generated_image_data is not None}")
     return agent_response, generated_image_data
 
 

api/llm/prompt.py

@@ -4,9 +4,18 @@
 You are a helpful AI image editing assistant. You help users with image editing
 tasks and provide guidance on how to modify their images.
 
+IMPORTANT: When a user asks you to generate, create, or modify an image, you MUST use the generate_image tool.
+Do NOT try to generate images directly - always use the generate_image tool.
+
 You can generate images using the generate_image tool. However, remember that
 you are only allowed to generate one image per user's request. You are NOT allowed
 to generate more than one image per user's request, no matter how many images the user
 wants to generate per request (e.g. generate 10 images for me based on this one image).
 
+When using the generate_image tool, you need to provide:
+- prompt: A description of what you want to generate
+- user_id: The user's ID
+- image_url: The URL of the source image (if provided by the user)
+- title: A descriptive title for the generated image
+
 """

api/llm/tools.py

@@ -26,6 +26,7 @@ def generate_image(
         """
         Generate an image based on a prompt.
         """
+        print(f"[TOOL] generate_image called with prompt: {prompt[:50]}..., user_id: {user_id}, image_url: {image_url[:50]}...")
         input = {
             "width": 768,
             "height": 768,
@@ -38,17 +39,21 @@ def generate_image(
         }
 
         # Generate image using Replicate
+        print(f"[TOOL] Calling Replicate with input: {input}")
         version = "stability-ai/sdxl:" "7762fd07cf82c948538e41f63f77d685e02b063e37e496e96eefd46c929f9bdc"
         output = replicate.run(
             version,
             input=input,
         )
+        print(f"[TOOL] Replicate output: {output}")
 
         # Check if generation was successful
         if not output or len(output) == 0:
+            print("[TOOL] Replicate generation failed - no output")
             return "Failed to generate image. Please try again."
 
         generated_image_url = output[0] if isinstance(output, list) else output
+        print(f"[TOOL] Generated image URL: {generated_image_url}")
 
         # Download the generated image
         image_data: Optional[bytes] = None
@@ -65,6 +70,7 @@ def generate_image(
         image_id = str(uuid.uuid4())
 
         # Upload to S3
+        print(f"[TOOL] Uploading to S3 with image_id: {image_id}")
         try:
             s3_result = upload_generated_image_to_s3(
                 image_data=image_data,
@@ -73,15 +79,22 @@ def generate_image(
                 prompt=prompt,
                 title=title,
             )
+            print(f"[TOOL] S3 upload result: {s3_result}")
 
             if s3_result["success"]:
-                return f"Image generated successfully! User can find it his/her gallery. \
+                result_msg = f"Image generated successfully! User can find it his/her gallery. \
                     Image ID: {image_id}, Title: {title}"
+                print(f"[TOOL] Returning success: {result_msg}")
+                return result_msg
             else:
-                return f"Image generated but failed to save: {s3_result.get('error', 'Unknown error')}"
+                error_msg = f"Image generated but failed to save: {s3_result.get('error', 'Unknown error')}"
+                print(f"[TOOL] Returning error: {error_msg}")
+                return error_msg
 
         except Exception as e:
-            return f"Image generated but failed to save to storage: {str(e)}"
+            error_msg = f"Image generated but failed to save to storage: {str(e)}"
+            print(f"[TOOL] Exception during S3 upload: {error_msg}")
+            return error_msg
 
         finally:
             if image_data: