Administrator
86386baa1a
Features: - 🤖 AI image generation with Gemini models - 🎬 Video planning and storyboard creation - ✨ Advanced prompt optimization and analysis - 🔧 Smart API key configuration and management - 📊 Multiple model support (Gemini 2.5 Flash/Pro) - 🛡️ Robust error handling and rate limiting - 📖 Comprehensive documentation and quick start guide Created by FloatFloat-chan with ❤️ for the ComfyUI community
208 lines
7.4 KiB
Python
208 lines
7.4 KiB
Python
import os
|
|
import asyncio
|
|
import aiohttp
|
|
import backoff
|
|
import logging
|
|
from typing import Optional, Dict, Any, List
|
|
from dataclasses import dataclass
|
|
|
|
from utils.error_handling import APIError, RateLimitError, error_handler
|
|
|
|
@dataclass
|
|
class APIConfig:
|
|
api_key: str
|
|
base_url: str = "https://generativelanguage.googleapis.com"
|
|
timeout: int = 120
|
|
max_retries: int = 3
|
|
|
|
class GoogleAIStudioClient:
|
|
def __init__(self, api_key: Optional[str] = None):
|
|
self.api_key = api_key or os.getenv("GOOGLE_API_KEY") or os.getenv("GEMINI_API_KEY")
|
|
if not self.api_key:
|
|
raise ValueError("Google AI Studio API key required. Set GOOGLE_API_KEY or GEMINI_API_KEY environment variable.")
|
|
|
|
self.config = APIConfig(api_key=self.api_key)
|
|
self.session: Optional[aiohttp.ClientSession] = None
|
|
self.logger = logging.getLogger(__name__)
|
|
|
|
async def __aenter__(self):
|
|
timeout = aiohttp.ClientTimeout(total=self.config.timeout)
|
|
self.session = aiohttp.ClientSession(
|
|
timeout=timeout,
|
|
headers={
|
|
"Content-Type": "application/json",
|
|
"x-goog-api-key": self.api_key
|
|
}
|
|
)
|
|
return self
|
|
|
|
async def __aexit__(self, exc_type, exc_val, exc_tb):
|
|
if self.session:
|
|
await self.session.close()
|
|
|
|
@backoff.on_exception(
|
|
backoff.expo,
|
|
(aiohttp.ClientError, asyncio.TimeoutError),
|
|
max_tries=3,
|
|
max_time=300
|
|
)
|
|
async def generate_content(self, prompt: str, model: str = "gemini-2.5-flash",
|
|
**kwargs) -> Dict[str, Any]:
|
|
"""Generate content using Google AI models"""
|
|
endpoint = f"/v1beta/models/{model}:generateContent"
|
|
|
|
payload = {
|
|
"contents": [{
|
|
"role": "user",
|
|
"parts": [{"text": prompt}]
|
|
}],
|
|
"generationConfig": {
|
|
"temperature": kwargs.get("temperature", 0.7),
|
|
"maxOutputTokens": kwargs.get("max_tokens", 8192),
|
|
"topP": kwargs.get("top_p", 0.95),
|
|
"topK": kwargs.get("top_k", 40)
|
|
}
|
|
}
|
|
|
|
# Add safety settings
|
|
payload["safetySettings"] = [
|
|
{
|
|
"category": "HARM_CATEGORY_HARASSMENT",
|
|
"threshold": "BLOCK_MEDIUM_AND_ABOVE"
|
|
},
|
|
{
|
|
"category": "HARM_CATEGORY_HATE_SPEECH",
|
|
"threshold": "BLOCK_MEDIUM_AND_ABOVE"
|
|
},
|
|
{
|
|
"category": "HARM_CATEGORY_SEXUALLY_EXPLICIT",
|
|
"threshold": "BLOCK_MEDIUM_AND_ABOVE"
|
|
},
|
|
{
|
|
"category": "HARM_CATEGORY_DANGEROUS_CONTENT",
|
|
"threshold": "BLOCK_MEDIUM_AND_ABOVE"
|
|
}
|
|
]
|
|
|
|
return await self._make_request(endpoint, payload)
|
|
|
|
async def generate_image(self, prompt: str, model: str = "gemini-2.5-flash",
|
|
**kwargs) -> Dict[str, Any]:
|
|
"""Generate image using Google AI models"""
|
|
# For image generation, we modify the prompt to be more specific
|
|
image_prompt = f"Generate an image: {prompt}"
|
|
|
|
# Add image-specific parameters
|
|
enhanced_kwargs = {
|
|
**kwargs,
|
|
"temperature": kwargs.get("temperature", 0.7),
|
|
"max_tokens": kwargs.get("max_tokens", 1000)
|
|
}
|
|
|
|
return await self.generate_content(image_prompt, model, **enhanced_kwargs)
|
|
|
|
async def generate_video(self, prompt: str, model: str = "gemini-2.5-flash",
|
|
duration: float = 5.0, **kwargs) -> Dict[str, Any]:
|
|
"""Generate video description using available models"""
|
|
video_prompt = f"Create a detailed video description for: {prompt}. Duration: {duration} seconds."
|
|
|
|
enhanced_kwargs = {
|
|
**kwargs,
|
|
"temperature": kwargs.get("temperature", 0.8),
|
|
"max_tokens": kwargs.get("max_tokens", 2000)
|
|
}
|
|
|
|
return await self.generate_content(video_prompt, model, **enhanced_kwargs)
|
|
|
|
@error_handler()
|
|
async def _make_request(self, endpoint: str, payload: Dict[str, Any]) -> Dict[str, Any]:
|
|
"""Make HTTP request to Google AI Studio API"""
|
|
if not self.session:
|
|
raise RuntimeError("Client session not initialized. Use 'async with' context manager.")
|
|
|
|
url = f"{self.config.base_url}{endpoint}"
|
|
self.logger.debug(f"Making request to: {url}")
|
|
|
|
try:
|
|
async with self.session.post(url, json=payload) as response:
|
|
if response.status == 429:
|
|
# Rate limit handling
|
|
retry_after = int(response.headers.get("retry-after", "60"))
|
|
error_msg = f"Rate limited. Retry after {retry_after}s"
|
|
self.logger.warning(error_msg)
|
|
raise RateLimitError(error_msg, retry_after)
|
|
|
|
if response.status == 403:
|
|
error_msg = "API key invalid or insufficient permissions"
|
|
raise APIError(error_msg)
|
|
|
|
if response.status == 400:
|
|
error_text = await response.text()
|
|
error_msg = f"Bad request: {error_text}"
|
|
raise APIError(error_msg)
|
|
|
|
response.raise_for_status()
|
|
result = await response.json()
|
|
|
|
self.logger.debug(f"API Response: {result}")
|
|
return result
|
|
|
|
except aiohttp.ClientError as e:
|
|
raise APIError(f"Network error: {str(e)}")
|
|
except asyncio.TimeoutError:
|
|
raise APIError("Request timeout")
|
|
|
|
def validate_model(self, model: str) -> bool:
|
|
"""Validate if model is supported"""
|
|
supported_models = [
|
|
"gemini-2.5-flash",
|
|
"gemini-2.5-pro",
|
|
"gemini-1.5-flash",
|
|
"gemini-1.5-pro"
|
|
]
|
|
return model in supported_models
|
|
|
|
def estimate_tokens(self, text: str) -> int:
|
|
"""Rough estimation of token count"""
|
|
# Simple estimation: ~4 characters per token
|
|
return len(text) // 4
|
|
|
|
class RateLimiter:
|
|
"""Rate limiter for Google AI Studio API calls"""
|
|
|
|
def __init__(self):
|
|
self.requests = {}
|
|
self.model_limits = {
|
|
"gemini-2.5-flash": {"rpm": 1000, "tpm": 2000000}, # Tier 1
|
|
"gemini-2.5-pro": {"rpm": 50, "tpm": 2000000},
|
|
"gemini-1.5-flash": {"rpm": 1000, "tpm": 1000000},
|
|
"gemini-1.5-pro": {"rpm": 50, "tpm": 32000}
|
|
}
|
|
|
|
async def acquire(self, model: str, estimated_tokens: int = 1000) -> None:
|
|
"""Acquire rate limit permission before making request"""
|
|
if model not in self.model_limits:
|
|
return # No limits defined for this model
|
|
|
|
now = asyncio.get_event_loop().time()
|
|
|
|
# Initialize tracking for this model if needed
|
|
if model not in self.requests:
|
|
self.requests[model] = []
|
|
|
|
# Clean old requests (older than 1 minute)
|
|
self.requests[model] = [
|
|
req_time for req_time in self.requests[model]
|
|
if now - req_time <= 60
|
|
]
|
|
|
|
limits = self.model_limits[model]
|
|
|
|
# Check RPM limit
|
|
if len(self.requests[model]) >= limits.get("rpm", float('inf')):
|
|
wait_time = 60 - (now - self.requests[model][0])
|
|
if wait_time > 0:
|
|
await asyncio.sleep(wait_time)
|
|
|
|
# Record this request
|
|
self.requests[model].append(now) |