@soleur/opencode 0.0.0-stage → 1.0.0

package/skills/gemini-imagegen/scripts/multi_turn_chat.py

@@ -0,0 +1,223 @@
+#!/usr/bin/env python3
+"""
+Interactive multi-turn image generation and refinement using Gemini API.
+
+Usage:
+    python multi_turn_chat.py [--model MODEL] [--output-dir DIR]
+
+This starts an interactive session where you can:
+- Generate images from prompts
+- Iteratively refine images through conversation
+- Load existing images for editing
+- Save images at any point
+
+Commands:
+    /save [filename]  - Save current image
+    /load <path>      - Load an image into the conversation
+    /clear            - Start fresh conversation
+    /quit             - Exit
+
+Environment:
+    GEMINI_API_KEY - Required API key
+"""
+
+import argparse
+import os
+import sys
+from datetime import datetime
+from pathlib import Path
+
+from PIL import Image
+from google import genai
+from google.genai import errors, types
+
+from _error_handling import NoImageError, handle_api_error, parse_image_response
+
+
+class ImageChat:
+    """Interactive chat session for image generation and refinement."""
+    
+    def __init__(
+        self,
+        model: str = "gemini-2.5-flash-image",
+        output_dir: str = ".",
+    ):
+        api_key = os.environ.get("GEMINI_API_KEY")
+        if not api_key:
+            raise EnvironmentError("GEMINI_API_KEY environment variable not set")
+        
+        self.client = genai.Client(api_key=api_key)
+        self.model = model
+        self.output_dir = Path(output_dir)
+        self.output_dir.mkdir(parents=True, exist_ok=True)
+        
+        self.chat = None
+        self.current_image = None
+        self.image_count = 0
+        
+        self._init_chat()
+    
+    def _init_chat(self):
+        """Initialize or reset the chat session."""
+        config = types.GenerateContentConfig(
+            response_modalities=["TEXT", "IMAGE"]
+        )
+        self.chat = self.client.chats.create(
+            model=self.model,
+            config=config,
+        )
+        self.current_image = None
+    
+    def send_message(self, message: str, image: Image.Image | None = None) -> tuple[str | None, Image.Image | None]:
+        """Send a message and optionally an image, return response text and image."""
+        contents = []
+        if message:
+            contents.append(message)
+        if image:
+            contents.append(image)
+        
+        if not contents:
+            return None, None
+
+        try:
+            response = self.chat.send_message(contents)
+        except errors.APIError as e:
+            handle_api_error(e)
+
+        try:
+            image_response, text_response = parse_image_response(response)
+            self.current_image = image_response
+        except NoImageError:
+            # In interactive chat, text-only responses are normal
+            text_response = None
+            if response.parts:
+                for part in response.parts:
+                    if part.text is not None:
+                        text_response = part.text
+            image_response = None
+
+        return text_response, image_response
+    
+    def save_image(self, filename: str | None = None) -> str | None:
+        """Save the current image to a file."""
+        if self.current_image is None:
+            return None
+        
+        if filename is None:
+            self.image_count += 1
+            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+            filename = f"image_{timestamp}_{self.image_count}.png"
+        
+        filepath = self.output_dir / filename
+        self.current_image.save(filepath)
+        return str(filepath)
+    
+    def load_image(self, path: str) -> Image.Image:
+        """Load an image from disk."""
+        img = Image.open(path)
+        self.current_image = img
+        return img
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="Interactive multi-turn image generation",
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+        epilog=__doc__
+    )
+    parser.add_argument(
+        "--model", "-m",
+        default="gemini-2.5-flash-image",
+        choices=["gemini-2.5-flash-image", "gemini-3-pro-image-preview"],
+        help="Model to use"
+    )
+    parser.add_argument(
+        "--output-dir", "-o",
+        default=".",
+        help="Directory to save images"
+    )
+    
+    args = parser.parse_args()
+    
+    try:
+        chat = ImageChat(model=args.model, output_dir=args.output_dir)
+    except Exception as e:
+        print(f"Error initializing: {e}", file=sys.stderr)
+        sys.exit(1)
+    
+    print(f"Gemini Image Chat ({args.model})")
+    print("Commands: /save [name], /load <path>, /clear, /quit")
+    print("-" * 50)
+    
+    while True:
+        try:
+            user_input = input("\nYou: ").strip()
+        except (EOFError, KeyboardInterrupt):
+            print("\nGoodbye!")
+            break
+        
+        if not user_input:
+            continue
+        
+        # Handle commands
+        if user_input.startswith("/"):
+            parts = user_input.split(maxsplit=1)
+            cmd = parts[0].lower()
+            arg = parts[1] if len(parts) > 1 else None
+            
+            if cmd == "/quit":
+                print("Goodbye!")
+                break
+            
+            elif cmd == "/clear":
+                chat._init_chat()
+                print("Conversation cleared.")
+                continue
+            
+            elif cmd == "/save":
+                path = chat.save_image(arg)
+                if path:
+                    print(f"Image saved to: {path}")
+                else:
+                    print("No image to save.")
+                continue
+            
+            elif cmd == "/load":
+                if not arg:
+                    print("Usage: /load <path>")
+                    continue
+                try:
+                    chat.load_image(arg)
+                    print(f"Loaded: {arg}")
+                    print("You can now describe edits to make.")
+                except Exception as e:
+                    print(f"Error loading image: {e}")
+                continue
+            
+            else:
+                print(f"Unknown command: {cmd}")
+                continue
+        
+        # Send message to model
+        try:
+            # If we have a loaded image and this is first message, include it
+            image_to_send = None
+            if chat.current_image and not chat.chat.history:
+                image_to_send = chat.current_image
+            
+            text, image = chat.send_message(user_input, image_to_send)
+            
+            if text:
+                print(f"\nGemini: {text}")
+            
+            if image:
+                # Auto-save
+                path = chat.save_image()
+                print(f"\n[Image generated: {path}]")
+            
+        except Exception as e:
+            print(f"\nError: {e}")
+
+
+if __name__ == "__main__":
+    main()

