From 7249b7b44b90be9675e4e5d164bbc9f7367c4ae4 Mon Sep 17 00:00:00 2001 From: xds Date: Tue, 17 Mar 2026 18:02:22 +0300 Subject: [PATCH] models + refactor --- .env | 5 +++- adapters/ai_proxy_adapter.py | 38 ++++++++++++++++++++++++++----- api/service/generation_service.py | 35 +++++++++++++++++++--------- config.py | 2 +- scheduler/daily_scheduler.py | 12 +++++----- 5 files changed, 67 insertions(+), 25 deletions(-) diff --git a/.env b/.env index 6a2082b..317b643 100644 --- a/.env +++ b/.env @@ -1,4 +1,5 @@ BOT_TOKEN=8495170789:AAHyjjhHwwVtd9_ROnjHqPHRdnmyVr1aeaY + # BOT_TOKEN=8011562605:AAF3kyzrZJgii0Jx-H8Sum5Njbo0BdbsiAo GEMINI_API_KEY=AIzaSyAHzDYhgjOqZZnvOnOFRGaSkKu4OAN3kZE MONGO_HOST=mongodb://admin:super_secure_password@31.59.58.220:27017/ @@ -9,5 +10,7 @@ MINIO_SECRET_KEY=SuperSecretPassword123! MINIO_BUCKET=ai-char MODE=production EXTERNAL_API_SECRET=Gt9TyQ8OAYhcELh2YCbKjdHLflZGufKHJZcG338MQDW +#PROXY_URL=http://82.22.174.14:8001 +PROXY_URL=http://localhost:8001 PROXY_SECRET_SALT=AbVJUkwTPcUWJWhPzmjXb5p4SYyKmYC5m1uVW7Dhi7o -SCHEDULER_CHARACTER_ID=69931c10721fbd539804589b \ No newline at end of file +SCHEDULER_CHARACTER_ID=69931c10721fbd539804589b diff --git a/adapters/ai_proxy_adapter.py b/adapters/ai_proxy_adapter.py index b886241..6f2bcb0 100644 --- a/adapters/ai_proxy_adapter.py +++ b/adapters/ai_proxy_adapter.py @@ -10,9 +10,14 @@ from config import settings logger = logging.getLogger(__name__) +class AIProxyException(Exception): + def __init__(self, message: str, error_code: str | None = None): + super().__init__(message) + self.error_code = error_code + class AIProxyAdapter: - def __init__(self, base_url: str = "http://82.22.174.14:8001", salt: str = None): - self.base_url = base_url.rstrip("/") + def __init__(self, base_url: str = None, salt: str = None): + self.base_url = (base_url or settings.PROXY_URL).rstrip("/") self.salt = salt or settings.PROXY_SECRET_SALT def _generate_headers(self) -> Dict[str, str]: @@ -25,6 +30,23 @@ class AIProxyAdapter: "X-Signature": signature } + def _handle_http_error(self, e: httpx.HTTPStatusError, context: str): + error_code = None + message = str(e) + try: + error_data = e.response.json() + detail = error_data.get("detail") + if isinstance(detail, dict): + error_code = detail.get("error_code") + message = detail.get("message", message) + elif isinstance(detail, str): + message = detail + except Exception: + pass + + logger.error(f"{context} Error: {message} (code: {error_code})") + raise AIProxyException(message, error_code=error_code) + async def generate_text(self, prompt: str, model: str = "gemini-3.1-pro-preview", asset_urls: List[str] | None = None) -> str: """ Generates text using the AI Proxy with signature verification. @@ -49,9 +71,11 @@ class AIProxyAdapter: logger.warning(f"AI Proxy generation finished with reason: {data.get('finish_reason')}") return data.get("response") or "" + except httpx.HTTPStatusError as e: + self._handle_http_error(e, "AI Proxy Text") except Exception as e: - logger.error(f"AI Proxy Text Error: {e}") - raise Exception(f"AI Proxy Text Error: {e}") + logger.error(f"AI Proxy Text General Error: {e}") + raise AIProxyException(f"AI Proxy Text Error: {e}") async def generate_image( self, @@ -95,6 +119,8 @@ class AIProxyAdapter: } return [byte_arr], metrics + except httpx.HTTPStatusError as e: + self._handle_http_error(e, "AI Proxy Image") except Exception as e: - logger.error(f"AI Proxy Image Error: {e}") - raise Exception(f"AI Proxy Image Error: {e}") + logger.error(f"AI Proxy Image General Error: {e}") + raise AIProxyException(f"AI Proxy Image Error: {e}") diff --git a/api/service/generation_service.py b/api/service/generation_service.py index b8afd9a..e9af766 100644 --- a/api/service/generation_service.py +++ b/api/service/generation_service.py @@ -12,7 +12,7 @@ from aiogram.types import BufferedInputFile from adapters.Exception import GoogleGenerationException from adapters.google_adapter import GoogleAdapter -from adapters.ai_proxy_adapter import AIProxyAdapter +from adapters.ai_proxy_adapter import AIProxyAdapter, AIProxyException from adapters.s3_adapter import S3Adapter from api.models import ( FinancialReport, UsageStats, UsageByEntity, @@ -180,13 +180,21 @@ class GenerationService: settings = await self.dao.settings.get_settings() if settings.use_ai_proxy: asset_urls = await self._prepare_asset_urls(asset_ids) if asset_ids else None - generated_images_io, metrics = await self.ai_proxy.generate_image( - prompt=generation_prompt, - aspect_ratio=generation.aspect_ratio, - quality=generation.quality, - model=generation.model or "gemini-3-pro-image-preview", - asset_urls=asset_urls - ) + try: + generated_images_io, metrics = await self.ai_proxy.generate_image( + prompt=generation_prompt, + aspect_ratio=generation.aspect_ratio, + quality=generation.quality, + model=generation.model or "gemini-3-pro-image-preview", + asset_urls=asset_urls + ) + except AIProxyException as e: + await self._handle_generation_failure(generation, e) + raise + except Exception as e: + await self._handle_generation_failure(generation, e) + raise + generated_bytes_list = [] for img_io in generated_images_io: img_io.seek(0) @@ -402,11 +410,16 @@ class GenerationService: generation.input_token_usage = metrics.get("input_token_usage") generation.output_token_usage = metrics.get("output_token_usage") - async def _handle_generation_failure(self, generation: Generation, error: Optional[Exception]): + async def _handle_generation_failure(self, generation: Generation, error: Optional[Exception], error_code: Optional[str] = None): logger.error(f"Generation {generation.id} failed: {error}") generation.status = GenerationStatus.FAILED - # Don't overwrite if reason is already set, unless a new error is provided - if error: + + # Priority: explicit error_code, then AIProxyException's code, then exception message + if error_code: + generation.failed_reason = error_code + elif isinstance(error, AIProxyException) and error.error_code: + generation.failed_reason = error.error_code + elif error: generation.failed_reason = str(error) elif not generation.failed_reason: generation.failed_reason = "Unknown error" diff --git a/config.py b/config.py index e92074a..1b75cce 100644 --- a/config.py +++ b/config.py @@ -33,7 +33,7 @@ class Settings(BaseSettings): # AI Proxy Security PROXY_SECRET_SALT: str = "AbVJUkwTPcUWJWhPzmjXb5p4SYyKmYC5m1uVW7Dhi7o" - + PROXY_URL:str = "http://localhost:8001" # JWT Security SECRET_KEY: str = "CHANGE_ME_TO_A_SUPER_SECRET_KEY" ALGORITHM: str = "HS256" diff --git a/scheduler/daily_scheduler.py b/scheduler/daily_scheduler.py index 5b80a9a..0218dd1 100644 --- a/scheduler/daily_scheduler.py +++ b/scheduler/daily_scheduler.py @@ -7,7 +7,7 @@ from aiogram import Bot from aiogram.types import BufferedInputFile, InlineKeyboardButton, InlineKeyboardMarkup from adapters.google_adapter import GoogleAdapter -from adapters.ai_proxy_adapter import AIProxyAdapter +from adapters.ai_proxy_adapter import AIProxyAdapter, AIProxyException from adapters.s3_adapter import S3Adapter from api.service.generation_service import GenerationService from models.Asset import Asset @@ -272,11 +272,11 @@ class DailyScheduler: try: # Delegate all heavy lifting to the existing service await self.generation_service.create_generation(generation) - except Exception: - # create_generation doesn't mark FAILED itself — the caller (_queued_generation_runner) does. - # So we need to handle failure here. - await self.generation_service._handle_generation_failure(generation, Exception("Image generation failed")) - raise + except AIProxyException as e: + # error code is already saved by create_generation + raise ValueError(f"AI Proxy Error: {e.error_code or str(e)}") + except Exception as e: + raise ValueError(f"Image generation failed: {str(e)}") # After create_generation, generation.result_list is populated if not generation.result_list: