xinyin025/one-api
1
0
Fork 0
mirror of https://github.com/songquanpeng/one-api.git synced 2025-11-14 20:23:46 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

✨ add: add images api

Browse Source
This commit is contained in:
Martial BE
2023-12-01 17:20:22 +08:00
parent 5b70ee3407
commit 9dd92bbddd
19 changed files with 296 additions and 12 deletions
Download Patch File Download Diff File Expand all files Collapse all files

7
providers/azure/base.go
View File

@@ -20,10 +20,13 @@ func CreateAzureProvider(c *gin.Context) *AzureProvider {
Completions: "/completions",
ChatCompletions: "/chat/completions",
Embeddings: "/embeddings",
AudioSpeech: "/audio/speech",
AudioTranscriptions: "/audio/transcriptions",
AudioTranslations: "/audio/translations",
Context: c,
ImagesGenerations: "/images/generations",
// ImagesEdit: "/images/edit",
// ImagesVariations: "/images/variations",
Context: c,
// AudioSpeech: "/audio/speech",
},
IsAzure: true,
},

102
providers/azure/image_generations.go Normal file
View File

@@ -0,0 +1,102 @@
package azure
import (
"errors"
"fmt"
"net/http"
"one-api/common"
"one-api/providers/openai"
"one-api/types"
"time"
)
func (c *ImageAzureResponse) ResponseHandler(resp *http.Response) (OpenAIResponse any, errWithCode *types.OpenAIErrorWithStatusCode) {
if c.Status == "canceled" || c.Status == "failed" {
errWithCode = &types.OpenAIErrorWithStatusCode{
OpenAIError: types.OpenAIError{
Message: c.Error.Message,
Type: "one_api_error",
Code: c.Error.Code,
},
StatusCode: resp.StatusCode,
}
return
}
operation_location := resp.Header.Get("operation-location")
if operation_location == "" {
return nil, types.ErrorWrapper(errors.New("image url is empty"), "get_images_url_failed", http.StatusInternalServerError)
}
client := common.NewClient()
req, err := client.NewRequest("GET", operation_location, common.WithHeader(c.Header))
if err != nil {
return nil, types.ErrorWrapper(err, "get_images_request_failed", http.StatusInternalServerError)
}
getImageAzureResponse := ImageAzureResponse{}
for i := 0; i < 3; i++ {
// 休眠 2 秒
time.Sleep(2 * time.Second)
_, errWithCode = common.SendRequest(req, &getImageAzureResponse, false)
fmt.Println("getImageAzureResponse", getImageAzureResponse)
if errWithCode != nil {
return
}
if getImageAzureResponse.Status == "canceled" || getImageAzureResponse.Status == "failed" {
return nil, &types.OpenAIErrorWithStatusCode{
OpenAIError: types.OpenAIError{
Message: c.Error.Message,
Type: "get_images_request_failed",
Code: c.Error.Code,
},
StatusCode: resp.StatusCode,
}
}
if getImageAzureResponse.Status == "succeeded" {
return getImageAzureResponse.Result, nil
}
}
return nil, types.ErrorWrapper(errors.New("get image Timeout"), "get_images_url_failed", http.StatusInternalServerError)
}
func (p *AzureProvider) ImageGenerationsAction(request *types.ImageRequest, isModelMapped bool, promptTokens int) (usage *types.Usage, errWithCode *types.OpenAIErrorWithStatusCode) {
requestBody, err := p.GetRequestBody(&request, isModelMapped)
if err != nil {
return nil, types.ErrorWrapper(err, "json_marshal_failed", http.StatusInternalServerError)
}
fullRequestURL := p.GetFullRequestURL(p.ImagesGenerations, request.Model)
headers := p.GetRequestHeaders()
client := common.NewClient()
req, err := client.NewRequest(p.Context.Request.Method, fullRequestURL, common.WithBody(requestBody), common.WithHeader(headers))
if err != nil {
return nil, types.ErrorWrapper(err, "new_request_failed", http.StatusInternalServerError)
}
if request.Model == "dall-e-2" {
imageAzureResponse := &ImageAzureResponse{
Header: headers,
}
errWithCode = p.SendRequest(req, imageAzureResponse, false)
} else {
openAIProviderImageResponseResponse := &openai.OpenAIProviderImageResponseResponse{}
errWithCode = p.SendRequest(req, openAIProviderImageResponseResponse, true)
}
if errWithCode != nil {
return
}
usage = &types.Usage{
PromptTokens: promptTokens,
CompletionTokens: 0,
TotalTokens: promptTokens,
}
return
}

21
providers/azure/type.go Normal file
View File

