added untested funcitonal tool for agent

Author: codinglabsong

api/.env.example

@@ -1,3 +1,8 @@
 REPLICATE_API_TOKEN=YOUR_REPLICATE_API_TOKEN
 GOOGLE_API_KEY=YOUR_GOOGLE_API_KEY
 DATABASE_URL=YOUR_DATABASE_URL
+
+AWS_REGION=us-east-1
+AWS_ACCESS_KEY_ID=your_access_key_here
+AWS_SECRET_ACCESS_KEY=your_secret_key_here
+AWS_S3_BUCKET_NAME=your-bucket-name-here

api/.pre-commit-config.yaml

@@ -15,6 +15,8 @@ repos:
     rev: v1.10.0
     hooks:
       - id: mypy
+        additional_dependencies:
+          - types-requests
 
   - repo: https://github.com/pre-commit/pre-commit-hooks
     rev: v4.6.0

api/llm/agent.py

@@ -9,6 +9,7 @@
 from langgraph.prebuilt import create_react_agent
 
 from llm.prompt import system_message
+from llm.tools import initialize_tools
 
 load_dotenv()
 
@@ -53,10 +54,13 @@ def get_agent():
         # Build LLM
         llm = ChatGoogleGenerativeAI(model="gemini-2.5-flash")
 
+        # Build tools
+        tools = initialize_tools()
+
         # Create agent
         _agent_executor = create_react_agent(
             llm,
-            tools=[],
+            tools=tools,
             prompt=system_message,
             checkpointer=get_checkpointer(),
         )

api/llm/tools.py

@@ -0,0 +1,89 @@
+import uuid
+from typing import Optional
+
+import replicate
+import requests
+from dotenv import load_dotenv
+from langchain_core.tools import tool
+
+from llm.utils import upload_generated_image_to_s3
+
+load_dotenv()
+
+
+def initialize_tools():
+    """Initialize the tools for the agent."""
+
+    @tool(
+        description="Generate an image based on a prompt",
+    )
+    def generate_image(
+        prompt: str,
+        user_id: str,
+        image_url: str,
+    ) -> str:
+        """
+        Generate an image based on a prompt.
+        """
+        input = {
+            "width": 768,
+            "height": 768,
+            "prompt": prompt,
+            "refine": "expert_ensemble_refiner",
+            "apply_watermark": False,
+            "num_inference_steps": 25,
+            "prompt_strength": 0.5,
+            "image": image_url,
+        }
+
+        # Generate image using Replicate
+        output = replicate.run(
+            "stability-ai/sdxl:7762fd07cf82c948538e41f63f77d6 \
+                85e02b063e37e496e96eefd46c929f9bdc",
+            input=input,
+        )
+
+        # Check if generation was successful
+        if not output or len(output) == 0:
+            return "Failed to generate image. Please try again."
+
+        generated_image_url = output[0] if isinstance(output, list) else output
+
+        # Download the generated image
+        image_data: Optional[bytes] = None
+        try:
+            response = requests.get(generated_image_url)
+            response.raise_for_status()
+            image_data = response.content  # get the actual image bytes in content into memory
+            # Close the response to free up resources
+            response.close()
+        except Exception as e:
+            return f"Failed to download generated image: {str(e)}"
+
+        # Generate unique ID for the image
+        image_id = str(uuid.uuid4())
+
+        # Upload to S3
+        try:
+            s3_result = upload_generated_image_to_s3(
+                image_data=image_data, image_id=image_id, user_id=user_id, prompt=prompt
+            )
+
+            if s3_result["success"]:
+                return f"Image generated successfully! User can find it his/her gallery. \
+                    Image ID: {image_id}"
+            else:
+                return (
+                    f"Image generated but failed to save: {s3_result.get('error', 'Unknown error')}"
+                )
+
+        except Exception as e:
+            # Clear image data from memory even if upload fails
+            return f"Image generated but failed to save to storage: {str(e)}"
+
+        finally:
+            if image_data:
+                # Clear image data from memory
+                del image_data
+
+    return [generate_image]

api/llm/utils.py

@@ -0,0 +1,68 @@
+import os
+from datetime import datetime
+from typing import Any, Dict
+
+import boto3
+from botocore.exceptions import ClientError
+
+
+def upload_generated_image_to_s3(
+    image_data: bytes, image_id: str, user_id: str, prompt: str
+) -> Dict[str, Any]:
+    """
+    Upload a generated image to S3.
+
+    Args:
+        image_data: The image data as bytes
+        image_id: Unique identifier for the image
+        user_id: User identifier
+        prompt: The prompt used to generate the image
+
+    Returns:
+        Dict with success status and URL or error message
+    """
+    try:
+        # Initialize S3 client
+        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"),
+        )
+
+        # Generate S3 key with userId and imageId for organization
+        key = f"users/{user_id}/images/{image_id}"
+        bucket_name = os.environ.get("AWS_S3_BUCKET_NAME")
+
+        if not bucket_name:
+            return {"success": False, "error": "AWS_S3_BUCKET_NAME environment variable is not set"}
+
+        # Upload to S3
+        s3_client.put_object(
+            Bucket=bucket_name,
+            Key=key,
+            Body=image_data,
+            ContentType="image/png",
+            Metadata={
+                "title": "Generated Image",  # TODO: add title to the image provided by agent
+                "imageId": image_id,
+                "userId": user_id,
+                "uploadedAt": datetime.now().isoformat(),
+                "type": "generated",
+                "generationPrompt": prompt,
+            },
+        )
+
+        # Generate presigned URL for reading the uploaded file (valid for 2 hours)
+        presigned_url = s3_client.generate_presigned_url(
+            "get_object",
+            Params={"Bucket": bucket_name, "Key": key},
+            ExpiresIn=7200,  # 2 hours
+        )
+
+        return {"success": True, "url": presigned_url, "image_id": image_id}
+
+    except ClientError as e:
+        return {"success": False, "error": str(e)}
+    except Exception as e:
+        return {"success": False, "error": str(e)}

api/pyproject.toml

@@ -21,6 +21,8 @@ dependencies = [
     "langchain[google-genai]",
     "langgraph-checkpoint-postgres>=0.2.0",
     "psycopg[binary]>=3.1.18",
+    "boto3",
+    "requests",
 ]
 
 [project.optional-dependencies]
@@ -30,6 +32,7 @@ dev = [
     "black>=24.4.0",
     "mypy>=1.10.0",
     "pre-commit>=3.7.0",
+    "types-requests",
 ]
 
 [tool.setuptools.packages.find]

src/lib/actions.ts

@@ -96,10 +96,11 @@ export async function uploadImageToS3(
       Body: buffer,
       ContentType: file.type,
       Metadata: {
-        originalName: file.name,
+        title: file.name,
         imageId: imageId,
         userId: userId,
         uploadedAt: new Date().toISOString(),
+        type: "uploaded",
       },
     });