models + refactor

This commit is contained in:
xds
2026-03-17 18:02:22 +03:00
parent 14f9e7b7e9
commit 7249b7b44b
5 changed files with 67 additions and 25 deletions

3
.env
View File

@@ -1,4 +1,5 @@
BOT_TOKEN=8495170789:AAHyjjhHwwVtd9_ROnjHqPHRdnmyVr1aeaY BOT_TOKEN=8495170789:AAHyjjhHwwVtd9_ROnjHqPHRdnmyVr1aeaY
# BOT_TOKEN=8011562605:AAF3kyzrZJgii0Jx-H8Sum5Njbo0BdbsiAo # BOT_TOKEN=8011562605:AAF3kyzrZJgii0Jx-H8Sum5Njbo0BdbsiAo
GEMINI_API_KEY=AIzaSyAHzDYhgjOqZZnvOnOFRGaSkKu4OAN3kZE GEMINI_API_KEY=AIzaSyAHzDYhgjOqZZnvOnOFRGaSkKu4OAN3kZE
MONGO_HOST=mongodb://admin:super_secure_password@31.59.58.220:27017/ MONGO_HOST=mongodb://admin:super_secure_password@31.59.58.220:27017/
@@ -9,5 +10,7 @@ MINIO_SECRET_KEY=SuperSecretPassword123!
MINIO_BUCKET=ai-char MINIO_BUCKET=ai-char
MODE=production MODE=production
EXTERNAL_API_SECRET=Gt9TyQ8OAYhcELh2YCbKjdHLflZGufKHJZcG338MQDW EXTERNAL_API_SECRET=Gt9TyQ8OAYhcELh2YCbKjdHLflZGufKHJZcG338MQDW
#PROXY_URL=http://82.22.174.14:8001
PROXY_URL=http://localhost:8001
PROXY_SECRET_SALT=AbVJUkwTPcUWJWhPzmjXb5p4SYyKmYC5m1uVW7Dhi7o PROXY_SECRET_SALT=AbVJUkwTPcUWJWhPzmjXb5p4SYyKmYC5m1uVW7Dhi7o
SCHEDULER_CHARACTER_ID=69931c10721fbd539804589b SCHEDULER_CHARACTER_ID=69931c10721fbd539804589b

View File

@@ -10,9 +10,14 @@ from config import settings
logger = logging.getLogger(__name__) 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: class AIProxyAdapter:
def __init__(self, base_url: str = "http://82.22.174.14:8001", salt: str = None): def __init__(self, base_url: str = None, salt: str = None):
self.base_url = base_url.rstrip("/") self.base_url = (base_url or settings.PROXY_URL).rstrip("/")
self.salt = salt or settings.PROXY_SECRET_SALT self.salt = salt or settings.PROXY_SECRET_SALT
def _generate_headers(self) -> Dict[str, str]: def _generate_headers(self) -> Dict[str, str]:
@@ -25,6 +30,23 @@ class AIProxyAdapter:
"X-Signature": signature "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: 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. 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')}") logger.warning(f"AI Proxy generation finished with reason: {data.get('finish_reason')}")
return data.get("response") or "" return data.get("response") or ""
except httpx.HTTPStatusError as e:
self._handle_http_error(e, "AI Proxy Text")
except Exception as e: except Exception as e:
logger.error(f"AI Proxy Text Error: {e}") logger.error(f"AI Proxy Text General Error: {e}")
raise Exception(f"AI Proxy Text Error: {e}") raise AIProxyException(f"AI Proxy Text Error: {e}")
async def generate_image( async def generate_image(
self, self,
@@ -95,6 +119,8 @@ class AIProxyAdapter:
} }
return [byte_arr], metrics return [byte_arr], metrics
except httpx.HTTPStatusError as e:
self._handle_http_error(e, "AI Proxy Image")
except Exception as e: except Exception as e:
logger.error(f"AI Proxy Image Error: {e}") logger.error(f"AI Proxy Image General Error: {e}")
raise Exception(f"AI Proxy Image Error: {e}") raise AIProxyException(f"AI Proxy Image Error: {e}")

