from flask import Flask, request, jsonify
import requests
import os

app = Flask(__name__)
GEMINI_URL = "https://generativelanguage.googleapis.com/v1beta/models/gemini-pro:generateContent"

@app.route('/proxy/chat', methods=['POST'])
def reverse_proxy():
    try:
        # 验证 API 密钥
        api_key = os.getenv("GEMINI_API_KEY")
        if not api_key:
            return jsonify({"error": "Missing API key"}), 500

        # 转发请求到 Gemini
        headers = {
            "Content-Type": "application/json",
            "x-goog-api-key": api_key
        }
        
        # 添加速率限制或自定义逻辑
        payload = {
            "contents": [{
                "parts": [{
                    "text": request.json.get("prompt", "请自我介绍")
                }]
            }]
        }

        response = requests.post(
            GEMINI_URL,
            headers=headers,
            json=payload,
            timeout=30
        )
        
        return jsonify(response.json()), response.status_code

    except requests.exceptions.Timeout:
        return jsonify({"error": "Request timeout"}), 504
    except Exception as e:
        return jsonify({"error": str(e)}), 500

if __name__ == '__main__':
    app.run(host='0.0.0.0', port=7860)