feat: cloudflare support native openai api (#1596)

This commit is contained in:
zijiren 2024-07-06 13:12:30 +08:00 committed by GitHub
parent d7a78f3397
commit efd30a40b3
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
3 changed files with 55 additions and 81 deletions

View File

@ -10,6 +10,7 @@ import (
"github.com/songquanpeng/one-api/relay/adaptor" "github.com/songquanpeng/one-api/relay/adaptor"
"github.com/songquanpeng/one-api/relay/meta" "github.com/songquanpeng/one-api/relay/meta"
"github.com/songquanpeng/one-api/relay/model" "github.com/songquanpeng/one-api/relay/model"
"github.com/songquanpeng/one-api/relay/relaymode"
) )
type Adaptor struct { type Adaptor struct {
@ -28,7 +29,14 @@ func (a *Adaptor) Init(meta *meta.Meta) {
} }
func (a *Adaptor) GetRequestURL(meta *meta.Meta) (string, error) { func (a *Adaptor) GetRequestURL(meta *meta.Meta) (string, error) {
return fmt.Sprintf("%s/client/v4/accounts/%s/ai/run/%s", meta.BaseURL, meta.Config.UserID, meta.ActualModelName), nil switch meta.Mode {
case relaymode.ChatCompletions:
return fmt.Sprintf("%s/client/v4/accounts/%s/ai/v1/chat/completions", meta.BaseURL, meta.Config.UserID), nil
case relaymode.Embeddings:
return fmt.Sprintf("%s/client/v4/accounts/%s/ai/v1/embeddings", meta.BaseURL, meta.Config.UserID), nil
default:
return fmt.Sprintf("%s/client/v4/accounts/%s/ai/run/%s", meta.BaseURL, meta.Config.UserID, meta.ActualModelName), nil
}
} }
func (a *Adaptor) SetupRequestHeader(c *gin.Context, req *http.Request, meta *meta.Meta) error { func (a *Adaptor) SetupRequestHeader(c *gin.Context, req *http.Request, meta *meta.Meta) error {
@ -41,7 +49,14 @@ func (a *Adaptor) ConvertRequest(c *gin.Context, relayMode int, request *model.G
if request == nil { if request == nil {
return nil, errors.New("request is nil") return nil, errors.New("request is nil")
} }
return ConvertRequest(*request), nil switch relayMode {
case relaymode.Completions:
return ConvertCompletionsRequest(*request), nil
case relaymode.ChatCompletions, relaymode.Embeddings:
return request, nil
default:
return nil, errors.New("not implemented")
}
} }
func (a *Adaptor) DoRequest(c *gin.Context, meta *meta.Meta, requestBody io.Reader) (*http.Response, error) { func (a *Adaptor) DoRequest(c *gin.Context, meta *meta.Meta, requestBody io.Reader) (*http.Response, error) {

View File

@ -3,11 +3,13 @@ package cloudflare
import ( import (
"bufio" "bufio"
"encoding/json" "encoding/json"
"github.com/songquanpeng/one-api/common/render"
"io" "io"
"net/http" "net/http"
"strings" "strings"
"github.com/songquanpeng/one-api/common/ctxkey"
"github.com/songquanpeng/one-api/common/render"
"github.com/gin-gonic/gin" "github.com/gin-gonic/gin"
"github.com/songquanpeng/one-api/common" "github.com/songquanpeng/one-api/common"
"github.com/songquanpeng/one-api/common/helper" "github.com/songquanpeng/one-api/common/helper"
@ -16,57 +18,23 @@ import (
"github.com/songquanpeng/one-api/relay/model" "github.com/songquanpeng/one-api/relay/model"
) )
func ConvertRequest(textRequest model.GeneralOpenAIRequest) *Request { func ConvertCompletionsRequest(textRequest model.GeneralOpenAIRequest) *Request {
var promptBuilder strings.Builder p, _ := textRequest.Prompt.(string)
for _, message := range textRequest.Messages {
promptBuilder.WriteString(message.StringContent())
promptBuilder.WriteString("\n") // 添加换行符来分隔每个消息
}
return &Request{ return &Request{
Prompt: p,
MaxTokens: textRequest.MaxTokens, MaxTokens: textRequest.MaxTokens,
Prompt: promptBuilder.String(),
Stream: textRequest.Stream, Stream: textRequest.Stream,
Temperature: textRequest.Temperature, Temperature: textRequest.Temperature,
} }
} }
func ResponseCloudflare2OpenAI(cloudflareResponse *Response) *openai.TextResponse {
choice := openai.TextResponseChoice{
Index: 0,
Message: model.Message{
Role: "assistant",
Content: cloudflareResponse.Result.Response,
},
FinishReason: "stop",
}
fullTextResponse := openai.TextResponse{
Object: "chat.completion",
Created: helper.GetTimestamp(),
Choices: []openai.TextResponseChoice{choice},
}
return &fullTextResponse
}
func StreamResponseCloudflare2OpenAI(cloudflareResponse *StreamResponse) *openai.ChatCompletionsStreamResponse {
var choice openai.ChatCompletionsStreamResponseChoice
choice.Delta.Content = cloudflareResponse.Response
choice.Delta.Role = "assistant"
openaiResponse := openai.ChatCompletionsStreamResponse{
Object: "chat.completion.chunk",
Choices: []openai.ChatCompletionsStreamResponseChoice{choice},
Created: helper.GetTimestamp(),
}
return &openaiResponse
}
func StreamHandler(c *gin.Context, resp *http.Response, promptTokens int, modelName string) (*model.ErrorWithStatusCode, *model.Usage) { func StreamHandler(c *gin.Context, resp *http.Response, promptTokens int, modelName string) (*model.ErrorWithStatusCode, *model.Usage) {
scanner := bufio.NewScanner(resp.Body) scanner := bufio.NewScanner(resp.Body)
scanner.Split(bufio.ScanLines) scanner.Split(bufio.ScanLines)
common.SetEventStreamHeaders(c) common.SetEventStreamHeaders(c)
id := helper.GetResponseID(c) id := helper.GetResponseID(c)
responseModel := c.GetString("original_model") responseModel := c.GetString(ctxkey.OriginalModel)
var responseText string var responseText string
for scanner.Scan() { for scanner.Scan() {
@ -77,22 +45,22 @@ func StreamHandler(c *gin.Context, resp *http.Response, promptTokens int, modelN
data = strings.TrimPrefix(data, "data: ") data = strings.TrimPrefix(data, "data: ")
data = strings.TrimSuffix(data, "\r") data = strings.TrimSuffix(data, "\r")
var cloudflareResponse StreamResponse if data == "[DONE]" {
err := json.Unmarshal([]byte(data), &cloudflareResponse) break
}
var response openai.ChatCompletionsStreamResponse
err := json.Unmarshal([]byte(data), &response)
if err != nil { if err != nil {
logger.SysError("error unmarshalling stream response: " + err.Error()) logger.SysError("error unmarshalling stream response: " + err.Error())
continue continue
} }
for _, v := range response.Choices {
response := StreamResponseCloudflare2OpenAI(&cloudflareResponse) v.Delta.Role = "assistant"
if response == nil { responseText += v.Delta.StringContent()
continue
} }
responseText += cloudflareResponse.Response
response.Id = id response.Id = id
response.Model = responseModel response.Model = modelName
err = render.ObjectData(c, response) err = render.ObjectData(c, response)
if err != nil { if err != nil {
logger.SysError(err.Error()) logger.SysError(err.Error())
@ -123,22 +91,25 @@ func Handler(c *gin.Context, resp *http.Response, promptTokens int, modelName st
if err != nil { if err != nil {
return openai.ErrorWrapper(err, "close_response_body_failed", http.StatusInternalServerError), nil return openai.ErrorWrapper(err, "close_response_body_failed", http.StatusInternalServerError), nil
} }
var cloudflareResponse Response var response openai.TextResponse
err = json.Unmarshal(responseBody, &cloudflareResponse) err = json.Unmarshal(responseBody, &response)
if err != nil { if err != nil {
return openai.ErrorWrapper(err, "unmarshal_response_body_failed", http.StatusInternalServerError), nil return openai.ErrorWrapper(err, "unmarshal_response_body_failed", http.StatusInternalServerError), nil
} }
fullTextResponse := ResponseCloudflare2OpenAI(&cloudflareResponse) response.Model = modelName
fullTextResponse.Model = modelName var responseText string
usage := openai.ResponseText2Usage(cloudflareResponse.Result.Response, modelName, promptTokens) for _, v := range response.Choices {
fullTextResponse.Usage = *usage responseText += v.Message.Content.(string)
fullTextResponse.Id = helper.GetResponseID(c) }
jsonResponse, err := json.Marshal(fullTextResponse) usage := openai.ResponseText2Usage(responseText, modelName, promptTokens)
response.Usage = *usage
response.Id = helper.GetResponseID(c)
jsonResponse, err := json.Marshal(response)
if err != nil { if err != nil {
return openai.ErrorWrapper(err, "marshal_response_body_failed", http.StatusInternalServerError), nil return openai.ErrorWrapper(err, "marshal_response_body_failed", http.StatusInternalServerError), nil
} }
c.Writer.Header().Set("Content-Type", "application/json") c.Writer.Header().Set("Content-Type", "application/json")
c.Writer.WriteHeader(resp.StatusCode) c.Writer.WriteHeader(resp.StatusCode)
_, err = c.Writer.Write(jsonResponse) _, _ = c.Writer.Write(jsonResponse)
return nil, usage return nil, usage
} }

View File

@ -1,25 +1,13 @@
package cloudflare package cloudflare
import "github.com/songquanpeng/one-api/relay/model"
type Request struct { type Request struct {
Lora string `json:"lora,omitempty"` Messages []model.Message `json:"messages,omitempty"`
MaxTokens int `json:"max_tokens,omitempty"` Lora string `json:"lora,omitempty"`
Prompt string `json:"prompt,omitempty"` MaxTokens int `json:"max_tokens,omitempty"`
Raw bool `json:"raw,omitempty"` Prompt string `json:"prompt,omitempty"`
Stream bool `json:"stream,omitempty"` Raw bool `json:"raw,omitempty"`
Temperature float64 `json:"temperature,omitempty"` Stream bool `json:"stream,omitempty"`
} Temperature float64 `json:"temperature,omitempty"`
type Result struct {
Response string `json:"response"`
}
type Response struct {
Result Result `json:"result"`
Success bool `json:"success"`
Errors []string `json:"errors"`
Messages []string `json:"messages"`
}
type StreamResponse struct {
Response string `json:"response"`
} }