package com.aliyun.webdemo.controller;

import com.alibaba.dashscope.app.Application;
import com.alibaba.dashscope.app.ApplicationParam;
import com.alibaba.dashscope.app.ApplicationResult;
import com.alibaba.dashscope.exception.InputRequiredException;
import com.alibaba.dashscope.exception.NoApiKeyException;
import com.google.gson.Gson;
import com.google.gson.JsonObject;
import com.google.gson.JsonParser;
import io.reactivex.Flowable;
import org.springframework.web.bind.annotation.*;
import org.springframework.web.servlet.mvc.method.annotation.ResponseBodyEmitter;

import java.util.concurrent.atomic.AtomicInteger;

/**
 * pom 配置依賴項
 * <dependency>
 *   <groupId>com.alibaba</groupId>
 *   <artifactId>dashscope-sdk-java</artifactId>
 *   <!-- 請將 'the-latest-version' 替換為查詢到的最新版本號：https://mvnrepository.com/artifact/com.alibaba/dashscope-sdk-java -->
 *   <version>the-latest-version</version>
 * </dependency>
 *
 * 將API Key配置到環境變量 http://m.bestwisewords.com/zh/model-studio/developer-reference/configure-api-key-through-environment-variables
 */
@RestController
@RequestMapping("/chatbot")
public class ChatbotController {

    /**
     * 實現 chat 接口，支持流式返回數據
     *
     * @param query
     * @return
     */
    @RequestMapping(value = "/chat", method = RequestMethod.POST)
    public ResponseBodyEmitter streamData(@RequestBody String query) {
        ResponseBodyEmitter emitter = new ResponseBodyEmitter(180000L);
        JsonObject jsonObject = new JsonParser().parse(query).getAsJsonObject();
        try {
            streamCall(emitter, jsonObject.get("prompt").getAsString());
        } catch (NoApiKeyException | InputRequiredException e) {
            e.printStackTrace();
        }

        return emitter;
    }

    /**
     * 調用百煉應用，封裝流式返回數據
     * 返回數據格式
     * id:1
     * event:result
     * :HTTP_STATUS/200
     * data:{"output":{"session_id":"xxx","finish_reason":"null","text":"相關的問題"}}
     *
     * @param emitter
     * @param query
     * @throws NoApiKeyException
     * @throws InputRequiredException
     */
    public void streamCall(ResponseBodyEmitter emitter, String query) throws NoApiKeyException, InputRequiredException {
        // appId 填入百煉應用 ID
        ApplicationParam param = ApplicationParam.builder()
                .appId("xxx")
                .prompt(query)
                .incrementalOutput(true)
                .build();

        Application application = new Application();
        Flowable<ApplicationResult> result = application.streamCall(param);
        AtomicInteger counter = new AtomicInteger(0);
        result.blockingForEach(data -> {
            int newValue = counter.incrementAndGet();
            String resData = "id:" + newValue + "\nevent:result\n:HTTP_STATUS/200\ndata:" + new Gson().toJson(data) + "\n\n";
            emitter.send(resData.getBytes(java.nio.charset.StandardCharsets.UTF_8));

            if ("stop".equals(data.getOutput().getFinishReason())) {
                emitter.complete();
            }
        });

    }
}

from flask import Flask, request, render_template, Response
from http import HTTPStatus
from dashscope import Application
import json

app = Flask(__name__)


@app.route("/")
def home(path=""):
    """
    首頁
    """
    return render_template("index.html")


def call_with_stream(query):
    """
    流式調用百煉 API
    需要將API Key配置到環境變量
    http://m.bestwisewords.com/zh/model-studio/developer-reference/configure-api-key-through-environment-variables
    返回數據格式
    id:1
    event:result
    :HTTP_STATUS/200
    data:{"output":{"session_id":"xxx","finish_reason":"null","text":"相關的問題"}}
    """
    # appId 填入百煉應用 ID
    responses = Application.call(
        app_id="xxx",
        prompt=query,
        stream=True,
        incremental_output=True,
    )
    i = 0
    for response in responses:
        if response.status_code != HTTPStatus.OK:
            print(
                "request_id=%s\n output=%s\n usage=%s\n"
                % (response.request_id, response.output, response.usage)
            )
        else:
            data = (
                "id:"
                + str(i)
                + "\nevent:result\n:HTTP_STATUS/200\ndata:"
                + json.stringfy(response)
                + "\n\n"
            )
            yield data


@app.route("/chat", methods=["POST"])
def create():
    raw_data = request.data
    data = json.loads(raw_data.decode("utf-8"))

    return Response(
        call_with_stream(data["prompt"]), mimetype="application/octet-stream"
    )


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