99 lines
4.4 KiB
Python
99 lines
4.4 KiB
Python
import os
|
||
import io
|
||
import logging
|
||
from datetime import datetime
|
||
from typing import List, Union, Dict, Any
|
||
from PIL import Image
|
||
|
||
# Импортируем из нового SDK
|
||
from google import genai
|
||
from google.genai import types
|
||
|
||
# Для настройки логгера
|
||
logger = logging.getLogger(__name__)
|
||
|
||
|
||
class GoogleAdapter:
|
||
def __init__(self, api_key: str):
|
||
if not api_key:
|
||
raise ValueError("API Key for Gemini is missing")
|
||
self.client = genai.Client(api_key=api_key)
|
||
# Укажите актуальную модель.
|
||
# Если gemini-3-pro-image-preview недоступна, используйте gemini-2.0-flash-exp
|
||
self.model_name = "gemini-3-pro-preview"
|
||
|
||
def generate(
|
||
self,
|
||
prompt: str,
|
||
image_bytes: bytes = None,
|
||
generate_image: bool = False
|
||
) -> Dict[str, Any]:
|
||
"""
|
||
Универсальный метод:
|
||
- Если generate_image=True: просим модель вернуть картинку (Image Generation).
|
||
- Если image_bytes переданы + generate_image=False: это Vision (описание фото).
|
||
- Если image_bytes + generate_image=True: это Image-to-Image (редактирование).
|
||
"""
|
||
if generate_image:
|
||
self.model_name = "gemini-3-pro-image-preview"
|
||
else :
|
||
self.model_name = "gemini-3-pro-preview"
|
||
contents = [prompt]
|
||
|
||
# Если есть входное изображение (для Vision или для редактирования)
|
||
if image_bytes:
|
||
try:
|
||
image = Image.open(io.BytesIO(image_bytes))
|
||
contents.append(image)
|
||
except Exception as e:
|
||
logger.error(f"Error processing input image: {e}")
|
||
return {"error": "Не удалось обработать входящее изображение."}
|
||
|
||
# Настраиваем конфигурацию
|
||
# Для генерации картинок добавляем 'IMAGE' в response_modalities
|
||
modalities = ['TEXT']
|
||
if generate_image:
|
||
modalities.append('IMAGE')
|
||
|
||
try:
|
||
# Вызов API (синхронный метод в обертке, но aiogram вызывает его в треде,
|
||
# либо используйте client.aio для асинхронности если поддерживается версией SDK)
|
||
# В google-genai v0.3+ есть асинхронный клиент, но для простоты здесь стандартный вызов.
|
||
# Чтобы не блокировать event loop, в main.py мы обернем это в to_thread при необходимости,
|
||
# но пока используем стандартный вызов.
|
||
|
||
response = self.client.models.generate_content(
|
||
model=self.model_name,
|
||
contents=contents,
|
||
config=types.GenerateContentConfig(
|
||
response_modalities=modalities,
|
||
temperature=0.7 if not generate_image else 1.0,
|
||
)
|
||
)
|
||
|
||
result = {"text": "", "images": []}
|
||
|
||
# Парсим ответ (Text или Inline Data)
|
||
if response.parts:
|
||
for part in response.parts:
|
||
if part.text:
|
||
result["text"] += part.text
|
||
|
||
# Проверяем наличие сгенерированного изображения
|
||
if part.inline_data:
|
||
# ИСПРАВЛЕНИЕ: Берем "сырые" байты напрямую из ответа
|
||
# Это работает быстрее и не вызывает ошибку с PIL
|
||
|
||
# part.inline_data.data — это уже bytes
|
||
byte_arr = io.BytesIO(part.inline_data.data)
|
||
now = datetime.now()
|
||
# Имя файла для телеграма (формально)
|
||
byte_arr.name = f'{now.timestamp()}.png'
|
||
|
||
result["images"].append(byte_arr)
|
||
|
||
return result
|
||
|
||
except Exception as e:
|
||
logger.error(f"Gemini API Error: {e}")
|
||
return {"error": f"Ошибка API: {str(e)}"} |