feat: add console streaming for debugging

server.py

@@ -0,0 +1,184 @@
+from fastapi import FastAPI, UploadFile, File, Form, HTTPException, BackgroundTasks
+from fastapi.responses import StreamingResponse
+from fastapi.middleware.cors import CORSMiddleware
+import asyncio
+import os
+from dotenv import load_dotenv
+from typing import Optional, List, Dict, Any
+import json
+from pydantic import BaseModel
+
+# Import project modules
+from src.agents.image_edit_agent import image_edit_agent, ImageEditDeps
+from src.agents.generic_agent import generic_agent
+from src.hopter.client import Hopter, Environment
+from src.services.generate_mask import GenerateMaskService
+from src.utils import upload_file_to_base64, upload_image
+
+# Load environment variables
+load_dotenv()
+
+app = FastAPI(title="Image Edit API")
+
+# Add CORS middleware
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],  # Allows all origins
+    allow_credentials=True,
+    allow_methods=["*"],  # Allows all methods
+    allow_headers=["*"],  # Allows all headers
+)
+
+class EditRequest(BaseModel):
+    edit_instruction: str
+    image_url: Optional[str] = None
+
+class MessageContent(BaseModel):
+    type: str
+    text: Optional[str] = None
+    image_url: Optional[Dict[str, str]] = None
+
+class Message(BaseModel):
+    content: List[MessageContent]
+
+
+@app.post("/test/stream")
+async def test(query: str):
+    async def stream_messages():
+        async with generic_agent.run_stream(query) as result:
+            async for message in result.stream(debounce_by=0.01):
+                yield json.dumps(message) + "\n"
+    
+    return StreamingResponse(stream_messages(), media_type="text/plain")
+
+@app.post("/edit")
+async def edit_image(request: EditRequest):
+    """
+    Edit an image based on the provided instruction.
+    Returns the URL of the edited image.
+    """
+    try:
+        # Initialize services
+        hopter = Hopter(
+            api_key=os.environ.get("HOPTER_API_KEY"),
+            environment=Environment.STAGING
+        )
+        mask_service = GenerateMaskService(hopter=hopter)
+
+        # Initialize dependencies
+        deps = ImageEditDeps(
+            edit_instruction=request.edit_instruction,
+            image_url=request.image_url,
+            hopter_client=hopter,
+            mask_service=mask_service
+        )
+
+        # Create messages
+        messages = [
+            {
+                "type": "text",
+                "text": request.edit_instruction
+            }
+        ]
+        
+        if request.image_url:
+            messages.append({
+                "type": "image_url",
+                "image_url": {
+                    "url": request.image_url
+                }
+            })
+
+        # Run the agent
+        result = await image_edit_agent.run(messages, deps=deps)
+        
+        # Return the result
+        return {"edited_image_url": result.edited_image_url}
+    
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+
+@app.post("/edit/stream")
+async def edit_image_stream(request: EditRequest):
+    """
+    Edit an image based on the provided instruction.
+    Streams the agent's responses back to the client.
+    """
+    try:
+        # Initialize services
+        hopter = Hopter(
+            api_key=os.environ.get("HOPTER_API_KEY"),
+            environment=Environment.STAGING
+        )
+        mask_service = GenerateMaskService(hopter=hopter)
+
+        # Initialize dependencies
+        deps = ImageEditDeps(
+            edit_instruction=request.edit_instruction,
+            image_url=request.image_url,
+            hopter_client=hopter,
+            mask_service=mask_service
+        )
+
+        # Create messages
+        messages = [
+            {
+                "type": "text",
+                "text": request.edit_instruction
+            }
+        ]
+        
+        if request.image_url:
+            messages.append({
+                "type": "image_url",
+                "image_url": {
+                    "url": request.image_url
+                }
+            })
+
+        async def stream_generator():
+            async with image_edit_agent.run_stream(messages, deps=deps) as result:
+                async for message in result.stream():
+                    # Convert message to JSON and yield
+                    yield json.dumps(message) + "\n"
+
+        return StreamingResponse(
+            stream_generator(),
+            media_type="application/x-ndjson"
+        )
+    
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+
+@app.post("/upload")
+async def upload_image_file(file: UploadFile = File(...)):
+    """
+    Upload an image file and return its URL.
+    """
+    try:
+        # Save the uploaded file to a temporary location
+        temp_file_path = f"/tmp/{file.filename}"
+        with open(temp_file_path, "wb") as buffer:
+            buffer.write(await file.read())
+        
+        # Upload the image to Google Cloud Storage
+        image_url = upload_image(temp_file_path)
+        
+        # Remove the temporary file
+        os.remove(temp_file_path)
+        
+        return {"image_url": image_url}
+    
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+
+@app.get("/health")
+async def health_check():
+    """
+    Health check endpoint.
+    """
+    return {"status": "ok"}
+
+if __name__ == "__main__":
+    import uvicorn
+    uvicorn.run(app, host="0.0.0.0", port=8000)

