200 lines
6.3 KiB
Python
200 lines
6.3 KiB
Python
"""
|
|
Google Gemini provider implementation.
|
|
Supports chat and image generation (Nano Banana / Nano Banana Pro).
|
|
"""
|
|
import aiohttp
|
|
from typing import Any
|
|
import torch
|
|
|
|
from .base import BaseProvider, ChatResponse
|
|
from ..image_utils import tensor_to_base64, base64_to_tensor
|
|
|
|
|
|
class GeminiProvider(BaseProvider):
|
|
"""Google Gemini API provider with image generation support."""
|
|
|
|
@property
|
|
def default_base_url(self) -> str:
|
|
return "https://generativelanguage.googleapis.com/v1beta"
|
|
|
|
def _build_contents(
|
|
self,
|
|
messages: list[dict[str, Any]],
|
|
images: list[torch.Tensor] | None = None,
|
|
) -> tuple[str | None, list[dict[str, Any]]]:
|
|
"""
|
|
Build Gemini-format contents.
|
|
|
|
Returns:
|
|
Tuple of (system_instruction, contents)
|
|
"""
|
|
system = None
|
|
contents = []
|
|
|
|
for i, msg in enumerate(messages):
|
|
role = msg["role"]
|
|
content = msg["content"]
|
|
|
|
# Extract system message
|
|
if role == "system":
|
|
system = content
|
|
continue
|
|
|
|
# Map roles
|
|
gemini_role = "user" if role == "user" else "model"
|
|
|
|
# Build parts
|
|
parts = []
|
|
|
|
# Add images to last user message
|
|
if role == "user" and images and i == len(messages) - 1:
|
|
for img in images:
|
|
b64 = tensor_to_base64(img)
|
|
parts.append({
|
|
"inline_data": {
|
|
"mime_type": "image/png",
|
|
"data": b64,
|
|
}
|
|
})
|
|
|
|
parts.append({"text": content})
|
|
contents.append({"role": gemini_role, "parts": parts})
|
|
|
|
return system, contents
|
|
|
|
async def chat(
|
|
self,
|
|
messages: list[dict[str, Any]],
|
|
model: str,
|
|
temperature: float = 1.0,
|
|
max_tokens: int = 2048,
|
|
images: list[torch.Tensor] | None = None,
|
|
enable_image_generation: bool = False,
|
|
**kwargs,
|
|
) -> ChatResponse:
|
|
"""Send chat request to Gemini API."""
|
|
|
|
url = f"{self.base_url}/models/{model}:generateContent"
|
|
|
|
# Add API key as query parameter
|
|
if "?" in url:
|
|
url += f"&key={self.api_key}"
|
|
else:
|
|
url += f"?key={self.api_key}"
|
|
|
|
headers = {"Content-Type": "application/json"}
|
|
|
|
system, contents = self._build_contents(messages, images)
|
|
|
|
payload = {
|
|
"contents": contents,
|
|
"generationConfig": {
|
|
"temperature": temperature,
|
|
"maxOutputTokens": max_tokens,
|
|
}
|
|
}
|
|
|
|
if system:
|
|
payload["systemInstruction"] = {"parts": [{"text": system}]}
|
|
|
|
# Enable image generation if requested
|
|
if enable_image_generation:
|
|
payload["generationConfig"]["responseModalities"] = ["TEXT", "IMAGE"]
|
|
|
|
async with aiohttp.ClientSession() as session:
|
|
async with session.post(url, headers=headers, json=payload) as resp:
|
|
if resp.status != 200:
|
|
error_text = await resp.text()
|
|
raise RuntimeError(f"Gemini API error {resp.status}: {error_text}")
|
|
|
|
data = await resp.json()
|
|
|
|
# Extract text and images from response
|
|
text = ""
|
|
result_image = None
|
|
|
|
candidates = data.get("candidates", [])
|
|
if candidates:
|
|
parts = candidates[0].get("content", {}).get("parts", [])
|
|
for part in parts:
|
|
if "text" in part:
|
|
text += part["text"]
|
|
elif "inlineData" in part:
|
|
inline = part["inlineData"]
|
|
if inline.get("mimeType", "").startswith("image/"):
|
|
result_image = base64_to_tensor(inline["data"])
|
|
|
|
return ChatResponse(text=text, image=result_image, raw_response=data)
|
|
|
|
async def generate_image(
|
|
self,
|
|
prompt: str,
|
|
model: str,
|
|
reference_image: torch.Tensor | None = None,
|
|
aspect_ratio: str = "1:1",
|
|
size: str = "1K",
|
|
**kwargs,
|
|
) -> ChatResponse:
|
|
"""Generate or edit image using Gemini."""
|
|
|
|
url = f"{self.base_url}/models/{model}:generateContent?key={self.api_key}"
|
|
headers = {"Content-Type": "application/json"}
|
|
|
|
# Build parts
|
|
parts = []
|
|
|
|
# Add reference image if provided
|
|
if reference_image is not None:
|
|
b64 = tensor_to_base64(reference_image)
|
|
parts.append({
|
|
"inline_data": {
|
|
"mime_type": "image/png",
|
|
"data": b64,
|
|
}
|
|
})
|
|
|
|
parts.append({"text": prompt})
|
|
|
|
payload = {
|
|
"contents": [{"parts": parts}],
|
|
"generationConfig": {
|
|
"responseModalities": ["TEXT", "IMAGE"],
|
|
}
|
|
}
|
|
|
|
# Add image config if supported
|
|
if aspect_ratio or size:
|
|
payload["generationConfig"]["imageConfig"] = {}
|
|
if aspect_ratio:
|
|
payload["generationConfig"]["imageConfig"]["aspectRatio"] = aspect_ratio
|
|
if size:
|
|
payload["generationConfig"]["imageConfig"]["imageSize"] = size
|
|
|
|
async with aiohttp.ClientSession() as session:
|
|
async with session.post(url, headers=headers, json=payload) as resp:
|
|
if resp.status != 200:
|
|
error_text = await resp.text()
|
|
raise RuntimeError(f"Gemini API error {resp.status}: {error_text}")
|
|
|
|
data = await resp.json()
|
|
|
|
# Extract text and image
|
|
text = ""
|
|
result_image = None
|
|
|
|
candidates = data.get("candidates", [])
|
|
if candidates:
|
|
parts = candidates[0].get("content", {}).get("parts", [])
|
|
for part in parts:
|
|
if "text" in part:
|
|
text += part["text"]
|
|
elif "inlineData" in part:
|
|
inline = part["inlineData"]
|
|
if inline.get("mimeType", "").startswith("image/"):
|
|
result_image = base64_to_tensor(inline["data"])
|
|
|
|
if result_image is None:
|
|
raise RuntimeError("Gemini did not return an image. Try a different prompt or model.")
|
|
|
|
return ChatResponse(text=text, image=result_image, raw_response=data)
|