win_monitor/backend/app/ai_gemini.py

from __future__ import annotations

from typing import Any
from urllib.parse import quote

import httpx


DEFAULT_GEMINI_BASE_URL = "https://generativelanguage.googleapis.com/v1beta"


def normalize_base_url(base_url: str | None) -> str:
    return (base_url or DEFAULT_GEMINI_BASE_URL).rstrip("/")


def chat(provider: dict[str, Any], model: dict[str, Any], prompt: str, temperature: float) -> dict[str, Any]:
    base_url = normalize_base_url(provider.get("base_url"))
    headers = {"Content-Type": "application/json"}
    api_key = provider.get("api_key")
    if api_key:
        headers["x-goog-api-key"] = api_key

    payload = {
        "contents": [{"parts": [{"text": prompt}]}],
        "generationConfig": {"temperature": temperature},
    }
    model_name = quote(model["name"], safe="")
    with httpx.Client(timeout=120) as client:
        response = client.post(f"{base_url}/models/{model_name}:generateContent", json=payload, headers=headers)
        response.raise_for_status()
        data = response.json()

    parts = data.get("candidates", [{}])[0].get("content", {}).get("parts", [])
    content = "".join(part.get("text", "") for part in parts)
    return {"content": content, "raw_response": data}


def chat_with_images(
    provider: dict[str, Any],
    model: dict[str, Any],
    prompt: str,
    images: list[dict[str, str]],
    temperature: float,
) -> dict[str, Any]:
    base_url = normalize_base_url(provider.get("base_url"))
    headers = {"Content-Type": "application/json"}
    api_key = provider.get("api_key")
    if api_key:
        headers["x-goog-api-key"] = api_key

    parts: list[dict[str, Any]] = [{"text": prompt}]
    for image in images:
        parts.append(
            {
                "inlineData": {
                    "mimeType": image["mime_type"],
                    "data": image["base64"],
                }
            }
        )

    payload = {
        "contents": [{"parts": parts}],
        "generationConfig": {"temperature": temperature},
    }
    model_name = quote(model["name"], safe="")
    with httpx.Client(timeout=180) as client:
        response = client.post(f"{base_url}/models/{model_name}:generateContent", json=payload, headers=headers)
        response.raise_for_status()
        data = response.json()

    response_parts = data.get("candidates", [{}])[0].get("content", {}).get("parts", [])
    content = "".join(part.get("text", "") for part in response_parts)
    return {"content": content, "raw_response": data}