src/agents/generic_agent.py

@@ -0,0 +1,23 @@
+from pydantic_ai import Agent, RunContext
+from pydantic_ai.models.openai import OpenAIModel
+from dotenv import load_dotenv
+import os
+
+load_dotenv()
+
+model = OpenAIModel(
+    "gpt-4o",
+    api_key=os.environ.get("OPENAI_API_KEY")
+)
+
+system_prompt = """
+You are a helpful assistant that can answer questions and help with tasks.
+"""
+
+generic_agent = Agent(
+    model=model,
+    system_prompt=system_prompt,
+    tools=[],
+)
+
+

stream_utils.py

@@ -0,0 +1,145 @@
+import json
+import requests
+from rich.console import Console
+from rich.live import Live
+from rich.markdown import Markdown
+from rich.panel import Panel
+from rich.text import Text
+
+class StreamResponseHandler:
+    """
+    A utility class for handling streaming responses from API endpoints.
+    Provides rich formatting and real-time updates of the response content.
+    """
+    
+    def __init__(self, console=None):
+        """
+        Initialize the stream response handler.
+        
+        Args:
+            console (Console, optional): A Rich console instance. If not provided, a new one will be created.
+        """
+        self.console = console or Console()
+    
+    def check_server_health(self, health_url="http://localhost:8000/health"):
+        """
+        Check if the server is running and accessible.
+        
+        Args:
+            health_url (str, optional): The URL to check server health. Defaults to "http://localhost:8000/health".
+            
+        Returns:
+            bool: True if the server is running and accessible, False otherwise.
+        """
+        try:
+            self.console.print("Checking server health...", style="bold yellow")
+            response = requests.get(health_url)
+            if response.status_code == 200:
+                self.console.print("[bold green]✓ Server is running and accessible.[/]")
+                return True
+            else:
+                self.console.print(f"[bold red]✗ Server health check failed[/] with status code: {response.status_code}")
+                return False
+        except requests.exceptions.ConnectionError:
+            self.console.print("[bold red]✗ Error:[/] Could not connect to the server. Make sure it's running.")
+            return False
+        except Exception as e:
+            self.console.print(f"[bold red]✗ Error checking server health:[/] {e}")
+            return False
+    
+    def stream_response(self, url, payload=None, params=None, method="POST", title="AI Response"):
+        """
+        Send a request to an endpoint and stream the output to the terminal.
+        
+        Args:
+            url (str): The URL of the endpoint to send the request to.
+            payload (dict, optional): The JSON payload to send in the request body. Defaults to None.
+            params (dict, optional): The query parameters to send in the request. Defaults to None.
+            method (str, optional): The HTTP method to use. Defaults to "POST".
+            title (str, optional): The title to display in the panel. Defaults to "AI Response".
+            
+        Returns:
+            bool: True if the streaming was successful, False otherwise.
+        """
+        # Display request information
+        self.console.print(f"Sending request to [bold cyan]{url}[/]")
+        if payload:
+            self.console.print("Payload:", style="bold")
+            self.console.print(json.dumps(payload, indent=2))
+        if params:
+            self.console.print("Parameters:", style="bold")
+            self.console.print(json.dumps(params, indent=2))
+        
+        try:
+            # Prepare the request
+            request_kwargs = {
+                "stream": True
+            }
+            if payload:
+                request_kwargs["json"] = payload
+            if params:
+                request_kwargs["params"] = params
+            
+            # Make the request
+            with getattr(requests, method.lower())(url, **request_kwargs) as response:
+                # Check if the request was successful
+                if response.status_code != 200:
+                    self.console.print(f"[bold red]Error:[/] Received status code {response.status_code}")
+                    self.console.print(f"Response: {response.text}")
+                    return False
+                
+                # Initialize an empty response text
+                full_response = ""
+                
+                # Use Rich's Live display to update the content in place
+                with Live(Panel("Waiting for response...", title=title, border_style="blue"), refresh_per_second=10) as live:
+                    # Process the streaming response
+                    for line in response.iter_lines():
+                        if line:
+                            # Decode the line and parse it as JSON
+                            decoded_line = line.decode('utf-8')
+                            try:
+                                # Parse the JSON
+                                data = json.loads(decoded_line)
+                                
+                                # Extract and display the content
+                                if isinstance(data, dict):
+                                    if "content" in data:
+                                        for content in data["content"]:
+                                            if content.get("type") == "text":
+                                                text_content = content.get("text", "")
+                                                # Append to the full response
+                                                full_response += text_content
+                                                # Update the live display with the current full response
+                                                live.update(Panel(Markdown(full_response), title=title, border_style="green"))
+                                            elif content.get("type") == "image_url":
+                                                image_url = content.get("image_url", {}).get("url", "")
+                                                # Add a note about the image URL
+                                                image_note = f"\n\n[Image URL: {image_url}]"
+                                                full_response += image_note
+                                                live.update(Panel(Markdown(full_response), title=title, border_style="green"))
+                                    elif "edited_image_url" in data:
+                                        # Handle edited image URL from edit endpoint
+                                        image_url = data.get("edited_image_url", "")
+                                        image_note = f"\n\n[Edited Image URL: {image_url}]"
+                                        full_response += image_note
+                                        live.update(Panel(Markdown(full_response), title=title, border_style="green"))
+                                    else:
+                                        # For other types of data, just show the JSON
+                                        live.update(Panel(Text(json.dumps(data, indent=2)), title="Raw JSON Response", border_style="yellow"))
+                                else:
+                                    live.update(Panel(Text(decoded_line), title="Raw Response", border_style="yellow"))
+                            except json.JSONDecodeError:
+                                # If it's not valid JSON, just show the raw line
+                                live.update(Panel(Text(f"Raw response: {decoded_line}"), title="Invalid JSON", border_style="red"))
+
+                self.console.print("[bold green]Stream completed.[/]")
+                return True
+        
+        except requests.exceptions.ConnectionError:
+            self.console.print(f"[bold red]Error:[/] Could not connect to the server at {url}", style="red")
+            self.console.print("Make sure the server is running and accessible.", style="red")
+            return False
+        except requests.exceptions.RequestException as e:
+            self.console.print(f"[bold red]Error:[/] {e}", style="red")
+            return False

