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

View File

@@ -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}")