package/skills/gemini-imagegen/scripts/test_error_handling.py

@@ -0,0 +1,179 @@
+"""Tests for _error_handling module."""
+
+import unittest
+from unittest.mock import MagicMock
+
+from google.genai import errors
+
+from _error_handling import (
+    NoImageError,
+    PermissionDeniedError,
+    QuotaExhaustedError,
+    SafetyFilterError,
+    check_quota,
+    handle_api_error,
+    parse_image_response,
+)
+
+
+def _make_client_error(code: int, message: str = "test error"):
+    """Create a ClientError with the given code using the real constructor."""
+    response_json = {"error": {"message": message, "code": code}}
+    return errors.ClientError(code, response_json)
+
+
+def _make_server_error(code: int = 500, message: str = "server error"):
+    """Create a ServerError with the given code using the real constructor."""
+    response_json = {"error": {"message": message, "code": code}}
+    return errors.ServerError(code, response_json)
+
+
+def _make_response(parts=None):
+    """Create a mock response with the given parts."""
+    response = MagicMock()
+    response.parts = parts
+    return response
+
+
+def _make_text_part(text: str):
+    """Create a mock text part."""
+    part = MagicMock()
+    part.text = text
+    part.inline_data = None
+    return part
+
+
+def _make_image_part():
+    """Create a mock image part that returns a saveable image."""
+    part = MagicMock()
+    part.text = None
+    part.inline_data = MagicMock()
+    mock_image = MagicMock()
+    part.as_image.return_value = mock_image
+    return part
+
+
+class TestHandleApiError(unittest.TestCase):
+    def test_quota_exhausted_on_429(self):
+        err = _make_client_error(429, "Rate limit exceeded")
+        with self.assertRaises(QuotaExhaustedError) as ctx:
+            handle_api_error(err)
+        self.assertIn("QUOTA EXHAUSTED", str(ctx.exception))
+
+    def test_permission_denied_on_403(self):
+        err = _make_client_error(403, "Forbidden")
+        with self.assertRaises(PermissionDeniedError) as ctx:
+            handle_api_error(err)
+        self.assertIn("PERMISSION DENIED", str(ctx.exception))
+
+    def test_generic_client_error_on_400(self):
+        err = _make_client_error(400, "Invalid request")
+        with self.assertRaises(RuntimeError) as ctx:
+            handle_api_error(err)
+        self.assertIn("API CLIENT ERROR (400)", str(ctx.exception))
+
+    def test_server_error_on_500(self):
+        err = _make_server_error(500, "Internal server error")
+        with self.assertRaises(RuntimeError) as ctx:
+            handle_api_error(err)
+        self.assertIn("API SERVER ERROR (500)", str(ctx.exception))
+
+    def test_chains_original_exception(self):
+        err = _make_client_error(429, "Rate limit")
+        with self.assertRaises(QuotaExhaustedError) as ctx:
+            handle_api_error(err)
+        self.assertIs(ctx.exception.__cause__, err)
+
+
+class TestParseImageResponse(unittest.TestCase):
+    def test_raises_on_empty_parts(self):
+        response = _make_response(parts=None)
+        with self.assertRaises(NoImageError) as ctx:
+            parse_image_response(response)
+        self.assertIn("empty response", str(ctx.exception))
+
+    def test_raises_on_empty_list(self):
+        response = _make_response(parts=[])
+        with self.assertRaises(NoImageError):
+            parse_image_response(response)
+
+    def test_returns_image_and_text(self):
+        text_part = _make_text_part("Here's your image")
+        image_part = _make_image_part()
+        response = _make_response(parts=[text_part, image_part])
+
+        image, text = parse_image_response(response)
+
+        self.assertEqual(text, "Here's your image")
+        self.assertIsNotNone(image)
+
+    def test_safety_filter_detected_via_text(self):
+        text_part = _make_text_part("The image was blocked by safety policy")
+        response = _make_response(parts=[text_part])
+
+        with self.assertRaises(SafetyFilterError) as ctx:
+            parse_image_response(response)
+        self.assertIn("SAFETY FILTER", str(ctx.exception))
+
+    def test_no_image_text_only(self):
+        text_part = _make_text_part("Here is some text about your request")
+        response = _make_response(parts=[text_part])
+
+        with self.assertRaises(NoImageError) as ctx:
+            parse_image_response(response)
+        self.assertIn("text only", str(ctx.exception))
+
+    def test_image_only_no_text(self):
+        image_part = _make_image_part()
+        response = _make_response(parts=[image_part])
+
+        image, text = parse_image_response(response)
+
+        self.assertIsNone(text)
+        self.assertIsNotNone(image)
+
+    def test_safety_keywords_case_insensitive(self):
+        text_part = _make_text_part("This content was BLOCKED by our POLICY")
+        response = _make_response(parts=[text_part])
+
+        with self.assertRaises(SafetyFilterError):
+            parse_image_response(response)
+
+    def test_no_false_positive_on_unrelated_text(self):
+        text_part = _make_text_part("Just text, no special keywords present")
+        response = _make_response(parts=[text_part])
+
+        with self.assertRaises(NoImageError):
+            parse_image_response(response)
+
+
+class TestCheckQuota(unittest.TestCase):
+    def test_quota_available(self):
+        client = MagicMock()
+        image_part = _make_image_part()
+        response = _make_response(parts=[image_part])
+        client.models.generate_content.return_value = response
+
+        # Should not raise
+        check_quota(client, model="gemini-2.5-flash-image")
+
+    def test_quota_exhausted_via_api_error(self):
+        client = MagicMock()
+        err = _make_client_error(429, "Quota exceeded")
+        client.models.generate_content.side_effect = err
+
+        with self.assertRaises(QuotaExhaustedError):
+            check_quota(client, model="gemini-2.5-flash-image")
+
+    def test_no_image_in_response_raises(self):
+        client = MagicMock()
+        text_part = _make_text_part("no image here")
+        response = _make_response(parts=[text_part])
+        client.models.generate_content.return_value = response
+
+        with self.assertRaises(NoImageError):
+            check_quota(client, model="gemini-2.5-flash-image")
+
+
+if __name__ == "__main__":
+    unittest.main()