fix some issues

gradio_app.py

@@ -12,7 +12,6 @@ from pydantic_ai.messages import (
 from src.agents.mask_generation_agent import EditImageResult
 from pydantic_ai.agent import Agent 
 from pydantic_ai.models.openai import OpenAIModel
-
 model = OpenAIModel(
     "gpt-4o",
     api_key=os.environ.get("OPENAI_API_KEY"),
@@ -45,37 +44,65 @@ def build_user_message(chat_input):
         ])
     return messages
 
-async def stream_from_agent(chat_input, chatbot, past_messages):
+def build_messages_for_agent(chat_input, past_messages):
+    # filter out image messages from past messages to save on tokens
+    messages = [msg for msg in past_messages 
+                if not (isinstance(msg, dict) 
+                        and msg.get("type") == "image_url")]
+    
+    # add the user's text message
+    if chat_input["text"]:
+        messages.append({
+            "type": "text",
+            "text": chat_input["text"]
+        })
+
+    # add the user's image message
+    files = chat_input.get("files", [])
+    image_url = image_path_to_uri(files[0]) if files else None
+    if image_url:
+        messages.append({
+            "type": "image_url",
+            "image_url": {"url": image_url}
+        })
+
+    return messages
+    
+
+async def stream_from_agent(chat_input, chatbot, past_messages, current_image):
+    # Prepare messages for the UI
     chatbot.extend(build_user_message(chat_input))
-    # Clear the input immediately after submission
-    yield {"text": "", "files": []}, chatbot, gr.skip
+    yield {"text": "", "files": []}, chatbot, gr.skip, gr.skip()
 
-    # for agent
+    # Prepare messages for the agent
     text = chat_input["text"]
-    images = [image_path_to_uri(image) for image in chat_input["files"]]
+    files = chat_input.get("files", [])
+    image_url = image_path_to_uri(files[0]) if files else None
     messages = [
         {
             "type": "text",
             "text": text
         },
     ]
-    if images:
-        messages.extend([
-            {"type": "image_url", "image_url": {"url": image}}
-            for image in images
-        ])
+    if image_url:
+        messages.append(
+            {"type": "image_url", "image_url": {"url": image_url}}
+        )
+        current_image = image_url
 
+    # Dependencies
     hopter = Hopter(os.environ.get("HOPTER_API_KEY"), environment=Environment.STAGING)
     mask_service = GenerateMaskService(hopter=hopter)
     deps = ImageEditDeps(
         edit_instruction=text,
-        image_url=images[0],
+        image_url=current_image,
         hopter_client=hopter,
         mask_service=mask_service
     )
+    # Run the agent
     async with mask_generation_agent.run_stream(
         messages,
-        deps=deps
+        deps=deps,
     ) as result:
         for message in result.new_messages():
             for call in message.parts:
@@ -88,9 +115,12 @@ async def stream_from_agent(chat_input, chatbot, past_messages):
                     metadata = {
                         'title': f'🛠️ Using {call.tool_name}',
                     }
+                    # set the tool call id so that when the tool returns
+                    # we can find this message and update with the result
                     if call.tool_call_id is not None:
                         metadata['id'] = call.tool_call_id
 