@@ -0,0 +1,21 @@
package azure
import "one-api/types"
type ImageAzureResponse struct {
ID string `json:"id,omitempty"`
Created int64 `json:"created,omitempty"`
Expires int64 `json:"expires,omitempty"`
Result types.ImageResponse `json:"result,omitempty"`
Status string `json:"status,omitempty"`
Error ImageAzureError `json:"error,omitempty"`
Header map[string]string `json:"header,omitempty"`
}
type ImageAzureError struct {
Code string `json:"code,omitempty"`
Target string `json:"target,omitempty"`
Message string `json:"message,omitempty"`
Details []string `json:"details,omitempty"`
InnerError any `json:"innererror,omitempty"`
}

9
providers/base/common.go
View File

@@ -23,6 +23,9 @@ type BaseProvider struct {
Moderation string
AudioTranscriptions string
AudioTranslations string
ImagesGenerations string
ImagesEdit string
ImagesVariations string
Proxy string
Context *gin.Context
}
@@ -141,6 +144,12 @@ func (p *BaseProvider) SupportAPI(relayMode int) bool {
return p.AudioTranslations != ""
case common.RelayModeModerations:
return p.Moderation != ""
case common.RelayModeImagesGenerations:
return p.ImagesGenerations != ""
case common.RelayModeImagesEdit:
return p.ImagesEdit != ""
case common.RelayModeImagesVariations:
return p.ImagesVariations != ""
default:
return false
}

6
providers/base/interface.go
View File

