refactor: clean up frontend files

README.md

@@ -5,6 +5,6 @@ colorFrom: yellow
 colorTo: purple
 sdk: gradio
 sdk_version: 5.16.1
-app_file: gradio_app.py
+app_file: app.py
 pinned: false
 ---

agent.py

@@ -1,158 +0,0 @@
-import streamlit as st
-from typing import TypedDict, Literal
-from pydantic_ai.messages import (
-    ModelRequest,
-    ModelResponse,
-    UserPromptPart,
-    TextPart,
-    ToolCallPart,
-    ToolReturnPart,
-)
-import asyncio
-from src.agents.mask_generation_agent import mask_generation_agent, ImageEditDeps
-from src.hopter.client import Hopter, Environment
-import os
-from src.services.generate_mask import GenerateMaskService
-from dotenv import load_dotenv
-from src.utils import image_path_to_uri
-load_dotenv()
-
-st.set_page_config(
-    page_title="Conversational Image Editor",
-    page_icon="🧊",
-    layout="wide",
-    initial_sidebar_state="collapsed"
-)
-
-hopter = Hopter(
-    api_key=os.getenv("HOPTER_API_KEY"),
-    environment=Environment.STAGING
-)
-mask_service = GenerateMaskService(hopter=hopter)
-user_msg_input_key = "input_user_msg"
-
-class ChatMessage(TypedDict):
-    """Format of messages sent to the browser/API."""
-
-    role: Literal['user', 'model']
-    timestamp: str
-    content: str
-
-
-def display_message_part(part):
-    """
-    Display a single part of a message in the Streamlit UI.
-    Customize how you display system prompts, user prompts,
-    tool calls, tool returns, etc.
-    """
-    # system-prompt
-    if part.part_kind == 'system-prompt':
-        with st.chat_message("system"):
-            st.markdown(f"**System**: {part.content}")
-    # user-prompt
-    elif part.part_kind == 'user-prompt':
-        with st.chat_message("user"):
-            st.markdown(part.content)
-    # text
-    elif part.part_kind == 'text':
-        with st.chat_message("assistant"):
-            st.markdown(part.content) 
-
-    # tool call
-    elif part.part_kind == 'tool-call':
-        with st.chat_message("assistant"):
-            st.markdown(f"**{part.tool_name}**: {part.args}")
-
-    # tool return
-    elif part.part_kind == 'tool-return':
-        with st.chat_message("assistant"):
-            st.markdown(f"**{part.tool_name}**: {part.content}")
-
-async def run_agent(user_input: str, image_b64: str):
-    messages = [
-        {
-            "type": "text",
-            "text": user_input
-        },
-        {
-            "type": "image_url",
-            "image_url": {
-                "url": image_b64
-            }
-        }
-    ]
-    deps = ImageEditDeps(
-        edit_instruction=user_input,
-        image_url=image_b64,
-        hopter_client=hopter,
-        mask_service=mask_service
-    )
-    async with mask_generation_agent.run_stream(
-        messages,
-        deps=deps
-    ) as result:
-        partial_text = ""
-        message_placeholder = st.empty()
-
-        # Render partial text as it arrives
-        async for chunk in result.stream_text(delta=True):
-            partial_text += chunk
-            message_placeholder.markdown(partial_text)
-
-        # Now that the stream is finished, we have a final result.
-        # Add new messages from this run, excluding user-prompt messages
-        filtered_messages = [msg for msg in result.new_messages() 
-                            if not (hasattr(msg, 'parts') and 
-                                    any(part.part_kind == 'user-prompt' for part in msg.parts))]
-        st.session_state.messages.extend(filtered_messages)
-
-        # Add the final response to the messages
-        st.session_state.messages.append(
-            ModelResponse(parts=[TextPart(content=partial_text)])
-        )
-    st.rerun()
-
-async def main():
-    st.title("Conversational Image Editor")
-
-    if "openai_model" not in st.session_state:
-        st.session_state["openai_model"] = "gpt-4o"
-
-    if "messages" not in st.session_state:
-        st.session_state.messages = []
-
-    if "image" not in st.session_state:
-        st.session_state.image = None
-
-    chat_col, image_col = st.columns(2)
-    with chat_col:
-        # Display all messages from the conversation so far
-        # Each message is either a ModelRequest or ModelResponse.
-        # We iterate over their parts to decide how to display them.
-        for msg in st.session_state.messages:
-            if isinstance(msg, ModelRequest) or isinstance(msg, ModelResponse) or isinstance(msg, ToolCallPart) or isinstance(msg, ToolReturnPart):
-                for part in msg.parts:
-                    display_message_part(part)
-
-    with image_col:
-        st.session_state.image = st.file_uploader("Upload an image", type=["png", "jpg", "jpeg"])
-        if st.session_state.image:
-            st.image(st.session_state.image)
-        else:
-            st.write("Upload an image to get started")
-
-    # Chat input for the user
-    user_input = st.chat_input("What would you like to edit your image?", disabled=not st.session_state.image)
-    if user_input and st.session_state.image:
-        st.session_state.messages.append(
-            ModelRequest(parts=[UserPromptPart(content=user_input)])
-        )
-
-        # Display the assistant's partial response while streaming
-        with st.chat_message("assistant"):
-            # Actually run the agent now, streaming the text
-            image_url = image_path_to_uri(st.session_state.image)
-            await run_agent(user_input, image_url)
-
-if __name__ == "__main__":
-    asyncio.run(main())