test_edit_stream.py

@@ -0,0 +1,104 @@
+import argparse
+import os
+import sys
+import requests
+import json
+from dotenv import load_dotenv
+from stream_utils import StreamResponseHandler
+
+# Load environment variables
+load_dotenv()
+
+def get_default_image():
+    """Get the default image path and convert it to a data URI."""
+    image_path = "./assets/lakeview.jpg"
+    if os.path.exists(image_path):
+        try:
+            from src.utils import image_path_to_uri
+            image_uri = image_path_to_uri(image_path)
+            print(f"Using default image: {image_path}")
+            return image_uri
+        except Exception as e:
+            print(f"Error converting image to URI: {e}")
+            return None
+    else:
+        print(f"Warning: Default image not found at {image_path}")
+        return None
+
+def upload_image(handler, image_path):
+    """
+    Upload an image to the server.
+    
+    Args:
+        handler (StreamResponseHandler): The stream response handler.
+        image_path (str): Path to the image file to upload.
+        
+    Returns:
+        str: The URL of the uploaded image, or None if upload failed.
+    """
+    if not os.path.exists(image_path):
+        handler.console.print(f"[bold red]Error:[/] Image file not found at {image_path}")
+        return None
+    
+    try:
+        handler.console.print(f"Uploading image: [bold]{image_path}[/]")
+        with open(image_path, 'rb') as f:
+            files = {'file': (os.path.basename(image_path), f)}
+            response = requests.post("http://localhost:8000/upload", files=files)
+            if response.status_code == 200:
+                image_url = response.json().get("image_url")
+                handler.console.print(f"Image uploaded successfully. URL: [bold green]{image_url}[/]")
+                return image_url
+            else:
+                handler.console.print(f"[bold red]Failed to upload image.[/] Status code: {response.status_code}")
+                handler.console.print(f"Response: {response.text}")
+                return None
+    except Exception as e:
+        handler.console.print(f"[bold red]Error uploading image:[/] {e}")
+        return None
+
+def main():
+    # Create a stream response handler
+    handler = StreamResponseHandler()
+    
+    # Parse command line arguments
+    parser = argparse.ArgumentParser(description="Test the image edit streaming API.")
+    parser.add_argument("--instruction", "-i", required=True, help="The edit instruction.")
+    parser.add_argument("--image", "-img", help="The URL of the image to edit.")
+    parser.add_argument("--upload", "-u", help="Path to an image file to upload first.")
+    
+    args = parser.parse_args()
+    
+    # Check if the server is running
+    if not handler.check_server_health():
+        sys.exit(1)
+    
+    image_url = args.image
+    
+    # If upload is specified, upload the image first
+    if args.upload:
+        image_url = upload_image(handler, args.upload)
+        if not image_url:
+            handler.console.print("[yellow]Warning:[/] Failed to upload image. Continuing without image URL.")
+    
+    # Use the default image if no image URL is provided
+    if not image_url:
+        image_url = get_default_image()
+        if not image_url:
+            handler.console.print("[yellow]No image URL provided and default image not available.[/]")
+            handler.console.print("The agent may ask for an image if needed.")
+    
+    # Prepare the payload for the edit request
+    payload = {
+        "edit_instruction": args.instruction
+    }
+    
+    if image_url:
+        payload["image_url"] = image_url
+    
+    # Stream the edit request
+    endpoint_url = "http://localhost:8000/edit/stream"
+    handler.stream_response(endpoint_url, payload=payload, title="Image Edit Response")
+
+if __name__ == "__main__":
+    main()

test_generic_stream.py

@@ -0,0 +1,29 @@
+import argparse
+import sys
+from dotenv import load_dotenv
+from stream_utils import StreamResponseHandler
+
+# Load environment variables
+load_dotenv()
+
+def main():
+    # Create a console for rich output
+    handler = StreamResponseHandler()
+    
+    # Parse command line arguments
+    parser = argparse.ArgumentParser(description="Test the generic agent streaming API.")
+    parser.add_argument("--query", "-q", required=True, help="The query or message to send to the generic agent.")
+    
+    args = parser.parse_args()
+    
+    # Check if the server is running
+    if not handler.check_server_health():
+        sys.exit(1)
+    
+    # Stream the generic request
+    endpoint_url = "http://localhost:8000/test/stream"
+    params = {"query": args.query}
+    handler.stream_response(endpoint_url, params=params, title="Generic Agent Response")
+
+if __name__ == "__main__":
+    main()