fixed, optimized, and organized api code for sending image meta data back to client

Author: codinglabsong

api/llm/agent.py

@@ -1,6 +1,5 @@
 import atexit
 import os
-import re
 from datetime import datetime
 from functools import lru_cache
 from typing import List, Optional
@@ -12,12 +11,15 @@
 
 from llm.prompt import system_message
 from llm.tools import initialize_tools
+from llm.utils import cleanup_old_tool_results, get_tool_result
 
 load_dotenv()
 
 # Global agent instance
 _agent_executor = None
 _checkpointer = None
+# Counter for periodic cleanup
+_request_count = 0
 
 
 @lru_cache
@@ -70,6 +72,125 @@ def get_agent():
     return _agent_executor
 
 
+def _build_message_with_context(message: str, selected_images: Optional[List[dict]]) -> str:
+    """Build the full message with image context if provided."""
+    if not selected_images or len(selected_images) == 0:
+        return message
+
+    image_context = "\n\nSelected Images:\n"
+    for i, img in enumerate(selected_images, 1):
+        image_context += f"{i}. {img.get('title', 'Untitled')} (ID: {img.get('id', 'unknown')})\n"
+        image_context += f"   Type: {img.get('type', 'unknown')}\n"
+        image_context += f"   Description: {img.get('description', 'No description')}\n"
+        if img.get("url"):
+            image_context += f"   URL: {img.get('url')}\n"
+        image_context += "\n"
+
+    return message + image_context
+
+
+def _extract_agent_response(response) -> str:
+    """Extract the agent's response text from the response object."""
+    if not response or "messages" not in response or len(response["messages"]) == 0:
+        return "I'm sorry, I couldn't process your request. Please try again."
+
+    last_message = response["messages"][-1]
+
+    # Handle None or unexpected message types
+    if last_message is None:
+        return "I'm sorry, I couldn't process your request. Please try again."
+
+    # Handle both AIMessage objects and dictionaries
+    if hasattr(last_message, "content"):
+        content = last_message.content
+        if content is None:
+            return "I'm sorry, I couldn't process your request. Please try again."
+        return content
+    elif isinstance(last_message, dict) and "content" in last_message:
+        content = last_message["content"]
+        if content is None:
+            return "I'm sorry, I couldn't process your request. Please try again."
+        return content
+
+    return "I'm sorry, I couldn't process your request. Please try again."
+
+
+def _generate_presigned_url(user_id: str, image_id: str) -> Optional[str]:
+    """Generate a presigned URL for an image."""
+    import boto3
+
+    s3_client = boto3.client(
+        "s3",
+        region_name=os.environ.get("AWS_REGION", "us-east-1"),
+        aws_access_key_id=os.environ.get("AWS_ACCESS_KEY_ID"),
+        aws_secret_access_key=os.environ.get("AWS_SECRET_ACCESS_KEY"),
+    )
+
+    bucket_name = os.environ.get("AWS_S3_BUCKET_NAME")
+    if not bucket_name:
+        print("[AGENT] AWS_S3_BUCKET_NAME not set")
+        return None
+
+    try:
+        s3_key = f"users/{user_id}/images/{image_id}"
+        print(f"[AGENT] Generating presigned URL for S3 key: {s3_key}")
+
+        presigned_url = s3_client.generate_presigned_url(
+            "get_object",
+            Params={"Bucket": bucket_name, "Key": s3_key},
+            ExpiresIn=7200,  # 2 hours
+        )
+        print(f"[AGENT] Generated presigned URL: {presigned_url[:50]}...")
+        return presigned_url
+
+    except Exception as e:
+        print(f"[AGENT] Error generating presigned URL: {e}")
+        return None
+
+
+def _process_generated_image(user_id: str, tool_result: dict) -> Optional[dict]:
+    """Process a generated image tool result and return image data."""
+    image_id = tool_result.get("image_id")
+    title = tool_result.get("title", "Generated Image")
+    prompt = tool_result.get("prompt", "Based on your request")
+
+    if not image_id:
+        print("[AGENT] No image_id found in tool result")
+        return None
+
+    print(f"[AGENT] Processing generated image with ID: {image_id}")
+
+    # Generate presigned URL
+    presigned_url = _generate_presigned_url(user_id, image_id)
+    if not presigned_url:
+        return None
+
+    # Create image data structure using data from tool result
+    generated_image_data = {
+        "id": image_id,
+        "url": presigned_url,
+        "title": title,
+        "description": f"AI-generated image: {prompt}",
+        "timestamp": datetime.now().isoformat(),
+        "type": "generated",
+    }
+
+    print(f"[AGENT] Created generated_image_data: {generated_image_data}")
+    return generated_image_data
+
+
+def _process_tool_results(user_id: str) -> Optional[dict]:
+    """Process any tool results for the user and return generated image data if found."""
+    print(f"[AGENT] Checking for tool results for user {user_id}")
+    tool_result = get_tool_result(user_id, "generate_image")
+
+    if tool_result:
+        print(f"[AGENT] Found tool result: {tool_result}")
+        return _process_generated_image(user_id, tool_result)
+
+    return None
+
+
 def chat_with_agent(message: str, user_id: str = "default", selected_images: Optional[List[dict]] = None) -> tuple[str, Optional[dict]]:
     """
     Send a message to the agent and get a response.
@@ -82,21 +203,19 @@ def chat_with_agent(message: str, user_id: str = "default", selected_images: Opt
     Returns:
         Tuple of (agent_response, generated_image_data)
     """
