101 lines
3.5 KiB
Python
101 lines
3.5 KiB
Python
import logging
|
|
import io
|
|
import httpx
|
|
import hashlib
|
|
import time
|
|
from typing import List, Tuple, Dict, Any, Optional
|
|
from datetime import datetime
|
|
from models.enums import AspectRatios, Quality
|
|
from config import settings
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
class AIProxyAdapter:
|
|
def __init__(self, base_url: str = "http://82.22.174.14:8001", salt: str = None):
|
|
self.base_url = base_url.rstrip("/")
|
|
self.salt = salt or settings.PROXY_SECRET_SALT
|
|
|
|
def _generate_headers(self) -> Dict[str, str]:
|
|
timestamp = int(time.time())
|
|
hash_input = f"{timestamp}{self.salt}".encode()
|
|
signature = hashlib.sha256(hash_input).hexdigest()
|
|
|
|
return {
|
|
"X-Timestamp": str(timestamp),
|
|
"X-Signature": signature
|
|
}
|
|
|
|
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.
|
|
"""
|
|
url = f"{self.base_url}/generate_text"
|
|
|
|
messages = [{"role": "user", "content": prompt}]
|
|
payload = {
|
|
"messages": messages,
|
|
"asset_urls": asset_urls
|
|
}
|
|
|
|
headers = self._generate_headers()
|
|
|
|
async with httpx.AsyncClient() as client:
|
|
try:
|
|
response = await client.post(url, json=payload, headers=headers, timeout=60.0)
|
|
response.raise_for_status()
|
|
data = response.json()
|
|
|
|
if data.get("finish_reason") != "STOP":
|
|
logger.warning(f"AI Proxy generation finished with reason: {data.get('finish_reason')}")
|
|
|
|
return data.get("response") or ""
|
|
except Exception as e:
|
|
logger.error(f"AI Proxy Text Error: {e}")
|
|
raise Exception(f"AI Proxy Text Error: {e}")
|
|
|
|
async def generate_image(
|
|
self,
|
|
prompt: str,
|
|
aspect_ratio: AspectRatios,
|
|
quality: Quality,
|
|
model: str = "gemini-3-pro-image-preview",
|
|
asset_urls: List[str] | None = None
|
|
) -> Tuple[List[io.BytesIO], Dict[str, Any]]:
|
|
"""
|
|
Generates image using the AI Proxy with signature verification.
|
|
"""
|
|
url = f"{self.base_url}/generate_image"
|
|
|
|
payload = {
|
|
"prompt": prompt,
|
|
"asset_urls": asset_urls
|
|
}
|
|
|
|
headers = self._generate_headers()
|
|
|
|
start_time = datetime.now()
|
|
async with httpx.AsyncClient() as client:
|
|
try:
|
|
response = await client.post(url, json=payload, headers=headers, timeout=120.0)
|
|
response.raise_for_status()
|
|
|
|
image_bytes = response.content
|
|
byte_arr = io.BytesIO(image_bytes)
|
|
byte_arr.name = f"{datetime.now().timestamp()}.png"
|
|
byte_arr.seek(0)
|
|
|
|
end_time = datetime.now()
|
|
api_duration = (end_time - start_time).total_seconds()
|
|
|
|
metrics = {
|
|
"api_execution_time_seconds": api_duration,
|
|
"token_usage": 0,
|
|
"input_token_usage": 0,
|
|
"output_token_usage": 0
|
|
}
|
|
|
|
return [byte_arr], metrics
|
|
except Exception as e:
|
|
logger.error(f"AI Proxy Image Error: {e}")
|
|
raise Exception(f"AI Proxy Image Error: {e}")
|