Update server.py

server.py

@@ -1,107 +1,123 @@
 import os
 import random
 import logging
-from flask import Flask, request, jsonify
+import base64
+from io import BytesIO
+from typing import Optional, List
+
+from fastapi import FastAPI, HTTPException, UploadFile, File, Form, Body
+from fastapi.middleware.cors import CORSMiddleware
+from pydantic import BaseModel
 import google.generativeai as genai
-from flask_cors import CORS
 from dotenv import load_dotenv
+from PIL import Image
 
-# Load environment variables (dùng cho chạy local, trên HF nó tự load từ Settings)
+# --- LOAD BIẾN MÔI TRƯỜNG ---
 load_dotenv()
 
-app = Flask(__name__)
-CORS(app)
+# --- CẤU HÌNH SERVER ---
+app = FastAPI(title="HT MATH UNIFIED SERVER")
 
-# Cấu hình Logging
-logging.basicConfig(level=logging.INFO)
-logger = logging.getLogger(__name__)
+# Cấu hình CORS (Cho phép cả Web và Desktop App gọi vào)
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],  # Cho phép tất cả các nguồn (Web, App Desktop)
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
 
-# --- CẤU HÌNH TỪ BIẾN MÔI TRƯỜNG ---
+# --- PHẦN 1: CẤU HÌNH CHO DESKTOP APP (Key Rotation & Models) ---
 
-# 1. Lấy danh sách API KEYS
-# Trên Hugging Face, bạn vào Settings -> Variables/Secrets tạo biến tên GEMINI_API_KEYS
-# Giá trị là các key cách nhau bằng dấu phẩy: key1,key2,key3
+# 1. Lấy danh sách API KEYS từ biến môi trường
+# Format trên HuggingFace: Key1,Key2,Key3
 api_keys_env = os.getenv("GEMINI_API_KEYS", "")
-API_KEYS = [k.strip() for k in api_keys_env.split(",") if k.strip()]
-
-if not API_KEYS:
-    logger.warning("CHÚ Ý: Chưa tìm thấy GEMINI_API_KEYS trong biến môi trường! Server sẽ không hoạt động đúng.")
+DESKTOP_API_KEYS = [k.strip() for k in api_keys_env.split(",") if k.strip()]
 
 # 2. Lấy danh sách MODELS
-# Tạo biến GEMINI_MODELS, giá trị ví dụ: gemini-2.5-flash,gemini-1.5-pro
 models_env = os.getenv("GEMINI_MODELS", "gemini-2.0-flash-exp,gemini-1.5-pro,gemini-1.5-flash")
-GEMINI_MODELS = [m.strip() for m in models_env.split(",") if m.strip()]
-
-@app.route('/')
-def home():
-    return jsonify({
-        "status": "online",
-        "message": "HT MATH Server is running",
-        "models_available": GEMINI_MODELS,
-        "keys_count": len(API_KEYS)
-    })
+DESKTOP_MODELS = [m.strip() for m in models_env.split(",") if m.strip()]
 
-@app.route('/api/models', methods=['GET'])
-def get_models():
-    """API trả về danh sách model để Client (main.py) cập nhật vào ComboBox"""
-    return jsonify({"models": GEMINI_MODELS})
+# Model dữ liệu nhận từ Desktop App
+class DesktopGenerateRequest(BaseModel):
+    prompt: str
+    model: Optional[str] = "gemini-1.5-flash"
+    image: Optional[str] = None # Base64 string
 
-@app.route('/api/generate', methods=['POST'])
-def generate_content():
-    if not API_KEYS:
-        return jsonify({'error': 'Server chưa cấu hình API Key'}), 500
+@app.get("/")
+async def root():
+    return {
+        "status": "online",
+        "server": "HT MATH UNIFIED SERVER",
+        "desktop_keys_loaded": len(DESKTOP_API_KEYS),
+        "web_support": "Active"
+    }
+
+# --- API 1: Lấy danh sách Model (Cho Desktop App) ---
+@app.get("/api/models")
+async def get_models_desktop():
+    return {"models": DESKTOP_MODELS}
+
+# --- API 2: Xử lý AI (Cho Desktop App - Có xoay vòng Key) ---
+@app.post("/api/generate")
+async def generate_content_desktop(req: DesktopGenerateRequest):
+    if not DESKTOP_API_KEYS:
+        raise HTTPException(status_code=500, detail="Server chưa cấu hình GEMINI_API_KEYS")
 
     try:
