package ai_model_cli

import (
	"context"
	"goskeleton/app/global/variable"
	"strings"

	"github.com/baidubce/bce-qianfan-sdk/go/qianfan"
	"github.com/gin-gonic/gin"
)

func RequestQianFan(cont *gin.Context) (r bool, c interface{}) {
	// 使用安全认证AK/SK鉴权，通过环境变量初始化；
	qianfan.GetConfig().AccessKey = variable.ConfigYml.GetString("BaiduCE.QianFanAccessKey")
	qianfan.GetConfig().SecretKey = variable.ConfigYml.GetString("BaiduCE.QianFanSecretKey")

	// 指定特定模型
	chat := qianfan.NewChatCompletion(
		qianfan.WithModel("ERNIE-4.0-8K"),
	)

	message := ""
	t := strings.TrimSpace(cont.PostForm("type"))
	b := strings.TrimSpace(cont.PostForm("background"))
	d := strings.TrimSpace(cont.PostForm("doc"))
	switch t {
	case "summary":
		message = "请在这个背景下：" + b + "\n对下文进行概括，不要改变原有语言：" + d
	case "decoration":
		message = "请在这个背景下：" + b + "\n对下文进行润色，不要改变原有语言：" + d
	case "extension":
		message = "请在这个背景下：" + b + "\n对下文进行续写，不要改变原有语言：" + d
	case "correction":
		message = "请在这个背景下：" + b + "\n对下文进行改错，不要改变原有语言：" + d
	case "translation":
		message = "请在这个背景下：" + b + "\n对下文进行翻译：" + d
	}
	resp, _ := chat.Do(
		context.TODO(),
		&qianfan.ChatCompletionRequest{
			Messages: []qianfan.ChatCompletionMessage{
				qianfan.ChatCompletionUserMessage(message),
			},
		},
	)
	return true, gin.H{"new_doc": resp.Result}
}
