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}