diff --git a/adapter/adapter.go b/adapter/adapter.go index 0f75bf5c..af48fb0b 100644 --- a/adapter/adapter.go +++ b/adapter/adapter.go @@ -1,6 +1,7 @@ package adapter import ( + "chat/adapter/bing" "chat/adapter/chatgpt" "chat/adapter/claude" "chat/adapter/palm2" @@ -61,6 +62,11 @@ func NewChatRequest(props *ChatProps, hook globals.Hook) error { return slack.NewChatInstanceFromConfig().CreateStreamChatRequest(&slack.ChatProps{ Message: props.Message, }, hook) + } else if globals.IsBingModel(props.Model) { + return bing.NewChatInstanceFromConfig().CreateStreamChatRequest(&bing.ChatProps{ + Model: props.Model, + Message: props.Message, + }, hook) } return nil diff --git a/adapter/bing/chat.go b/adapter/bing/chat.go new file mode 100644 index 00000000..9bc06dff --- /dev/null +++ b/adapter/bing/chat.go @@ -0,0 +1,56 @@ +package bing + +import ( + "chat/globals" + "chat/utils" + "fmt" + "strings" +) + +type ChatProps struct { + Message []globals.Message + Model string +} + +func (c *ChatInstance) CreateStreamChatRequest(props *ChatProps, hook globals.Hook) error { + var conn *utils.WebSocket + if conn = utils.NewWebsocketClient(c.GetEndpoint()); conn == nil { + return fmt.Errorf("sparkdesk error: websocket connection failed") + } + defer conn.DeferClose() + + model, _ := strings.CutPrefix(props.Model, "bing-") + if err := conn.SendJSON(&ChatRequest{ + Prompt: props.Message[len(props.Message)-1].Content, + Cookies: c.Cookies, + Model: model, + }); err != nil { + return err + } + + for { + form := utils.ReadForm[ChatResponse](conn) + if form == nil { + return nil + } + + if form.Error != "" && form.End { + return fmt.Errorf("bing error: %s", form.Error) + } + + if err := hook(form.Response); err != nil { + return err + } + + if len(form.Suggested) > 0 { + message := "" + for _, suggested := range form.Suggested { + message += fmt.Sprintf("- %s\n", suggested) + } + + if err := hook(fmt.Sprintf("\n\n%s", message)); err != nil { + return err + } + } + } +} diff --git a/adapter/bing/struct.go b/adapter/bing/struct.go index 2135f884..c323ac3e 100644 --- a/adapter/bing/struct.go +++ b/adapter/bing/struct.go @@ -1,14 +1,28 @@ package bing import ( - "chat/globals" + "chat/utils" + "fmt" + "github.com/spf13/viper" ) type ChatInstance struct { Endpoint string - Cookies map[string]string + Cookies *map[string]interface{} } -type ChatProps struct { - Message []globals.Message +func (c *ChatInstance) GetEndpoint() string { + return fmt.Sprintf("%s/chat", c.Endpoint) +} + +func NewChatInstance(endpoint, cookies string) *ChatInstance { + form := utils.UnmarshalForm[map[string]interface{}](cookies) + return &ChatInstance{ + Endpoint: endpoint, + Cookies: form, + } +} + +func NewChatInstanceFromConfig() *ChatInstance { + return NewChatInstance(viper.GetString("bing.endpoint"), viper.GetString("bing.cookies")) } diff --git a/adapter/bing/types.go b/adapter/bing/types.go new file mode 100644 index 00000000..6491f58e --- /dev/null +++ b/adapter/bing/types.go @@ -0,0 +1,16 @@ +package bing + +// see https://github.com/Deeptrain-Community/chatnio-bing-service + +type ChatRequest struct { + Prompt string `json:"prompt"` + Cookies *map[string]interface{} `json:"cookies"` + Model string `json:"model"` +} + +type ChatResponse struct { + Response string `json:"response"` + Suggested []string `json:"suggested"` + Error string `json:"error"` + End bool `json:"end"` +} diff --git a/globals/variables.go b/globals/variables.go index bc4b9eb0..d10e16dd 100644 --- a/globals/variables.go +++ b/globals/variables.go @@ -32,6 +32,9 @@ const ( ClaudeSlack = "claude-slack" SparkDesk = "spark-desk" ChatBison001 = "chat-bison-001" + BingCreative = "bing-creative" + BingBalanced = "bing-balanced" + BingPrecise = "bing-precise" ) var GPT3TurboArray = []string{ @@ -63,6 +66,12 @@ var ClaudeModelArray = []string{ Claude2100k, } +var BingModelArray = []string{ + BingCreative, + BingBalanced, + BingPrecise, +} + var LongContextModelArray = []string{ GPT3Turbo16k, GPT3Turbo16k0613, @@ -119,6 +128,10 @@ func IsPalm2Model(model string) bool { return model == ChatBison001 } +func IsBingModel(model string) bool { + return in(model, BingModelArray) +} + func IsLongContextModel(model string) bool { return in(model, LongContextModelArray) } diff --git a/manager/transhipment.go b/manager/transhipment.go index c5e5bc98..3642051f 100644 --- a/manager/transhipment.go +++ b/manager/transhipment.go @@ -164,9 +164,9 @@ func getStreamTranshipmentForm(id string, created int64, form TranshipmentForm, }, }, Usage: Usage{ - PromptTokens: buffer.CountInputToken(), - CompletionTokens: buffer.CountOutputToken(), - TotalTokens: buffer.CountToken(), + PromptTokens: utils.MultiF(end, func() int { return buffer.CountInputToken() }, 0), + CompletionTokens: utils.MultiF(end, func() int { return buffer.CountOutputToken() }, 0), + TotalTokens: utils.MultiF(end, func() int { return buffer.CountToken() }, 0), }, Quota: buffer.GetQuota(), } diff --git a/utils/base.go b/utils/base.go index dc02eb8a..aadb97a4 100644 --- a/utils/base.go +++ b/utils/base.go @@ -101,10 +101,10 @@ func Multi[T comparable](condition bool, tval, fval T) T { } } -func MultiF[T comparable](condition bool, tval, fval func() T) T { +func MultiF[T comparable](condition bool, tval func() T, fval T) T { if condition { return tval() } else { - return fval() + return fval } }