+                    # Create a tool call message to show on the UI
                     gr_message = {
                         'role': 'assistant',
                         'content': 'Parameters: ' + call_args,
@@ -99,10 +129,11 @@ async def stream_from_agent(chat_input, chatbot, past_messages):
                     chatbot.append(gr_message)
                 if isinstance(call, ToolReturnPart):
                     for gr_message in chatbot:
-                        if (
-                            gr_message.get('metadata', {}).get('id', '')
-                            == call.tool_call_id
-                        ):
+                        # Skip messages without metadata
+                        if not gr_message.get('metadata'):
+                            continue
+
+                        if gr_message['metadata'].get('id', '') == call.tool_call_id:
                             if isinstance(call.content, EditImageResult):
                                 chatbot.append({
                                     "role": "assistant",
@@ -113,15 +144,15 @@ async def stream_from_agent(chat_input, chatbot, past_messages):
                                 gr_message['content'] += (
                                     f'\nOutput: {call.content}'
                                 )
-                yield gr.skip(), chatbot, gr.skip()
+                yield gr.skip(), chatbot, gr.skip(), gr.skip()
 
         chatbot.append({'role': 'assistant', 'content': ''})
         async for message in result.stream_text():
             chatbot[-1]['content'] = message
-            yield gr.skip(), chatbot, gr.skip()
+            yield gr.skip(), chatbot, gr.skip(), gr.skip()
         past_messages = result.all_messages()
 
-        yield gr.Textbox(interactive=True), gr.skip(), past_messages
+        yield gr.Textbox(interactive=True), gr.skip(), past_messages, current_image
 
 with gr.Blocks() as demo:
     gr.HTML(
@@ -138,6 +169,7 @@ with gr.Blocks() as demo:
 """
     )
 
+    current_image = gr.State(None)
     past_messages = gr.State([])
     chatbot = gr.Chatbot(
         label='Image Editing Assistant',
@@ -147,15 +179,15 @@ with gr.Blocks() as demo:
     with gr.Row():
         chat_input = gr.MultimodalTextbox(
             interactive=True,
-            file_count="multiple",
+            file_count="single",
             show_label=False,
             placeholder='How would you like to edit this image?',
-            sources=["upload", "microphone"]
+            sources=["upload"]
         )
     generation = chat_input.submit(
         stream_from_agent,
-        inputs=[chat_input, chatbot, past_messages],
-        outputs=[chat_input, chatbot, past_messages],
+        inputs=[chat_input, chatbot, past_messages, current_image],
+        outputs=[chat_input, chatbot, past_messages, current_image],
     )
 
 if __name__ == '__main__':

src/agents/mask_generation_agent.py

@@ -4,6 +4,7 @@ from dotenv import load_dotenv
 import os
 import asyncio
 from dataclasses import dataclass
+from typing import Optional
 import logfire
 from src.services.generate_mask import GenerateMaskService
 from src.hopter.client import Hopter, Environment, MagicReplaceInput, SuperResolutionInput
@@ -18,14 +19,15 @@ logfire.instrument_openai()
 system_prompt = """
 I will give you an editing instruction of the image. 
 if the edit instruction involved modifying parts of the image, please generate a mask for it.
+if images are not provided, ask the user to provide an image.
 """
 
 @dataclass
 class ImageEditDeps:
     edit_instruction: str
-    image_url: str
     hopter_client: Hopter
     mask_service: GenerateMaskService
+    image_url: Optional[str] = None
 
 model = OpenAIModel(
     "gpt-4o",
@@ -47,25 +49,16 @@ mask_generation_agent = Agent(
     deps_type=ImageEditDeps
 )
 
-# @mask_generation_agent.tool
-# async def generate_mask(ctx: RunContext[ImageEditDeps]) -> MaskGenerationResult:
-#     """
-#     Generate a mask for the image editing instruction.
-#     """
-#     print("Invoking generate_mask tool")
-#     service = GenerateMaskService()
-#     mask_instruction = await service.get_mask_generation_instruction(ctx.deps.edit_instruction, ctx.deps.image_url)
-#     response = mask_instruction.model_dump_json(indent=4)
-#     print(f"generate_mask tool response: {response}")
-
-#     mask = await service.generate_mask(mask_instruction, ctx.deps.image_url)
-#     print("Exiting generate_mask tool")
-#     return MaskGenerationResult(mask_image_base64=mask)
-
 @mask_generation_agent.tool
 async def edit_object(ctx: RunContext[ImageEditDeps]) -> EditImageResult:
     """
-    Edit an object in the image.
+    Use this tool to edit an object in the image. for example:
+    - remove the pole
+    - replace the dog with a cat
+    - change the background to a beach
+    - remove the person in the image
+    - change the hair color to red
+    - change the hat to a cap
     """
     edit_instruction = ctx.deps.edit_instruction
     image_url = ctx.deps.image_url