@@ -50,11 +50,17 @@ type TranscriptionsInterface interface {
TranscriptionsAction(request *types.AudioRequest, isModelMapped bool, promptTokens int) (usage *types.Usage, errWithCode *types.OpenAIErrorWithStatusCode)
}
// 语音翻译接口
type TranslationInterface interface {
ProviderInterface
TranslationAction(request *types.AudioRequest, isModelMapped bool, promptTokens int) (usage *types.Usage, errWithCode *types.OpenAIErrorWithStatusCode)
}
type ImageGenerationsInterface interface {
ProviderInterface
ImageGenerationsAction(request *types.ImageRequest, isModelMapped bool, promptTokens int) (usage *types.Usage, errWithCode *types.OpenAIErrorWithStatusCode)
}
// 余额接口
type BalanceInterface interface {
BalanceAction(channel *model.Channel) (float64, error)

13
providers/openai/base.go
View File

@@ -38,6 +38,9 @@ func CreateOpenAIProvider(c *gin.Context, baseURL string) *OpenAIProvider {
AudioSpeech: "/v1/audio/speech",
AudioTranscriptions: "/v1/audio/transcriptions",
AudioTranslations: "/v1/audio/translations",
ImagesGenerations: "/v1/images/generations",
ImagesEdit: "/v1/images/edit",
ImagesVariations: "/v1/images/variations",
Context: c,
},
IsAzure: false,
@@ -50,7 +53,13 @@ func (p *OpenAIProvider) GetFullRequestURL(requestURL string, modelName string)
if p.IsAzure {
apiVersion := p.Context.GetString("api_version")
requestURL = fmt.Sprintf("/openai/deployments/%s%s?api-version=%s", modelName, requestURL, apiVersion)
if modelName == "dall-e-2" {
// 因为dall-e-3需要api-version=2023-12-01-preview，但是该版本
// 已经没有dall-e-2了，所以暂时写死
requestURL = fmt.Sprintf("/openai/%s:submit?api-version=2023-09-01-preview", requestURL)
} else {
requestURL = fmt.Sprintf("/openai/deployments/%s%s?api-version=%s", modelName, requestURL, apiVersion)
}
}
if strings.HasPrefix(baseURL, "https://gateway.ai.cloudflare.com") {
@@ -78,7 +87,7 @@ func (p *OpenAIProvider) GetRequestHeaders() (headers map[string]string) {
}
// 获取请求体
func (p *OpenAIProvider) getRequestBody(request any, isModelMapped bool) (requestBody io.Reader, err error) {
func (p *OpenAIProvider) GetRequestBody(request any, isModelMapped bool) (requestBody io.Reader, err error) {
if isModelMapped {
jsonStr, err := json.Marshal(request)
if err != nil {

2
providers/openai/chat.go
View File

@@ -26,7 +26,7 @@ func (c *OpenAIProviderChatStreamResponse) responseStreamHandler() (responseText
}
func (p *OpenAIProvider) ChatAction(request *types.ChatCompletionRequest, isModelMapped bool, promptTokens int) (usage *types.Usage, errWithCode *types.OpenAIErrorWithStatusCode) {
requestBody, err := p.getRequestBody(&request, isModelMapped)
requestBody, err := p.GetRequestBody(&request, isModelMapped)
if err != nil {
return nil, types.ErrorWrapper(err, "json_marshal_failed", http.StatusInternalServerError)
}

2
providers/openai/completion.go
View File

@@ -26,7 +26,7 @@ func (c *OpenAIProviderCompletionResponse) responseStreamHandler() (responseText
}
func (p *OpenAIProvider) CompleteAction(request *types.CompletionRequest, isModelMapped bool, promptTokens int) (usage *types.Usage, errWithCode *types.OpenAIErrorWithStatusCode) {
requestBody, err := p.getRequestBody(&request, isModelMapped)
requestBody, err := p.GetRequestBody(&request, isModelMapped)
if err != nil {
return nil, types.ErrorWrapper(err, "json_marshal_failed", http.StatusInternalServerError)
}

2
providers/openai/embeddings.go
View File

@@ -19,7 +19,7 @@ func (c *OpenAIProviderEmbeddingsResponse) ResponseHandler(resp *http.Response)
func (p *OpenAIProvider) EmbeddingsAction(request *types.EmbeddingRequest, isModelMapped bool, promptTokens int) (usage *types.Usage, errWithCode *types.OpenAIErrorWithStatusCode) {
requestBody, err := p.getRequestBody(&request, isModelMapped)
requestBody, err := p.GetRequestBody(&request, isModelMapped)
if err != nil {
return nil, types.ErrorWrapper(err, "json_marshal_failed", http.StatusInternalServerError)
}

49
providers/openai/image_generations.go Normal file
View File

@@ -0,0 +1,49 @@
package openai
import (
"net/http"
"one-api/common"
"one-api/types"
)
func (c *OpenAIProviderImageResponseResponse) ResponseHandler(resp *http.Response) (OpenAIResponse any, errWithCode *types.OpenAIErrorWithStatusCode) {
if c.Error.Type != "" {
errWithCode = &types.OpenAIErrorWithStatusCode{
OpenAIError: c.Error,
StatusCode: resp.StatusCode,
}
return
}
return nil, nil
}
func (p *OpenAIProvider) ImageGenerationsAction(request *types.ImageRequest, isModelMapped bool, promptTokens int) (usage *types.Usage, errWithCode *types.OpenAIErrorWithStatusCode) {
requestBody, err := p.GetRequestBody(&request, isModelMapped)
if err != nil {
return nil, types.ErrorWrapper(err, "json_marshal_failed", http.StatusInternalServerError)
}
fullRequestURL := p.GetFullRequestURL(p.ImagesGenerations, request.Model)
headers := p.GetRequestHeaders()
client := common.NewClient()
req, err := client.NewRequest(p.Context.Request.Method, fullRequestURL, common.WithBody(requestBody), common.WithHeader(headers))
if err != nil {
return nil, types.ErrorWrapper(err, "new_request_failed", http.StatusInternalServerError)
}
openAIProviderImageResponseResponse := &OpenAIProviderImageResponseResponse{}
errWithCode = p.SendRequest(req, openAIProviderImageResponseResponse, true)
if errWithCode != nil {
return
}
usage = &types.Usage{
PromptTokens: promptTokens,
CompletionTokens: 0,
TotalTokens: promptTokens,
}
return
}

2
providers/openai/moderation.go
View File

@@ -19,7 +19,7 @@ func (c *OpenAIProviderModerationResponse) ResponseHandler(resp *http.Response)
func (p *OpenAIProvider) ModerationAction(request *types.ModerationRequest, isModelMapped bool, promptTokens int) (usage *types.Usage, errWithCode *types.OpenAIErrorWithStatusCode) {
requestBody, err := p.getRequestBody(&request, isModelMapped)
requestBody, err := p.GetRequestBody(&request, isModelMapped)
if err != nil {
return nil, types.ErrorWrapper(err, "json_marshal_failed", http.StatusInternalServerError)
}

2
providers/openai/speech.go
View File

@@ -8,7 +8,7 @@ import (
func (p *OpenAIProvider) SpeechAction(request *types.SpeechAudioRequest, isModelMapped bool, promptTokens int) (usage *types.Usage, errWithCode *types.OpenAIErrorWithStatusCode) {
requestBody, err := p.getRequestBody(&request, isModelMapped)
requestBody, err := p.GetRequestBody(&request, isModelMapped)
if err != nil {
return nil, types.ErrorWrapper(err, "json_marshal_failed", http.StatusInternalServerError)
}

5
providers/openai/type.go
View File

@@ -37,3 +37,8 @@ type OpenAIProviderTranscriptionsTextResponse string
func (a *OpenAIProviderTranscriptionsTextResponse) GetString() *string {
return (*string)(a)
}
type OpenAIProviderImageResponseResponse struct {
types.ImageResponse
types.OpenAIErrorResponse
}