app.py

import requests
from flask import Flask, request, jsonify
import logging
import re

app = Flask(__name__)
logging.basicConfig(level=logging.DEBUG)

def preprocess_prompt(prompt):
    # 使用非捕获组 (?:...) 来匹配任意字符
    pattern = r'@startmindmap\n((?:.*?\n)*?)@endmindmap'
    match = re.search(pattern, prompt, re.DOTALL)
    if match:
        mindmap_content = match.group(1)
        processed_prompt = f"@startmindmap\n{mindmap_content}\n@endmindmap"
    else:
        processed_prompt = prompt
    return processed_prompt

@app.route('/v1/images/generations', methods=['POST'])
def generate_image():
    # 获取前端请求中的数据
    data = request.json
    logging.debug(f"Received request data: {data}")

    # 预处理 prompt
    processed_prompt = preprocess_prompt(data['prompt'])
    logging.debug(f"Processed prompt: {processed_prompt}")

    # # 将数据发送到 https://plantumll.azurewebsites.net/coder 接口
    # response = requests.post('https://plantumll.azurewebsites.net/coder', data=processed_prompt.encode('utf-8'))

    # 将数据发送到 https://plantuml-server-jetty.onrender.com/coder 接口
    response = requests.post('https://mistpe-plantuml.hf.space/coder', data=processed_prompt.encode('utf-8'))
    logging.debug(f"Response status code: {response.status_code}")

    # 检查响应状态码
    if response.status_code != 200:
        return jsonify({'error': 'Failed to generate image'}), response.status_code

    # 获取完整的响应内容
    response_data = response.content
    logging.debug(f"Response content: {response_data}")

    created = 1589478378 # 假设这是一个固定值
    # url = f"https://plantumll.azurewebsites.net/png/{response_data.decode('utf-8')}"
    url = f"https://mistpe-plantuml.hf.space/png/{response_data.decode('utf-8')}"

    # 构造返回给前端的数据
    result = {
        "created": created,
        "data": [
            {
                "url": url
            }
        ]
    }
    return jsonify(result)

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