update main.py

This commit is contained in:
Hossein 2025-11-30 18:47:46 +03:30
parent 329ef99282
commit 111aa883e4

31
main.py
View File

@ -10,12 +10,14 @@ from fastapi.middleware.cors import CORSMiddleware
from dotenv import load_dotenv
from google import genai
from google.genai.types import Part
from google.genai import types as genai_types
from PIL import Image
import httpx # ✅ for proxy support
# --- Load .env (for local development) ---
# --- Load .env (for local development only; in Docker envs are already set) ---
load_dotenv()
# ---- Gemini SDK Setup ----
# ---- Gemini SDK + Proxy Setup ----
GEMINI_API_KEY = os.getenv("GEMINI_API_KEY")
# Strip whitespace (common issue with .env files)
@ -25,11 +27,30 @@ if GEMINI_API_KEY:
if not GEMINI_API_KEY:
raise RuntimeError("GEMINI_API_KEY not set")
client = genai.Client(api_key=GEMINI_API_KEY)
IMAGE_MODEL = "gemini-2.5-flash-image"
AUTH_KEY = os.getenv("AIFIT_AUTH_KEY")
# 🔌 Proxy URL, e.g. "socks5://xray:10808"
PROXY_URL = os.getenv("AIFIT_PROXY_URL")
http_options = None
if PROXY_URL:
# Use httpx transports so ALL Gemini calls go through Xray
http_options = genai_types.HttpOptions(
client_args={
"transport": httpx.HTTPTransport(proxy=PROXY_URL),
},
async_client_args={
"transport": httpx.AsyncHTTPTransport(proxy=PROXY_URL),
},
)
client = genai.Client(
api_key=GEMINI_API_KEY,
http_options=http_options, # ✅ None locally, proxy in server
)
app = FastAPI()
# (Optional) CORS if you call it from frontend directly
@ -368,7 +389,7 @@ SECOND image = outfit_photo (outfit reference).
cand = response.candidates[0]
finish_reason = getattr(cand, "finish_reason", None)
print(f"[DEBUG] Finish reason: {finish_reason}")
print(f"[DEBUG] Safety ratings: {getattr(cand, 'safety_ratings', None)}")
print(f"[DEBUG] Safety ratings: {getattr(cand, "safety_ratings", None)}")
if hasattr(cand, "content") and getattr(cand.content, "parts", None):
print(f"[DEBUG] Parts in first candidate ({len(cand.content.parts)} total):")
@ -415,7 +436,7 @@ SECOND image = outfit_photo (outfit reference).
safety_ratings = cand.safety_ratings
if safety_ratings:
blocked = [
f"{getattr(r, 'category', 'unknown')}"
f"{getattr(r, "category", "unknown")}"
for r in safety_ratings
if hasattr(r, "blocked") and getattr(r, "blocked", False)
]