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

app = Flask(__name__)

# 请替换为你的 DeepSeek API 地址和 API Key
API_URL = "https://api.deepseek.com"
API_KEY = os.getenv('DeepSeek_API_KEY')

def call_deepseek(prompt, max_new_tokens=150):
    """
    调用 DeepSeek API 生成文本
    :param prompt: 输入的提示文本
    :param max_new_tokens: 生成的最大新 token 数
    :return: 生成的文本
    """
    headers = {
        "Content-Type": "application/json",
        "Authorization": f"Bearer {API_KEY}"
    }
    data = {
        "prompt": prompt,
        "max_new_tokens": max_new_tokens
    }
    try:
        response = requests.post(API_URL, headers=headers, data=json.dumps(data))
        response.raise_for_status()
        result = response.json().get("generated_text")
        return result
    except requests.RequestException as e:
        print(f"调用 DeepSeek API 时请求出错: {e}")
    except (KeyError, json.JSONDecodeError) as e:
        print(f"解析 DeepSeek API 响应时出错: {e}")
    return None

@app.route('/generate', methods=['POST'])
def generate_text():
    data = request.get_json()
    prompt = data.get('prompt')
    max_new_tokens = data.get('max_new_tokens', 150)

    if not prompt:
        return jsonify({"error": "Prompt is required"}), 400

    response = call_deepseek(prompt, max_new_tokens)
    if response:
        return jsonify({"generated_text": response})
    else:
        return jsonify({"error": "Failed to generate text"}), 500

if __name__ == "__main__":
    app.run(debug=True)