image_mcp_server.py

image_mcp_server.py

import json
import logging
import os
import subprocess
from typing import Annotated
from pathlib import Path

import aiohttp
from azure.identity import AzureDeveloperCliCredential, ManagedIdentityCredential
from azure.search.documents.aio import SearchClient
from azure.search.documents.models import VectorizableTextQuery
from dotenv import load_dotenv
from fastmcp import FastMCP
from fastmcp.tools.tool import ToolResult
from fastmcp.utilities.types import File

logger = logging.getLogger(__name__)
logging.basicConfig(level=logging.WARNING)
logger.setLevel(logging.INFO)
mcp = FastMCP(
    name="ImageSearchServer",
    instructions="Search for images using natural language queries. Returns matching images from an Azure AI Search index.",
)

# Global search client (initialized on first use)
_search_client: SearchClient | None = None


def load_azd_env():
    """Get path to current azd env file and load file using python-dotenv"""
    result = subprocess.run("azd env list -o json", shell=True, capture_output=True, text=True)
    if result.returncode != 0:
        raise Exception("Error loading azd env")
    env_json = json.loads(result.stdout)
    env_file_path = None
    for entry in env_json:
        if entry["IsDefault"]:
            env_file_path = entry["DotEnvPath"]
    if not env_file_path:
        raise Exception("No default azd env file found")
    logger.info(f"Loading azd env from {env_file_path}")
    load_dotenv(env_file_path, override=True)


def get_search_client() -> SearchClient:
    """Get or create the Azure Search client."""
    global _search_client
    if _search_client is None:
        if not os.getenv("RUNNING_IN_PRODUCTION"):
            load_azd_env()

        AZURE_SEARCH_SERVICE = os.environ["AZURE_SEARCH_SERVICE"]
        AZURE_SEARCH_INDEX = os.environ["AZURE_SEARCH_INDEX"]

        if os.getenv("RUNNING_IN_PRODUCTION"):
            credential = ManagedIdentityCredential(client_id=os.environ["AZURE_CLIENT_ID"])
        else:
            credential = AzureDeveloperCliCredential(tenant_id=os.environ["AZURE_TENANT_ID"])

        _search_client = SearchClient(
            endpoint=f"https://{AZURE_SEARCH_SERVICE}.search.windows.net",
            index_name=AZURE_SEARCH_INDEX,
            credential=credential,
        )
    return _search_client


def get_image_format(url: str) -> str:
    """Extract image format from URL path."""
    url_lower = url.lower()
    if ".png" in url_lower:
        return "png"
    elif ".gif" in url_lower:
        return "gif"
    elif ".webp" in url_lower:
        return "webp"
    return "jpeg"  # Default to JPEG


async def fetch_image_bytes(url: str) -> bytes:
    """Fetch image bytes from a URL."""
    async with aiohttp.ClientSession() as session:
        async with session.get(url) as response:
            response.raise_for_status()
            return await response.read()


@mcp.tool(annotations={"readOnlyHint": True})
async def image_search(
    query: Annotated[str, "Text description of images to find (e.g., 'red dress', 'blue shirt')"],
    max_results: Annotated[int, "Maximum number of images to return (1-20)"] = 5,
) -> ToolResult:
    """
    Search for images matching a natural language query.

    Uses Azure AI Search with vector embeddings to find images that match
    the semantic meaning of your query. Returns the actual image data.
    """
    # Clamp max_results to reasonable bounds
    max_results = max(1, min(20, max_results))

    search_client = get_search_client()

    results = await search_client.search(
        search_text=None,
        top=max_results,
        vector_queries=[VectorizableTextQuery(k_nearest_neighbors=max_results, fields="embedding", text=query)],
        select="metadata_storage_path",
    )

    files: list[File] = []
    image_paths: list[str] = []
    result_index = 0
    async for result in results:
        result_index += 1
        url = result["metadata_storage_path"]
        try:
            image_bytes = await fetch_image_bytes(url)
            image_format = get_image_format(url)
            filename = os.path.basename(url.split("?", maxsplit=1)[0])
            if not filename:
                filename = f"image-{result_index}.{image_format}"
            file_basename = Path(filename).stem
            files.append(File(data=image_bytes, format=image_format, name=file_basename))
            image_paths.append(url)
            logger.info(f"Fetched image from {url} ({len(image_bytes)} bytes, format={image_format})")
        except Exception as e:
            logger.error(f"Failed to fetch image from {url}: {e}")
            continue

    return ToolResult(
        content=files,
        structured_content={
            "query": query,
            "paths": image_paths,
        },
    )


if __name__ == "__main__":
    mcp.run(transport="http", host="0.0.0.0", port=8001)