+    global _request_count
+
+    # Periodic cleanup every 10 requests
+    _request_count += 1
+    if _request_count % 10 == 0:
+        print(f"[AGENT] Running periodic cleanup (request #{_request_count})")
+        cleanup_old_tool_results()
+
     print(f"[AGENT] Starting chat_with_agent - user_id: {user_id}, message: {message[:100]}...")
     agent = get_agent()
 
     # Prepare the message with context
-    full_message = message
-    if selected_images and len(selected_images) > 0:
-        image_context = "\n\nSelected Images:\n"
-        for i, img in enumerate(selected_images, 1):
-            image_context += f"{i}. {img.get('title', 'Untitled')} (ID: {img.get('id', 'unknown')})\n"
-            image_context += f"   Type: {img.get('type', 'unknown')}\n"
-            image_context += f"   Description: {img.get('description', 'No description')}\n"
-            if img.get("url"):
-                image_context += f"   URL: {img.get('url')}\n"
-            image_context += "\n"
-        full_message = message + image_context
+    full_message = _build_message_with_context(message, selected_images)
 
     # Configure thread ID for conversation continuity
     config = {"configurable": {"thread_id": user_id}}
@@ -106,100 +225,12 @@ def chat_with_agent(message: str, user_id: str = "default", selected_images: Opt
     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"]:
-            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
-                    title = "Generated Image"
-
-                    # Pattern 1: "Image ID: uuid"
-                    match = re.search(r"Image ID: ([a-f0-9-]+)", tool_result)
-                    if match:
-                        image_id = match.group(1)
-
-                    # Pattern 2: "ID: uuid"
-                    if not image_id:
-                        match = re.search(r"ID: ([a-f0-9-]+)", tool_result)
-                        if match:
-                            image_id = match.group(1)
-
-                    # Extract title if present
-                    title_match = re.search(r"Title: (.+?)(?:\n|$)", tool_result)
-                    if title_match:
-                        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
-
-                        s3_client = boto3.client(
-                            "s3",
-                            region_name=os.environ.get("AWS_REGION", "us-east-1"),
-                            aws_access_key_id=os.environ.get("AWS_ACCESS_KEY_ID"),
-                            aws_secret_access_key=os.environ.get("AWS_SECRET_ACCESS_KEY"),
-                        )
-
-                        bucket_name = os.environ.get("AWS_S3_BUCKET_NAME")
-                        print(f"[AGENT] Using bucket: {bucket_name}")
-                        if bucket_name:
-                            try:
-                                # Get metadata from S3
-                                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,
-                                    "url": presigned_url,
-                                    "title": metadata.get("title", title),
-                                    "description": f"AI-generated image: {metadata.get('generationPrompt', 'Based on your request')}",
-                                    "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"[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."
+    # Extract the agent's response
+    agent_response = _extract_agent_response(response)
+    print(f"[AGENT] Extracted agent response: {agent_response[:100]}...")
+
+    # Check for tool results and process generated images
+    generated_image_data = _process_tool_results(user_id)
 
     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/tools.py

@@ -6,7 +6,7 @@
 from dotenv import load_dotenv
 from langchain_core.tools import tool
 
-from llm.utils import upload_generated_image_to_s3
+from llm.utils import store_tool_result, upload_generated_image_to_s3
 
 load_dotenv()
 
@@ -82,6 +82,10 @@ def generate_image(
             print(f"[TOOL] S3 upload result: {s3_result}")
 
             if s3_result["success"]:
+                # Store structured result for the agent to retrieve
+                tool_result = {"image_id": image_id, "title": title, "prompt": prompt, "success": True}
+                store_tool_result(user_id, "generate_image", tool_result)
+
                 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}")
@@ -111,7 +115,8 @@ def generate_image(
         {
             "prompt": "A woman in a beautiful sunset over a calm ocean",
             "user_id": "123",
-            "image_url": "https://replicate.delivery/pbxt/N55l5TWGh8mSlNzW8usReoaNhGbFwvLeZR3TX1NL4pd2Wtfv/replicate-prediction-f2d25rg6gnrma0cq257vdw2n4c.png",
+            "image_url": "https://example.com/image.jpg",
+            "title": "Test Image",
         }
     )
     print(output)