From 72c55d8b1da7a7ef6eed8d1ef3120bdb60f3442a Mon Sep 17 00:00:00 2001 From: Zhang Minghan Date: Sun, 18 Feb 2024 21:52:34 +0800 Subject: [PATCH] fix: auto get role --- globals/constant.go | 1 + manager/chat_completions.go | 13 +++++++++++++ 2 files changed, 14 insertions(+) diff --git a/globals/constant.go b/globals/constant.go index 2f7bb213..4e8199bf 100644 --- a/globals/constant.go +++ b/globals/constant.go @@ -5,6 +5,7 @@ const ( User = "user" Assistant = "assistant" Tool = "tool" + Function = "function" ) const ( diff --git a/manager/chat_completions.go b/manager/chat_completions.go index c9b19075..a5d4b4c2 100644 --- a/manager/chat_completions.go +++ b/manager/chat_completions.go @@ -156,6 +156,18 @@ func getFinishReason(buffer *utils.Buffer, end bool) interface{} { return ReasonStop } +func getRole(data *globals.Chunk) string { + if data.Content != "" { + return globals.Assistant + } else if data.ToolCall != nil { + return globals.Tool + } else if data.FunctionCall != nil { + return globals.Function + } + + return "" +} + func getStreamTranshipmentForm(id string, created int64, form RelayForm, data *globals.Chunk, buffer *utils.Buffer, end bool, err error) RelayStreamResponse { return RelayStreamResponse{ Id: fmt.Sprintf("chatcmpl-%s", id), @@ -166,6 +178,7 @@ func getStreamTranshipmentForm(id string, created int64, form RelayForm, data *g { Index: 0, Delta: Message{ + Role: getRole(data), Content: data.Content, ToolCalls: data.ToolCall, FunctionCall: data.FunctionCall,