models + refactor

adapters/google_adapter.py

@@ -8,7 +8,7 @@ from google import genai
 from google.genai import types
 
 from adapters.Exception import GoogleGenerationException
-from models.enums import AspectRatios, Quality
+from models.enums import AspectRatios, Quality, TextModel, ImageModel
 
 logger = logging.getLogger(__name__)
 
@@ -19,10 +19,6 @@ class GoogleAdapter:
             raise ValueError("API Key for Gemini is missing")
         self.client = genai.Client(api_key=api_key)
 
-        # Константы моделей
-        self.TEXT_MODEL = "gemini-3.1-pro-preview"
-        self.IMAGE_MODEL = "gemini-3-pro-image-preview"
-
     def _prepare_contents(self, prompt: str, images_list: List[bytes] | None = None) -> tuple:
         """Вспомогательный метод для подготовки контента (текст + картинки).
         Returns (contents, opened_images) — caller MUST close opened_images after use."""
@@ -41,16 +37,19 @@ class GoogleAdapter:
             logger.info("Preparing content with no images")
         return contents, opened_images
 
-    def generate_text(self, prompt: str, images_list: List[bytes] | None = None) -> str:
+    def generate_text(self, prompt: str, model: str = "gemini-3.1-pro-preview", images_list: List[bytes] | None = None) -> str:
         """
         Генерация текста (Чат или Vision).
         Возвращает строку с ответом.
         """
+        if model not in [m.value for m in TextModel]:
+            raise ValueError(f"Invalid model for text generation: {model}. Expected one of: {[m.value for m in TextModel]}")
+
         contents, opened_images = self._prepare_contents(prompt, images_list)
-        logger.info(f"Generating text: {prompt}")
+        logger.info(f"Generating text: {prompt} with model: {model}")
         try:
             response = self.client.models.generate_content(
-                model=self.TEXT_MODEL,
+                model=model,
                 contents=contents,
                 config=types.GenerateContentConfig(
                     response_modalities=['TEXT'],
@@ -74,21 +73,23 @@ class GoogleAdapter:
             for img in opened_images:
                 img.close()
 
-    def generate_image(self, prompt: str, aspect_ratio: AspectRatios, quality: Quality, images_list: List[bytes] | None = None, ) -> Tuple[List[io.BytesIO], Dict[str, Any]]:
+    def generate_image(self, prompt: str, aspect_ratio: AspectRatios, quality: Quality, model: str = "gemini-3-pro-image-preview", images_list: List[bytes] | None = None, ) -> Tuple[List[io.BytesIO], Dict[str, Any]]:
         """
         Генерация изображений (Text-to-Image или Image-to-Image).
         Возвращает список байтовых потоков (готовых к отправке).
         """
+        if model not in [m.value for m in ImageModel]:
+            raise ValueError(f"Invalid model for image generation: {model}. Expected one of: {[m.value for m in ImageModel]}")
 
         contents, opened_images = self._prepare_contents(prompt, images_list)
-        logger.info(f"Generating image. Prompt length: {len(prompt)}, Ratio: {aspect_ratio}, Quality: {quality}")
+        logger.info(f"Generating image. Prompt length: {len(prompt)}, Ratio: {aspect_ratio}, Quality: {quality}, Model: {model}")
 
         start_time = datetime.now()
         token_usage = 0
         
         try:
             response = self.client.models.generate_content(
-                model=self.IMAGE_MODEL,
+                model=model,
                 contents=contents,
                 config=types.GenerateContentConfig(
                     response_modalities=['IMAGE'],