View File

@@ -12,7 +12,7 @@ from aiogram.types import BufferedInputFile
from adapters.Exception import GoogleGenerationException from adapters.Exception import GoogleGenerationException
from adapters.google_adapter import GoogleAdapter 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 adapters.s3_adapter import S3Adapter
from api.models import ( from api.models import (
FinancialReport, UsageStats, UsageByEntity, FinancialReport, UsageStats, UsageByEntity,
@@ -180,13 +180,21 @@ class GenerationService:
settings = await self.dao.settings.get_settings() settings = await self.dao.settings.get_settings()
if settings.use_ai_proxy: if settings.use_ai_proxy:
asset_urls = await self._prepare_asset_urls(asset_ids) if asset_ids else None asset_urls = await self._prepare_asset_urls(asset_ids) if asset_ids else None
generated_images_io, metrics = await self.ai_proxy.generate_image( try:
prompt=generation_prompt, generated_images_io, metrics = await self.ai_proxy.generate_image(
aspect_ratio=generation.aspect_ratio, prompt=generation_prompt,
quality=generation.quality, aspect_ratio=generation.aspect_ratio,
model=generation.model or "gemini-3-pro-image-preview", quality=generation.quality,
asset_urls=asset_urls 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 = [] generated_bytes_list = []
for img_io in generated_images_io: for img_io in generated_images_io:
img_io.seek(0) img_io.seek(0)
@@ -402,11 +410,16 @@ class GenerationService:
generation.input_token_usage = metrics.get("input_token_usage") generation.input_token_usage = metrics.get("input_token_usage")
generation.output_token_usage = metrics.get("output_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}") logger.error(f"Generation {generation.id} failed: {error}")
generation.status = GenerationStatus.FAILED 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) generation.failed_reason = str(error)
elif not generation.failed_reason: elif not generation.failed_reason:
generation.failed_reason = "Unknown error" generation.failed_reason = "Unknown error"

View File

@@ -33,7 +33,7 @@ class Settings(BaseSettings):
# AI Proxy Security # AI Proxy Security
PROXY_SECRET_SALT: str = "AbVJUkwTPcUWJWhPzmjXb5p4SYyKmYC5m1uVW7Dhi7o" PROXY_SECRET_SALT: str = "AbVJUkwTPcUWJWhPzmjXb5p4SYyKmYC5m1uVW7Dhi7o"
PROXY_URL:str = "http://localhost:8001"
# JWT Security # JWT Security
SECRET_KEY: str = "CHANGE_ME_TO_A_SUPER_SECRET_KEY" SECRET_KEY: str = "CHANGE_ME_TO_A_SUPER_SECRET_KEY"
ALGORITHM: str = "HS256" ALGORITHM: str = "HS256"

View File

@@ -7,7 +7,7 @@ from aiogram import Bot
from aiogram.types import BufferedInputFile, InlineKeyboardButton, InlineKeyboardMarkup from aiogram.types import BufferedInputFile, InlineKeyboardButton, InlineKeyboardMarkup
from adapters.google_adapter import GoogleAdapter 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 adapters.s3_adapter import S3Adapter
from api.service.generation_service import GenerationService from api.service.generation_service import GenerationService
from models.Asset import Asset from models.Asset import Asset
@@ -272,11 +272,11 @@ class DailyScheduler:
try: try:
# Delegate all heavy lifting to the existing service # Delegate all heavy lifting to the existing service
await self.generation_service.create_generation(generation) await self.generation_service.create_generation(generation)
except Exception: except AIProxyException as e:
# create_generation doesn't mark FAILED itself — the caller (_queued_generation_runner) does. # error code is already saved by create_generation
# So we need to handle failure here. raise ValueError(f"AI Proxy Error: {e.error_code or str(e)}")
await self.generation_service._handle_generation_failure(generation, Exception("Image generation failed")) except Exception as e:
raise raise ValueError(f"Image generation failed: {str(e)}")
# After create_generation, generation.result_list is populated # After create_generation, generation.result_list is populated
if not generation.result_list: if not generation.result_list: