xinyin025/one-api
1
0
Fork 0
mirror of https://github.com/songquanpeng/one-api.git synced 2025-11-17 13:43:42 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

🐛 fix: stream mode delay issue (#53)

Browse Source
This commit is contained in:
Buer
2024-01-25 11:56:31 +08:00
committed by GitHub
parent 705804e6dd
commit d7193b8e46
20 changed files with 291 additions and 262 deletions
Download Patch File Download Diff File Expand all files Collapse all files

34
providers/openai/chat.go
View File

@@ -2,16 +2,19 @@ package openai
import (
"encoding/json"
"io"
"net/http"
"one-api/common"
"one-api/common/requester"
"one-api/types"
"strings"
"time"
)
type OpenAIStreamHandler struct {
Usage *types.Usage
ModelName string
isAzure bool
}
func (p *OpenAIProvider) CreateChatCompletion(request *types.ChatCompletionRequest) (openaiResponse *types.ChatCompletionResponse, errWithCode *types.OpenAIErrorWithStatusCode) {
@@ -43,7 +46,7 @@ func (p *OpenAIProvider) CreateChatCompletion(request *types.ChatCompletionReque
return &response.ChatCompletionResponse, nil
}
func (p *OpenAIProvider) CreateChatCompletionStream(request *types.ChatCompletionRequest) (requester.StreamReaderInterface[types.ChatCompletionStreamResponse], *types.OpenAIErrorWithStatusCode) {
func (p *OpenAIProvider) CreateChatCompletionStream(request *types.ChatCompletionRequest) (requester.StreamReaderInterface[string], *types.OpenAIErrorWithStatusCode) {
req, errWithCode := p.GetRequestTextBody(common.RelayModeChatCompletions, request.Model, request)
if errWithCode != nil {
return nil, errWithCode
@@ -59,16 +62,17 @@ func (p *OpenAIProvider) CreateChatCompletionStream(request *types.ChatCompletio
chatHandler := OpenAIStreamHandler{
Usage: p.Usage,
ModelName: request.Model,
isAzure: p.IsAzure,
}
return requester.RequestStream[types.ChatCompletionStreamResponse](p.Requester, resp, chatHandler.HandlerChatStream)
return requester.RequestStream[string](p.Requester, resp, chatHandler.HandlerChatStream)
}
func (h *OpenAIStreamHandler) HandlerChatStream(rawLine *[]byte, isFinished *bool, response *[]types.ChatCompletionStreamResponse) error {
func (h *OpenAIStreamHandler) HandlerChatStream(rawLine *[]byte, dataChan chan string, errChan chan error) {
// 如果rawLine 前缀不为data:，则直接返回
if !strings.HasPrefix(string(*rawLine), "data: ") {
*rawLine = nil
return nil
return
}
// 去除前缀
@@ -76,26 +80,32 @@ func (h *OpenAIStreamHandler) HandlerChatStream(rawLine *[]byte, isFinished *boo
// 如果等于 DONE 则结束
if string(*rawLine) == "[DONE]" {
*isFinished = true
return nil
errChan <- io.EOF
*rawLine = requester.StreamClosed
return
}
var openaiResponse OpenAIProviderChatStreamResponse
err := json.Unmarshal(*rawLine, &openaiResponse)
if err != nil {
return common.ErrorToOpenAIError(err)
errChan <- common.ErrorToOpenAIError(err)
return
}
error := ErrorHandle(&openaiResponse.OpenAIErrorResponse)
if error != nil {
return error
errChan <- error
return
}
dataChan <- string(*rawLine)
if h.isAzure {
// 阻塞 20ms
time.Sleep(20 * time.Millisecond)
}
countTokenText := common.CountTokenText(openaiResponse.getResponseText(), h.ModelName)
h.Usage.CompletionTokens += countTokenText
h.Usage.TotalTokens += countTokenText
*response = append(*response, openaiResponse.ChatCompletionStreamResponse)
return nil
}

26
providers/openai/completion.go
View File

@@ -2,6 +2,7 @@ package openai
import (
"encoding/json"
"io"
"net/http"
"one-api/common"
"one-api/common/requester"
@@ -38,7 +39,7 @@ func (p *OpenAIProvider) CreateCompletion(request *types.CompletionRequest) (ope
return &response.CompletionResponse, nil
}
func (p *OpenAIProvider) CreateCompletionStream(request *types.CompletionRequest) (stream requester.StreamReaderInterface[types.CompletionResponse], errWithCode *types.OpenAIErrorWithStatusCode) {
func (p *OpenAIProvider) CreateCompletionStream(request *types.CompletionRequest) (stream requester.StreamReaderInterface[string], errWithCode *types.OpenAIErrorWithStatusCode) {
req, errWithCode := p.GetRequestTextBody(common.RelayModeChatCompletions, request.Model, request)
if errWithCode != nil {
return nil, errWithCode
@@ -56,14 +57,14 @@ func (p *OpenAIProvider) CreateCompletionStream(request *types.CompletionRequest
ModelName: request.Model,
}
return requester.RequestStream[types.CompletionResponse](p.Requester, resp, chatHandler.handlerCompletionStream)
return requester.RequestStream[string](p.Requester, resp, chatHandler.handlerCompletionStream)
}
func (h *OpenAIStreamHandler) handlerCompletionStream(rawLine *[]byte, isFinished *bool, response *[]types.CompletionResponse) error {
func (h *OpenAIStreamHandler) handlerCompletionStream(rawLine *[]byte, dataChan chan string, errChan chan error) {
// 如果rawLine 前缀不为data:，则直接返回
if !strings.HasPrefix(string(*rawLine), "data: ") {
*rawLine = nil
return nil
return
}
// 去除前缀
@@ -71,26 +72,27 @@ func (h *OpenAIStreamHandler) handlerCompletionStream(rawLine *[]byte, isFinishe
// 如果等于 DONE 则结束
if string(*rawLine) == "[DONE]" {
*isFinished = true
return nil
errChan <- io.EOF
*rawLine = requester.StreamClosed
return
}
var openaiResponse OpenAIProviderCompletionResponse
err := json.Unmarshal(*rawLine, &openaiResponse)
if err != nil {
return common.ErrorToOpenAIError(err)
errChan <- common.ErrorToOpenAIError(err)
return
}
error := ErrorHandle(&openaiResponse.OpenAIErrorResponse)
if error != nil {
return error
errChan <- error
return
}
dataChan <- string(*rawLine)
countTokenText := common.CountTokenText(openaiResponse.getResponseText(), h.ModelName)
h.Usage.CompletionTokens += countTokenText
h.Usage.TotalTokens += countTokenText
*response = append(*response, openaiResponse.CompletionResponse)
return nil
}