extensions/visual_qa.py

import base64
import mimetypes
import os
import uuid

import requests
from dotenv import load_dotenv

from smolagents import tool


load_dotenv(override=True)

# Function to encode the image
def encode_image(image_path):
    if image_path.startswith("http"):
        user_agent = "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/119.0.0.0 Safari/537.36 Edg/119.0.0.0"
        request_kwargs = {
            "headers": {"User-Agent": user_agent},
            "stream": True,
        }

        # Send a HTTP request to the URL
        response = requests.get(image_path, **request_kwargs)
        response.raise_for_status()
        content_type = response.headers.get("content-type", "")

        extension = mimetypes.guess_extension(content_type)
        if extension is None:
            extension = ".download"

        fname = str(uuid.uuid4()) + extension
        download_path = os.path.abspath(os.path.join("downloads", fname))

        with open(download_path, "wb") as fh:
            for chunk in response.iter_content(chunk_size=512):
                fh.write(chunk)

        image_path = download_path

    with open(image_path, "rb") as image_file:
        return base64.b64encode(image_file.read()).decode("utf-8")


@tool
def visualizer(image_path: str, question: str | None = None) -> str:
    """A tool that can answer questions about attached images.

    Args:
        image_path: The path to the image on which to answer the question. This should be a local path to downloaded image.
        question: The question to answer.
    """
    import mimetypes
    import os

    import requests

    from .visual_qa import encode_image

    add_note = False
    if not question:
        add_note = True
        question = "Please write a detailed caption for this image."
    if not isinstance(image_path, str):
        raise Exception("You should provide at least `image_path` string argument to this tool!")

    api_key = os.getenv("GEMINI_API_KEY")
    if not api_key:
        raise Exception("Google API key not found. Please set the GEMINI_API_KEY environment variable.")

    mime_type, _ = mimetypes.guess_type(image_path)
    base64_image = encode_image(image_path)

    payload = {
        "contents": [
            {
                "parts": [
                    {"text": question},
                    {
                        "inline_data": {
                            "mime_type": mime_type,
                            "data": base64_image,
                        }
                    },
                ],
            }
        ],
        "generationConfig": {"maxOutputTokens": 2048},
    }
    headers = {"Content-Type": "application/json"}
    url = f"https://generativelanguage.googleapis.com/v1beta/models/gemini-2.5-flash:generateContent?key={api_key}"
    response = requests.post(url, headers=headers, json=payload)
    
    if response.status_code != 200:
        raise Exception(f"Request failed with status {response.status_code}: {response.text}")

    print(response.json())
    try:
        response_data = response.json()
        candidate = response_data["candidates"][0]
        
        # Improved error handling for specific API responses
        finish_reason = candidate.get("finishReason")
        if finish_reason == "MAX_TOKENS":
            raise Exception("The model's response was truncated because it reached the maximum token limit. The returned content may be incomplete.")
        
        if "parts" not in candidate["content"]:
            raise Exception(f"The model returned empty content. Finish Reason: {finish_reason}. Full response: {response_data}")

        output = candidate["content"]["parts"][0]["text"]
    except (KeyError, IndexError) as e:
        # Fallback for any other unexpected format
        raise Exception(f"Response format unexpected: {response.json()}") from e

    if add_note:
        output = f"You did not provide a particular question, so here is a detailed caption for the image: {output}"

    return output