app.py

@@ -1,59 +1,11 @@
-from openai import OpenAI
-import streamlit as st
-from src.utils import image_path_to_uri
+import gradio as gr
+import image_edit_demo
+import image_edit_chat
 
-st.title("ChatGPT-like clone")
+with gr.Blocks() as demo:
+    image_edit_chat.demo.render()
+with demo.route("PicEdit"):
+    image_edit_demo.demo.render()
 
-client = OpenAI(api_key=st.secrets["OPENAI_API_KEY"])
-image = st.file_uploader("Upload an image", type=["png", "jpg", "jpeg"])
-
-if "openai_model" not in st.session_state:
-    st.session_state["openai_model"] = "gpt-4o"
-
-if "messages" not in st.session_state:
-    st.session_state.messages = []
-
-if "image" not in st.session_state:
-    st.session_state.image = None
-
-for message in st.session_state.messages:
-    with st.chat_message(message["role"]):
-        for content in message["content"]:
-            if content["type"] == "text":
-                st.markdown(content["text"])
-            elif content["type"] == "image_url":
-                st.image(content["image_url"]["url"])
-
-if prompt := st.chat_input("What is up?"):
-    # Add image to the message if it exists
-    if image is not None:
-        st.session_state.image = image
-    if st.session_state.image:
-        image_url = image_path_to_uri(st.session_state.image)
-        st.session_state.messages.append({"role": "user", "content": [
-            {"type": "text", "text": prompt},
-            {"type": "image_url", "image_url": {"url": image_url}}
-        ]})
-    else:
-        st.session_state.messages.append({"role": "user", "content": [
-            {"type": "text", "text": prompt}
-        ]})
-
-    with st.chat_message("user"):
-        if st.session_state.image:
-            st.image(st.session_state.image)
-        st.markdown(prompt)
-
-    with st.chat_message("assistant"):
-        stream = client.chat.completions.create(
-            model=st.session_state["openai_model"],
-            messages=[
-                {"role": m["role"], "content": m["content"]}
-                for m in st.session_state.messages
-            ],
-            stream=True,
-        )
-        response = st.write_stream(stream)
-    st.session_state.messages.append({"role": "assistant", "content": [
-        {"type": "text", "text": response}
-    ]})
+if __name__ == "__main__":
+    demo.launch()

gradio_app.py

@@ -1,11 +0,0 @@
-import gradio as gr
-import gradio_demo
-import gradio_chat
-
-with gr.Blocks() as demo:
-    gradio_chat.demo.render()
-with demo.route("PicEdit"):
-    gradio_demo.demo.render()
-
-if __name__ == "__main__":
-    demo.launch()

gradio_chat.py → image_edit_chat.py

@@ -4,7 +4,7 @@ import os
 from src.hopter.client import Hopter, Environment
 from src.services.generate_mask import GenerateMaskService
 from dotenv import load_dotenv
-from src.utils import image_path_to_uri, upload_image
+from src.utils import upload_image
 from pydantic_ai.messages import (
     ToolCallPart,
     ToolReturnPart