models + refactor

adapters/ai_proxy_adapter.py

@@ -0,0 +1,100 @@
+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}")