update main.py
This commit is contained in:
parent
329ef99282
commit
111aa883e4
31
main.py
31
main.py
@ -10,12 +10,14 @@ from fastapi.middleware.cors import CORSMiddleware
|
|||||||
from dotenv import load_dotenv
|
from dotenv import load_dotenv
|
||||||
from google import genai
|
from google import genai
|
||||||
from google.genai.types import Part
|
from google.genai.types import Part
|
||||||
|
from google.genai import types as genai_types
|
||||||
from PIL import Image
|
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()
|
load_dotenv()
|
||||||
|
|
||||||
# ---- Gemini SDK Setup ----
|
# ---- Gemini SDK + Proxy Setup ----
|
||||||
GEMINI_API_KEY = os.getenv("GEMINI_API_KEY")
|
GEMINI_API_KEY = os.getenv("GEMINI_API_KEY")
|
||||||
|
|
||||||
# Strip whitespace (common issue with .env files)
|
# Strip whitespace (common issue with .env files)
|
||||||
@ -25,11 +27,30 @@ if GEMINI_API_KEY:
|
|||||||
if not GEMINI_API_KEY:
|
if not GEMINI_API_KEY:
|
||||||
raise RuntimeError("GEMINI_API_KEY not set")
|
raise RuntimeError("GEMINI_API_KEY not set")
|
||||||
|
|
||||||
client = genai.Client(api_key=GEMINI_API_KEY)
|
|
||||||
IMAGE_MODEL = "gemini-2.5-flash-image"
|
IMAGE_MODEL = "gemini-2.5-flash-image"
|
||||||
|
|
||||||
AUTH_KEY = os.getenv("AIFIT_AUTH_KEY")
|
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()
|
app = FastAPI()
|
||||||
|
|
||||||
# (Optional) CORS if you call it from frontend directly
|
# (Optional) CORS if you call it from frontend directly
|
||||||
@ -368,7 +389,7 @@ SECOND image = outfit_photo (outfit reference).
|
|||||||
cand = response.candidates[0]
|
cand = response.candidates[0]
|
||||||
finish_reason = getattr(cand, "finish_reason", None)
|
finish_reason = getattr(cand, "finish_reason", None)
|
||||||
print(f"[DEBUG] Finish reason: {finish_reason}")
|
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):
|
if hasattr(cand, "content") and getattr(cand.content, "parts", None):
|
||||||
print(f"[DEBUG] Parts in first candidate ({len(cand.content.parts)} total):")
|
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
|
safety_ratings = cand.safety_ratings
|
||||||
if safety_ratings:
|
if safety_ratings:
|
||||||
blocked = [
|
blocked = [
|
||||||
f"{getattr(r, 'category', 'unknown')}"
|
f"{getattr(r, "category", "unknown")}"
|
||||||
for r in safety_ratings
|
for r in safety_ratings
|
||||||
if hasattr(r, "blocked") and getattr(r, "blocked", False)
|
if hasattr(r, "blocked") and getattr(r, "blocked", False)
|
||||||
]
|
]
|
||||||
|
|||||||
Loading…
Reference in New Issue
Block a user