saw-go/service/doubao.go

package service

import (
	"StuAcaWorksAI/dao"
	"StuAcaWorksAI/proto"
	"StuAcaWorksAI/worker"
	"context"
	"encoding/json"
	"fmt"
	"github.com/ollama/ollama/api"
	"github.com/volcengine/volcengine-go-sdk/service/arkruntime"
	"github.com/volcengine/volcengine-go-sdk/service/arkruntime/model"
	"github.com/volcengine/volcengine-go-sdk/service/arkruntime/utils"
	"github.com/volcengine/volcengine-go-sdk/volcengine"
	"io"
	"log"
	"time"
)

func doubao(modelParam proto.ModelParam, question, channel string, SessionID, UserID int, modelID int) {
	doubaoClient := arkruntime.NewClientWithApiKey(
		modelParam.APIKey,
		arkruntime.WithBaseUrl(modelParam.Url),
	)
	ctx := context.Background()
	messages := make([]model.ChatCompletionMessage, 0)
	err := GetDouBaoSessionHistoryMsg(SessionID, "", &messages, question)
	if err != nil {
		log.Println("get doubao session history message error:", err)
	}
	log.Println("doubao messages:", messages)

	req := model.CreateChatCompletionRequest{
		Model: modelParam.Model,
		Messages: []*model.ChatCompletionMessage{
			{
				Role: model.ChatMessageRoleSystem,
				Content: &model.ChatCompletionMessageContent{
					StringValue: volcengine.String("你是豆包，是由字节跳动开发的 AI 人工智能助手"),
				},
			},
			{
				Role: model.ChatMessageRoleUser,
				Content: &model.ChatCompletionMessageContent{
					StringValue: volcengine.String(question),
				},
			},
		},
	}

	stream, err := doubaoClient.CreateChatCompletionStream(ctx, req)
	if err != nil {
		fmt.Printf("stream chat error: %v\n", err)
		return
	}
	defer func(stream *utils.ChatCompletionStreamReader) {
		err2 := stream.Close()
		if err2 != nil {
			log.Println("close stream error:", err2)
		}
	}(stream)

	answer := ""
	for {
		recv, err3 := stream.Recv()
		if err3 == io.EOF {
			log.Println("doubao stream end:", err3)
			return
		}
		if err3 != nil {
			fmt.Printf("doubao Stream chat error: %v\n", err)
			return
		}
		go CreateAIStreamMsg(UserID, modelID, SessionID, recv, modelParam.Model)
		doubaoToGeneralMassageAndSendMsgQueue(&recv, channel, SessionID, UserID)
		choices := recv.Choices[0]
		if choices.FinishReason == proto.FinishReasonStop {
			answer += choices.Delta.Content
			break
		} else {
			answer += choices.Delta.Content
		}
	}
	//将消息存入数据库
	err, _ = dao.CreateMessage(SessionID, modelID, UserID, proto.ModelToUserMsgType, proto.MsgHasRead, answer)
	if err != nil {
		log.Println("spark create message error:", err)
	}
}

func doubaoToGeneralMassageAndSendMsgQueue(data *model.ChatCompletionStreamResponse, channel string, SessionID, UserID int) {
	var aiMsg proto.AIQueueMessage
	var wsMsg proto.WSMessage
	var apiMsg api.GenerateResponse

	apiMsg.Model = "doubao"
	choices := data.Choices[0]
	apiMsg.Response = choices.Delta.Content
	if choices.FinishReason == proto.FinishReasonStop {
		apiMsg.Done = true
	} else {
		apiMsg.Done = false
	}
	//消息队列部分
	aiMsg.Type = "doubao"
	aiMsg.Msg = apiMsg
	//ws发送消息部分
	wsMsg.Msg = aiMsg
	wsMsg.SessionID = SessionID
	wsMsg.ToID = UserID
	wsMsg.Type = "doubao"

	//发送消息
	wsMsgStr, _ := json.Marshal(wsMsg)
	worker.Publish(channel, string(wsMsgStr), time.Second*60)
}

func GetDouBaoSessionHistoryMsg(sessionID int, systemPrompt string, messages *[]model.ChatCompletionMessage, prompt string) error {
	var err error
	//添加系统
	if systemPrompt != "" {
		var message model.ChatCompletionMessage
		message.Role = model.ChatMessageRoleSystem
		message.Content = &model.ChatCompletionMessageContent{
			StringValue: volcengine.String(systemPrompt),
		}
		*messages = append(*messages, message)
	}
	//获取该会话的历史消息
	msgs := dao.FindMessageBySessionIDV2(sessionID)
	for _, v := range msgs {
		var message model.ChatCompletionMessage
		if v.Type == proto.UserToModelMsgType {
			message.Role = model.ChatMessageRoleUser
		} else {
			message.Role = model.ChatMessageRoleAssistant
		}
		message.Content = &model.ChatCompletionMessageContent{
			StringValue: volcengine.String(v.Msg),
		}
		*messages = append(*messages, message)
	}
	//添加本次请求消息
	var message model.ChatCompletionMessage
	message.Role = model.ChatMessageRoleUser
	message.Content = &model.ChatCompletionMessageContent{
		StringValue: volcengine.String(prompt),
	}
	return err
}