111 lines
2.9 KiB
Go
111 lines
2.9 KiB
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()
|
||
|
|
|
||
|
|
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.UserAndModelMsgType, 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)
|
||
|
|
}
|