one-api/providers/moonshot/chat.go
2024-05-21 03:10:23 +08:00

73 lines
2.1 KiB
Go

package moonshot
import (
"net/http"
"one-api/common"
"one-api/common/requester"
"one-api/providers/openai"
"one-api/types"
)
func (p *MoonshotProvider) CreateChatCompletion(request *types.ChatCompletionRequest) (openaiResponse *types.ChatCompletionResponse, errWithCode *types.OpenAIErrorWithStatusCode) {
request.ClearEmptyMessages()
req, errWithCode := p.GetRequestTextBody(common.RelayModeChatCompletions, request.Model, request)
if errWithCode != nil {
return nil, errWithCode
}
defer req.Body.Close()
response := &openai.OpenAIProviderChatResponse{}
// 发送请求
_, errWithCode = p.Requester.SendRequest(req, response, false)
if errWithCode != nil {
return nil, errWithCode
}
// 检测是否错误
openaiErr := openai.ErrorHandle(&response.OpenAIErrorResponse)
if openaiErr != nil {
errWithCode = &types.OpenAIErrorWithStatusCode{
OpenAIError: *openaiErr,
StatusCode: http.StatusBadRequest,
}
return nil, errWithCode
}
if response.Usage == nil {
response.Usage = &types.Usage{
PromptTokens: p.Usage.PromptTokens,
CompletionTokens: 0,
TotalTokens: 0,
}
// 那么需要计算
response.Usage.CompletionTokens = common.CountTokenText(response.GetContent(), request.Model)
response.Usage.TotalTokens = response.Usage.PromptTokens + response.Usage.CompletionTokens
}
*p.Usage = *response.Usage
return &response.ChatCompletionResponse, nil
}
func (p *MoonshotProvider) CreateChatCompletionStream(request *types.ChatCompletionRequest) (requester.StreamReaderInterface[string], *types.OpenAIErrorWithStatusCode) {
request.ClearEmptyMessages()
req, errWithCode := p.GetRequestTextBody(common.RelayModeChatCompletions, request.Model, request)
if errWithCode != nil {
return nil, errWithCode
}
defer req.Body.Close()
// 发送请求
resp, errWithCode := p.Requester.SendRequestRaw(req)
if errWithCode != nil {
return nil, errWithCode
}
chatHandler := openai.OpenAIStreamHandler{
Usage: p.Usage,
ModelName: request.Model,
}
return requester.RequestStream(p.Requester, resp, chatHandler.HandlerChatStream)
}