Update app.py

app.py

@@ -16,14 +16,25 @@ from PIL import Image
 import os
 import uuid
 import logging  # Add this line
+import boto3
+from botocore.exceptions import NoCredentialsError, PartialCredentialsError
 
 # Set up logging
 logging.basicConfig(level=logging.DEBUG)
 
+# Initialize the S3 client
+s3_client = boto3.client('s3')
+BUCKET_NAME = "wt-video-dl"  # Replace with your bucket name
+
+os.environ['AWS_ACCESS_KEY_ID'] = 'AKIA2YB4Z26VQAUNACVC'  # Replace with your AWS access key
+os.environ['AWS_SECRET_ACCESS_KEY'] = 'r3GmeE3WslFlaqOrPKBpz0pddAuXH9SDIpOLRGbv'  # Replace with your AWS secret key
+os.environ['AWS_DEFAULT_REGION'] = 'us-east-2'  # Replace with your region (e.g., 'us-east-1', 'us-west-2')
+
+
 #IMAGE_DIR = "./images" 
 IMAGE_DIR = "/tmp"
 os.makedirs(IMAGE_DIR, exist_ok=True)  
-IMGUR_CLIENT_ID = "429c0410bdece6a"
+
 
 GITHUB_API_URL = "https://api.github.com/search/repositories"
 ACCESS_TOKEN = os.getenv("github_pat")
@@ -50,7 +61,25 @@ embeddings = topic_data["embeddings"]
 
 discovered_repos = []  
 
-
+# Function to upload image to S3
+def upload_image_to_s3(image_data, filename):
+    try:
+        # Upload the image data to S3
+        s3_client.put_object(
+            Bucket=BUCKET_NAME,
+            Key=filename,
+            Body=image_data,
+            ContentType='image/png',
+            ACL='public-read'  # Make the image publicly readable
+        )
+        # Generate the S3 URL
+        s3_url = f"https://{BUCKET_NAME}.s3.amazonaws.com/{filename}"
+        return s3_url
+    except (NoCredentialsError, PartialCredentialsError) as e:
+        return f"Error with AWS credentials: {str(e)}"
+    except Exception as e:
+        return f"Error uploading image to S3: {str(e)}"
+		
 def search_similar_topics(input_text):
     if not input_text.strip():
         return "Enter topics to see suggestions."
@@ -484,65 +513,34 @@ def validate_and_fix_dot_code(dot_code):
     return dot_code
 
 def render_dot_code(dot_code, filename=None):
-    """
-    Renders Graphviz DOT code and uploads it to Imgur.
-
-    Args:
-        dot_code (str): The DOT code to render.
-        filename (str): Optional name for the output file (used for debugging purposes).
-
-    Returns:
-        str: Imgur URL of the uploaded image, or error message.
-    """
+    
     if not filename:
-        filename = f"diagram_{uuid.uuid4().hex}"
-	# Add /tmp directory to the filename
-    filename = os.path.join("/tmp", filename)
-	
-    print(f"Using IMGUR_CLIENT_ID: {IMGUR_CLIENT_ID}")
-
+        filename = f"diagram_{uuid.uuid4().hex}.png"  # Generate a unique filename
+		
     try:
         # Render the DOT code to an in-memory PNG
         src = Source(dot_code, format="png")
-        rendered_png = src.pipe()
-		
-		# Ensure IMGUR_CLIENT_ID is set
-        if not IMGUR_CLIENT_ID:
-            raise ValueError("IMGUR_CLIENT_ID is not set")
-
-        # Upload the rendered PNG to Imgur
-        headers = {"Authorization": f"Client-ID {IMGUR_CLIENT_ID}"}
-        #files = {"image": rendered_png}
-        files = {"image": ("image.png", rendered_png, "image/png")}
-        logging.debug(f"Sending request to Imgur with headers: {headers}")
-        response = requests.post("https://api.imgur.com/3/image", headers=headers, files=files)
-		
-        logging.debug(f"Imgur response status code: {response.status_code}")
-        logging.debug(f"Imgur response content: {response.text}")
-
-        if response.status_code == 200:
-            imgur_url = response.json()["data"]["link"]
-            logging.info(f"Diagram successfully uploaded to Imgur: {imgur_url}")
-            print(f"Diagram successfully uploaded to Imgur: {imgur_url}")
-            return imgur_url
-        else:
-            error_message = response.json().get("data", {}).get("error", "Unknown error")
-            raise Exception(f"Imgur upload failed: {error_message}. Response: {response.text}")
+        rendered_png = src.pipe()  # In-memory PNG image data
+
+        # Upload the rendered PNG to S3
+        s3_url = upload_image_to_s3(rendered_png, filename)
+
+        # Return the S3 URL
+        return s3_url
 
     except Exception as e:
-        logging.error(f"Error rendering or uploading diagram: {e}", exc_info=True)
         return f"Error rendering or uploading diagram: {str(e)}"
-
+		
+		
 import time
 
 def handle_generate_diagram(code_content, diagram_type, retries=5, wait_time=1):
     
-    imgur_url = render_dot_code(generate_dot_code_from_code(code_content, diagram_type))
-    if imgur_url.startswith("http"):  # Check if the response is a valid URL
-        return f'<img src="{imgur_url}" alt="Generated Diagram" style="max-width: 100%; height: auto;">'
+    s3_url = render_dot_code(generate_dot_code_from_code(code_content, diagram_type))
+    if s3_url.startswith("http"):  # Check if the response is a valid URL
+        return f'<img src="{s3_url}" alt="Generated Diagram" style="max-width: 100%; height: auto;">'
     else:
-        return f"<p>Error: {imgur_url}</p>"  # Return the error message in HTML format
-
+        return f"<p>Error: {s3_url}</p>"  # Return the error message in HTML format
     
 # Gradio Interface
 with gr.Blocks() as demo:
@@ -692,7 +690,6 @@ with gr.Blocks() as demo:
         )
         generate_diagram_button = gr.Button("Generate Diagram")
       with gr.Row():
-        
         diagram_output = gr.HTML(
           label="Generated Diagram",