Update app.py

app.py

@@ -5,6 +5,41 @@ import streamlit as st
 
 from graph import run_graph
 
+
+# -------------------------------------------------------
+# HELPER FUNCTIONS
+# -------------------------------------------------------
+
+def format_items(items):
+    """Safely convert lists containing strings/dicts into text."""
+
+    if not items:
+        return ""
+
+    output = []
+
+    for item in items:
+
+        if isinstance(item, dict):
+
+            if "name" in item:
+                output.append(str(item["name"]))
+
+            elif "label" in item:
+                output.append(str(item["label"]))
+
+            elif "text" in item:
+                output.append(str(item["text"]))
+
+            else:
+                output.append(str(item))
+
+        else:
+            output.append(str(item))
+
+    return ", ".join(output)
+
+
 # -------------------------------------------------------
 # PAGE CONFIG
 # -------------------------------------------------------
@@ -16,6 +51,7 @@ st.set_page_config(
 
 st.title("🧠 AI Image Studio (Agents + LangGraph)")
 
+
 # -------------------------------------------------------
 # SESSION STATE
 # -------------------------------------------------------
@@ -32,6 +68,7 @@ if "analysis" not in st.session_state:
 if "last_uploaded" not in st.session_state:
     st.session_state.last_uploaded = []
 
+
 # -------------------------------------------------------
 # SIDEBAR
 # -------------------------------------------------------
@@ -44,6 +81,7 @@ uploaded_files = st.sidebar.file_uploader(
     accept_multiple_files=True
 )
 
+
 # -------------------------------------------------------
 # SAVE FILES
 # -------------------------------------------------------
@@ -56,20 +94,16 @@ if uploaded_files:
 
     for file in uploaded_files:
 
-        file_path = os.path.join("uploads", file.name)
+        path = os.path.join("uploads", file.name)
 
-        with open(file_path, "wb") as f:
+        with open(path, "wb") as f:
             f.write(file.getbuffer())
 
-        paths.append(file_path)
+        paths.append(path)
 
-    # Save paths
     st.session_state.uploaded_paths = paths
 
-    # ---------------------------------------------------
-    # AUTO ANALYZE ONLY IF NEW IMAGE
-    # ---------------------------------------------------
-
+    # Only analyze if a different image is uploaded
     if paths != st.session_state.last_uploaded:
 
         st.session_state.last_uploaded = paths
@@ -87,10 +121,7 @@ if uploaded_files:
 
         if isinstance(analysis, dict):
 
-            description = analysis.get(
-                "description",
-                "Image analyzed."
-            )
+            description = analysis.get("description", "")
 
             resolution = analysis.get("resolution", "")
 
@@ -98,48 +129,77 @@ if uploaded_files:
 
             lighting = analysis.get("lighting", "")
 
-            objects = analysis.get("objects", [])
+            people = analysis.get("people", "")
 
-            message = f"""📷 **Image Analysis**
+            background = analysis.get("background", "")
 
-**Description**
+            quality = analysis.get("quality", "")
 
-{description}
-"""
+            camera = analysis.get("camera_angle", "")
+
+            objects = format_items(
+                analysis.get("objects", [])
+            )
+
+            colors = format_items(
+                analysis.get("colors", [])
+            )
+
+            suggestions = analysis.get(
+                "suggestions",
+                []
+            )
+
+            message = "📷 **Image Analysis**\n\n"
+
+            if description:
+                message += f"### Description\n{description}\n\n"
 
             if resolution:
-                message += f"\n**Resolution:** {resolution}"
+                message += f"**Resolution:** {resolution}\n\n"
 
             if style:
-                message += f"\n\n**Style:** {style}"
+                message += f"**Style:** {style}\n\n"
 
             if lighting:
-                message += f"\n\n**Lighting:** {lighting}"
+                message += f"**Lighting:** {lighting}\n\n"
 
-            if objects:
-                message += f"\n\n**Objects:** {', '.join(objects)}"
+            if people:
+                message += f"**People:** {people}\n\n"
 
-            message += """
+            if background:
+                message += f"**Background:** {background}\n\n"
 
----
+            if camera:
+                message += f"**Camera Angle:** {camera}\n\n"
 
-### What would you like to do?
+            if quality:
+                message += f"**Quality:** {quality}\n\n"
 
-• Remove background
+            if colors:
+                message += f"**Colors:** {colors}\n\n"
 
-• Make it Anime
+            if objects:
+                message += f"**Objects:** {objects}\n\n"
 
-• Make it Pixar style
+            if suggestions:
 
-• Enhance quality
+                message += "### Suggested Actions\n"
 
-• Replace objects
+                for s in suggestions:
+                    message += f"• {s}\n"
 
-• Change colors
+            else:
 
-• Add new objects
+                message += """
+### Suggested Actions
 
-• Extract text
+• Remove background
+• Make it Anime
+• Pixar Style
+• Change Background
+• Replace Object
+• Enhance Quality
 """
 
             st.session_state.messages.append(
@@ -158,6 +218,7 @@ if uploaded_files:
                 }
             )
 
+
 # -------------------------------------------------------
 # SIDEBAR PREVIEW
 # -------------------------------------------------------
@@ -166,19 +227,20 @@ if st.session_state.uploaded_paths:
 
     st.sidebar.subheader("🖼 Preview")
 
-    for img in st.session_state.uploaded_paths:
+    for path in st.session_state.uploaded_paths:
 
         st.sidebar.image(
-            img,
-            caption=os.path.basename(img),
+            path,
+            caption=os.path.basename(path),
             use_container_width=True
         )
 
+
 # -------------------------------------------------------
-# CHAT
+# CHAT HISTORY
 # -------------------------------------------------------
 
-st.subheader("💬 Chat with AI Agent")
+st.subheader("💬 Chat")
 
 for msg in st.session_state.messages:
 
@@ -186,14 +248,16 @@ for msg in st.session_state.messages:
 
         st.markdown(msg["content"])
 
+
 # -------------------------------------------------------
-# USER INPUT
+# CHAT INPUT
 # -------------------------------------------------------
 
 user_input = st.chat_input(
-    "Generate / Edit / Analyze image..."
+    "Ask AI to edit, analyze or generate..."
 )
 
+
 # -------------------------------------------------------
 # RUN AGENT
 # -------------------------------------------------------
@@ -213,7 +277,7 @@ if user_input:
 
     with st.chat_message("assistant"):
 
-        with st.spinner("🤖 AI Agent thinking..."):
+        with st.spinner("🤖 Thinking..."):
 
             result = run_graph(
                 user_input,
@@ -222,59 +286,60 @@ if user_input:
 
             output = result.get("result")
 
-            assistant_message = ""
+            assistant_text = ""
 
-            # ------------------------------------------------
-            # IMAGE OUTPUT
-            # ------------------------------------------------
+            # ---------------- IMAGE ----------------
 
             if isinstance(output, Image.Image):
 
                 st.image(
                     output,
-                    caption="Generated Image",
                     use_container_width=True
                 )
 
-                buffer = io.BytesIO()
+                buf = io.BytesIO()
 
-                output.save(buffer, format="PNG")
+                output.save(
+                    buf,
+                    format="PNG"
+                )
 
                 st.download_button(
-                    "⬇ Download Image",
-                    buffer.getvalue(),
-                    "result.png",
+                    "⬇ Download",
+                    buf.getvalue(),
+                    "edited_image.png",
                     "image/png"
                 )
 
-                assistant_message = "✅ Image generated successfully."
+                assistant_text = "✅ Image edited successfully."
 
-            # ------------------------------------------------
-            # DICTIONARY OUTPUT
-            # ------------------------------------------------
+            # ---------------- JSON ----------------
 
             elif isinstance(output, dict):
 
-                st.json(output)
+                if "description" in output:
 
-                assistant_message = output.get(
-                    "description",
-                    str(output)
-                )
+                    st.markdown(output["description"])
+
+                    assistant_text = output["description"]
+
+                else:
+
+                    st.json(output)
+
+                    assistant_text = str(output)
 
-            # ------------------------------------------------
-            # TEXT OUTPUT
-            # ------------------------------------------------
+            # ---------------- TEXT ----------------
 
             else:
 
                 st.markdown(str(output))
 
-                assistant_message = str(output)
+                assistant_text = str(output)
 
     st.session_state.messages.append(
         {
             "role": "assistant",
-            "content": assistant_message
+            "content": assistant_text
         }
     )