-        data = request.json
-        if not data:
-            return jsonify({'error': 'No data provided'}), 400
-
-        prompt = data.get('prompt', '')
-        image_data = data.get('image')  # Base64 string
-        model_name = data.get('model', GEMINI_MODELS[0]) # Mặc định lấy model đầu tiên nếu client không gửi
-
-        if not prompt and not image_data:
-            return jsonify({'error': 'Cần cung cấp nội dung hoặc hình ảnh'}), 400
-
-        # --- RANDOM KEY SELECTION ---
-        # Chọn ngẫu nhiên 1 key để tránh rate limit
-        current_key = random.choice(API_KEYS)
+        # 1. Chọn ngẫu nhiên 1 Key (Load Balancing)
+        current_key = random.choice(DESKTOP_API_KEYS)
         genai.configure(api_key=current_key)
         
-        # Log (ẩn bớt key để bảo mật)
-        masked_key = current_key[:5] + "..." + current_key[-3:]
-        logger.info(f"Processing request with Model: {model_name} using Key: {masked_key}")
-
-        # Khởi tạo model
+        # 2. Chọn Model
+        # Nếu model client gửi lên không có trong danh sách hỗ trợ, dùng model mặc định
+        model_name = req.model if req.model in DESKTOP_MODELS else DESKTOP_MODELS[0]
         model = genai.GenerativeModel(model_name)
 
-        # Xử lý input (Text + Image hoặc Text only)
-        content_parts = [prompt]
-        
-        if image_data:
+        # 3. Chuẩn bị nội dung gửi đi
+        content_parts = [req.prompt]
+
+        # 4. Xử lý ảnh (nếu có)
+        if req.image:
             try:
-                import base64
-                from io import BytesIO
-                from PIL import Image
-                
-                image_bytes = base64.b64decode(image_data)
+                # Desktop App gửi ảnh dạng Base64 string
+                image_bytes = base64.b64decode(req.image)
                 image = Image.open(BytesIO(image_bytes))
                 content_parts.append(image)
             except Exception as e:
-                logger.error(f"Lỗi xử lý ảnh: {str(e)}")
-                return jsonify({'error': f'Lỗi xử lý ảnh: {str(e)}'}), 400
+                raise HTTPException(status_code=400, detail=f"Lỗi xử lý ảnh base64: {str(e)}")
 
-        # Gọi Google Gemini
+        # 5. Gọi Google Gemini
         response = model.generate_content(content_parts)
         
         if response.text:
-            return jsonify({'result': response.text})
+            return {"result": response.text}
         else:
-            return jsonify({'result': 'Không có phản hồi từ Gemini'}), 500
+            raise HTTPException(status_code=500, detail="Gemini không trả về nội dung text.")
 
     except Exception as e:
-        logger.error(f"Lỗi Server: {str(e)}")
-        return jsonify({'error': str(e)}), 500
+        print(f"Lỗi Desktop API: {e}")
+        raise HTTPException(status_code=500, detail=str(e))
+
+
+# ==============================================================================
+# --- PHẦN 2: KHU VỰC CỦA WEB APP (COPY CODE CŨ VÀO DƯỚI ĐÂY) ---
+# ==============================================================================
+
+# Bạn hãy dán các hàm xử lý của Web App vào đây.
+# Ví dụ: import supabase, các hàm convert_file, login, signup...
+# Đảm bảo giữ nguyên logic cũ của Web App.
+
+# Ví dụ mẫu (Nếu bạn dùng Supabase):
+# from supabase import create_client, Client
+# url: str = os.environ.get("SUPABASE_URL")
+# key: str = os.environ.get("SUPABASE_KEY")
+# supabase: Client = create_client(url, key)
+
+# @app.post("/convert")
+# async def convert_file(...):
+#     ... code cũ của bạn ...
 
-if __name__ == '__main__':
-    app.run(host='0.0.0.0', port=7860)
+if __name__ == "__main__":
+    import uvicorn
+    uvicorn.run(app, host="0.0.0.0", port=7860)