solace-agent-mesh 1.11.2__py3-none-any.whl

solace_agent_mesh/agent/tools/image_tools.py

@@ -0,0 +1,1185 @@
+"""
+Collection of Python tools for image generation, manipulation, and multimodal content analysis.
+Includes tools for image description and audio description using vision and audio APIs.
+"""
+
+import logging
+import asyncio
+import base64
+import inspect
+import json
+import os
+import uuid
+from datetime import datetime, timezone
+from typing import Any, Dict, Optional
+
+import httpx
+from google.adk.tools import ToolContext
+
+from ..utils.artifact_helpers import (
+    save_artifact_with_metadata,
+    DEFAULT_SCHEMA_MAX_KEYS,
+)
+from ..utils.context_helpers import get_original_session_id
+
+from google.genai import types as adk_types
+from .tool_definition import BuiltinTool
+from .registry import tool_registry
+
+log = logging.getLogger(__name__)
+
+async def create_image_from_description(
+    image_description: str,
+    output_filename: Optional[str] = None,
+    tool_context: ToolContext = None,
+    tool_config: Optional[Dict[str, Any]] = None,
+) -> Dict[str, Any]:
+    """
+    Generates an image based on a textual description using LiteLLM and saves it as a PNG artifact.
+    Configuration for LiteLLM (model, api_key, etc.) is expected in `tool_config`.
+
+    Args:
+        image_description: The textual prompt to use for image generation.
+        output_filename: Optional. The desired filename for the output PNG image.
+                         If not provided, a unique name like 'generated_image_<uuid>.png' will be used.
+        tool_context: The context provided by the ADK framework.
+        tool_config: Optional dictionary containing specific configuration for this tool.
+
+    Returns:
+        A dictionary containing:
+        - "status": "success" or "error".
+        - "message": A descriptive message about the outcome.
+        - "output_filename": The name of the saved image artifact (if successful).
+        - "output_version": The version of the saved image artifact (if successful).
+        - "result_preview": A brief preview message (if successful).
+    """
+    log_identifier = f"[ImageTools:create_image_from_description]"
+    if not tool_context:
+        log.error(f"{log_identifier} ToolContext is missing.")
+        return {"status": "error", "message": "ToolContext is missing."}
+
+    try:
+        inv_context = tool_context._invocation_context
+        if not inv_context:
+            raise ValueError("InvocationContext is not available.")
+
+        app_name = getattr(inv_context, "app_name", None)
+        user_id = getattr(inv_context, "user_id", None)
+        session_id = get_original_session_id(inv_context)
+        artifact_service = getattr(inv_context, "artifact_service", None)
+
+        if not all([app_name, user_id, session_id, artifact_service]):
+            missing_parts = [
+                part
+                for part, val in [
+                    ("app_name", app_name),
+                    ("user_id", user_id),
+                    ("session_id", session_id),
+                    ("artifact_service", artifact_service),
+                ]
+                if not val
+            ]
+            raise ValueError(
+                f"Missing required context parts: {', '.join(missing_parts)}"
+            )
+
+        log.info(f"{log_identifier} Processing request for session {session_id}.")
+
+        current_tool_config = tool_config if tool_config is not None else {}
+
+        if not current_tool_config:
+            log.warning(
+                f"{log_identifier} Tool-specific configuration (tool_config) is empty."
+            )
+
+        model_name = current_tool_config.get("model")
+        api_key = current_tool_config.get("api_key")
+        api_base = current_tool_config.get("api_base")
+        extra_params = current_tool_config.get("extra_params", {})
+
+        if not model_name:
+            raise ValueError("'model' configuration is missing in tool_config.")
+        if not api_key:
+            raise ValueError("'api_key' configuration is missing in tool_config.")
+        if not api_base:
+            raise ValueError("'api_base' configuration is missing in tool_config.")
+
+        if "/" in model_name:
+            original_model_name = model_name
+            model_name = model_name.split("/", 1)[-1]
+            log.debug(
+                f"{log_identifier} Original model name '{original_model_name}' processed to '{model_name}' for API call."
+            )
+
+        log.debug(
+            f"{log_identifier} Using image generation model: {model_name} via direct API call to: {api_base}"
+        )
+
+        api_url = f"{api_base.rstrip('/')}/v1/images/generations"
+        headers = {
+            "Content-Type": "application/json",
+            "Authorization": f"Bearer {api_key}",
+        }
+        payload = {"model": model_name, "prompt": image_description, **extra_params}
+
+        log.debug(
+            f"{log_identifier} Calling image generation API with prompt: '{image_description[:100]}...' and payload: {json.dumps(payload)}"
+        )
+
+        try:
+            async with httpx.AsyncClient(timeout=60.0) as client:
+                http_response = await client.post(
+                    api_url, headers=headers, json=payload
+                )
+                http_response.raise_for_status()
+                response_data = http_response.json()
+        except httpx.HTTPStatusError as hse:
+            log.error(
+                f"{log_identifier} HTTP error calling image generation API {hse.request.url}: {hse.response.status_code} - {hse.response.text}"
+            )
+            return {
+                "status": "error",
+                "message": f"API error generating image: {hse.response.status_code} - {hse.response.text}",
+            }
+        except httpx.RequestError as re:
+            log.error(
+                f"{log_identifier} Request error calling image generation API {re.request.url}: {re}"
+            )
+            return {
+                "status": "error",
+                "message": f"Request error generating image: {re}",
+            }
+        except Exception as e:
+            log.error(f"{log_identifier} Error calling image generation API: {e}")
+            return {"status": "error", "message": f"Error generating image: {e}"}
+
+        log.debug(f"{log_identifier} Image generation API response received.")
+
+        if (
+            not response_data
+            or not response_data.get("data")
+            or not response_data["data"][0]
+        ):
+            log.error(
+                f"{log_identifier} API did not return valid image data. Response: {json.dumps(response_data)}"
+            )
+            raise ValueError("Image generation API did not return valid image data.")
+
+        image_data_item = response_data["data"][0]
+        image_bytes = None
+
+        if image_data_item.get("url"):
+            image_url = image_data_item["url"]
+            log.info(f"{log_identifier} Fetching image from URL: {image_url}")
+            async with httpx.AsyncClient() as client:
+                http_response = await client.get(image_url, timeout=30.0)
+                http_response.raise_for_status()
+                image_bytes = http_response.content
+            log.info(f"{log_identifier} Image fetched successfully from URL.")
+        elif image_data_item.get("b64_json"):
+            log.info(f"{log_identifier} Decoding image from b64_json.")
+            image_bytes = base64.b64decode(image_data_item["b64_json"])
+            log.info(f"{log_identifier} Image decoded successfully from b64_json.")
+        else:
+            raise ValueError(
+                "No valid image data (URL or b64_json) found in LiteLLM response."
+            )
+
+        if not image_bytes:
+            raise ValueError("Failed to retrieve image bytes.")
+
+        final_output_filename = ""
+        if output_filename:
+            if not output_filename.lower().endswith(".png"):
+                final_output_filename = f"{output_filename}.png"
+            else:
+                final_output_filename = output_filename
+        else:
+            final_output_filename = f"generated_image_{uuid.uuid4()}.png"
+        log.debug(
+            f"{log_identifier} Determined output filename: {final_output_filename}"
+        )
+
+        output_mime_type = "image/png"
+        current_timestamp_iso = datetime.now(timezone.utc).isoformat()
+
+        metadata_dict = {
+            "description": f"Image generated from prompt: {image_description}",
+            "source_prompt": image_description,
+            "generation_tool": "direct_api",
+            "generation_model": model_name,
+            "request_timestamp": current_timestamp_iso,
+            "original_requested_filename": (
+                output_filename if output_filename else "N/A"
+            ),
+        }
+        if extra_params:
+            metadata_dict["api_request_params"] = json.dumps(extra_params)
+
+        log.info(
+            f"{log_identifier} Saving artifact '{final_output_filename}' with mime_type '{output_mime_type}'."
+        )
+        save_result = await save_artifact_with_metadata(
+            artifact_service=artifact_service,
+            app_name=app_name,
+            user_id=user_id,
+            session_id=session_id,
+            filename=final_output_filename,
+            content_bytes=image_bytes,
+            mime_type=output_mime_type,
+            metadata_dict=metadata_dict,
+            timestamp=datetime.now(timezone.utc),
+            schema_max_keys=DEFAULT_SCHEMA_MAX_KEYS,
+            tool_context=tool_context,
+        )
+
+        if save_result.get("status") == "error":
+            raise IOError(
+                f"Failed to save image artifact: {save_result.get('message', 'Unknown error')}"
+            )
+
+        log.info(
+            f"{log_identifier} Artifact '{final_output_filename}' v{save_result['data_version']} saved successfully."
+        )
+
+        return {
+            "status": "success",
+            "message": "Image generated and saved successfully.",
+            "output_filename": final_output_filename,
+            "output_version": save_result["data_version"],
+            "result_preview": f"Image '{final_output_filename}' (v{save_result['data_version']}) created from prompt: \"{image_description[:50]}...\"",
+        }
+
+    except ValueError as ve:
+        log.error(f"{log_identifier} Value error: {ve}")
+        return {"status": "error", "message": str(ve)}
+    except httpx.HTTPStatusError as hse:
+        log.error(
+            f"{log_identifier} HTTP error fetching image from URL {hse.request.url}: {hse.response.status_code} - {hse.response.text}"
+        )
+        return {
+            "status": "error",
+            "message": f"HTTP error fetching image: {hse.response.status_code}",
+        }
+    except httpx.RequestError as re:
+        log.error(
+            f"{log_identifier} Request error fetching image from URL {re.request.url}: {re}"
+        )
+        return {"status": "error", "message": f"Request error fetching image: {re}"}
+    except IOError as ioe:
+        log.error(f"{log_identifier} IO error: {ioe}")
+        return {"status": "error", "message": str(ioe)}
+    except Exception as e:
+        log.exception(
+            f"{log_identifier} Unexpected error in create_image_from_description: {e}"
+        )
+        return {"status": "error", "message": f"An unexpected error occurred: {e}"}
+
+
+def _get_image_mime_type(filename: str) -> str:
+    """Get MIME type from file extension."""
+    ext = os.path.splitext(filename)[1].lower()
+    mime_mapping = {
+        ".jpg": "image/jpeg",
+        ".jpeg": "image/jpeg",
+        ".png": "image/png",
+        ".webp": "image/webp",
+        ".gif": "image/gif",
+    }
+    return mime_mapping.get(ext, "application/octet-stream")
+
+
+def _is_supported_image_format(filename: str) -> bool:
+    """Check if the image format is supported."""
+    ext = os.path.splitext(filename)[1].lower()
+    supported_formats = {".jpg", ".jpeg", ".png", ".webp", ".gif"}
+    return ext in supported_formats
+
+
+def _create_data_url(image_bytes: bytes, mime_type: str) -> str:
+    """Create base64 data URL from image bytes."""
+    base64_data = base64.b64encode(image_bytes).decode("utf-8")
+    return f"data:{mime_type};base64,{base64_data}"
+
+
+async def describe_image(
+    image_filename: str,
+    prompt: str = "What is in this image?",
+    tool_context: ToolContext = None,
+    tool_config: Optional[Dict[str, Any]] = None,
+) -> Dict[str, Any]:
+    """
+    Describes an image using an OpenAI-compatible vision API.
+
+    Args:
+        image_filename: The filename (and optional :version) of the input image artifact.
+        prompt: Custom prompt for image analysis (default: "What is in this image?").
+        tool_context: The context provided by the ADK framework.
+        tool_config: Configuration dictionary containing model, api_base, api_key.
+
+    Returns:
+        A dictionary containing:
+        - "status": "success" or "error".
+        - "message": A descriptive message about the outcome.
+        - "description": The image description from the API (if successful).
+        - "image_filename": The name of the input image artifact (if successful).
+        - "image_version": The version of the input image artifact (if successful).
+        - "tokens_used": Token usage information from the API (if available).
+    """
+    log_identifier = f"[ImageTools:describe_image:{image_filename}]"
+    if not tool_context:
+        log.error(f"{log_identifier} ToolContext is missing.")
+        return {"status": "error", "message": "ToolContext is missing."}
+
+    try:
+        inv_context = tool_context._invocation_context
+        if not inv_context:
+            raise ValueError("InvocationContext is not available.")
+
+        app_name = getattr(inv_context, "app_name", None)
+        user_id = getattr(inv_context, "user_id", None)
+        session_id = get_original_session_id(inv_context)
+        artifact_service = getattr(inv_context, "artifact_service", None)
+
+        if not all([app_name, user_id, session_id, artifact_service]):
+            missing_parts = [
+                part
+                for part, val in [
+                    ("app_name", app_name),
+                    ("user_id", user_id),
+                    ("session_id", session_id),
+                    ("artifact_service", artifact_service),
+                ]
+                if not val
+            ]
+            raise ValueError(
+                f"Missing required context parts: {', '.join(missing_parts)}"
+            )
+
+        log.info(f"{log_identifier} Processing request for session {session_id}.")
+
+        current_tool_config = tool_config if tool_config is not None else {}
+
+        if not current_tool_config:
+            log.warning(
+                f"{log_identifier} Tool-specific configuration (tool_config) is empty."
+            )
+
+        model_name = current_tool_config.get("model")
+        api_key = current_tool_config.get("api_key")
+        api_base = current_tool_config.get("api_base")
+
+        if not model_name:
+            raise ValueError("'model' configuration is missing in tool_config.")
+        if not api_key:
+            raise ValueError("'api_key' configuration is missing in tool_config.")
+        if not api_base:
+            raise ValueError("'api_base' configuration is missing in tool_config.")
+
+        log.debug(f"{log_identifier} Using model: {model_name}, API base: {api_base}")
+
+        parts = image_filename.rsplit(":", 1)
+        filename_base_for_load = parts[0]
+        version_str = parts[1] if len(parts) > 1 else None
+        version_to_load = int(version_str) if version_str else None
+
+        if not _is_supported_image_format(filename_base_for_load):
+            raise ValueError(
+                f"Unsupported image format. Supported formats: .png, .jpg, .jpeg, .webp, .gif"
+            )
+
+        if version_to_load is None:
+            list_versions_method = getattr(artifact_service, "list_versions")
+            if inspect.iscoroutinefunction(list_versions_method):
+                versions = await list_versions_method(
+                    app_name=app_name,
+                    user_id=user_id,
+                    session_id=session_id,
+                    filename=filename_base_for_load,
+                )
+            else:
+                versions = await asyncio.to_thread(
+                    list_versions_method,
+                    app_name=app_name,
+                    user_id=user_id,
+                    session_id=session_id,
+                    filename=filename_base_for_load,
+                )
+            if not versions:
+                raise FileNotFoundError(
+                    f"Image artifact '{filename_base_for_load}' not found."
+                )
+            version_to_load = max(versions)
+            log.debug(
+                f"{log_identifier} Using latest version for input: {version_to_load}"
+            )
+
+        load_artifact_method = getattr(artifact_service, "load_artifact")
+        if inspect.iscoroutinefunction(load_artifact_method):
+            image_artifact_part = await load_artifact_method(
+                app_name=app_name,
+                user_id=user_id,
+                session_id=session_id,
+                filename=filename_base_for_load,
+                version=version_to_load,
+            )
+        else:
+            image_artifact_part = await asyncio.to_thread(
+                load_artifact_method,
+                app_name=app_name,
+                user_id=user_id,
+                session_id=session_id,
+                filename=filename_base_for_load,
+                version=version_to_load,
+            )
+
+        if not image_artifact_part or not image_artifact_part.inline_data:
+            raise FileNotFoundError(
+                f"Content for image artifact '{filename_base_for_load}' v{version_to_load} not found."
+            )
+
+        image_bytes = image_artifact_part.inline_data.data
+        log.debug(f"{log_identifier} Loaded image artifact: {len(image_bytes)} bytes")
+
+        mime_type = _get_image_mime_type(filename_base_for_load)
+        data_url = _create_data_url(image_bytes, mime_type)
+        log.debug(f"{log_identifier} Created data URL with MIME type: {mime_type}")
+
+        api_url = f"{api_base.rstrip('/')}/v1/chat/completions"
+        headers = {
+            "Content-Type": "application/json",
+            "Authorization": f"Bearer {api_key}",
+        }
+
+        request_data = {
+            "model": model_name,
+            "messages": [
+                {
+                    "role": "user",
+                    "content": [
+                        {"type": "text", "text": prompt},
+                        {"type": "image_url", "image_url": {"url": data_url}},
+                    ],
+                }
+            ],
+        }
+
+        log.debug(
+            f"{log_identifier} Calling vision API with prompt: '{prompt[:100]}...'"
+        )
+
+        async with httpx.AsyncClient(timeout=60.0) as client:
+            response = await client.post(api_url, headers=headers, json=request_data)
+            response.raise_for_status()
+            response_data = response.json()
+        log.debug(f"{log_identifier} Vision API response received.")
+
+        if not response_data.get("choices") or not response_data["choices"]:
+            raise ValueError("API response does not contain valid choices.")
+
+        choice = response_data["choices"][0]
+        if not choice.get("message") or not choice["message"].get("content"):
+            raise ValueError("API response does not contain valid message content.")
+
+        description = choice["message"]["content"]
+
+        tokens_used = response_data.get("usage", {})
+
+        log.info(
+            f"{log_identifier} Image described successfully. Description length: {len(description)} characters"
+        )
+
+        return {
+            "status": "success",
+            "message": "Image described successfully",
+            "description": description,
+            "image_filename": filename_base_for_load,
+            "image_version": version_to_load,
+            "tokens_used": tokens_used,
+        }
+
+    except FileNotFoundError as e:
+        log.warning(f"{log_identifier} File not found error: {e}")
+        return {"status": "error", "message": str(e)}
+    except ValueError as ve:
+        log.error(f"{log_identifier} Value error: {ve}")
+        return {"status": "error", "message": str(ve)}
+    except httpx.HTTPStatusError as hse:
+        log.error(
+            f"{log_identifier} HTTP error calling vision API: {hse.response.status_code} - {hse.response.text}"
+        )
+        return {"status": "error", "message": f"API error: {hse.response.status_code}"}
+    except httpx.RequestError as re:
+        log.error(f"{log_identifier} Request error calling vision API: {re}")
+        return {"status": "error", "message": f"Request error: {re}"}
+    except json.JSONDecodeError as jde:
+        log.error(f"{log_identifier} JSON decode error: {jde}")
+        return {"status": "error", "message": "Invalid JSON response from API"}
+    except Exception as e:
+        log.exception(f"{log_identifier} Unexpected error in describe_image: {e}")
+        return {"status": "error", "message": f"An unexpected error occurred: {e}"}
+
+
+def _get_audio_format(filename: str) -> str:
+    """Get audio format from file extension."""
+    ext = os.path.splitext(filename)[1].lower()
+    format_mapping = {".wav": "wav", ".mp3": "mp3"}
+    return format_mapping.get(ext, "wav")
+
+
+def _is_supported_audio_format(filename: str) -> bool:
+    """Check if the audio format is supported."""
+    ext = os.path.splitext(filename)[1].lower()
+    supported_formats = {".wav", ".mp3"}
+    return ext in supported_formats
+
+
+def _encode_audio_to_base64(audio_bytes: bytes) -> str:
+    """Encode audio bytes to base64 string."""
+    return base64.b64encode(audio_bytes).decode("utf-8")
+
+
+async def describe_audio(
+    audio_filename: str,
+    prompt: str = "What is in this recording?",
+    tool_context: ToolContext = None,
+    tool_config: Optional[Dict[str, Any]] = None,
+) -> Dict[str, Any]:
+    """
+    Describes an audio recording using an OpenAI-compatible audio API.
+
+    Args:
+        audio_filename: The filename (and optional :version) of the input audio artifact.
+        prompt: Custom prompt for audio analysis (default: "What is in this recording?").
+        tool_context: The context provided by the ADK framework.
+        tool_config: Configuration dictionary containing model, api_base, api_key.
+
+    Returns:
+        A dictionary containing:
+        - "status": "success" or "error".
+        - "message": A descriptive message about the outcome.
+        - "description": The audio description from the API (if successful).
+        - "audio_filename": The name of the input audio artifact (if successful).
+        - "audio_version": The version of the input audio artifact (if successful).
+        - "tokens_used": Token usage information from the API (if available).
+    """
+    log_identifier = f"[ImageTools:describe_audio:{audio_filename}]"
+    if not tool_context:
+        log.error(f"{log_identifier} ToolContext is missing.")
+        return {"status": "error", "message": "ToolContext is missing."}
+
+    try:
+        inv_context = tool_context._invocation_context
+        if not inv_context:
+            raise ValueError("InvocationContext is not available.")
+
+        app_name = getattr(inv_context, "app_name", None)
+        user_id = getattr(inv_context, "user_id", None)
+        session_id = get_original_session_id(inv_context)
+        artifact_service = getattr(inv_context, "artifact_service", None)
+
+        if not all([app_name, user_id, session_id, artifact_service]):
+            missing_parts = [
+                part
+                for part, val in [
+                    ("app_name", app_name),
+                    ("user_id", user_id),
+                    ("session_id", session_id),
+                    ("artifact_service", artifact_service),
+                ]
+                if not val
+            ]
+            raise ValueError(
+                f"Missing required context parts: {', '.join(missing_parts)}"
+            )
+
+        log.info(f"{log_identifier} Processing request for session {session_id}.")
+
+        current_tool_config = tool_config if tool_config is not None else {}
+
+        if not current_tool_config:
+            log.warning(
+                f"{log_identifier} Tool-specific configuration (tool_config) is empty."
+            )
+
+        model_name = current_tool_config.get("model")
+        api_key = current_tool_config.get("api_key")
+        api_base = current_tool_config.get("api_base")
+
+        if not model_name:
+            raise ValueError("'model' configuration is missing in tool_config.")
+        if not api_key:
+            raise ValueError("'api_key' configuration is missing in tool_config.")
+        if not api_base:
+            raise ValueError("'api_base' configuration is missing in tool_config.")
+
+        log.debug(f"{log_identifier} Using model: {model_name}, API base: {api_base}")
+
+        parts = audio_filename.rsplit(":", 1)
+        filename_base_for_load = parts[0]
+        version_str = parts[1] if len(parts) > 1 else None
+        version_to_load = int(version_str) if version_str else None
+
+        if not _is_supported_audio_format(filename_base_for_load):
+            raise ValueError(f"Unsupported audio format. Supported formats: .wav, .mp3")
+
+        if version_to_load is None:
+            list_versions_method = getattr(artifact_service, "list_versions")
+            if inspect.iscoroutinefunction(list_versions_method):
+                versions = await list_versions_method(
+                    app_name=app_name,
+                    user_id=user_id,
+                    session_id=session_id,
+                    filename=filename_base_for_load,
+                )
+            else:
+                versions = await asyncio.to_thread(
+                    list_versions_method,
+                    app_name=app_name,
+                    user_id=user_id,
+                    session_id=session_id,
+                    filename=filename_base_for_load,
+                )
+            if not versions:
+                raise FileNotFoundError(
+                    f"Audio artifact '{filename_base_for_load}' not found."
+                )
+            version_to_load = max(versions)
+            log.debug(
+                f"{log_identifier} Using latest version for input: {version_to_load}"
+            )
+
+        load_artifact_method = getattr(artifact_service, "load_artifact")
+        if inspect.iscoroutinefunction(load_artifact_method):
+            audio_artifact_part = await load_artifact_method(
+                app_name=app_name,
+                user_id=user_id,
+                session_id=session_id,
+                filename=filename_base_for_load,
+                version=version_to_load,
+            )
+        else:
+            audio_artifact_part = await asyncio.to_thread(
+                load_artifact_method,
+                app_name=app_name,
+                user_id=user_id,
+                session_id=session_id,
+                filename=filename_base_for_load,
+                version=version_to_load,
+            )
+
+        if not audio_artifact_part or not audio_artifact_part.inline_data:
+            raise FileNotFoundError(
+                f"Content for audio artifact '{filename_base_for_load}' v{version_to_load} not found."
+            )
+
+        audio_bytes = audio_artifact_part.inline_data.data
+        log.debug(f"{log_identifier} Loaded audio artifact: {len(audio_bytes)} bytes")
+
+        audio_format = _get_audio_format(filename_base_for_load)
+        base64_audio = _encode_audio_to_base64(audio_bytes)
+        log.debug(
+            f"{log_identifier} Encoded audio to base64 with format: {audio_format}"
+        )
+
+        api_url = f"{api_base.rstrip('/')}/v1/chat/completions"
+        headers = {
+            "Content-Type": "application/json",
+            "Authorization": f"Bearer {api_key}",
+        }
+
+        request_data = {
+            "model": model_name,
+            "modalities": ["audio", "text"],
+            "audio": {"voice": "alloy", "format": audio_format},
+            "messages": [
+                {
+                    "role": "user",
+                    "content": [
+                        {"type": "text", "text": prompt},
+                        {
+                            "type": "input_audio",
+                            "input_audio": {
+                                "data": base64_audio,
+                                "format": audio_format,
+                            },
+                        },
+                    ],
+                }
+            ],
+        }
+
+        log.debug(
+            f"{log_identifier} Calling audio API with prompt: '{prompt[:100]}...'"
+        )
+
+        async with httpx.AsyncClient(timeout=60.0) as client:
+            response = await client.post(api_url, headers=headers, json=request_data)
+            response.raise_for_status()
+            response_data = response.json()
+
+        log.debug(f"{log_identifier} Audio API response received.")
+
+        if not response_data.get("choices") or not response_data["choices"]:
+            raise ValueError("API response does not contain valid choices.")
+
+        choice = response_data["choices"][0]
+        if not choice.get("message") or not choice["message"].get("content"):
+            raise ValueError("API response does not contain valid message content.")
+
+        description = choice["message"]["content"]
+
+        tokens_used = response_data.get("usage", {})
+
+        log.info(
+            f"{log_identifier} Audio described successfully. Description length: {len(description)} characters"
+        )
+
+        return {
+            "status": "success",
+            "message": "Audio described successfully",
+            "description": description,
+            "audio_filename": filename_base_for_load,
+            "audio_version": version_to_load,
+            "tokens_used": tokens_used,
+        }
+
+    except FileNotFoundError as e:
+        log.warning(f"{log_identifier} File not found error: {e}")
+        return {"status": "error", "message": str(e)}
+    except ValueError as ve:
+        log.error(f"{log_identifier} Value error: {ve}")
+        return {"status": "error", "message": str(ve)}
+    except httpx.HTTPStatusError as hse:
+        log.error(
+            f"{log_identifier} HTTP error calling audio API: {hse.response.status_code} - {hse.response.text}"
+        )
+        return {"status": "error", "message": f"API error: {hse.response.status_code}"}
+    except httpx.RequestError as re:
+        log.error(f"{log_identifier} Request error calling audio API: {re}")
+        return {"status": "error", "message": f"Request error: {re}"}
+    except json.JSONDecodeError as jde:
+        log.error(f"{log_identifier} JSON decode error: {jde}")
+        return {"status": "error", "message": "Invalid JSON response from API"}
+    except Exception as e:
+        log.exception(f"{log_identifier} Unexpected error in describe_audio: {e}")
+        return {"status": "error", "message": f"An unexpected error occurred: {e}"}
+
+
+async def edit_image_with_gemini(
+    image_filename: str,
+    edit_prompt: str,
+    output_filename: Optional[str] = None,
+    tool_context: ToolContext = None,
+    tool_config: Optional[Dict[str, Any]] = None,
+) -> Dict[str, Any]:
+    """
+    Edits an existing image based on a text prompt using Google's Gemini 2.0 Flash Preview Image Generation model.
+
+    Args:
+        image_filename: The filename (and optional :version) of the input image artifact.
+        edit_prompt: Text description of the desired edits to apply to the image.
+        output_filename: Optional. The desired filename for the output edited image.
+                        If not provided, a unique name like 'edited_image_<uuid>.jpg' will be used.
+        tool_context: The context provided by the ADK framework.
+        tool_config: Configuration dictionary containing gemini_api_key and model.
+
+    Returns:
+        A dictionary containing:
+        - "status": "success" or "error".
+        - "message": A descriptive message about the outcome.
+        - "output_filename": The name of the saved edited image artifact (if successful).
+        - "output_version": The version of the saved edited image artifact (if successful).
+        - "result_preview": A brief preview message (if successful).
+        - "original_filename": The name of the input image artifact (if successful).
+        - "original_version": The version of the input image artifact (if successful).
+    """
+    log_identifier = f"[ImageTools:edit_image_with_gemini:{image_filename}]"
+    if not tool_context:
+        log.error(f"{log_identifier} ToolContext is missing.")
+        return {"status": "error", "message": "ToolContext is missing."}
+
+    try:
+        try:
+            from google import genai
+            from google.genai import types
+            from PIL import Image as PILImage
+            from io import BytesIO
+        except ImportError as ie:
+            log.error(f"{log_identifier} Required dependencies not available: {ie}")
+            return {
+                "status": "error",
+                "message": f"Required dependencies not available: {ie}",
+            }
+
+        inv_context = tool_context._invocation_context
+        if not inv_context:
+            raise ValueError("InvocationContext is not available.")
+
+        app_name = getattr(inv_context, "app_name", None)
+        user_id = getattr(inv_context, "user_id", None)
+        session_id = get_original_session_id(inv_context)
+        artifact_service = getattr(inv_context, "artifact_service", None)
+
+        if not all([app_name, user_id, session_id, artifact_service]):
+            missing_parts = [
+                part
+                for part, val in [
+                    ("app_name", app_name),
+                    ("user_id", user_id),
+                    ("session_id", session_id),
+                    ("artifact_service", artifact_service),
+                ]
+                if not val
+            ]
+            raise ValueError(
+                f"Missing required context parts: {', '.join(missing_parts)}"
+            )
+
+        log.info(
+            f"{log_identifier} Processing image edit request for session {session_id}."
+        )
+
+        current_tool_config = tool_config if tool_config is not None else {}
+
+        if not current_tool_config:
+            log.warning(
+                f"{log_identifier} Tool-specific configuration (tool_config) is empty."
+            )
+
+        gemini_api_key = current_tool_config.get("gemini_api_key")
+        model_name = current_tool_config.get(
+            "model", "gemini-2.0-flash-preview-image-generation"
+        )
+
+        if not gemini_api_key:
+            raise ValueError(
+                "'gemini_api_key' configuration is missing in tool_config."
+            )
+
+        log.debug(f"{log_identifier} Using Gemini model: {model_name}")
+
+        parts = image_filename.rsplit(":", 1)
+        filename_base_for_load = parts[0]
+        version_str = parts[1] if len(parts) > 1 else None
+        version_to_load = int(version_str) if version_str else None
+
+        if not _is_supported_image_format(filename_base_for_load):
+            raise ValueError(
+                f"Unsupported image format. Supported formats: .png, .jpg, .jpeg, .webp, .gif"
+            )
+
+        if version_to_load is None:
+            list_versions_method = getattr(artifact_service, "list_versions")
+            if inspect.iscoroutinefunction(list_versions_method):
+                versions = await list_versions_method(
+                    app_name=app_name,
+                    user_id=user_id,
+                    session_id=session_id,
+                    filename=filename_base_for_load,
+                )
+            else:
+                versions = await asyncio.to_thread(
+                    list_versions_method,
+                    app_name=app_name,
+                    user_id=user_id,
+                    session_id=session_id,
+                    filename=filename_base_for_load,
+                )
+            if not versions:
+                raise FileNotFoundError(
+                    f"Image artifact '{filename_base_for_load}' not found."
+                )
+            version_to_load = max(versions)
+            log.debug(
+                f"{log_identifier} Using latest version for input: {version_to_load}"
+            )
+
+        load_artifact_method = getattr(artifact_service, "load_artifact")
+        if inspect.iscoroutinefunction(load_artifact_method):
+            image_artifact_part = await load_artifact_method(
+                app_name=app_name,
+                user_id=user_id,
+                session_id=session_id,
+                filename=filename_base_for_load,
+                version=version_to_load,
+            )
+        else:
+            image_artifact_part = await asyncio.to_thread(
+                load_artifact_method,
+                app_name=app_name,
+                user_id=user_id,
+                session_id=session_id,
+                filename=filename_base_for_load,
+                version=version_to_load,
+            )
+
+        if not image_artifact_part or not image_artifact_part.inline_data:
+            raise FileNotFoundError(
+                f"Content for image artifact '{filename_base_for_load}' v{version_to_load} not found."
+            )
+
+        image_bytes = image_artifact_part.inline_data.data
+        log.debug(f"{log_identifier} Loaded image artifact: {len(image_bytes)} bytes")
+
+        try:
+            from PIL import UnidentifiedImageError
+
+            pil_image = PILImage.open(BytesIO(image_bytes))
+            log.debug(
+                f"{log_identifier} Converted to PIL Image: {pil_image.size}, mode: {pil_image.mode}"
+            )
+        except UnidentifiedImageError as e:
+            log.error(f"{log_identifier} Unidentified image error: {e}")
+            raise ValueError(f"Cannot identify image file: {e}")
+        except IOError as e:
+            log.error(f"{log_identifier} IO error: {e}")
+            raise ValueError(f"Cannot identify image file: {e}")
+        except Exception as e:
+            raise ValueError(f"Failed to process image data: {e}")
+
+        try:
+            client = genai.Client(api_key=gemini_api_key)
+            log.debug(f"{log_identifier} Initialized Gemini client")
+        except Exception as e:
+            raise ValueError(f"Failed to initialize Gemini client: {e}")
+
+        text_input = (edit_prompt,)
+
+        log.debug(
+            f"{log_identifier} Calling Gemini API with edit prompt: '{edit_prompt[:100]}...'"
+        )
+
+        try:
+            response = await asyncio.to_thread(
+                client.models.generate_content,
+                model=model_name,
+                contents=[text_input, pil_image],
+                config=types.GenerateContentConfig(
+                    response_modalities=["TEXT", "IMAGE"]
+                ),
+            )
+            log.debug(f"{log_identifier} Gemini API response received.")
+        except Exception as e:
+            raise ValueError(f"Gemini API call failed: {e}")
+
+        edited_image_bytes = None
+        response_text = None
+
+        if not response.candidates or not response.candidates[0].content.parts:
+            raise ValueError("Gemini API did not return valid content.")
+
+        for part in response.candidates[0].content.parts:
+            if part.text is not None:
+                response_text = part.text
+                log.debug(
+                    f"{log_identifier} Received text response: {response_text[:100]}..."
+                )
+            elif part.inline_data is not None:
+                edited_pil_image = PILImage.open(BytesIO(part.inline_data.data))
+                output_buffer = BytesIO()
+                if edited_pil_image.mode == "RGBA":
+                    rgb_image = PILImage.new(
+                        "RGB", edited_pil_image.size, (255, 255, 255)
+                    )
+                    rgb_image.paste(edited_pil_image, mask=edited_pil_image.split()[-1])
+                    rgb_image.save(output_buffer, format="JPEG", quality=95)
+                else:
+                    edited_pil_image.save(output_buffer, format="JPEG", quality=95)
+                edited_image_bytes = output_buffer.getvalue()
+                log.debug(
+                    f"{log_identifier} Processed edited image: {len(edited_image_bytes)} bytes"
+                )
+
+        if not edited_image_bytes:
+            raise ValueError("No edited image data received from Gemini API.")
+
+        final_output_filename = ""
+        if output_filename:
+            sane_filename = os.path.basename(output_filename)
+            if not sane_filename.lower().endswith((".jpg", ".jpeg")):
+                final_output_filename = f"{sane_filename}.jpg"
+            else:
+                final_output_filename = sane_filename
+        else:
+            base_name = os.path.splitext(filename_base_for_load)[0]
+            final_output_filename = f"edited_{base_name}_{uuid.uuid4().hex[:8]}.jpg"
+
+        log.debug(
+            f"{log_identifier} Determined output filename: {final_output_filename}"
+        )
+
+        output_mime_type = "image/jpeg"
+        current_timestamp_iso = datetime.now(timezone.utc).isoformat()
+
+        metadata_dict = {
+            "description": f"Image edited with prompt: {edit_prompt}",
+            "original_image": filename_base_for_load,
+            "original_version": version_to_load,
+            "edit_prompt": edit_prompt,
+            "editing_tool": "gemini",
+            "editing_model": model_name,
+            "request_timestamp": current_timestamp_iso,
+            "original_requested_filename": (
+                output_filename if output_filename else "N/A"
+            ),
+        }
+        if response_text:
+            metadata_dict["gemini_response_text"] = response_text
+
+        log.info(
+            f"{log_identifier} Saving edited image artifact '{final_output_filename}' with mime_type '{output_mime_type}'."
+        )
+        save_result = await save_artifact_with_metadata(
+            artifact_service=artifact_service,
+            app_name=app_name,
+            user_id=user_id,
+            session_id=session_id,
+            filename=final_output_filename,
+            content_bytes=edited_image_bytes,
+            mime_type=output_mime_type,
+            metadata_dict=metadata_dict,
+            timestamp=datetime.now(timezone.utc),
+            schema_max_keys=DEFAULT_SCHEMA_MAX_KEYS,
+            tool_context=tool_context,
+        )
+
+        if save_result.get("status") == "error":
+            raise IOError(
+                f"Failed to save edited image artifact: {save_result.get('message', 'Unknown error')}"
+            )
+
+        log.info(
+            f"{log_identifier} Edited image artifact '{final_output_filename}' v{save_result['data_version']} saved successfully."
+        )
+
+        return {
+            "status": "success",
+            "message": "Image edited and saved successfully.",
+            "output_filename": final_output_filename,
+            "output_version": save_result["data_version"],
+            "result_preview": f"Edited image '{final_output_filename}' (v{save_result['data_version']}) created from '{filename_base_for_load}' with prompt: \"{edit_prompt[:50]}...\"",
+            "original_filename": filename_base_for_load,
+            "original_version": version_to_load,
+        }
+
+    except FileNotFoundError as e:
+        log.warning(f"{log_identifier} File not found error: {e}")
+        return {"status": "error", "message": str(e)}
+    except ValueError as ve:
+        log.error(f"{log_identifier} Value error: {ve}")
+        return {"status": "error", "message": str(ve)}
+    except IOError as ioe:
+        log.error(f"{log_identifier} IO error: {ioe}")
+        return {"status": "error", "message": str(ioe)}
+    except Exception as e:
+        log.exception(
+            f"{log_identifier} Unexpected error in edit_image_with_gemini: {e}"
+        )
+        return {"status": "error", "message": f"An unexpected error occurred: {e}"}
+
+
+create_image_from_description_tool_def = BuiltinTool(
+    name="create_image_from_description",
+    implementation=create_image_from_description,
+    description="Generates an image based on a textual description using a configured image generation model (e.g., via LiteLLM) and saves it as a PNG artifact.",
+    category="image",
+    required_scopes=["tool:image:create"],
+    parameters=adk_types.Schema(
+        type=adk_types.Type.OBJECT,
+        properties={
+            "image_description": adk_types.Schema(
+                type=adk_types.Type.STRING,
+                description="The textual prompt to use for image generation.",
+            ),
+            "output_filename": adk_types.Schema(
+                type=adk_types.Type.STRING,
+                description="Optional. The desired filename for the output PNG image.",
+                nullable=True,
+            ),
+        },
+        required=["image_description"],
+    ),
+    examples=[],
+)
+
+describe_image_tool_def = BuiltinTool(
+    name="describe_image",
+    implementation=describe_image,
+    description="Describes an image using an OpenAI-compatible vision API.",
+    category="image",
+    required_scopes=["tool:image:describe"],
+    parameters=adk_types.Schema(
+        type=adk_types.Type.OBJECT,
+        properties={
+            "image_filename": adk_types.Schema(
+                type=adk_types.Type.STRING,
+                description="The filename (and optional :version) of the input image artifact.",
+            ),
+            "prompt": adk_types.Schema(
+                type=adk_types.Type.STRING,
+                description="Custom prompt for image analysis.",
+                nullable=True,
+            ),
+        },
+        required=["image_filename"],
+    ),
+    examples=[],
+)
+
+describe_audio_tool_def = BuiltinTool(
+    name="describe_audio",
+    implementation=describe_audio,
+    description="Describes an audio recording using a multimodal API.",
+    category="image",
+    required_scopes=["tool:audio:describe"],
+    parameters=adk_types.Schema(
+        type=adk_types.Type.OBJECT,
+        properties={
+            "audio_filename": adk_types.Schema(
+                type=adk_types.Type.STRING,
+                description="The filename (and optional :version) of the input audio artifact.",
+            ),
+            "prompt": adk_types.Schema(
+                type=adk_types.Type.STRING,
+                description="Custom prompt for audio analysis.",
+                nullable=True,
+            ),
+        },
+        required=["audio_filename"],
+    ),
+    examples=[],
+)
+
+edit_image_with_gemini_tool_def = BuiltinTool(
+    name="edit_image_with_gemini",
+    implementation=edit_image_with_gemini,
+    description="Edits an existing image based on a text prompt using Google's Gemini 2.0 Flash Preview Image Generation model.",
+    category="image",
+    required_scopes=["tool:image:edit"],
+    parameters=adk_types.Schema(
+        type=adk_types.Type.OBJECT,
+        properties={
+            "image_filename": adk_types.Schema(
+                type=adk_types.Type.STRING,
+                description="The filename (and optional :version) of the input image artifact.",
+            ),
+            "edit_prompt": adk_types.Schema(
+                type=adk_types.Type.STRING,
+                description="Text description of the desired edits to apply to the image.",
+            ),
+            "output_filename": adk_types.Schema(
+                type=adk_types.Type.STRING,
+                description="Optional. The desired filename for the output edited image.",
+                nullable=True,
+            ),
+        },
+        required=["image_filename", "edit_prompt"],
+    ),
+    examples=[],
+)
+
+tool_registry.register(create_image_from_description_tool_def)
+tool_registry.register(describe_image_tool_def)
+tool_registry.register(describe_audio_tool_def)
+tool_registry.register(edit_image_with_gemini_tool_def)