feat: publish v1.0.0-rc.36-custom.1
Publish Docker image (Multi-arch) / Build & push (amd64) (push) Failing after 9m20s
Release (Linux, macOS, Windows) / Linux Release (push) Failing after 55m7s
Publish Docker image (Multi-arch) / Build & push (arm64) (push) Canceled after 0s
Release (Linux, macOS, Windows) / macOS Release (push) Canceled after 0s
Release (Linux, macOS, Windows) / Windows Release (push) Canceled after 0s
Publish Docker image (Multi-arch) / Create multi-arch manifests (push) Canceled after 0s

This commit is contained in:
Codex committed 2026-09-11 05:07:48 +08:00
commit bb7968b612
2381 files changed
+551494

No files matched your search

+138
View File
@@ -0,0 +1,138 @@
package relay
import (
"errors"
"fmt"
"io"
"net/http"
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/constant"
"github.com/QuantumNous/new-api/logger"
relaycommon "github.com/QuantumNous/new-api/relay/common"
"github.com/QuantumNous/new-api/relay/helper"
"github.com/QuantumNous/new-api/relaykit/dto"
"github.com/QuantumNous/new-api/relaykit/types"
"github.com/QuantumNous/new-api/service"
"github.com/gin-gonic/gin"
)
func AlphaSearchHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError *types.NewAPIError) {
info.InitChannelMeta(c)
switch info.ChannelType {
case constant.ChannelTypeSub2API,
constant.ChannelTypeNewAPI,
constant.ChannelTypeCodex,
constant.ChannelTypeAdvancedCustom:
default:
// Allow retry onto another channel that may support this endpoint.
return types.NewError(
errors.New("channel does not support /v1/alpha/search"),
types.ErrorCodeInvalidRequest,
)
}
request, ok := info.Request.(*dto.AlphaSearchRequest)
if !ok {
return types.NewErrorWithStatusCode(
fmt.Errorf("invalid request type, expected *dto.AlphaSearchRequest, got %T", info.Request),
types.ErrorCodeInvalidRequest,
http.StatusBadRequest,
types.ErrOptionWithSkipRetry(),
)
}
err := helper.ModelMappedHelper(c, info, request)
if err != nil {
return types.NewError(err, types.ErrorCodeChannelModelMappedError, types.ErrOptionWithSkipRetry())
}
jsonData, err := buildAlphaSearchRequestBody(request.RawBody, info.OriginModelName, info.UpstreamModelName)
if err != nil {
return types.NewError(err, types.ErrorCodeConvertRequestFailed, types.ErrOptionWithSkipRetry())
}
if len(info.ParamOverride) > 0 {
jsonData, err = relaycommon.ApplyParamOverrideWithRelayInfo(jsonData, info)
if err != nil {
return newAPIErrorFromParamOverride(err)
}
}
logger.LogDebug(c, "requestBody: %s", jsonData)
body, closer, err := relaycommon.NewOutboundJSONBody(jsonData)
if err != nil {
return types.NewError(err, types.ErrorCodeConvertRequestFailed, types.ErrOptionWithSkipRetry())
}
defer closer.Close()
adaptor := GetAdaptor(info.ApiType)
if adaptor == nil {
return types.NewError(fmt.Errorf("invalid api type: %d", info.ApiType), types.ErrorCodeInvalidApiType, types.ErrOptionWithSkipRetry())
}
adaptor.Init(info)
resp, err := adaptor.DoRequest(c, info, body)
if err != nil {
return types.NewOpenAIError(err, types.ErrorCodeDoRequestFailed, http.StatusInternalServerError)
}
statusCodeMappingStr := c.GetString("status_code_mapping")
httpResp, ok := resp.(*http.Response)
if !ok || httpResp == nil {
return types.NewOpenAIError(errors.New("invalid http response"), types.ErrorCodeDoRequestFailed, http.StatusInternalServerError)
}
defer httpResp.Body.Close()
if httpResp.StatusCode < 200 || httpResp.StatusCode >= 300 {
newAPIError = service.RelayErrorHandler(c.Request.Context(), httpResp, false)
service.ResetStatusCode(newAPIError, statusCodeMappingStr)
return newAPIError
}
if contentType := httpResp.Header.Get("Content-Type"); contentType != "" {
c.Writer.Header().Set("Content-Type", contentType)
}
c.Writer.WriteHeader(httpResp.StatusCode)
if _, err := io.Copy(c.Writer, httpResp.Body); err != nil {
return types.NewError(err, types.ErrorCodeDoRequestFailed, types.ErrOptionWithSkipRetry())
}
// Upstream alpha search returns no usage; bill one web_search_preview call.
if info.ResponsesUsageInfo == nil {
info.ResponsesUsageInfo = &relaycommon.ResponsesUsageInfo{
BuiltInTools: make(map[string]*relaycommon.BuildInToolInfo),
}
}
if info.ResponsesUsageInfo.BuiltInTools == nil {
info.ResponsesUsageInfo.BuiltInTools = make(map[string]*relaycommon.BuildInToolInfo)
}
info.ResponsesUsageInfo.BuiltInTools[dto.BuildInToolWebSearchPreview] = &relaycommon.BuildInToolInfo{
ToolName: dto.BuildInToolWebSearchPreview,
CallCount: 1,
}
usage := &dto.Usage{}
service.PostTextConsumeQuota(c, info, usage, nil)
return nil
}
// buildAlphaSearchRequestBody returns RawBody unchanged unless the model was
// mapped, in which case only the "model" field is rewritten so unknown fields
// are preserved.
func buildAlphaSearchRequestBody(rawBody []byte, originModel, upstreamModel string) ([]byte, error) {
if len(rawBody) == 0 {
return nil, errors.New("empty alpha search request body")
}
if upstreamModel == "" || upstreamModel == originModel {
return rawBody, nil
}
var body map[string]any
if err := common.Unmarshal(rawBody, &body); err != nil {
return nil, err
}
body["model"] = upstreamModel
return common.Marshal(body)
}
+47
View File
@@ -0,0 +1,47 @@
package relay
import (
"testing"
"github.com/QuantumNous/new-api/common"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
func TestBuildAlphaSearchRequestBodyPreservesUnknownFields(t *testing.T) {
raw := []byte(`{
"id":"req_1",
"model":"gpt-5.1",
"input":[{"role":"user","content":"hi"}],
"commands":{"search_query":[{"q":"weather","recency":1}]},
"settings":{"locale":"en"},
"future_field":{"nested":true}
}`)
out, err := buildAlphaSearchRequestBody(raw, "gpt-5.1", "gpt-5.1-mapped")
require.NoError(t, err)
var body map[string]any
require.NoError(t, common.Unmarshal(out, &body))
assert.Equal(t, "gpt-5.1-mapped", body["model"])
assert.Equal(t, "req_1", body["id"])
require.Contains(t, body, "commands")
require.Contains(t, body, "settings")
require.Contains(t, body, "future_field")
require.Contains(t, body, "input")
commands, ok := body["commands"].(map[string]any)
require.True(t, ok)
require.Contains(t, commands, "search_query")
future, ok := body["future_field"].(map[string]any)
require.True(t, ok)
assert.Equal(t, true, future["nested"])
}
func TestBuildAlphaSearchRequestBodyNoMappingKeepsRawBytes(t *testing.T) {
raw := []byte(`{"model":"gpt-5.1","commands":{"search_query":[{"q":"x"}]},"future_field":1}`)
out, err := buildAlphaSearchRequestBody(raw, "gpt-5.1", "gpt-5.1")
require.NoError(t, err)
assert.Equal(t, raw, out)
}
+77
View File
@@ -0,0 +1,77 @@
package relay
import (
"errors"
"fmt"
"net/http"
"github.com/QuantumNous/new-api/common"
relaycommon "github.com/QuantumNous/new-api/relay/common"
"github.com/QuantumNous/new-api/relay/helper"
"github.com/QuantumNous/new-api/relaykit/dto"
"github.com/QuantumNous/new-api/relaykit/types"
"github.com/QuantumNous/new-api/service"
"github.com/gin-gonic/gin"
)
func AudioHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError *types.NewAPIError) {
info.InitChannelMeta(c)
audioReq, ok := info.Request.(*dto.AudioRequest)
if !ok {
return types.NewError(errors.New("invalid request type"), types.ErrorCodeInvalidRequest, types.ErrOptionWithSkipRetry())
}
request, err := common.DeepCopy(audioReq)
if err != nil {
return types.NewError(fmt.Errorf("failed to copy request to AudioRequest: %w", err), types.ErrorCodeInvalidRequest, types.ErrOptionWithSkipRetry())
}
err = helper.ModelMappedHelper(c, info, request)
if err != nil {
return types.NewError(err, types.ErrorCodeChannelModelMappedError, types.ErrOptionWithSkipRetry())
}
adaptor := GetAdaptor(info.ApiType)
if adaptor == nil {
return types.NewError(fmt.Errorf("invalid api type: %d", info.ApiType), types.ErrorCodeInvalidApiType, types.ErrOptionWithSkipRetry())
}
adaptor.Init(info)
ioReader, err := adaptor.ConvertAudioRequest(c, info, *request)
if err != nil {
return types.NewError(err, types.ErrorCodeConvertRequestFailed, types.ErrOptionWithSkipRetry())
}
resp, err := adaptor.DoRequest(c, info, ioReader)
if err != nil {
return types.NewOpenAIError(err, types.ErrorCodeDoRequestFailed, http.StatusInternalServerError)
}
statusCodeMappingStr := c.GetString("status_code_mapping")
var httpResp *http.Response
if resp != nil {
httpResp = resp.(*http.Response)
if httpResp.StatusCode != http.StatusOK {
newAPIError = service.RelayErrorHandler(c.Request.Context(), httpResp, false)
// reset status code 重置状态码
service.ResetStatusCode(newAPIError, statusCodeMappingStr)
return newAPIError
}
}
usage, newAPIError := adaptor.DoResponse(c, httpResp, info)
if newAPIError != nil {
// reset status code 重置状态码
service.ResetStatusCode(newAPIError, statusCodeMappingStr)
return newAPIError
}
if usage.(*dto.Usage).CompletionTokenDetails.AudioTokens > 0 || usage.(*dto.Usage).PromptTokensDetails.AudioTokens > 0 {
service.PostAudioConsumeQuota(c, info, usage.(*dto.Usage), "")
} else {
service.PostTextConsumeQuota(c, info, usage.(*dto.Usage), nil)
}
return nil
}
+135
View File
@@ -0,0 +1,135 @@
package channel
import (
"io"
"net/http"
taskdto "github.com/QuantumNous/new-api/dto"
"github.com/QuantumNous/new-api/model"
relaycommon "github.com/QuantumNous/new-api/relay/common"
"github.com/QuantumNous/new-api/relaykit/dto"
"github.com/QuantumNous/new-api/relaykit/types"
hosttypes "github.com/QuantumNous/new-api/types"
"github.com/gin-gonic/gin"
)
type Adaptor interface {
// Init IsStream bool
Init(info *relaycommon.RelayInfo)
GetRequestURL(info *relaycommon.RelayInfo) (string, error)
SetupRequestHeader(c *gin.Context, req *http.Header, info *relaycommon.RelayInfo) error
ConvertOpenAIRequest(c *gin.Context, info *relaycommon.RelayInfo, request *dto.GeneralOpenAIRequest) (any, error)
ConvertRerankRequest(c *gin.Context, relayMode int, request dto.RerankRequest) (any, error)
ConvertEmbeddingRequest(c *gin.Context, info *relaycommon.RelayInfo, request dto.EmbeddingRequest) (any, error)
ConvertAudioRequest(c *gin.Context, info *relaycommon.RelayInfo, request dto.AudioRequest) (io.Reader, error)
ConvertImageRequest(c *gin.Context, info *relaycommon.RelayInfo, request dto.ImageRequest) (any, error)
ConvertOpenAIResponsesRequest(c *gin.Context, info *relaycommon.RelayInfo, request dto.OpenAIResponsesRequest) (any, error)
DoRequest(c *gin.Context, info *relaycommon.RelayInfo, requestBody io.Reader) (any, error)
DoResponse(c *gin.Context, resp *http.Response, info *relaycommon.RelayInfo) (usage any, err *types.NewAPIError)
GetModelList() []string
GetChannelName() string
ConvertClaudeRequest(c *gin.Context, info *relaycommon.RelayInfo, request *dto.ClaudeRequest) (any, error)
ConvertGeminiRequest(c *gin.Context, info *relaycommon.RelayInfo, request *dto.GeminiChatRequest) (any, error)
}
type TaskAdaptor interface {
Init(info *relaycommon.RelayInfo)
ValidateRequestAndSetAction(c *gin.Context, info *relaycommon.RelayInfo) *taskdto.TaskError
// ── Billing ──────────────────────────────────────────────────────
// EstimateBilling returns OtherRatios for pre-charge based on user request.
// Called after ValidateRequestAndSetAction, before price calculation.
// Adaptors should extract duration, resolution, etc. from the parsed request
// and return them as ratio multipliers (e.g. {"seconds": 5, "size": 1.666}).
// Return nil to use the base model price without extra ratios.
EstimateBilling(c *gin.Context, info *relaycommon.RelayInfo) map[string]float64
// AdjustBillingOnSubmit returns adjusted OtherRatios from the upstream
// submit response. Called after a successful ParseResponse.
// If the upstream returned actual parameters that differ from the estimate
// (e.g. actual seconds), return updated ratios so the caller can recalculate
// the quota and settle the delta with the pre-charge.
// Return nil if no adjustment is needed.
AdjustBillingOnSubmit(info *relaycommon.RelayInfo, taskData []byte) map[string]float64
// AdjustBillingOnComplete returns the actual quota when a task reaches a
// terminal state (success/failure) during polling.
// Called by the polling loop after ParseTaskResult.
// Return a positive value to trigger delta settlement (supplement / refund).
// Return 0 to keep the pre-charged amount unchanged.
AdjustBillingOnComplete(task *model.Task, taskResult *relaycommon.TaskInfo) int
// ── Request / Response ───────────────────────────────────────────
BuildRequestURL(info *relaycommon.RelayInfo) (string, error)
BuildRequestHeader(c *gin.Context, req *http.Request, info *relaycommon.RelayInfo) error
BuildRequestBody(c *gin.Context, info *relaycommon.RelayInfo) (io.Reader, error)
DoRequest(c *gin.Context, info *relaycommon.RelayInfo, requestBody io.Reader) (*http.Response, error)
ParseResponse(c *gin.Context, resp *http.Response, info *relaycommon.RelayInfo) (*TaskSubmitResponse, *taskdto.TaskError)
GetModelList() []string
GetChannelName() string
// ── Polling ──────────────────────────────────────────────────────
FetchTask(baseUrl, key string, task *model.Task, proxy string) (*http.Response, error)
ParseTaskResult(task *model.Task, resp *http.Response, respBody []byte) (*relaycommon.TaskInfo, error)
}
// TaskSubmitResponse is the transport-independent result of parsing an
// upstream task submission. Parsing must not write to the client response.
type TaskSubmitResponse struct {
UpstreamTaskID string
TaskData []byte
ClientResponse any
Immediate *relaycommon.TaskInfo
PluginState []byte
}
type OpenAIVideoConverter interface {
ConvertToOpenAIVideo(originTask *model.Task) ([]byte, error)
}
type TaskArtifact = hosttypes.TaskArtifact
type TaskArtifactClientRequest struct {
Method string `json:"method"`
Headers map[string]string `json:"headers,omitempty"`
}
type TaskArtifactProvider interface {
ListArtifacts(task *model.Task) ([]TaskArtifact, error)
}
type TaskContentRequest struct {
URL string
Method string
Headers map[string]string
Body []byte
Credentialless bool
}
type TaskContentRequestProvider interface {
BuildContentRequest(task *model.Task, artifactKey string, clientRequest TaskArtifactClientRequest) (*TaskContentRequest, error)
}
type TaskUsageFactsProvider interface {
ExtractUsageFacts(c *gin.Context, info *relaycommon.RelayInfo) map[string]any
}
// TaskValidatedBillingProvider lets an adaptor reject invalid usage facts at
// the existing estimate point, after model mapping and before quota
// multiplication. Non-plugin task adaptors keep using EstimateBilling.
type TaskValidatedBillingProvider interface {
EstimateBillingValidated(c *gin.Context, info *relaycommon.RelayInfo) (map[string]float64, error)
}
// TaskValidatedUsageFactsProvider is the tiered-billing counterpart to
// TaskValidatedBillingProvider.
type TaskValidatedUsageFactsProvider interface {
ExtractUsageFactsValidated(c *gin.Context, info *relaycommon.RelayInfo) (map[string]any, error)
}
+575
View File
@@ -0,0 +1,575 @@
package advancedcustom
import (
"errors"
"fmt"
"io"
"net/http"
"net/url"
"strings"
"github.com/QuantumNous/new-api/constant"
"github.com/QuantumNous/new-api/relay/channel"
"github.com/QuantumNous/new-api/relay/channel/claude"
"github.com/QuantumNous/new-api/relay/channel/gemini"
"github.com/QuantumNous/new-api/relay/channel/openai"
relaycommon "github.com/QuantumNous/new-api/relay/common"
relayconstant "github.com/QuantumNous/new-api/relay/constant"
"github.com/QuantumNous/new-api/relaykit/dto"
"github.com/QuantumNous/new-api/relaykit/relayconvert"
"github.com/QuantumNous/new-api/relaykit/types"
"github.com/QuantumNous/new-api/service"
"github.com/gin-gonic/gin"
"github.com/samber/lo"
)
const ChannelName = "advanced_custom"
const advancedCustomModelPlaceholder = "{model}"
type Adaptor struct {
openaiAdaptor openai.Adaptor
claudeAdaptor claude.Adaptor
geminiAdaptor gemini.Adaptor
resolved bool
converted bool
route dto.AdvancedCustomRoute
converter string
}
func (a *Adaptor) Init(info *relaycommon.RelayInfo) {
a.openaiAdaptor.Init(info)
a.claudeAdaptor.Init(info)
a.geminiAdaptor.Init(info)
}
func (a *Adaptor) ConvertOpenAIRequest(c *gin.Context, info *relaycommon.RelayInfo, request *dto.GeneralOpenAIRequest) (any, error) {
converter, err := a.resolveForConversion(c, info)
if err != nil {
return nil, err
}
if converter == relayconvert.ConverterNone {
return a.convertOpenAICompatibleRequest(c, info, request)
}
switch converter {
case relayconvert.ConverterOpenAIChatToClaudeMessages,
relayconvert.ConverterOpenAIChatToOpenAIResponses,
relayconvert.ConverterOpenAIChatToGeminiContent:
result, err := service.ConvertRequestByID(c, info, converter, request)
if err != nil {
return nil, err
}
return result.Value, nil
default:
return nil, fmt.Errorf("converter %q does not support OpenAI chat completions requests", converter)
}
}
func (a *Adaptor) ConvertClaudeRequest(c *gin.Context, info *relaycommon.RelayInfo, request *dto.ClaudeRequest) (any, error) {
converter, err := a.resolveForConversion(c, info)
if err != nil {
return nil, err
}
switch converter {
case relayconvert.ConverterNone:
return a.claudeAdaptor.ConvertClaudeRequest(c, info, request)
case relayconvert.ConverterClaudeMessagesToOpenAIChat:
result, err := service.ConvertRequestByID(c, info, converter, request)
if err != nil {
return nil, err
}
chatRequest, ok := result.Value.(*dto.GeneralOpenAIRequest)
if !ok {
return nil, fmt.Errorf("expected OpenAI chat completions request, got %T", result.Value)
}
return a.convertOpenAICompatibleRequest(c, info, chatRequest)
default:
return nil, fmt.Errorf("converter %q does not support Anthropic Messages requests", converter)
}
}
func (a *Adaptor) ConvertGeminiRequest(c *gin.Context, info *relaycommon.RelayInfo, request *dto.GeminiChatRequest) (any, error) {
converter, err := a.resolveForConversion(c, info)
if err != nil {
return nil, err
}
switch converter {
case relayconvert.ConverterNone:
return a.geminiAdaptor.ConvertGeminiRequest(c, info, request)
case relayconvert.ConverterGeminiContentToOpenAIChat:
result, err := service.ConvertRequestByID(c, info, converter, request)
if err != nil {
return nil, err
}
chatRequest, ok := result.Value.(*dto.GeneralOpenAIRequest)
if !ok {
return nil, fmt.Errorf("expected OpenAI chat completions request, got %T", result.Value)
}
return a.convertOpenAICompatibleRequest(c, info, chatRequest)
default:
return nil, fmt.Errorf("converter %q does not support Gemini generateContent requests", converter)
}
}
func (a *Adaptor) ConvertOpenAIResponsesRequest(c *gin.Context, info *relaycommon.RelayInfo, request dto.OpenAIResponsesRequest) (any, error) {
converter, err := a.resolveForConversion(c, info)
if err != nil {
return nil, err
}
switch converter {
case relayconvert.ConverterNone:
return a.convertOpenAICompatibleResponsesRequest(c, info, request)
case relayconvert.ConverterOpenAIResponsesToOpenAIChat:
result, err := service.ConvertRequestByID(c, info, converter, request)
if err != nil {
return nil, err
}
chatRequest, ok := result.Value.(*dto.GeneralOpenAIRequest)
if !ok {
return nil, fmt.Errorf("expected OpenAI chat completions request, got %T", result.Value)
}
return a.convertOpenAICompatibleRequest(c, info, chatRequest)
case relayconvert.ConverterOpenAIResponsesToGemini:
result, err := service.ConvertRequestByID(c, info, converter, request)
if err != nil {
return nil, err
}
geminiRequest, ok := result.Value.(*dto.GeminiChatRequest)
if !ok {
return nil, fmt.Errorf("expected Gemini generateContent request, got %T", result.Value)
}
return geminiRequest, nil
default:
return nil, fmt.Errorf("converter %q does not support OpenAI Responses requests", converter)
}
}
func (a *Adaptor) ConvertEmbeddingRequest(c *gin.Context, info *relaycommon.RelayInfo, request dto.EmbeddingRequest) (any, error) {
converter, err := a.resolveForConversion(c, info)
if err != nil {
return nil, err
}
if converter != relayconvert.ConverterNone {
return nil, fmt.Errorf("converter %q does not support embedding requests", converter)
}
return a.convertOpenAICompatibleEmbeddingRequest(c, info, request)
}
func (a *Adaptor) ConvertAudioRequest(c *gin.Context, info *relaycommon.RelayInfo, request dto.AudioRequest) (io.Reader, error) {
converter, err := a.resolveForConversion(c, info)
if err != nil {
return nil, err
}
if converter != relayconvert.ConverterNone {
return nil, fmt.Errorf("converter %q does not support audio requests", converter)
}
return a.convertOpenAICompatibleAudioRequest(c, info, request)
}
func (a *Adaptor) ConvertImageRequest(c *gin.Context, info *relaycommon.RelayInfo, request dto.ImageRequest) (any, error) {
converter, err := a.resolveForConversion(c, info)
if err != nil {
return nil, err
}
if converter != relayconvert.ConverterNone {
return nil, fmt.Errorf("converter %q does not support image requests", converter)
}
return a.convertOpenAICompatibleImageRequest(c, info, request)
}
func (a *Adaptor) ConvertRerankRequest(c *gin.Context, relayMode int, request dto.RerankRequest) (any, error) {
a.converted = true
return a.openaiAdaptor.ConvertRerankRequest(c, relayMode, request)
}
func (a *Adaptor) GetRequestURL(info *relaycommon.RelayInfo) (string, error) {
if err := a.resolve(nil, info); err != nil {
return "", err
}
return a.routeURL(info)
}
func (a *Adaptor) BuildModelListRequest(info *relaycommon.RelayInfo) (string, http.Header, error) {
return a.buildManagementRequest(info, dto.AdvancedCustomModelListPath)
}
func (a *Adaptor) BuildBalanceRequest(info *relaycommon.RelayInfo) (string, http.Header, error) {
return a.buildManagementRequest(info, dto.AdvancedCustomBalancePath)
}
func (a *Adaptor) buildManagementRequest(info *relaycommon.RelayInfo, managementPath string) (string, http.Header, error) {
if info == nil {
return "", nil, errors.New("missing relay info")
}
config := info.ChannelOtherSettings.AdvancedCustom
if config == nil {
return "", nil, errors.New("advanced_custom is required")
}
if err := config.Validate(); err != nil {
return "", nil, err
}
var route dto.AdvancedCustomRoute
var ok bool
switch managementPath {
case dto.AdvancedCustomModelListPath:
route, ok = config.ModelListRoute()
case dto.AdvancedCustomBalancePath:
route, ok = config.BalanceRoute()
default:
return "", nil, fmt.Errorf("unsupported advanced custom management path: %s", managementPath)
}
if !ok {
return "", nil, fmt.Errorf("advanced custom channel does not configure a %s route", managementPath)
}
converter := strings.TrimSpace(route.Converter)
if converter == "" {
converter = relayconvert.ConverterNone
}
if converter != relayconvert.ConverterNone {
return "", nil, fmt.Errorf("converter %q does not support %s requests", converter, managementPath)
}
requestURL, err := buildRouteURL(route, converter, info)
if err != nil {
return "", nil, err
}
header := http.Header{}
auth := route.Auth
if auth == nil {
header.Set("Authorization", "Bearer "+info.ApiKey)
return requestURL, header, nil
}
switch strings.TrimSpace(auth.Type) {
case dto.AdvancedCustomAuthTypeNone, dto.AdvancedCustomAuthTypeQuery:
case dto.AdvancedCustomAuthTypeHeader:
header.Set(strings.TrimSpace(auth.Name), applyAuthTemplate(auth.Value, info.ApiKey))
default:
return "", nil, fmt.Errorf("invalid advanced custom auth type: %s", auth.Type)
}
return requestURL, header, nil
}
func (a *Adaptor) SetupRequestHeader(c *gin.Context, header *http.Header, info *relaycommon.RelayInfo) error {
if err := a.resolve(c, info); err != nil {
return err
}
channel.SetupApiRequestHeader(info, c, header)
auth := a.route.Auth
if auth == nil {
header.Set("Authorization", "Bearer "+info.ApiKey)
} else {
switch strings.TrimSpace(auth.Type) {
case dto.AdvancedCustomAuthTypeNone:
case dto.AdvancedCustomAuthTypeHeader:
header.Set(strings.TrimSpace(auth.Name), applyAuthTemplate(auth.Value, info.ApiKey))
case dto.AdvancedCustomAuthTypeQuery:
default:
return fmt.Errorf("invalid advanced custom auth type: %s", auth.Type)
}
}
if shouldApplyClaudeHeaders(a.converter, info) {
applyClaudeHeaders(c, header, info)
}
return nil
}
func (a *Adaptor) DoRequest(c *gin.Context, info *relaycommon.RelayInfo, requestBody io.Reader) (any, error) {
if err := a.resolve(c, info); err != nil {
return nil, err
}
if !a.converted && a.converter != relayconvert.ConverterNone {
return nil, errors.New("advanced custom converter routes cannot be used with pass-through request body")
}
if info.RelayMode == relayconstant.RelayModeAudioTranscription ||
info.RelayMode == relayconstant.RelayModeAudioTranslation ||
(info.RelayMode == relayconstant.RelayModeImagesEdits && !isJSONRequest(c)) {
return channel.DoFormRequest(a, c, info, requestBody)
}
if info.RelayMode == relayconstant.RelayModeRealtime {
return channel.DoWssRequest(a, c, info, requestBody)
}
return channel.DoApiRequest(a, c, info, requestBody)
}
func (a *Adaptor) DoResponse(c *gin.Context, resp *http.Response, info *relaycommon.RelayInfo) (usage any, err *types.NewAPIError) {
if err := a.resolve(c, info); err != nil {
return nil, types.NewOpenAIError(err, types.ErrorCodeInvalidRequest, http.StatusBadRequest, types.ErrOptionWithSkipRetry())
}
switch a.converter {
case relayconvert.ConverterNone:
return a.doNativeResponse(c, resp, info)
case relayconvert.ConverterClaudeMessagesToOpenAIChat,
relayconvert.ConverterGeminiContentToOpenAIChat:
return a.openaiAdaptor.DoResponse(c, resp, info)
case relayconvert.ConverterOpenAIChatToClaudeMessages:
return a.claudeAdaptor.DoResponse(c, resp, info)
case relayconvert.ConverterOpenAIChatToGeminiContent:
return a.geminiAdaptor.DoResponse(c, resp, info)
case relayconvert.ConverterOpenAIResponsesToGemini:
return a.geminiAdaptor.DoResponse(c, resp, info)
case relayconvert.ConverterOpenAIChatToOpenAIResponses:
if info.IsStream {
return openai.OaiResponsesToChatStreamHandler(c, info, resp)
}
return openai.OaiResponsesToChatHandler(c, info, resp)
case relayconvert.ConverterOpenAIResponsesToOpenAIChat:
if info.IsStream {
return openai.OaiChatToResponsesStreamHandler(c, info, resp)
}
return openai.OaiChatToResponsesHandler(c, info, resp)
default:
return nil, types.NewOpenAIError(fmt.Errorf("unsupported advanced custom converter: %s", a.converter), types.ErrorCodeInvalidRequest, http.StatusBadRequest, types.ErrOptionWithSkipRetry())
}
}
func (a *Adaptor) GetModelList() []string {
models := make([]string, 0, len(openai.ModelList)+len(claude.ModelList)+len(gemini.ModelList))
models = append(models, openai.ModelList...)
models = append(models, claude.ModelList...)
models = append(models, gemini.ModelList...)
return lo.Uniq(models)
}
func (a *Adaptor) GetChannelName() string {
return ChannelName
}
func (a *Adaptor) doNativeResponse(c *gin.Context, resp *http.Response, info *relaycommon.RelayInfo) (any, *types.NewAPIError) {
switch info.RelayFormat {
case types.RelayFormatClaude:
return a.claudeAdaptor.DoResponse(c, resp, info)
case types.RelayFormatGemini:
return a.geminiAdaptor.DoResponse(c, resp, info)
default:
return a.openaiAdaptor.DoResponse(c, resp, info)
}
}
func (a *Adaptor) resolveForConversion(c *gin.Context, info *relaycommon.RelayInfo) (string, error) {
if err := a.resolve(c, info); err != nil {
return "", err
}
a.converted = true
return a.converter, nil
}
func (a *Adaptor) resolve(c *gin.Context, info *relaycommon.RelayInfo) error {
if a.resolved {
return nil
}
if info == nil {
return errors.New("missing relay info")
}
config := info.ChannelOtherSettings.AdvancedCustom
if config == nil {
return errors.New("advanced_custom is required")
}
if err := config.Validate(); err != nil {
return err
}
incomingPath := incomingRequestPath(c, info)
route, ok := config.MatchPathForModel(incomingPath, info.OriginModelName)
if ok {
route.Converter = strings.TrimSpace(route.Converter)
if route.Converter == "" {
route.Converter = relayconvert.ConverterNone
}
a.route = route
a.converter = route.Converter
a.resolved = true
return nil
}
return fmt.Errorf("advanced custom channel does not support request path %s for model %s", incomingPath, info.OriginModelName)
}
func incomingRequestPath(c *gin.Context, info *relaycommon.RelayInfo) string {
if c != nil && c.Request != nil && c.Request.URL != nil {
return c.Request.URL.Path
}
if info == nil {
return ""
}
return strings.Split(info.RequestURLPath, "?")[0]
}
func (a *Adaptor) routeURL(info *relaycommon.RelayInfo) (string, error) {
return buildRouteURL(a.route, a.converter, info)
}
func buildRouteURL(route dto.AdvancedCustomRoute, converter string, info *relaycommon.RelayInfo) (string, error) {
parsedURL, err := resolveUpstreamTargetURL(applyUpstreamPathTemplate(strings.TrimSpace(route.UpstreamPath), info), info)
if err != nil {
return "", err
}
if shouldUseGeminiStreamURL(converter, info) {
useGeminiStreamGenerateContentURL(parsedURL)
}
if info != nil && info.RelayMode == relayconstant.RelayModeRealtime {
switch parsedURL.Scheme {
case "https":
parsedURL.Scheme = "wss"
case "http":
parsedURL.Scheme = "ws"
}
}
if route.Auth != nil && strings.TrimSpace(route.Auth.Type) == dto.AdvancedCustomAuthTypeQuery {
query := parsedURL.Query()
query.Set(strings.TrimSpace(route.Auth.Name), applyAuthTemplate(route.Auth.Value, info.ApiKey))
parsedURL.RawQuery = query.Encode()
}
return parsedURL.String(), nil
}
func resolveUpstreamTargetURL(upstreamPath string, info *relaycommon.RelayInfo) (*url.URL, error) {
if strings.HasPrefix(upstreamPath, "/") {
if strings.HasPrefix(upstreamPath, "//") {
return nil, errors.New("advanced custom upstream path must be a full URL or a path starting with /")
}
if info == nil || strings.TrimSpace(info.ChannelBaseUrl) == "" {
return nil, errors.New("channel base URL is required when advanced custom upstream path is relative")
}
return joinBaseURLAndUpstreamPath(info.ChannelBaseUrl, upstreamPath)
}
parsedURL, err := url.Parse(upstreamPath)
if err != nil {
return nil, err
}
if parsedURL.Scheme == "" || parsedURL.Host == "" {
return nil, errors.New("advanced custom upstream path must be a full URL or a path starting with /")
}
if !strings.EqualFold(parsedURL.Scheme, "http") && !strings.EqualFold(parsedURL.Scheme, "https") {
return nil, errors.New("advanced custom upstream path must use http or https")
}
return parsedURL, nil
}
func joinBaseURLAndUpstreamPath(baseURL string, upstreamPath string) (*url.URL, error) {
parsedBaseURL, err := url.Parse(strings.TrimSpace(baseURL))
if err != nil {
return nil, err
}
if parsedBaseURL.Scheme == "" || parsedBaseURL.Host == "" {
return nil, errors.New("channel base URL must be a full URL when advanced custom upstream path is relative")
}
if !strings.EqualFold(parsedBaseURL.Scheme, "http") && !strings.EqualFold(parsedBaseURL.Scheme, "https") {
return nil, errors.New("channel base URL must use http or https when advanced custom upstream path is relative")
}
parsedPath, err := url.Parse(upstreamPath)
if err != nil {
return nil, err
}
parsedBaseURL.Path = strings.TrimRight(parsedBaseURL.Path, "/") + "/" + strings.TrimLeft(parsedPath.Path, "/")
parsedBaseURL.RawPath = ""
parsedBaseURL.RawQuery = parsedPath.RawQuery
parsedBaseURL.Fragment = parsedPath.Fragment
return parsedBaseURL, nil
}
func applyUpstreamPathTemplate(upstreamPath string, info *relaycommon.RelayInfo) string {
if info == nil {
return upstreamPath
}
return strings.ReplaceAll(upstreamPath, advancedCustomModelPlaceholder, info.UpstreamModelName)
}
func shouldUseGeminiStreamURL(converter string, info *relaycommon.RelayInfo) bool {
return info != nil &&
info.IsStream &&
(converter == relayconvert.ConverterOpenAIChatToGeminiContent ||
converter == relayconvert.ConverterOpenAIResponsesToGemini)
}
func useGeminiStreamGenerateContentURL(parsedURL *url.URL) {
if strings.Contains(parsedURL.Path, ":generateContent") {
parsedURL.Path = strings.Replace(parsedURL.Path, ":generateContent", ":streamGenerateContent", 1)
}
if strings.Contains(parsedURL.Path, ":streamGenerateContent") {
query := parsedURL.Query()
query.Set("alt", "sse")
parsedURL.RawQuery = query.Encode()
}
}
func shouldApplyClaudeHeaders(converter string, info *relaycommon.RelayInfo) bool {
return converter == relayconvert.ConverterOpenAIChatToClaudeMessages ||
(converter == relayconvert.ConverterNone && info != nil && info.RelayFormat == types.RelayFormatClaude)
}
func applyClaudeHeaders(c *gin.Context, header *http.Header, info *relaycommon.RelayInfo) {
anthropicVersion := ""
if c != nil && c.Request != nil {
anthropicVersion = c.Request.Header.Get("anthropic-version")
}
if anthropicVersion == "" {
anthropicVersion = "2023-06-01"
}
header.Set("anthropic-version", anthropicVersion)
if c != nil {
claude.CommonClaudeHeadersOperation(c, header, info)
}
}
func applyAuthTemplate(template string, apiKey string) string {
return strings.ReplaceAll(template, "{api_key}", apiKey)
}
func isJSONRequest(c *gin.Context) bool {
if c == nil || c.Request == nil {
return false
}
return strings.Contains(strings.ToLower(c.Request.Header.Get("Content-Type")), "application/json")
}
func (a *Adaptor) convertOpenAICompatibleRequest(c *gin.Context, info *relaycommon.RelayInfo, request *dto.GeneralOpenAIRequest) (any, error) {
old := info.ChannelType
info.ChannelType = constant.ChannelTypeOpenAI
converted, err := a.openaiAdaptor.ConvertOpenAIRequest(c, info, request)
info.ChannelType = old
return converted, err
}
func (a *Adaptor) convertOpenAICompatibleResponsesRequest(c *gin.Context, info *relaycommon.RelayInfo, request dto.OpenAIResponsesRequest) (any, error) {
old := info.ChannelType
info.ChannelType = constant.ChannelTypeOpenAI
converted, err := a.openaiAdaptor.ConvertOpenAIResponsesRequest(c, info, request)
info.ChannelType = old
return converted, err
}
func (a *Adaptor) convertOpenAICompatibleEmbeddingRequest(c *gin.Context, info *relaycommon.RelayInfo, request dto.EmbeddingRequest) (any, error) {
old := info.ChannelType
info.ChannelType = constant.ChannelTypeOpenAI
converted, err := a.openaiAdaptor.ConvertEmbeddingRequest(c, info, request)
info.ChannelType = old
return converted, err
}
func (a *Adaptor) convertOpenAICompatibleAudioRequest(c *gin.Context, info *relaycommon.RelayInfo, request dto.AudioRequest) (io.Reader, error) {
old := info.ChannelType
info.ChannelType = constant.ChannelTypeOpenAI
converted, err := a.openaiAdaptor.ConvertAudioRequest(c, info, request)
info.ChannelType = old
return converted, err
}
func (a *Adaptor) convertOpenAICompatibleImageRequest(c *gin.Context, info *relaycommon.RelayInfo, request dto.ImageRequest) (any, error) {
old := info.ChannelType
info.ChannelType = constant.ChannelTypeOpenAI
converted, err := a.openaiAdaptor.ConvertImageRequest(c, info, request)
info.ChannelType = old
return converted, err
}
@@ -0,0 +1,863 @@
package advancedcustom
import (
"bytes"
"io"
"net/http"
"net/http/httptest"
"net/url"
"testing"
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/constant"
relaycommon "github.com/QuantumNous/new-api/relay/common"
relayconstant "github.com/QuantumNous/new-api/relay/constant"
"github.com/QuantumNous/new-api/relaykit/dto"
"github.com/QuantumNous/new-api/relaykit/relayconvert"
"github.com/QuantumNous/new-api/relaykit/types"
"github.com/QuantumNous/new-api/setting/model_setting"
"github.com/gin-gonic/gin"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
func TestAdaptorUsesExactRouteAndQueryAuth(t *testing.T) {
adaptor := &Adaptor{}
info := advancedCustomRelayInfo(&dto.AdvancedCustomConfig{
Routes: []dto.AdvancedCustomRoute{
{
IncomingPath: "/v1/messages",
UpstreamPath: "https://upstream.example/v1/chat/completions?existing=1",
Converter: relayconvert.ConverterClaudeMessagesToOpenAIChat,
Auth: &dto.AdvancedCustomRouteAuth{
Type: dto.AdvancedCustomAuthTypeQuery,
Name: "api_key",
Value: "{api_key}",
},
},
},
})
info.RequestURLPath = "/v1/messages?client=1"
requestURL, err := adaptor.GetRequestURL(info)
require.NoError(t, err)
parsedURL, err := url.Parse(requestURL)
require.NoError(t, err)
assert.Equal(t, "https", parsedURL.Scheme)
assert.Equal(t, "upstream.example", parsedURL.Host)
assert.Equal(t, "/v1/chat/completions", parsedURL.Path)
assert.Equal(t, "1", parsedURL.Query().Get("existing"))
assert.Equal(t, "sk-test", parsedURL.Query().Get("api_key"))
}
func TestAdaptorJoinsUpstreamPathWithChannelBaseURL(t *testing.T) {
adaptor := &Adaptor{}
info := advancedCustomRelayInfo(&dto.AdvancedCustomConfig{
Routes: []dto.AdvancedCustomRoute{
{
IncomingPath: "/v1/chat/completions",
UpstreamPath: "/proxy/v1/chat/completions?existing=1",
Converter: relayconvert.ConverterNone,
Auth: &dto.AdvancedCustomRouteAuth{
Type: dto.AdvancedCustomAuthTypeQuery,
Name: "api_key",
Value: "{api_key}",
},
},
},
})
info.ChannelBaseUrl = "https://gateway.example/base"
requestURL, err := adaptor.GetRequestURL(info)
require.NoError(t, err)
parsedURL, err := url.Parse(requestURL)
require.NoError(t, err)
assert.Equal(t, "https", parsedURL.Scheme)
assert.Equal(t, "gateway.example", parsedURL.Host)
assert.Equal(t, "/base/proxy/v1/chat/completions", parsedURL.Path)
assert.Equal(t, "1", parsedURL.Query().Get("existing"))
assert.Equal(t, "sk-test", parsedURL.Query().Get("api_key"))
}
func TestAdaptorReturnsErrorWhenUpstreamPathNeedsMissingBaseURL(t *testing.T) {
adaptor := &Adaptor{}
info := advancedCustomRelayInfo(&dto.AdvancedCustomConfig{
Routes: []dto.AdvancedCustomRoute{
{
IncomingPath: "/v1/chat/completions",
UpstreamPath: "/v1/chat/completions",
Converter: relayconvert.ConverterNone,
},
},
})
info.ChannelBaseUrl = ""
_, err := adaptor.GetRequestURL(info)
require.Error(t, err)
assert.Contains(t, err.Error(), "base URL is required")
}
func TestAdaptorSetupRequestHeaderUsesDefaultBearerAuth(t *testing.T) {
adaptor := &Adaptor{}
info := advancedCustomRelayInfo(&dto.AdvancedCustomConfig{
Routes: []dto.AdvancedCustomRoute{
{
IncomingPath: "/v1/chat/completions",
UpstreamPath: "https://upstream.example/v1/chat/completions",
Converter: relayconvert.ConverterNone,
},
},
})
c := advancedCustomGinContext("/v1/chat/completions")
header := http.Header{}
require.NoError(t, adaptor.SetupRequestHeader(c, &header, info))
assert.Equal(t, "Bearer sk-test", header.Get("Authorization"))
}
func TestAdaptorSetupRequestHeaderUsesConfiguredHeaderAuth(t *testing.T) {
adaptor := &Adaptor{}
info := advancedCustomRelayInfo(&dto.AdvancedCustomConfig{
Routes: []dto.AdvancedCustomRoute{
{
IncomingPath: "/v1/chat/completions",
UpstreamPath: "https://upstream.example/v1/chat/completions",
Converter: relayconvert.ConverterNone,
Auth: &dto.AdvancedCustomRouteAuth{
Type: dto.AdvancedCustomAuthTypeHeader,
Name: "x-api-key",
Value: "{api_key}",
},
},
},
})
c := advancedCustomGinContext("/v1/chat/completions")
header := http.Header{}
require.NoError(t, adaptor.SetupRequestHeader(c, &header, info))
assert.Empty(t, header.Get("Authorization"))
assert.Equal(t, "sk-test", header.Get("x-api-key"))
}
func TestAdaptorSetupRequestHeaderAddsClaudeDefaultHeaders(t *testing.T) {
adaptor := &Adaptor{}
info := advancedCustomRelayInfo(&dto.AdvancedCustomConfig{
Routes: []dto.AdvancedCustomRoute{
{
IncomingPath: "/v1/messages",
UpstreamPath: "https://api.anthropic.com/v1/messages",
Converter: relayconvert.ConverterNone,
Auth: &dto.AdvancedCustomRouteAuth{
Type: dto.AdvancedCustomAuthTypeHeader,
Name: "x-api-key",
Value: "{api_key}",
},
},
},
})
info.RelayFormat = types.RelayFormatClaude
c := advancedCustomGinContext("/v1/messages")
header := http.Header{}
require.NoError(t, adaptor.SetupRequestHeader(c, &header, info))
assert.Equal(t, "sk-test", header.Get("x-api-key"))
assert.Equal(t, "2023-06-01", header.Get("anthropic-version"))
}
func TestAdaptorReturnsErrorWhenNoRouteMatchesPath(t *testing.T) {
adaptor := &Adaptor{}
info := advancedCustomRelayInfo(&dto.AdvancedCustomConfig{
Routes: []dto.AdvancedCustomRoute{
{
IncomingPath: "/v1/messages",
UpstreamPath: "https://upstream.example/v1/chat/completions",
Converter: relayconvert.ConverterClaudeMessagesToOpenAIChat,
},
},
})
info.RequestURLPath = "/v1/chat/completions"
_, err := adaptor.GetRequestURL(info)
require.Error(t, err)
assert.Contains(t, err.Error(), "does not support request path")
}
func TestAdaptorReplacesModelPlaceholderInRouteURL(t *testing.T) {
adaptor := &Adaptor{}
info := advancedCustomRelayInfo(&dto.AdvancedCustomConfig{
Routes: []dto.AdvancedCustomRoute{
{
IncomingPath: "/v1/chat/completions",
UpstreamPath: "https://generativelanguage.googleapis.com/v1beta/models/{model}:generateContent",
Converter: relayconvert.ConverterOpenAIChatToGeminiContent,
Auth: &dto.AdvancedCustomRouteAuth{
Type: dto.AdvancedCustomAuthTypeQuery,
Name: "key",
Value: "{api_key}",
},
},
},
})
info.UpstreamModelName = "gemini-2.5-flash"
requestURL, err := adaptor.GetRequestURL(info)
require.NoError(t, err)
parsedURL, err := url.Parse(requestURL)
require.NoError(t, err)
assert.Equal(t, "/v1beta/models/gemini-2.5-flash:generateContent", parsedURL.Path)
assert.Equal(t, "sk-test", parsedURL.Query().Get("key"))
assert.Empty(t, parsedURL.Query().Get("alt"))
}
func TestAdaptorSwitchesGeminiGenerateContentURLForStream(t *testing.T) {
adaptor := &Adaptor{}
info := advancedCustomRelayInfo(&dto.AdvancedCustomConfig{
Routes: []dto.AdvancedCustomRoute{
{
IncomingPath: "/v1/chat/completions",
UpstreamPath: "https://generativelanguage.googleapis.com/v1beta/models/{model}:generateContent?existing=1",
Converter: relayconvert.ConverterOpenAIChatToGeminiContent,
Auth: &dto.AdvancedCustomRouteAuth{
Type: dto.AdvancedCustomAuthTypeQuery,
Name: "key",
Value: "{api_key}",
},
},
},
})
info.UpstreamModelName = "gemini-2.5-pro"
info.IsStream = true
requestURL, err := adaptor.GetRequestURL(info)
require.NoError(t, err)
parsedURL, err := url.Parse(requestURL)
require.NoError(t, err)
assert.Equal(t, "/v1beta/models/gemini-2.5-pro:streamGenerateContent", parsedURL.Path)
assert.Equal(t, "sse", parsedURL.Query().Get("alt"))
assert.Equal(t, "1", parsedURL.Query().Get("existing"))
assert.Equal(t, "sk-test", parsedURL.Query().Get("key"))
}
func TestAdaptorMatchesGeminiIncomingPathTemplate(t *testing.T) {
tests := []struct {
name string
requestURLPath string
wantRequestPath string
}{
{
name: "generate content",
requestURLPath: "/v1beta/models/gemini-2.5-flash:generateContent",
wantRequestPath: "/v1/chat/completions",
},
{
name: "stream generate content",
requestURLPath: "/v1beta/models/gemini-2.5-flash:streamGenerateContent?alt=sse",
wantRequestPath: "/v1/chat/completions",
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
adaptor := &Adaptor{}
info := advancedCustomRelayInfo(&dto.AdvancedCustomConfig{
Routes: []dto.AdvancedCustomRoute{
{
IncomingPath: "/v1beta/models/{model}:generateContent",
UpstreamPath: "https://upstream.example/v1/chat/completions",
Converter: relayconvert.ConverterGeminiContentToOpenAIChat,
},
},
})
info.RequestURLPath = tt.requestURLPath
requestURL, err := adaptor.GetRequestURL(info)
require.NoError(t, err)
parsedURL, err := url.Parse(requestURL)
require.NoError(t, err)
assert.Equal(t, tt.wantRequestPath, parsedURL.Path)
})
}
}
func TestAdaptorBuildModelListRequestUsesConfiguredRouteAuth(t *testing.T) {
adaptor := &Adaptor{}
info := advancedCustomRelayInfo(&dto.AdvancedCustomConfig{
Routes: []dto.AdvancedCustomRoute{
{
IncomingPath: "/v1/models",
UpstreamPath: "/provider/models",
Converter: relayconvert.ConverterNone,
Auth: &dto.AdvancedCustomRouteAuth{
Type: dto.AdvancedCustomAuthTypeHeader,
Name: "x-api-key",
Value: "token {api_key}",
},
},
},
})
info.RequestURLPath = "/v1/models"
requestURL, header, err := adaptor.BuildModelListRequest(info)
require.NoError(t, err)
parsedURL, err := url.Parse(requestURL)
require.NoError(t, err)
assert.Equal(t, "fallback.example", parsedURL.Host)
assert.Equal(t, "/provider/models", parsedURL.Path)
assert.Equal(t, "token sk-test", header.Get("x-api-key"))
assert.Empty(t, header.Get("Authorization"))
}
func TestAdaptorBuildModelListRequestUsesConfiguredQueryAuth(t *testing.T) {
adaptor := &Adaptor{}
info := advancedCustomRelayInfo(&dto.AdvancedCustomConfig{
Routes: []dto.AdvancedCustomRoute{
{
IncomingPath: "/v1/models",
UpstreamPath: "https://upstream.example/v1/models?existing=1",
Converter: relayconvert.ConverterNone,
Auth: &dto.AdvancedCustomRouteAuth{
Type: dto.AdvancedCustomAuthTypeQuery,
Name: "key",
Value: "{api_key}",
},
},
},
})
info.RequestURLPath = "/v1/models"
requestURL, header, err := adaptor.BuildModelListRequest(info)
require.NoError(t, err)
parsedURL, err := url.Parse(requestURL)
require.NoError(t, err)
assert.Equal(t, "upstream.example", parsedURL.Host)
assert.Equal(t, "/v1/models", parsedURL.Path)
assert.Equal(t, "1", parsedURL.Query().Get("existing"))
assert.Equal(t, "sk-test", parsedURL.Query().Get("key"))
assert.Empty(t, header.Get("Authorization"))
}
func TestAdaptorBuildModelListRequestDefaultAndNoAuth(t *testing.T) {
tests := []struct {
name string
auth *dto.AdvancedCustomRouteAuth
wantAuthorization string
}{
{
name: "default bearer",
wantAuthorization: "Bearer sk-test",
},
{
name: "no authentication",
auth: &dto.AdvancedCustomRouteAuth{
Type: dto.AdvancedCustomAuthTypeNone,
},
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
info := advancedCustomRelayInfo(&dto.AdvancedCustomConfig{
Routes: []dto.AdvancedCustomRoute{
{
IncomingPath: dto.AdvancedCustomModelListPath,
UpstreamPath: "/provider/models",
Auth: tt.auth,
},
},
})
info.RequestURLPath = "/unrelated/path"
requestURL, header, err := (&Adaptor{}).BuildModelListRequest(info)
require.NoError(t, err)
assert.Equal(t, "https://fallback.example/provider/models", requestURL)
assert.Equal(t, tt.wantAuthorization, header.Get("Authorization"))
})
}
}
func TestAdaptorBuildModelListRequestDoesNotReuseRelayRoute(t *testing.T) {
adaptor := &Adaptor{}
info := advancedCustomRelayInfo(&dto.AdvancedCustomConfig{
Routes: []dto.AdvancedCustomRoute{
{
IncomingPath: "/v1/chat/completions",
UpstreamPath: "/chat",
},
{
IncomingPath: dto.AdvancedCustomModelListPath,
UpstreamPath: "/provider/models",
},
},
})
chatURL, err := adaptor.GetRequestURL(info)
require.NoError(t, err)
assert.Equal(t, "https://fallback.example/chat", chatURL)
modelURL, header, err := adaptor.BuildModelListRequest(info)
require.NoError(t, err)
assert.Equal(t, "https://fallback.example/provider/models", modelURL)
assert.Equal(t, "Bearer sk-test", header.Get("Authorization"))
}
func TestAdaptorBuildModelListRequestRequiresConfiguredRoute(t *testing.T) {
info := advancedCustomRelayInfo(&dto.AdvancedCustomConfig{
Routes: []dto.AdvancedCustomRoute{
{
IncomingPath: "/v1/chat/completions",
UpstreamPath: "/v1/chat/completions",
},
},
})
_, _, err := (&Adaptor{}).BuildModelListRequest(info)
require.Error(t, err)
assert.Contains(t, err.Error(), "does not configure a /v1/models route")
}
func TestAdaptorBuildBalanceRequestUsesConfiguredRoute(t *testing.T) {
info := advancedCustomRelayInfo(&dto.AdvancedCustomConfig{
Routes: []dto.AdvancedCustomRoute{
{
IncomingPath: dto.AdvancedCustomModelListPath,
UpstreamPath: "/provider/models",
},
{
IncomingPath: dto.AdvancedCustomBalancePath,
UpstreamPath: "/provider/balance?existing=1",
Auth: &dto.AdvancedCustomRouteAuth{
Type: dto.AdvancedCustomAuthTypeQuery,
Name: "token",
Value: "prefix-{api_key}",
},
},
},
})
requestURL, header, err := (&Adaptor{}).BuildBalanceRequest(info)
require.NoError(t, err)
parsedURL, err := url.Parse(requestURL)
require.NoError(t, err)
assert.Equal(t, "/provider/balance", parsedURL.Path)
assert.Equal(t, "1", parsedURL.Query().Get("existing"))
assert.Equal(t, "prefix-sk-test", parsedURL.Query().Get("token"))
assert.Empty(t, header.Get("Authorization"))
}
func TestAdaptorBuildBalanceRequestRequiresConfiguredRoute(t *testing.T) {
info := advancedCustomRelayInfo(&dto.AdvancedCustomConfig{
Routes: []dto.AdvancedCustomRoute{{
IncomingPath: dto.AdvancedCustomModelListPath,
UpstreamPath: "/provider/models",
}},
})
_, _, err := (&Adaptor{}).BuildBalanceRequest(info)
require.Error(t, err)
assert.Contains(t, err.Error(), "does not configure a /v1/dashboard/billing/credit_grants route")
}
func TestAdaptorConvertsResponsesRequestToOpenAIChatUpstream(t *testing.T) {
adaptor := &Adaptor{}
info := advancedCustomRelayInfo(&dto.AdvancedCustomConfig{
Routes: []dto.AdvancedCustomRoute{
{
IncomingPath: "/v1/responses",
UpstreamPath: "/v1/chat/completions",
Converter: relayconvert.ConverterOpenAIResponsesToOpenAIChat,
},
},
})
info.RelayMode = relayconstant.RelayModeResponses
info.RequestURLPath = "/v1/responses"
gin.SetMode(gin.TestMode)
recorder := httptest.NewRecorder()
c, _ := gin.CreateTestContext(recorder)
c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", nil)
c.Request.Header.Set("Content-Type", "application/json")
converted, err := adaptor.ConvertOpenAIResponsesRequest(c, info, dto.OpenAIResponsesRequest{
Model: "gpt-test",
Instructions: mustAdvancedCustomRawMessage(t, "system rules"),
Input: mustAdvancedCustomRawMessage(t, "hello"),
})
require.NoError(t, err)
chatReq, ok := converted.(*dto.GeneralOpenAIRequest)
require.True(t, ok)
assert.Equal(t, "gpt-test", chatReq.Model)
require.Len(t, chatReq.Messages, 2)
assert.Equal(t, "system", chatReq.Messages[0].Role)
assert.Equal(t, "system rules", chatReq.Messages[0].StringContent())
assert.Equal(t, "user", chatReq.Messages[1].Role)
assert.Equal(t, "hello", chatReq.Messages[1].StringContent())
requestURL, err := adaptor.GetRequestURL(info)
require.NoError(t, err)
parsedURL, err := url.Parse(requestURL)
require.NoError(t, err)
assert.Equal(t, "/v1/chat/completions", parsedURL.Path)
}
func TestAdaptorSelectsDuplicateResponsesRoutesByModel(t *testing.T) {
config := &dto.AdvancedCustomConfig{
Routes: []dto.AdvancedCustomRoute{
{
IncomingPath: "/v1/responses",
UpstreamPath: "/v1/chat/completions",
Converter: relayconvert.ConverterOpenAIResponsesToOpenAIChat,
Models: []string{"gpt-test"},
},
{
IncomingPath: "/v1/responses",
UpstreamPath: "/v1beta/models/{model}:generateContent",
Converter: relayconvert.ConverterOpenAIResponsesToGemini,
Models: []string{"gemini-test"},
},
},
}
chatAdaptor := &Adaptor{}
chatInfo := advancedCustomRelayInfo(config)
chatInfo.RelayFormat = types.RelayFormatOpenAIResponses
chatInfo.RelayMode = relayconstant.RelayModeResponses
chatInfo.RequestURLPath = "/v1/responses"
chatInfo.OriginModelName = "gpt-test"
chatInfo.UpstreamModelName = "gpt-test"
chatConverted, err := chatAdaptor.ConvertOpenAIResponsesRequest(advancedCustomGinContext("/v1/responses"), chatInfo, dto.OpenAIResponsesRequest{
Model: "gpt-test",
Input: mustAdvancedCustomRawMessage(t, "hello"),
})
require.NoError(t, err)
_, ok := chatConverted.(*dto.GeneralOpenAIRequest)
require.True(t, ok)
geminiAdaptor := &Adaptor{}
geminiInfo := advancedCustomRelayInfo(config)
geminiInfo.RelayFormat = types.RelayFormatOpenAIResponses
geminiInfo.RelayMode = relayconstant.RelayModeResponses
geminiInfo.RequestURLPath = "/v1/responses"
geminiInfo.OriginModelName = "gemini-test"
geminiInfo.UpstreamModelName = "gemini-test"
geminiInfo.IsStream = true
geminiConverted, err := geminiAdaptor.ConvertOpenAIResponsesRequest(advancedCustomGinContext("/v1/responses"), geminiInfo, dto.OpenAIResponsesRequest{
Model: "gemini-test",
Input: mustAdvancedCustomRawMessage(t, "hello"),
})
require.NoError(t, err)
_, ok = geminiConverted.(*dto.GeminiChatRequest)
require.True(t, ok)
requestURL, err := geminiAdaptor.GetRequestURL(geminiInfo)
require.NoError(t, err)
parsedURL, err := url.Parse(requestURL)
require.NoError(t, err)
assert.Equal(t, "/v1beta/models/gemini-test:streamGenerateContent", parsedURL.Path)
assert.Equal(t, "sse", parsedURL.Query().Get("alt"))
}
func TestAdaptorResponsesToGeminiUsesResponsesBridge(t *testing.T) {
adaptor := &Adaptor{}
info := advancedCustomRelayInfo(&dto.AdvancedCustomConfig{
Routes: []dto.AdvancedCustomRoute{
{
IncomingPath: "/v1/responses",
UpstreamPath: "/v1beta/models/{model}:generateContent",
Converter: relayconvert.ConverterOpenAIResponsesToGemini,
Models: []string{"gemini-test"},
},
},
})
info.RelayFormat = types.RelayFormatOpenAIResponses
info.RelayMode = relayconstant.RelayModeResponses
info.RequestURLPath = "/v1/responses"
info.OriginModelName = "gemini-test"
info.UpstreamModelName = "gemini-test"
gin.SetMode(gin.TestMode)
recorder := httptest.NewRecorder()
c, _ := gin.CreateTestContext(recorder)
c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", nil)
c.Request.Header.Set("Content-Type", "application/json")
payload := dto.GeminiChatResponse{
Candidates: []dto.GeminiChatCandidate{
{
Content: dto.GeminiChatContent{
Role: "model",
Parts: []dto.GeminiPart{
{Text: "hello"},
},
},
},
},
UsageMetadata: dto.GeminiUsageMetadata{
PromptTokenCount: 2,
CandidatesTokenCount: 3,
TotalTokenCount: 5,
},
}
body, err := common.Marshal(payload)
require.NoError(t, err)
usage, newAPIError := adaptor.DoResponse(c, &http.Response{
Body: io.NopCloser(bytes.NewReader(body)),
}, info)
require.Nil(t, newAPIError)
require.NotNil(t, usage)
got := recorder.Body.String()
assert.Contains(t, got, `"object":"response"`)
assert.Contains(t, got, `"type":"output_text"`)
assert.Contains(t, got, `"text":"hello"`)
assert.NotContains(t, got, `"candidates"`)
}
func TestAdaptorResponsesToGeminiAddsThoughtSignatureForFunctionCallHistory(t *testing.T) {
geminiSettings := model_setting.GetGeminiSettings()
originalThoughtSignatureEnabled := geminiSettings.FunctionCallThoughtSignatureEnabled
geminiSettings.FunctionCallThoughtSignatureEnabled = true
t.Cleanup(func() {
geminiSettings.FunctionCallThoughtSignatureEnabled = originalThoughtSignatureEnabled
})
adaptor := &Adaptor{}
info := advancedCustomRelayInfo(&dto.AdvancedCustomConfig{
Routes: []dto.AdvancedCustomRoute{
{
IncomingPath: "/v1/responses",
UpstreamPath: "/v1beta/models/{model}:generateContent",
Converter: relayconvert.ConverterOpenAIResponsesToGemini,
Models: []string{"gemini-test"},
},
},
})
info.RelayFormat = types.RelayFormatOpenAIResponses
info.RelayMode = relayconstant.RelayModeResponses
info.RequestURLPath = "/v1/responses"
info.OriginModelName = "gemini-test"
info.UpstreamModelName = "gemini-test"
converted, err := adaptor.ConvertOpenAIResponsesRequest(advancedCustomGinContext("/v1/responses"), info, dto.OpenAIResponsesRequest{
Model: "gemini-test",
Input: mustAdvancedCustomRawMessage(t, []map[string]any{
{
"role": "user",
"content": "hi",
},
{
"type": "function_call",
"call_id": "call_1",
"name": "glob",
"arguments": map[string]any{"query": "*"},
},
{
"type": "function_call_output",
"call_id": "call_1",
"output": []map[string]any{{"path": "report.md"}},
},
}),
Tools: mustAdvancedCustomRawMessage(t, []map[string]any{
{"type": "function", "name": "glob", "parameters": map[string]any{"type": "object"}},
}),
})
require.NoError(t, err)
geminiReq, ok := converted.(*dto.GeminiChatRequest)
require.True(t, ok)
require.Len(t, geminiReq.Contents, 3)
require.Len(t, geminiReq.Contents[1].Parts, 1)
require.NotNil(t, geminiReq.Contents[1].Parts[0].FunctionCall)
assert.NotEmpty(t, geminiReq.Contents[1].Parts[0].ThoughtSignature)
require.Len(t, geminiReq.Contents[2].Parts, 1)
require.NotNil(t, geminiReq.Contents[2].Parts[0].FunctionResponse)
assert.Empty(t, geminiReq.Contents[2].Parts[0].ThoughtSignature)
}
func TestAdaptorConvertsOpenAIChatRequestToResponsesUpstream(t *testing.T) {
adaptor := &Adaptor{}
info := advancedCustomRelayInfo(&dto.AdvancedCustomConfig{
Routes: []dto.AdvancedCustomRoute{
{
IncomingPath: "/v1/chat/completions",
UpstreamPath: "/v1/responses",
Converter: relayconvert.ConverterOpenAIChatToOpenAIResponses,
},
},
})
c := advancedCustomGinContext("/v1/chat/completions")
converted, err := adaptor.ConvertOpenAIRequest(c, info, &dto.GeneralOpenAIRequest{
Model: "gpt-test",
Messages: []dto.Message{
{Role: "user", Content: "hello"},
},
})
require.NoError(t, err)
responsesReq, ok := converted.(*dto.OpenAIResponsesRequest)
require.True(t, ok)
assert.Equal(t, "gpt-test", responsesReq.Model)
assert.NotEmpty(t, responsesReq.Input)
}
func TestAdaptorConvertsOpenAIChatRequestToClaudeUpstream(t *testing.T) {
adaptor := &Adaptor{}
info := advancedCustomRelayInfo(&dto.AdvancedCustomConfig{
Routes: []dto.AdvancedCustomRoute{
{
IncomingPath: "/v1/chat/completions",
UpstreamPath: "/v1/messages",
Converter: relayconvert.ConverterOpenAIChatToClaudeMessages,
},
},
})
c := advancedCustomGinContext("/v1/chat/completions")
converted, err := adaptor.ConvertOpenAIRequest(c, info, &dto.GeneralOpenAIRequest{
Model: "claude-test",
Messages: []dto.Message{
{Role: "user", Content: "hello"},
},
})
require.NoError(t, err)
claudeReq, ok := converted.(*dto.ClaudeRequest)
require.True(t, ok)
assert.Equal(t, "claude-test", claudeReq.Model)
require.Len(t, claudeReq.Messages, 1)
assert.Equal(t, "user", claudeReq.Messages[0].Role)
}
func TestAdaptorConvertsOpenAIChatRequestToGeminiUpstream(t *testing.T) {
adaptor := &Adaptor{}
info := advancedCustomRelayInfo(&dto.AdvancedCustomConfig{
Routes: []dto.AdvancedCustomRoute{
{
IncomingPath: "/v1/chat/completions",
UpstreamPath: "/v1beta/models/{model}:generateContent",
Converter: relayconvert.ConverterOpenAIChatToGeminiContent,
},
},
})
info.UpstreamModelName = "gemini-2.5-flash"
c := advancedCustomGinContext("/v1/chat/completions")
converted, err := adaptor.ConvertOpenAIRequest(c, info, &dto.GeneralOpenAIRequest{
Model: "gemini-2.5-flash",
Messages: []dto.Message{
{Role: "user", Content: "hello"},
},
})
require.NoError(t, err)
geminiReq, ok := converted.(*dto.GeminiChatRequest)
require.True(t, ok)
require.Len(t, geminiReq.Contents, 1)
assert.Equal(t, "user", geminiReq.Contents[0].Role)
}
func TestAdaptorConvertsClaudeRequestToOpenAIChatUpstream(t *testing.T) {
adaptor := &Adaptor{}
info := advancedCustomRelayInfo(&dto.AdvancedCustomConfig{
Routes: []dto.AdvancedCustomRoute{
{
IncomingPath: "/v1/messages",
UpstreamPath: "/v1/chat/completions",
Converter: relayconvert.ConverterClaudeMessagesToOpenAIChat,
},
},
})
info.RelayFormat = types.RelayFormatClaude
info.RequestURLPath = "/v1/messages"
c := advancedCustomGinContext("/v1/messages")
converted, err := adaptor.ConvertClaudeRequest(c, info, &dto.ClaudeRequest{
Model: "gpt-test",
Messages: []dto.ClaudeMessage{
{Role: "user", Content: "hello"},
},
})
require.NoError(t, err)
chatReq, ok := converted.(*dto.GeneralOpenAIRequest)
require.True(t, ok)
assert.Equal(t, "gpt-test", chatReq.Model)
require.Len(t, chatReq.Messages, 1)
assert.Equal(t, "user", chatReq.Messages[0].Role)
}
func TestAdaptorConvertsGeminiRequestToOpenAIChatUpstream(t *testing.T) {
adaptor := &Adaptor{}
info := advancedCustomRelayInfo(&dto.AdvancedCustomConfig{
Routes: []dto.AdvancedCustomRoute{
{
IncomingPath: "/v1beta/models/{model}:generateContent",
UpstreamPath: "/v1/chat/completions",
Converter: relayconvert.ConverterGeminiContentToOpenAIChat,
},
},
})
info.RelayFormat = types.RelayFormatGemini
info.RequestURLPath = "/v1beta/models/gemini-2.5-flash:generateContent"
info.UpstreamModelName = "gpt-test"
c := advancedCustomGinContext("/v1beta/models/gemini-2.5-flash:generateContent")
converted, err := adaptor.ConvertGeminiRequest(c, info, &dto.GeminiChatRequest{
Contents: []dto.GeminiChatContent{
{
Role: "user",
Parts: []dto.GeminiPart{
{Text: "hello"},
},
},
},
})
require.NoError(t, err)
chatReq, ok := converted.(*dto.GeneralOpenAIRequest)
require.True(t, ok)
assert.Equal(t, "gpt-test", chatReq.Model)
require.Len(t, chatReq.Messages, 1)
assert.Equal(t, "user", chatReq.Messages[0].Role)
}
func advancedCustomRelayInfo(config *dto.AdvancedCustomConfig) *relaycommon.RelayInfo {
return &relaycommon.RelayInfo{
RelayFormat: types.RelayFormatOpenAI,
RelayMode: relayconstant.RelayModeChatCompletions,
RequestURLPath: "/v1/chat/completions",
OriginModelName: "gpt-test",
ChannelMeta: &relaycommon.ChannelMeta{
ApiKey: "sk-test",
ChannelBaseUrl: "https://fallback.example",
ChannelType: constant.ChannelTypeAdvancedCustom,
UpstreamModelName: "gpt-test",
ChannelOtherSettings: dto.ChannelOtherSettings{
AdvancedCustom: config,
},
},
}
}
func advancedCustomGinContext(path string) *gin.Context {
gin.SetMode(gin.TestMode)
c, _ := gin.CreateTestContext(httptest.NewRecorder())
c.Request = httptest.NewRequest(http.MethodPost, path, nil)
c.Request.Header.Set("Content-Type", "application/json")
return c
}
func mustAdvancedCustomRawMessage(t *testing.T, value any) []byte {
t.Helper()
raw, err := common.Marshal(value)
require.NoError(t, err)
return raw
}
+5
View File
@@ -0,0 +1,5 @@
package ai360
var ModelList []string
var ChannelName = "ai360"
+277
View File
@@ -0,0 +1,277 @@
package ali
import (
"errors"
"fmt"
"io"
"net/http"
"strings"
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/relay/channel"
"github.com/QuantumNous/new-api/relay/channel/claude"
"github.com/QuantumNous/new-api/relay/channel/openai"
relaycommon "github.com/QuantumNous/new-api/relay/common"
"github.com/QuantumNous/new-api/relay/constant"
"github.com/QuantumNous/new-api/relaykit/dto"
"github.com/QuantumNous/new-api/relaykit/types"
"github.com/QuantumNous/new-api/service"
"github.com/QuantumNous/new-api/setting/model_setting"
"github.com/gin-gonic/gin"
"github.com/samber/lo"
)
type Adaptor struct {
IsSyncImageModel bool
}
const aliAnthropicMessagesModelsEnv = "ALI_ANTHROPIC_MESSAGES_MODELS"
const defaultAliAnthropicMessagesModels = "qwen,deepseek-v4,kimi,glm,minimax-m"
/*
var syncModels = []string{
"z-image",
"qwen-image",
"wan2.6",
}
*/
func supportsAliAnthropicMessages(modelName string) bool {
normalizedModelName := strings.ToLower(strings.TrimSpace(modelName))
if normalizedModelName == "" {
return false
}
return lo.SomeBy(aliAnthropicMessagesModelPatterns(), func(pattern string) bool {
return strings.Contains(normalizedModelName, pattern)
})
}
func aliAnthropicMessagesModelPatterns() []string {
configuredModels := common.GetEnvOrDefaultString(aliAnthropicMessagesModelsEnv, defaultAliAnthropicMessagesModels)
return lo.FilterMap(strings.Split(configuredModels, ","), func(item string, _ int) (string, bool) {
pattern := strings.ToLower(strings.TrimSpace(item))
return pattern, pattern != ""
})
}
var syncModels = []string{
"z-image",
"qwen-image",
"wan2.6",
}
func isSyncImageModel(modelName string) bool {
return model_setting.IsSyncImageModel(modelName)
}
func (a *Adaptor) ConvertGeminiRequest(*gin.Context, *relaycommon.RelayInfo, *dto.GeminiChatRequest) (any, error) {
//TODO implement me
return nil, errors.New("not implemented")
}
func (a *Adaptor) ConvertClaudeRequest(c *gin.Context, info *relaycommon.RelayInfo, req *dto.ClaudeRequest) (any, error) {
if supportsAliAnthropicMessages(info.UpstreamModelName) {
return req, nil
}
result, err := service.ConvertRequest(c, info, types.RelayFormatOpenAI, req)
if err != nil {
return nil, err
}
oaiReq, ok := result.Value.(*dto.GeneralOpenAIRequest)
if !ok {
return nil, fmt.Errorf("expected OpenAI chat completions request, got %T", result.Value)
}
if info.SupportStreamOptions && info.IsStream {
oaiReq.StreamOptions = &dto.StreamOptions{IncludeUsage: true}
}
return a.ConvertOpenAIRequest(c, info, oaiReq)
}
func (a *Adaptor) Init(info *relaycommon.RelayInfo) {
}
func (a *Adaptor) GetRequestURL(info *relaycommon.RelayInfo) (string, error) {
var fullRequestURL string
switch info.RelayFormat {
case types.RelayFormatClaude:
if supportsAliAnthropicMessages(info.UpstreamModelName) {
fullRequestURL = fmt.Sprintf("%s/apps/anthropic/v1/messages", info.ChannelBaseUrl)
} else {
fullRequestURL = fmt.Sprintf("%s/compatible-mode/v1/chat/completions", info.ChannelBaseUrl)
}
default:
switch info.RelayMode {
case constant.RelayModeEmbeddings:
fullRequestURL = fmt.Sprintf("%s/compatible-mode/v1/embeddings", info.ChannelBaseUrl)
case constant.RelayModeRerank:
fullRequestURL = fmt.Sprintf("%s/api/v1/services/rerank/text-rerank/text-rerank", info.ChannelBaseUrl)
case constant.RelayModeResponses:
fullRequestURL = fmt.Sprintf("%s/api/v2/apps/protocols/compatible-mode/v1/responses", info.ChannelBaseUrl)
case constant.RelayModeImagesGenerations:
if isSyncImageModel(info.UpstreamModelName) {
fullRequestURL = fmt.Sprintf("%s/api/v1/services/aigc/multimodal-generation/generation", info.ChannelBaseUrl)
} else {
fullRequestURL = fmt.Sprintf("%s/api/v1/services/aigc/text2image/image-synthesis", info.ChannelBaseUrl)
}
case constant.RelayModeImagesEdits:
if isOldWanModel(info.UpstreamModelName) {
fullRequestURL = fmt.Sprintf("%s/api/v1/services/aigc/image2image/image-synthesis", info.ChannelBaseUrl)
} else if isWanModel(info.UpstreamModelName) {
fullRequestURL = fmt.Sprintf("%s/api/v1/services/aigc/image-generation/generation", info.ChannelBaseUrl)
} else {
fullRequestURL = fmt.Sprintf("%s/api/v1/services/aigc/multimodal-generation/generation", info.ChannelBaseUrl)
}
case constant.RelayModeCompletions:
fullRequestURL = fmt.Sprintf("%s/compatible-mode/v1/completions", info.ChannelBaseUrl)
default:
fullRequestURL = fmt.Sprintf("%s/compatible-mode/v1/chat/completions", info.ChannelBaseUrl)
}
}
return fullRequestURL, nil
}
func (a *Adaptor) SetupRequestHeader(c *gin.Context, req *http.Header, info *relaycommon.RelayInfo) error {
channel.SetupApiRequestHeader(info, c, req)
req.Set("Authorization", "Bearer "+info.ApiKey)
if info.IsStream {
req.Set("X-DashScope-SSE", "enable")
}
if c.GetString("plugin") != "" {
req.Set("X-DashScope-Plugin", c.GetString("plugin"))
}
if info.RelayMode == constant.RelayModeImagesGenerations {
if isSyncImageModel(info.UpstreamModelName) {
} else {
req.Set("X-DashScope-Async", "enable")
}
}
if info.RelayMode == constant.RelayModeImagesEdits {
if isWanModel(info.UpstreamModelName) {
req.Set("X-DashScope-Async", "enable")
}
req.Set("Content-Type", "application/json")
}
return nil
}
func (a *Adaptor) ConvertOpenAIRequest(c *gin.Context, info *relaycommon.RelayInfo, request *dto.GeneralOpenAIRequest) (any, error) {
if request == nil {
return nil, errors.New("request is nil")
}
// docs: https://bailian.console.aliyun.com/?tab=api#/api/?type=model&url=2712216
// fix: InternalError.Algo.InvalidParameter: The value of the enable_thinking parameter is restricted to True.
//if strings.Contains(request.Model, "thinking") {
// request.EnableThinking = true
// request.Stream = true
// info.IsStream = true
//}
//// fix: ali parameter.enable_thinking must be set to false for non-streaming calls
//if !info.IsStream {
// request.EnableThinking = false
//}
switch info.RelayMode {
default:
aliReq := requestOpenAI2Ali(*request, info.UpstreamModelName)
return aliReq, nil
}
}
func (a *Adaptor) ConvertImageRequest(c *gin.Context, info *relaycommon.RelayInfo, request dto.ImageRequest) (any, error) {
if info.RelayMode == constant.RelayModeImagesGenerations {
if isSyncImageModel(info.UpstreamModelName) {
a.IsSyncImageModel = true
}
aliRequest, err := oaiImage2AliImageRequest(info, request, a.IsSyncImageModel)
if err != nil {
return nil, fmt.Errorf("convert image request to async ali image request failed: %w", err)
}
return aliRequest, nil
} else if info.RelayMode == constant.RelayModeImagesEdits {
if isOldWanModel(info.UpstreamModelName) {
return oaiFormEdit2WanxImageEdit(c, info, request)
}
if isSyncImageModel(info.UpstreamModelName) {
if isWanModel(info.UpstreamModelName) {
a.IsSyncImageModel = false
} else {
a.IsSyncImageModel = true
}
}
// ali image edit https://bailian.console.aliyun.com/?tab=api#/api/?type=model&url=2976416
// 如果用户使用表单,则需要解析表单数据
if strings.Contains(c.Request.Header.Get("Content-Type"), "multipart/form-data") {
aliRequest, err := oaiFormEdit2AliImageEdit(c, info, request)
if err != nil {
return nil, fmt.Errorf("convert image edit form request failed: %w", err)
}
return aliRequest, nil
} else {
aliRequest, err := oaiImage2AliImageRequest(info, request, a.IsSyncImageModel)
if err != nil {
return nil, fmt.Errorf("convert image request to async ali image request failed: %w", err)
}
return aliRequest, nil
}
}
return nil, fmt.Errorf("unsupported image relay mode: %d", info.RelayMode)
}
func (a *Adaptor) ConvertRerankRequest(c *gin.Context, relayMode int, request dto.RerankRequest) (any, error) {
return ConvertRerankRequest(request), nil
}
func (a *Adaptor) ConvertEmbeddingRequest(c *gin.Context, info *relaycommon.RelayInfo, request dto.EmbeddingRequest) (any, error) {
return request, nil
}
func (a *Adaptor) ConvertAudioRequest(c *gin.Context, info *relaycommon.RelayInfo, request dto.AudioRequest) (io.Reader, error) {
//TODO implement me
return nil, errors.New("not implemented")
}
func (a *Adaptor) ConvertOpenAIResponsesRequest(c *gin.Context, info *relaycommon.RelayInfo, request dto.OpenAIResponsesRequest) (any, error) {
return request, nil
}
func (a *Adaptor) DoRequest(c *gin.Context, info *relaycommon.RelayInfo, requestBody io.Reader) (any, error) {
return channel.DoApiRequest(a, c, info, requestBody)
}
func (a *Adaptor) DoResponse(c *gin.Context, resp *http.Response, info *relaycommon.RelayInfo) (usage any, err *types.NewAPIError) {
switch info.RelayFormat {
case types.RelayFormatClaude:
if supportsAliAnthropicMessages(info.UpstreamModelName) {
adaptor := claude.Adaptor{}
return adaptor.DoResponse(c, resp, info)
}
adaptor := openai.Adaptor{}
return adaptor.DoResponse(c, resp, info)
default:
switch info.RelayMode {
case constant.RelayModeImagesGenerations:
err, usage = aliImageHandler(a, c, resp, info)
case constant.RelayModeImagesEdits:
err, usage = aliImageHandler(a, c, resp, info)
case constant.RelayModeRerank:
err, usage = RerankHandler(c, resp, info)
default:
adaptor := openai.Adaptor{}
usage, err = adaptor.DoResponse(c, resp, info)
}
return usage, err
}
}
func (a *Adaptor) GetModelList() []string {
return ModelList
}
func (a *Adaptor) GetChannelName() string {
return ChannelName
}
+248
View File
@@ -0,0 +1,248 @@
package ali
import (
"encoding/base64"
"encoding/json"
"fmt"
"io"
"net/http"
"net/http/httptest"
"strings"
"sync/atomic"
"testing"
"time"
"github.com/QuantumNous/new-api/common"
rootconstant "github.com/QuantumNous/new-api/constant"
relaycommon "github.com/QuantumNous/new-api/relay/common"
"github.com/QuantumNous/new-api/relay/constant"
relayhelper "github.com/QuantumNous/new-api/relay/helper"
"github.com/QuantumNous/new-api/relaykit/dto"
"github.com/QuantumNous/new-api/service"
"github.com/QuantumNous/new-api/setting/system_setting"
"github.com/gin-gonic/gin"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"github.com/tidwall/gjson"
)
func TestConvertOpenAIRequestFiltersThinkingBudgetByUpstreamModel(t *testing.T) {
tests := []struct {
name string
requestModel string
upstreamModel string
budget string
wantBudget bool
wantValue int64
}{
{
name: "qwen",
requestModel: "qwen-plus",
upstreamModel: "qwen-plus",
budget: "128",
wantBudget: true,
wantValue: 128,
},
{
name: "qwq explicit zero",
requestModel: "qwq-32b",
upstreamModel: "qwq-32b",
budget: "0",
wantBudget: true,
wantValue: 0,
},
{
name: "unsupported upstream overrides qwen request",
requestModel: "qwen-plus",
upstreamModel: "deepseek-r1",
budget: "128",
wantBudget: false,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
request := &dto.GeneralOpenAIRequest{
Model: tt.requestModel,
EnableThinking: json.RawMessage(`true`),
ThinkingBudget: json.RawMessage(tt.budget),
}
info := &relaycommon.RelayInfo{
ChannelMeta: &relaycommon.ChannelMeta{
UpstreamModelName: tt.upstreamModel,
},
}
convertedValue, err := (&Adaptor{}).ConvertOpenAIRequest(nil, info, request)
require.NoError(t, err)
converted, ok := convertedValue.(*dto.GeneralOpenAIRequest)
require.True(t, ok)
if tt.wantBudget {
assert.Equal(t, tt.budget, string(converted.ThinkingBudget))
} else {
assert.Nil(t, converted.ThinkingBudget)
}
encoded, err := common.Marshal(converted)
require.NoError(t, err)
assert.True(t, gjson.GetBytes(encoded, "enable_thinking").Bool())
value := gjson.GetBytes(encoded, "thinking_budget")
assert.Equal(t, tt.wantBudget, value.Exists())
if tt.wantBudget {
assert.Equal(t, tt.wantValue, value.Int())
}
})
}
}
func TestConvertOpenAIRequestPreservesExplicitZeroForMappedQwenModel(t *testing.T) {
const (
clientModel = "customer-model"
upstreamModel = "Qwen/Qwen3-235B-A22B-Thinking-2507"
)
c, _ := gin.CreateTestContext(httptest.NewRecorder())
c.Set("model_mapping", `{"customer-model":"Qwen/Qwen3-235B-A22B-Thinking-2507"}`)
request := &dto.GeneralOpenAIRequest{
Model: clientModel,
EnableThinking: json.RawMessage(`true`),
ThinkingBudget: json.RawMessage(`0`),
}
info := &relaycommon.RelayInfo{
OriginModelName: clientModel,
ChannelMeta: &relaycommon.ChannelMeta{
UpstreamModelName: clientModel,
},
}
err := relayhelper.ModelMappedHelper(c, info, request)
require.NoError(t, err)
assert.True(t, info.IsModelMapped)
assert.Equal(t, upstreamModel, info.UpstreamModelName)
assert.Equal(t, upstreamModel, request.Model)
convertedValue, err := (&Adaptor{}).ConvertOpenAIRequest(c, info, request)
require.NoError(t, err)
converted, ok := convertedValue.(*dto.GeneralOpenAIRequest)
require.True(t, ok)
assert.Equal(t, json.RawMessage(`0`), converted.ThinkingBudget)
encoded, err := common.Marshal(converted)
require.NoError(t, err)
value := gjson.GetBytes(encoded, "thinking_budget")
assert.True(t, value.Exists())
assert.Equal(t, int64(0), value.Int())
}
func TestMappedAliImageModelUsesUpstreamProtocol(t *testing.T) {
c, _ := gin.CreateTestContext(httptest.NewRecorder())
c.Request = httptest.NewRequest(http.MethodPost, "/v1/images/generations", nil)
info := &relaycommon.RelayInfo{
RelayMode: constant.RelayModeImagesGenerations,
OriginModelName: "customer-image-model",
ChannelMeta: &relaycommon.ChannelMeta{
ChannelBaseUrl: "https://dashscope.aliyuncs.com",
UpstreamModelName: "qwen-image-3.0-pro",
},
}
adaptor := &Adaptor{}
url, err := adaptor.GetRequestURL(info)
require.NoError(t, err)
assert.Equal(t, "https://dashscope.aliyuncs.com/api/v1/services/aigc/multimodal-generation/generation", url)
header := http.Header{}
require.NoError(t, adaptor.SetupRequestHeader(c, &header, info))
assert.Empty(t, header.Get("X-DashScope-Async"))
converted, err := adaptor.ConvertImageRequest(c, info, dto.ImageRequest{
Model: info.UpstreamModelName,
Prompt: "poster",
})
require.NoError(t, err)
assert.True(t, adaptor.IsSyncImageModel)
assert.IsType(t, &AliImageRequest{}, converted)
}
func TestAliImageHandlerHonorsRequestResponseFormat(t *testing.T) {
imageBytes := []byte("ali-image")
var downloads atomic.Int32
imageServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
downloads.Add(1)
w.Header().Set("Content-Type", "image/png")
_, _ = w.Write(imageBytes)
}))
t.Cleanup(imageServer.Close)
fetchSetting := system_setting.GetFetchSetting()
require.NotNil(t, fetchSetting)
originalFetchSetting := *fetchSetting
fetchSetting.EnableSSRFProtection = false
t.Cleanup(func() {
*fetchSetting = originalFetchSetting
})
originalMaxFileDownloadMB := rootconstant.MaxFileDownloadMB
rootconstant.MaxFileDownloadMB = 1
t.Cleanup(func() {
rootconstant.MaxFileDownloadMB = originalMaxFileDownloadMB
})
service.InitHttpClient()
tests := []struct {
name string
responseFormat string
wantBase64 string
wantDownloads int32
}{
{
name: "base64",
responseFormat: "b64_json",
wantBase64: base64.StdEncoding.EncodeToString(imageBytes),
wantDownloads: 1,
},
{
name: "url",
responseFormat: "url",
},
{
name: "default",
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
downloads.Store(0)
recorder := httptest.NewRecorder()
c, _ := gin.CreateTestContext(recorder)
info := &relaycommon.RelayInfo{
RelayMode: constant.RelayModeImagesGenerations,
StartTime: time.Unix(1, 0),
Request: &dto.ImageRequest{
ResponseFormat: tt.responseFormat,
},
}
responseBody := fmt.Sprintf(`{"output":{"results":[{"url":%q}]}}`, imageServer.URL)
resp := &http.Response{
StatusCode: http.StatusOK,
Header: http.Header{},
Body: io.NopCloser(strings.NewReader(responseBody)),
}
newAPIError, usage := aliImageHandler(&Adaptor{IsSyncImageModel: true}, c, resp, info)
require.Nil(t, newAPIError)
require.NotNil(t, usage)
var imageResponse dto.ImageResponse
require.NoError(t, common.Unmarshal(recorder.Body.Bytes(), &imageResponse))
require.Len(t, imageResponse.Data, 1)
assert.Equal(t, imageServer.URL, imageResponse.Data[0].Url)
assert.Equal(t, tt.wantBase64, imageResponse.Data[0].B64Json)
assert.Equal(t, tt.wantDownloads, downloads.Load())
})
}
}
+14
View File
@@ -0,0 +1,14 @@
package ali
var ModelList = []string{
"qwen-turbo",
"qwen-plus",
"qwen-max",
"qwen-max-longcontext",
"qwq-32b",
"qwen3-235b-a22b",
"text-embedding-v1",
"gte-rerank-v2",
}
var ChannelName = "ali"
+236
View File
@@ -0,0 +1,236 @@
package ali
import (
"strings"
"github.com/QuantumNous/new-api/logger"
"github.com/QuantumNous/new-api/relaykit/dto"
"github.com/QuantumNous/new-api/service"
"github.com/gin-gonic/gin"
)
type AliMessage struct {
Content any `json:"content"`
Role string `json:"role"`
}
type AliMediaContent struct {
Image string `json:"image,omitempty"`
Text string `json:"text,omitempty"`
}
type AliInput struct {
Prompt string `json:"prompt,omitempty"`
//History []AliMessage `json:"history,omitempty"`
Messages []AliMessage `json:"messages"`
}
type AliParameters struct {
TopP float64 `json:"top_p,omitempty"`
TopK int `json:"top_k,omitempty"`
Seed uint64 `json:"seed,omitempty"`
EnableSearch bool `json:"enable_search,omitempty"`
IncrementalOutput bool `json:"incremental_output,omitempty"`
}
type AliChatRequest struct {
Model string `json:"model"`
Input AliInput `json:"input"`
Parameters AliParameters `json:"parameters"`
}
type AliEmbeddingRequest struct {
Model string `json:"model"`
Input struct {
Texts []string `json:"texts"`
} `json:"input"`
Parameters *struct {
TextType string `json:"text_type,omitempty"`
} `json:"parameters,omitempty"`
}
type AliEmbedding struct {
Embedding []float64 `json:"embedding"`
TextIndex int `json:"text_index"`
}
type AliEmbeddingResponse struct {
Output struct {
Embeddings []AliEmbedding `json:"embeddings"`
} `json:"output"`
Usage AliUsage `json:"usage"`
AliError
}
type AliError struct {
Code string `json:"code"`
Message string `json:"message"`
RequestId string `json:"request_id"`
}
type AliUsage struct {
InputTokens int `json:"input_tokens"`
OutputTokens int `json:"output_tokens"`
TotalTokens int `json:"total_tokens"`
ImageCount int `json:"image_count,omitempty"`
}
type TaskResult struct {
B64Image string `json:"b64_image,omitempty"`
Url string `json:"url,omitempty"`
Code string `json:"code,omitempty"`
Message string `json:"message,omitempty"`
}
type AliOutput struct {
TaskId string `json:"task_id,omitempty"`
TaskStatus string `json:"task_status,omitempty"`
Text string `json:"text"`
FinishReason string `json:"finish_reason"`
Message string `json:"message,omitempty"`
Code string `json:"code,omitempty"`
Results []TaskResult `json:"results,omitempty"`
Choices []struct {
FinishReason string `json:"finish_reason,omitempty"`
Message struct {
Role string `json:"role,omitempty"`
Content []AliMediaContent `json:"content,omitempty"`
ReasoningContent string `json:"reasoning_content,omitempty"`
} `json:"message"`
} `json:"choices,omitempty"`
}
func (o *AliOutput) ChoicesToOpenAIImageDate(c *gin.Context, responseFormat string) []dto.ImageData {
var imageData []dto.ImageData
if len(o.Choices) > 0 {
for _, choice := range o.Choices {
var data dto.ImageData
for _, content := range choice.Message.Content {
if content.Image != "" {
if strings.HasPrefix(content.Image, "http") {
var b64Json string
if responseFormat == "b64_json" {
_, b64, err := service.GetImageFromUrl(content.Image)
if err != nil {
logger.LogError(c, "get_image_data_failed: "+err.Error())
continue
}
b64Json = b64
}
data.Url = content.Image
data.B64Json = b64Json
} else {
data.B64Json = content.Image
}
} else if content.Text != "" {
data.RevisedPrompt = content.Text
}
}
imageData = append(imageData, data)
}
}
return imageData
}
func (o *AliOutput) ResultToOpenAIImageDate(c *gin.Context, responseFormat string) []dto.ImageData {
var imageData []dto.ImageData
for _, data := range o.Results {
var b64Json string
if responseFormat == "b64_json" {
_, b64, err := service.GetImageFromUrl(data.Url)
if err != nil {
logger.LogError(c, "get_image_data_failed: "+err.Error())
continue
}
b64Json = b64
} else {
b64Json = data.B64Image
}
imageData = append(imageData, dto.ImageData{
Url: data.Url,
B64Json: b64Json,
RevisedPrompt: "",
})
}
return imageData
}
type AliResponse struct {
Output AliOutput `json:"output"`
Usage AliUsage `json:"usage"`
AliError
}
type AliImageRequest struct {
Model string `json:"model"`
Input any `json:"input"`
Parameters AliImageParameters `json:"parameters"`
ResponseFormat string `json:"response_format,omitempty"`
}
type AliImageParameters struct {
Size string `json:"size,omitempty"`
N int `json:"n,omitempty"`
Steps string `json:"steps,omitempty"`
Scale string `json:"scale,omitempty"`
Watermark *bool `json:"watermark,omitempty"`
PromptExtend *bool `json:"prompt_extend,omitempty"`
ThinkingMode *bool `json:"thinking_mode,omitempty"`
EnableSequential *bool `json:"enable_sequential,omitempty"`
BboxList any `json:"bbox_list,omitempty"`
ColorPalette any `json:"color_palette,omitempty"`
Seed *int `json:"seed,omitempty"`
}
func (p *AliImageParameters) PromptExtendValue() bool {
if p != nil && p.PromptExtend != nil {
return *p.PromptExtend
}
return false
}
type AliImageInput struct {
Prompt string `json:"prompt,omitempty"`
NegativePrompt string `json:"negative_prompt,omitempty"`
Messages []AliMessage `json:"messages,omitempty"`
}
type WanImageInput struct {
Prompt string `json:"prompt"` // 必需:文本提示词,描述生成图像中期望包含的元素和视觉特点
Images []string `json:"images"` // 必需:图像URL数组,长度不超过2,支持HTTP/HTTPS URL或Base64编码
NegativePrompt string `json:"negative_prompt,omitempty"` // 可选:反向提示词,描述不希望在画面中看到的内容
}
type WanImageParameters struct {
N int `json:"n,omitempty"` // 生成图片数量,取值范围1-4,默认4
Watermark *bool `json:"watermark,omitempty"` // 是否添加水印标识,默认false
Seed int `json:"seed,omitempty"` // 随机数种子,取值范围[0, 2147483647]
Strength float64 `json:"strength,omitempty"` // 修改幅度 0.0-1.0,默认0.5(部分模型支持)
}
type AliRerankParameters struct {
TopN *int `json:"top_n,omitempty"`
ReturnDocuments *bool `json:"return_documents,omitempty"`
}
type AliRerankInput struct {
Query string `json:"query"`
Documents []any `json:"documents"`
}
type AliRerankRequest struct {
Model string `json:"model"`
Input AliRerankInput `json:"input"`
Parameters AliRerankParameters `json:"parameters"`
}
type AliRerankResponse struct {
Output struct {
Results []dto.RerankResponseResult `json:"results"`
} `json:"output"`
Usage AliUsage `json:"usage"`
RequestId string `json:"request_id"`
AliError
}
+351
View File
@@ -0,0 +1,351 @@
package ali
import (
"encoding/base64"
"errors"
"fmt"
"io"
"mime/multipart"
"net/http"
"strings"
"time"
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/logger"
relaycommon "github.com/QuantumNous/new-api/relay/common"
"github.com/QuantumNous/new-api/relaykit/dto"
"github.com/QuantumNous/new-api/relaykit/types"
"github.com/QuantumNous/new-api/service"
"github.com/gin-gonic/gin"
"github.com/samber/lo"
)
func oaiImage2AliImageRequest(info *relaycommon.RelayInfo, request dto.ImageRequest, isSync bool) (*AliImageRequest, error) {
var imageRequest AliImageRequest
imageRequest.Model = request.Model
imageRequest.ResponseFormat = request.ResponseFormat
if request.Extra != nil {
if val, ok := request.Extra["parameters"]; ok {
err := common.Unmarshal(val, &imageRequest.Parameters)
if err != nil {
return nil, fmt.Errorf("invalid parameters field: %w", err)
}
} else {
// 兼容没有parameters字段的情况,从openai标准字段中提取参数
imageRequest.Parameters = AliImageParameters{
Size: strings.Replace(request.Size, "x", "*", -1),
N: int(lo.FromPtrOr(request.N, uint(1))),
Watermark: request.Watermark,
}
}
if val, ok := request.Extra["input"]; ok {
err := common.Unmarshal(val, &imageRequest.Input)
if err != nil {
return nil, fmt.Errorf("invalid input field: %w", err)
}
}
}
if strings.Contains(request.Model, "z-image") {
// z-image 开启prompt_extend后,按2倍计费
if imageRequest.Parameters.PromptExtendValue() {
info.PriceData.AddOtherRatio("prompt_extend", 2)
}
}
// Parameters may come from Extra["parameters"], bypassing the standard
// top-level n validation; enforce the same bound before it becomes a
// billing multiplier.
if imageRequest.Parameters.N < 0 || imageRequest.Parameters.N > dto.MaxImageN {
return nil, fmt.Errorf("parameters.n must be an integer between 1 and %d", dto.MaxImageN)
}
if imageRequest.Parameters.N != 0 {
info.PriceData.AddOtherRatio("n", float64(imageRequest.Parameters.N))
}
// 同步图片模型和异步图片模型请求格式不一样
if isSync {
if imageRequest.Input == nil {
imageRequest.Input = AliImageInput{
Messages: []AliMessage{
{
Role: "user",
Content: []AliMediaContent{
{
Text: request.Prompt,
},
},
},
},
}
}
} else {
if imageRequest.Input == nil {
imageRequest.Input = AliImageInput{
Prompt: request.Prompt,
}
}
}
return &imageRequest, nil
}
func getImageBase64sFromForm(c *gin.Context, fieldName string) ([]string, error) {
mf := c.Request.MultipartForm
if mf == nil {
if _, err := c.MultipartForm(); err != nil {
return nil, fmt.Errorf("failed to parse image edit form request: %w", err)
}
mf = c.Request.MultipartForm
}
var imageFiles []*multipart.FileHeader
var exists bool
// First check for standard "image" field
if imageFiles, exists = mf.File["image"]; !exists || len(imageFiles) == 0 {
// If not found, check for "image[]" field
if imageFiles, exists = mf.File["image[]"]; !exists || len(imageFiles) == 0 {
// If still not found, iterate through all fields to find any that start with "image["
foundArrayImages := false
for fieldName, files := range mf.File {
if strings.HasPrefix(fieldName, "image[") && len(files) > 0 {
foundArrayImages = true
imageFiles = append(imageFiles, files...)
}
}
// If no image fields found at all
if !foundArrayImages && (len(imageFiles) == 0) {
return nil, errors.New("image is required")
}
}
}
if len(imageFiles) == 0 {
return nil, errors.New("image is required")
}
//if len(imageFiles) > 1 {
// return nil, errors.New("only one image is supported for qwen edit")
//}
// 获取base64编码的图片
var imageBase64s []string
for _, file := range imageFiles {
image, err := file.Open()
if err != nil {
return nil, errors.New("failed to open image file")
}
// 读取文件内容
imageData, err := io.ReadAll(image)
if err != nil {
return nil, errors.New("failed to read image file")
}
// 获取MIME类型
mimeType := http.DetectContentType(imageData)
// 编码为base64
base64Data := base64.StdEncoding.EncodeToString(imageData)
// 构造data URL格式
dataURL := fmt.Sprintf("data:%s;base64,%s", mimeType, base64Data)
imageBase64s = append(imageBase64s, dataURL)
image.Close()
}
return imageBase64s, nil
}
func oaiFormEdit2AliImageEdit(c *gin.Context, info *relaycommon.RelayInfo, request dto.ImageRequest) (*AliImageRequest, error) {
var imageRequest AliImageRequest
imageRequest.Model = request.Model
imageRequest.ResponseFormat = request.ResponseFormat
imageBase64s, err := getImageBase64sFromForm(c, "image")
if err != nil {
return nil, fmt.Errorf("get image base64s from form failed: %w", err)
}
//dto.MediaContent{}
mediaContents := make([]AliMediaContent, len(imageBase64s))
for i, b64 := range imageBase64s {
mediaContents[i] = AliMediaContent{
Image: b64,
}
}
mediaContents = append(mediaContents, AliMediaContent{
Text: request.Prompt,
})
imageRequest.Input = AliImageInput{
Messages: []AliMessage{
{
Role: "user",
Content: mediaContents,
},
},
}
imageRequest.Parameters = AliImageParameters{
N: int(lo.FromPtrOr(request.N, uint(1))),
Watermark: request.Watermark,
}
return &imageRequest, nil
}
func updateTask(info *relaycommon.RelayInfo, taskID string) (*AliResponse, error, []byte) {
url := fmt.Sprintf("%s/api/v1/tasks/%s", info.ChannelBaseUrl, taskID)
var aliResponse AliResponse
req, err := http.NewRequest("GET", url, nil)
if err != nil {
return &aliResponse, err, nil
}
req.Header.Set("Authorization", "Bearer "+info.ApiKey)
client := &http.Client{}
resp, err := client.Do(req)
if err != nil {
common.SysLog("updateTask client.Do err: " + err.Error())
return &aliResponse, err, nil
}
defer resp.Body.Close()
responseBody, err := io.ReadAll(resp.Body)
var response AliResponse
err = common.Unmarshal(responseBody, &response)
if err != nil {
common.SysLog("updateTask NewDecoder err: " + err.Error())
return &aliResponse, err, nil
}
return &response, nil, responseBody
}
func asyncTaskWait(c *gin.Context, info *relaycommon.RelayInfo, taskID string) (*AliResponse, []byte, error) {
waitSeconds := 10
step := 0
maxStep := 20
var taskResponse AliResponse
var responseBody []byte
time.Sleep(time.Duration(5) * time.Second)
for {
logger.LogDebug(c, "asyncTaskWait step %d/%d, wait %d seconds", step, maxStep, waitSeconds)
step++
rsp, err, body := updateTask(info, taskID)
responseBody = body
if err != nil {
logger.LogWarn(c, "asyncTaskWait UpdateTask err: "+err.Error())
time.Sleep(time.Duration(waitSeconds) * time.Second)
continue
}
if rsp.Output.TaskStatus == "" {
return &taskResponse, responseBody, nil
}
switch rsp.Output.TaskStatus {
case "FAILED":
fallthrough
case "CANCELED":
fallthrough
case "SUCCEEDED":
fallthrough
case "UNKNOWN":
return rsp, responseBody, nil
}
if step >= maxStep {
break
}
time.Sleep(time.Duration(waitSeconds) * time.Second)
}
return nil, nil, fmt.Errorf("aliAsyncTaskWait timeout")
}
func responseAli2OpenAIImage(c *gin.Context, response *AliResponse, originBody []byte, info *relaycommon.RelayInfo, responseFormat string) *dto.ImageResponse {
imageResponse := dto.ImageResponse{
Created: info.StartTime.Unix(),
}
if len(response.Output.Results) > 0 {
imageResponse.Data = response.Output.ResultToOpenAIImageDate(c, responseFormat)
} else if len(response.Output.Choices) > 0 {
imageResponse.Data = response.Output.ChoicesToOpenAIImageDate(c, responseFormat)
}
imageResponse.Metadata = originBody
return &imageResponse
}
func aliImageHandler(a *Adaptor, c *gin.Context, resp *http.Response, info *relaycommon.RelayInfo) (*types.NewAPIError, *dto.Usage) {
responseFormat := ""
if imageReq, ok := info.Request.(*dto.ImageRequest); ok {
responseFormat = imageReq.ResponseFormat
}
var aliTaskResponse AliResponse
responseBody, err := io.ReadAll(resp.Body)
if err != nil {
return types.NewOpenAIError(err, types.ErrorCodeReadResponseBodyFailed, http.StatusInternalServerError), nil
}
service.CloseResponseBodyGracefully(resp)
err = common.Unmarshal(responseBody, &aliTaskResponse)
if err != nil {
return types.NewOpenAIError(err, types.ErrorCodeBadResponseBody, http.StatusInternalServerError), nil
}
if aliTaskResponse.Message != "" {
logger.LogError(c, "ali_async_task_failed: "+aliTaskResponse.Message)
return types.NewError(errors.New(aliTaskResponse.Message), types.ErrorCodeBadResponse), nil
}
var (
aliResponse *AliResponse
originRespBody []byte
)
if a.IsSyncImageModel {
aliResponse = &aliTaskResponse
originRespBody = responseBody
} else {
// 异步图片模型需要轮询任务结果
aliResponse, originRespBody, err = asyncTaskWait(c, info, aliTaskResponse.Output.TaskId)
if err != nil {
return types.NewError(err, types.ErrorCodeBadResponse), nil
}
if aliResponse.Output.TaskStatus != "SUCCEEDED" {
return types.WithOpenAIError(types.OpenAIError{
Message: aliResponse.Output.Message,
Type: "ali_error",
Param: "",
Code: aliResponse.Output.Code,
}, resp.StatusCode), nil
}
}
if a.IsSyncImageModel {
logger.LogDebug(c, "ali_sync_image_result: %s", originRespBody)
} else {
logger.LogDebug(c, "ali_async_image_result: %s", originRespBody)
}
imageResponses := responseAli2OpenAIImage(c, aliResponse, originRespBody, info, responseFormat)
if aliResponse.Usage.ImageCount != 0 {
info.PriceData.AddOtherRatio("n", float64(aliResponse.Usage.ImageCount))
} else if len(imageResponses.Data) != 0 {
info.PriceData.AddOtherRatio("n", float64(len(imageResponses.Data)))
}
jsonResponse, err := common.Marshal(imageResponses)
if err != nil {
return types.NewError(err, types.ErrorCodeBadResponseBody), nil
}
service.IOCopyBytesGracefully(c, resp, jsonResponse)
return nil, &dto.Usage{}
}
+49
View File
@@ -0,0 +1,49 @@
package ali
import (
"fmt"
"strings"
"github.com/QuantumNous/new-api/common"
relaycommon "github.com/QuantumNous/new-api/relay/common"
"github.com/QuantumNous/new-api/relaykit/dto"
"github.com/gin-gonic/gin"
"github.com/samber/lo"
)
func oaiFormEdit2WanxImageEdit(c *gin.Context, info *relaycommon.RelayInfo, request dto.ImageRequest) (*AliImageRequest, error) {
var err error
var imageRequest AliImageRequest
imageRequest.Model = request.Model
imageRequest.ResponseFormat = request.ResponseFormat
wanInput := WanImageInput{
Prompt: request.Prompt,
}
if err := common.UnmarshalBodyReusable(c, &wanInput); err != nil {
return nil, err
}
if wanInput.Images, err = getImageBase64sFromForm(c, "image"); err != nil {
return nil, fmt.Errorf("get image base64s from form failed: %w", err)
}
//wanParams := WanImageParameters{
// N: int(request.N),
//}
imageRequest.Input = wanInput
imageRequest.Parameters = AliImageParameters{
N: int(lo.FromPtrOr(request.N, uint(1))),
}
info.PriceData.AddOtherRatio("n", float64(imageRequest.Parameters.N))
return &imageRequest, nil
}
func isOldWanModel(modelName string) bool {
return strings.Contains(modelName, "wan") &&
!lo.SomeBy([]string{"wan2.6", "wan2.7"}, func(v string) bool { return strings.Contains(modelName, v) })
}
func isWanModel(modelName string) bool {
return strings.Contains(modelName, "wan")
}
+75
View File
@@ -0,0 +1,75 @@
package ali
import (
"encoding/json"
"io"
"net/http"
relaycommon "github.com/QuantumNous/new-api/relay/common"
"github.com/QuantumNous/new-api/relaykit/dto"
"github.com/QuantumNous/new-api/relaykit/types"
"github.com/QuantumNous/new-api/service"
"github.com/gin-gonic/gin"
)
func ConvertRerankRequest(request dto.RerankRequest) *AliRerankRequest {
returnDocuments := request.ReturnDocuments
if returnDocuments == nil {
t := true
returnDocuments = &t
}
return &AliRerankRequest{
Model: request.Model,
Input: AliRerankInput{
Query: request.Query,
Documents: request.Documents,
},
Parameters: AliRerankParameters{
TopN: request.TopN,
ReturnDocuments: returnDocuments,
},
}
}
func RerankHandler(c *gin.Context, resp *http.Response, info *relaycommon.RelayInfo) (*types.NewAPIError, *dto.Usage) {
responseBody, err := io.ReadAll(resp.Body)
if err != nil {
return types.NewOpenAIError(err, types.ErrorCodeReadResponseBodyFailed, http.StatusInternalServerError), nil
}
service.CloseResponseBodyGracefully(resp)
var aliResponse AliRerankResponse
err = json.Unmarshal(responseBody, &aliResponse)
if err != nil {
return types.NewOpenAIError(err, types.ErrorCodeBadResponseBody, http.StatusInternalServerError), nil
}
if aliResponse.Code != "" {
return types.WithOpenAIError(types.OpenAIError{
Message: aliResponse.Message,
Type: aliResponse.Code,
Param: aliResponse.RequestId,
Code: aliResponse.Code,
}, resp.StatusCode), nil
}
usage := dto.Usage{
PromptTokens: aliResponse.Usage.TotalTokens,
CompletionTokens: 0,
TotalTokens: aliResponse.Usage.TotalTokens,
}
rerankResponse := dto.RerankResponse{
Results: aliResponse.Output.Results,
Usage: usage,
}
jsonResponse, err := json.Marshal(rerankResponse)
if err != nil {
return types.NewError(err, types.ErrorCodeBadResponseBody), nil
}
c.Writer.Header().Set("Content-Type", "application/json")
c.Writer.WriteHeader(resp.StatusCode)
c.Writer.Write(jsonResponse)
return nil, &usage
}
+36
View File
@@ -0,0 +1,36 @@
package ali
import (
"github.com/QuantumNous/new-api/relaykit/dto"
"github.com/samber/lo"
)
// https://help.aliyun.com/document_detail/613695.html?spm=a2c4g.2399480.0.0.1adb778fAdzP9w#341800c0f8w0r
const EnableSearchModelSuffix = "-internet"
func requestOpenAI2Ali(request dto.GeneralOpenAIRequest, upstreamModelName string) *dto.GeneralOpenAIRequest {
modelName := upstreamModelName
if modelName == "" {
modelName = request.Model
}
if !dto.IsQwenThinkingBudgetModel(modelName) {
request.ThinkingBudget = nil
}
// DashScope rejects top_p at the 0 and 1 boundaries, so an explicit value is
// clamped into the open interval. The clamp stays at two decimals because
// some models on the platform reject a third decimal with
// "top_p参数非法:限制小数点[2]位".
//
// A request that omits top_p is left untouched: injecting a value would
// silently replace the model's own default with near-greedy decoding.
if request.TopP != nil {
if *request.TopP >= 1 {
request.TopP = lo.ToPtr(0.99)
} else if *request.TopP <= 0 {
request.TopP = lo.ToPtr(0.01)
}
}
return &request
}
+59
View File
@@ -0,0 +1,59 @@
package ali
import (
"testing"
"github.com/QuantumNous/new-api/relaykit/dto"
"github.com/samber/lo"
"github.com/stretchr/testify/assert"
)
func TestRequestOpenAI2AliTopP(t *testing.T) {
tests := []struct {
name string
topP *float64
want *float64
}{
{
name: "omitted top_p is not injected",
topP: nil,
want: nil,
},
{
name: "in-range top_p is preserved",
topP: lo.ToPtr(0.8),
want: lo.ToPtr(0.8),
},
{
name: "top_p of 1 is clamped to two decimals",
topP: lo.ToPtr(1.0),
want: lo.ToPtr(0.99),
},
{
name: "top_p above 1 is clamped to two decimals",
topP: lo.ToPtr(1.5),
want: lo.ToPtr(0.99),
},
{
name: "top_p of 0 is clamped to two decimals",
topP: lo.ToPtr(0.0),
want: lo.ToPtr(0.01),
},
{
name: "negative top_p is clamped to two decimals",
topP: lo.ToPtr(-0.3),
want: lo.ToPtr(0.01),
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
got := requestOpenAI2Ali(dto.GeneralOpenAIRequest{
Model: "qwen-plus",
TopP: tt.topP,
}, "qwen-plus")
assert.Equal(t, tt.want, got.TopP)
})
}
}
+596
View File
@@ -0,0 +1,596 @@
package channel
import (
"context"
"errors"
"fmt"
"io"
"net/http"
"regexp"
"strings"
"sync"
"time"
common2 "github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/logger"
"github.com/QuantumNous/new-api/relay/common"
"github.com/QuantumNous/new-api/relay/constant"
"github.com/QuantumNous/new-api/relay/helper"
"github.com/QuantumNous/new-api/relaykit/types"
"github.com/QuantumNous/new-api/service"
"github.com/QuantumNous/new-api/setting/operation_setting"
"github.com/bytedance/gopkg/util/gopool"
"github.com/gin-gonic/gin"
"github.com/gorilla/websocket"
)
// ApplyUpstreamBodyMetadata restores metadata that net/http cannot infer from
// a ReplayableBody. Callers must pass the original body because NewRequest
// hides its dynamic type behind req.Body's io.ReadCloser wrapper.
func ApplyUpstreamBodyMetadata(req *http.Request, body io.Reader) {
replayable, ok := body.(common2.ReplayableBody)
if !ok {
return
}
// BodyStorage structurally satisfies ReplayableBody, but it also exposes
// io.Closer. If a caller passes the storage directly instead of using
// NewReplayableBodyReader, hide Close before the transport takes ownership
// of req.Body so the shared replay source remains available to GetBody.
if _, rawStorage := body.(common2.BodyStorage); rawStorage {
req.Body = io.NopCloser(body)
}
req.ContentLength = replayable.Size()
if req.GetBody == nil {
req.GetBody = replayable.NewReader
}
}
func SetupApiRequestHeader(info *common.RelayInfo, c *gin.Context, req *http.Header) {
if info.RelayMode == constant.RelayModeAudioTranscription || info.RelayMode == constant.RelayModeAudioTranslation {
// multipart/form-data
} else if info.RelayMode == constant.RelayModeRealtime {
// websocket
} else {
req.Set("Content-Type", c.Request.Header.Get("Content-Type"))
req.Set("Accept", c.Request.Header.Get("Accept"))
if info.IsStream && c.Request.Header.Get("Accept") == "" {
req.Set("Accept", "text/event-stream")
}
}
}
const clientHeaderPlaceholderPrefix = "{client_header:"
const (
headerPassthroughAllKey = "*"
headerPassthroughRegexPrefix = "re:"
headerPassthroughRegexPrefixV2 = "regex:"
)
var passthroughSkipHeaderNamesLower = map[string]struct{}{
// RFC 7230 hop-by-hop headers.
"connection": {},
"keep-alive": {},
"proxy-authenticate": {},
"proxy-authorization": {},
"te": {},
"trailer": {},
"transfer-encoding": {},
"upgrade": {},
"cookie": {},
// Additional headers that should not be forwarded by name-matching passthrough rules.
"host": {},
"content-length": {},
"accept-encoding": {},
// Do not passthrough credentials by wildcard/regex.
"authorization": {},
"x-api-key": {},
"x-goog-api-key": {},
// WebSocket handshake headers are generated by the client/dialer.
"sec-websocket-key": {},
"sec-websocket-version": {},
"sec-websocket-extensions": {},
}
var headerPassthroughRegexCache sync.Map // map[string]*regexp.Regexp
func getHeaderPassthroughRegex(pattern string) (*regexp.Regexp, error) {
pattern = strings.TrimSpace(pattern)
if pattern == "" {
return nil, errors.New("empty regex pattern")
}
if v, ok := headerPassthroughRegexCache.Load(pattern); ok {
if re, ok := v.(*regexp.Regexp); ok {
return re, nil
}
headerPassthroughRegexCache.Delete(pattern)
}
compiled, err := regexp.Compile(pattern)
if err != nil {
return nil, err
}
actual, _ := headerPassthroughRegexCache.LoadOrStore(pattern, compiled)
if re, ok := actual.(*regexp.Regexp); ok {
return re, nil
}
return compiled, nil
}
func IsHeaderPassthroughRuleKey(key string) bool {
return isHeaderPassthroughRuleKey(key)
}
func isHeaderPassthroughRuleKey(key string) bool {
key = strings.TrimSpace(key)
if key == "" {
return false
}
if key == headerPassthroughAllKey {
return true
}
lower := strings.ToLower(key)
return strings.HasPrefix(lower, headerPassthroughRegexPrefix) || strings.HasPrefix(lower, headerPassthroughRegexPrefixV2)
}
func shouldSkipPassthroughHeader(name string) bool {
name = strings.TrimSpace(name)
if name == "" {
return true
}
lower := strings.ToLower(name)
if _, ok := passthroughSkipHeaderNamesLower[lower]; ok {
return true
}
return false
}
func applyHeaderOverridePlaceholders(template string, c *gin.Context, apiKey string) (string, bool, error) {
trimmed := strings.TrimSpace(template)
if strings.HasPrefix(trimmed, clientHeaderPlaceholderPrefix) {
afterPrefix := trimmed[len(clientHeaderPlaceholderPrefix):]
end := strings.Index(afterPrefix, "}")
if end < 0 || end != len(afterPrefix)-1 {
return "", false, fmt.Errorf("client_header placeholder must be the full value: %q", template)
}
name := strings.TrimSpace(afterPrefix[:end])
if name == "" {
return "", false, fmt.Errorf("client_header placeholder name is empty: %q", template)
}
if c == nil || c.Request == nil {
return "", false, fmt.Errorf("missing request context for client_header placeholder")
}
clientHeaderValue := c.Request.Header.Get(name)
if strings.TrimSpace(clientHeaderValue) == "" {
return "", false, nil
}
// Do not interpolate {api_key} inside client-supplied content.
return clientHeaderValue, true, nil
}
if strings.Contains(template, "{api_key}") {
template = strings.ReplaceAll(template, "{api_key}", apiKey)
}
if strings.TrimSpace(template) == "" {
return "", false, nil
}
return template, true, nil
}
// processHeaderOverride applies channel header overrides, with placeholder substitution.
// Supported placeholders:
// - {api_key}: resolved to the channel API key
// - {client_header:<name>}: resolved to the incoming request header value
//
// Header passthrough rules (keys only; values are ignored):
// - "*": passthrough all incoming headers by name (excluding unsafe headers)
// - "re:<regex>" / "regex:<regex>": passthrough headers whose names match the regex (Go regexp)
//
// Passthrough rules are applied first, then normal overrides are applied, so explicit overrides win.
func processHeaderOverride(info *common.RelayInfo, c *gin.Context) (map[string]string, error) {
headerOverride := make(map[string]string)
if info == nil {
return headerOverride, nil
}
headerOverrideSource := common.GetEffectiveHeaderOverride(info)
passAll := false
var passthroughRegex []*regexp.Regexp
if !info.IsChannelTest {
for k := range headerOverrideSource {
key := strings.TrimSpace(strings.ToLower(k))
if key == "" {
continue
}
if key == headerPassthroughAllKey {
passAll = true
continue
}
var pattern string
switch {
case strings.HasPrefix(key, headerPassthroughRegexPrefix):
pattern = strings.TrimSpace(key[len(headerPassthroughRegexPrefix):])
case strings.HasPrefix(key, headerPassthroughRegexPrefixV2):
pattern = strings.TrimSpace(key[len(headerPassthroughRegexPrefixV2):])
default:
continue
}
if pattern == "" {
return nil, types.NewError(fmt.Errorf("header passthrough regex pattern is empty: %q", k), types.ErrorCodeChannelHeaderOverrideInvalid)
}
compiled, err := getHeaderPassthroughRegex(pattern)
if err != nil {
return nil, types.NewError(err, types.ErrorCodeChannelHeaderOverrideInvalid)
}
passthroughRegex = append(passthroughRegex, compiled)
}
}
if passAll || len(passthroughRegex) > 0 {
if c == nil || c.Request == nil {
return nil, types.NewError(fmt.Errorf("missing request context for header passthrough"), types.ErrorCodeChannelHeaderOverrideInvalid)
}
for name := range c.Request.Header {
if shouldSkipPassthroughHeader(name) {
continue
}
if !passAll {
matched := false
for _, re := range passthroughRegex {
if re.MatchString(name) {
matched = true
break
}
}
if !matched {
continue
}
}
value := strings.TrimSpace(c.Request.Header.Get(name))
if value == "" {
continue
}
headerOverride[strings.ToLower(strings.TrimSpace(name))] = value
}
}
for k, v := range headerOverrideSource {
if isHeaderPassthroughRuleKey(k) {
continue
}
key := strings.TrimSpace(strings.ToLower(k))
if key == "" {
continue
}
str, ok := v.(string)
if !ok {
return nil, types.NewError(nil, types.ErrorCodeChannelHeaderOverrideInvalid)
}
if info.IsChannelTest && strings.HasPrefix(strings.TrimSpace(str), clientHeaderPlaceholderPrefix) {
continue
}
value, include, err := applyHeaderOverridePlaceholders(str, c, info.ApiKey)
if err != nil {
return nil, types.NewError(err, types.ErrorCodeChannelHeaderOverrideInvalid)
}
if !include {
continue
}
headerOverride[key] = value
}
return headerOverride, nil
}
func ResolveHeaderOverride(info *common.RelayInfo, c *gin.Context) (map[string]string, error) {
return processHeaderOverride(info, c)
}
func applyHeaderOverrideToRequest(req *http.Request, headerOverride map[string]string) {
if req == nil {
return
}
for key, value := range headerOverride {
req.Header.Set(key, value)
// set Host in req
if strings.EqualFold(key, "Host") {
req.Host = value
}
}
}
func DoApiRequest(a Adaptor, c *gin.Context, info *common.RelayInfo, requestBody io.Reader) (*http.Response, error) {
fullRequestURL, err := a.GetRequestURL(info)
if err != nil {
return nil, fmt.Errorf("get request url failed: %w", err)
}
logger.LogDebug(c, "fullRequestURL: %s", common.SanitizeURLForLog(fullRequestURL))
req, err := http.NewRequest(c.Request.Method, fullRequestURL, requestBody)
if err != nil {
return nil, fmt.Errorf("new request failed: %w", err)
}
ApplyUpstreamBodyMetadata(req, requestBody)
headers := req.Header
err = a.SetupRequestHeader(c, &headers, info)
if err != nil {
return nil, fmt.Errorf("setup request header failed: %w", err)
}
// 在 SetupRequestHeader 之后应用 Header Override,确保用户设置优先级最高
// 这样可以覆盖默认的 Authorization header 设置
headerOverride, err := processHeaderOverride(info, c)
if err != nil {
return nil, err
}
applyHeaderOverrideToRequest(req, headerOverride)
resp, err := doRequest(c, req, info)
if err != nil {
return nil, fmt.Errorf("do request failed: %w", err)
}
return resp, nil
}
func DoFormRequest(a Adaptor, c *gin.Context, info *common.RelayInfo, requestBody io.Reader) (*http.Response, error) {
fullRequestURL, err := a.GetRequestURL(info)
if err != nil {
return nil, fmt.Errorf("get request url failed: %w", err)
}
logger.LogDebug(c, "fullRequestURL: %s", common.SanitizeURLForLog(fullRequestURL))
req, err := http.NewRequest(c.Request.Method, fullRequestURL, requestBody)
if err != nil {
return nil, fmt.Errorf("new request failed: %w", err)
}
ApplyUpstreamBodyMetadata(req, requestBody)
// set form data
req.Header.Set("Content-Type", c.Request.Header.Get("Content-Type"))
headers := req.Header
err = a.SetupRequestHeader(c, &headers, info)
if err != nil {
return nil, fmt.Errorf("setup request header failed: %w", err)
}
// 在 SetupRequestHeader 之后应用 Header Override,确保用户设置优先级最高
// 这样可以覆盖默认的 Authorization header 设置
headerOverride, err := processHeaderOverride(info, c)
if err != nil {
return nil, err
}
applyHeaderOverrideToRequest(req, headerOverride)
resp, err := doRequest(c, req, info)
if err != nil {
return nil, fmt.Errorf("do request failed: %w", err)
}
return resp, nil
}
func DoWssRequest(a Adaptor, c *gin.Context, info *common.RelayInfo, requestBody io.Reader) (*websocket.Conn, error) {
fullRequestURL, err := a.GetRequestURL(info)
if err != nil {
return nil, fmt.Errorf("get request url failed: %w", err)
}
targetHeader := http.Header{}
err = a.SetupRequestHeader(c, &targetHeader, info)
if err != nil {
return nil, fmt.Errorf("setup request header failed: %w", err)
}
// 在 SetupRequestHeader 之后应用 Header Override,确保用户设置优先级最高
// 这样可以覆盖默认的 Authorization header 设置
headerOverride, err := processHeaderOverride(info, c)
if err != nil {
return nil, err
}
for key, value := range headerOverride {
targetHeader.Set(key, value)
}
targetHeader.Set("Content-Type", c.Request.Header.Get("Content-Type"))
targetConn, _, err := websocket.DefaultDialer.Dial(fullRequestURL, targetHeader)
if err != nil {
return nil, fmt.Errorf("dial failed to %s: %w", common.SanitizeURLForLog(fullRequestURL), err)
}
// send request body
//all, err := io.ReadAll(requestBody)
//err = service.WssString(c, targetConn, string(all))
return targetConn, nil
}
func startPingKeepAlive(c *gin.Context, pingInterval time.Duration) (context.CancelFunc, <-chan struct{}) {
pingerCtx, stopPinger := context.WithCancel(context.Background())
done := make(chan struct{})
gopool.Go(func() {
defer close(done)
defer func() {
// 增加panic恢复处理
if r := recover(); r != nil {
logger.LogDebug(c, "SSE ping goroutine panic recovered: %v", r)
}
logger.LogDebug(c, "SSE ping goroutine stopped")
}()
if pingInterval <= 0 {
pingInterval = helper.DefaultPingInterval
}
ticker := time.NewTicker(pingInterval)
// 确保在任何情况下都清理ticker
defer func() {
ticker.Stop()
logger.LogDebug(c, "SSE ping ticker stopped")
}()
var pingMutex sync.Mutex
logger.LogDebug(c, "SSE ping goroutine started")
// 增加超时控制,防止goroutine长时间运行
maxPingDuration := 120 * time.Minute // 最大ping持续时间
pingTimeout := time.NewTimer(maxPingDuration)
defer pingTimeout.Stop()
for {
select {
// 发送 ping 数据
case <-ticker.C:
if err := sendPingData(c, &pingMutex); err != nil {
logger.LogDebug(c, "SSE ping error, stopping goroutine: %s", err.Error())
return
}
// 收到退出信号
case <-pingerCtx.Done():
return
// request 结束
case <-c.Request.Context().Done():
return
// 超时保护,防止goroutine无限运行
case <-pingTimeout.C:
logger.LogDebug(c, "SSE ping goroutine timeout, stopping")
return
}
}
})
return stopPinger, done
}
func sendPingData(c *gin.Context, mutex *sync.Mutex) error {
mutex.Lock()
defer mutex.Unlock()
// Bound the write so a slow client cannot block this goroutine forever;
// doRequest's defer waits for the pinger to exit before returning.
helper.ExtendWriteDeadline(c)
err := helper.PingData(c)
if err != nil {
logger.LogError(c, "SSE ping error: "+err.Error())
return err
}
logger.LogDebug(c, "SSE ping data sent")
return nil
}
func DoRequest(c *gin.Context, req *http.Request, info *common.RelayInfo) (*http.Response, error) {
return doRequest(c, req, info)
}
// keepUpstreamRedirectResponse stops net/http from following redirects while
// returning the upstream 3xx response to the relay without an extra error.
func keepUpstreamRedirectResponse(_ *http.Request, _ []*http.Request) error {
return http.ErrUseLastResponse
}
func doRequest(c *gin.Context, req *http.Request, info *common.RelayInfo) (*http.Response, error) {
client, err := service.GetHttpClientWithProxySettings(info.ChannelSetting.Proxy, info.ChannelSetting)
if err != nil {
return nil, fmt.Errorf("new proxy http client failed: %w", err)
}
// Clients are cached and shared across channels, so override redirect
// behavior on a shallow copy instead of mutating the cached client. This
// still reuses its transport and connection pools, including HTTP/2's
// transparent stream retries.
relayClient := *client
relayClient.CheckRedirect = keepUpstreamRedirectResponse
if common2.DebugEnabled && req != nil && req.URL != nil {
policy := service.NormalizeHTTPTransportPolicy(info.ChannelSetting)
logger.LogDebug(c, fmt.Sprintf(
"http transport select: host=%s protocol=%s shards=%d policy=%s",
req.URL.Host,
policy.Protocol,
policy.Shards,
policy.String(),
))
}
var stopPinger context.CancelFunc
var pingerDone <-chan struct{}
if info.IsStream {
helper.SetEventStreamHeaders(c)
// 处理流式请求的 ping 保活
generalSettings := operation_setting.GetGeneralSetting()
if generalSettings.PingIntervalEnabled && !info.DisablePing {
pingInterval := time.Duration(generalSettings.PingIntervalSeconds) * time.Second
stopPinger, pingerDone = startPingKeepAlive(c, pingInterval)
// 使用defer确保在任何情况下都能停止ping goroutine
defer func() {
if stopPinger != nil {
stopPinger()
<-pingerDone
logger.LogDebug(c, "SSE ping goroutine stopped by defer")
}
}()
}
}
resp, err := relayClient.Do(req)
if err != nil {
logger.LogError(c, "do request failed: "+err.Error())
return nil, types.NewError(err, types.ErrorCodeDoRequestFailed, types.ErrOptionWithHideErrMsg("upstream error: do request failed"))
}
if resp == nil {
return nil, errors.New("resp is nil")
}
if common2.DebugEnabled {
policy := service.NormalizeHTTPTransportPolicy(info.ChannelSetting)
logger.LogDebug(c, fmt.Sprintf(
"http transport negotiated: host=%s protocol=%s shards=%d policy=%s negotiated=%s",
req.URL.Host,
policy.Protocol,
policy.Shards,
policy.String(),
resp.Proto,
))
}
if upID := resp.Header.Get(common2.RequestIdKey); upID != "" {
c.Set(common2.UpstreamRequestIdKey, upID)
}
_ = req.Body.Close()
_ = c.Request.Body.Close()
return resp, nil
}
func DoTaskApiRequest(a TaskAdaptor, c *gin.Context, info *common.RelayInfo, requestBody io.Reader) (*http.Response, error) {
fullRequestURL, err := a.BuildRequestURL(info)
if err != nil {
return nil, err
}
req, err := newTaskAPIRequest(c, fullRequestURL, requestBody)
if err != nil {
return nil, fmt.Errorf("new request failed: %w", err)
}
ApplyUpstreamBodyMetadata(req, requestBody)
// Do NOT wrap requestBody in a GetBody closure here: returning the same
// (already consumed) reader would make any transport-level retry silently
// replay an empty body. http.NewRequest already derives a correct,
// snapshot-based GetBody for *bytes.Reader/Buffer/strings.Reader bodies
// (which most task adaptors pass in); ApplyUpstreamBodyMetadata wires the
// same contract for bodies that explicitly implement ReplayableBody.
// Otherwise GetBody stays nil so the transport fails the retry instead of
// sending a corrupted request.
err = a.BuildRequestHeader(c, req, info)
if err != nil {
return nil, fmt.Errorf("setup request header failed: %w", err)
}
resp, err := doRequest(c, req, info)
if err != nil {
return nil, fmt.Errorf("do request failed: %w", err)
}
return resp, nil
}
func newTaskAPIRequest(c *gin.Context, fullRequestURL string, requestBody io.Reader) (*http.Request, error) {
if c == nil || c.Request == nil {
return nil, errors.New("task client request is missing")
}
return http.NewRequestWithContext(c.Request.Context(), c.Request.Method, fullRequestURL, requestBody)
}
+593
View File
@@ -0,0 +1,593 @@
package channel
import (
"bytes"
"context"
"crypto/tls"
"fmt"
"io"
"net"
"net/http"
"net/http/httptest"
"strings"
"testing"
"time"
"github.com/QuantumNous/new-api/common"
relaycommon "github.com/QuantumNous/new-api/relay/common"
"github.com/QuantumNous/new-api/service"
"github.com/gin-gonic/gin"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"golang.org/x/net/http2"
"golang.org/x/net/http2/hpack"
)
func TestApplyUpstreamBodyMetadataSetsReplayableMetadata(t *testing.T) {
t.Parallel()
payload := []byte(`{"model":"test-model","messages":[{"role":"user","content":"hi"}]}`)
body, closer, err := relaycommon.NewOutboundJSONBody(payload)
require.NoError(t, err)
defer closer.Close()
// NewRequest hides the body's dynamic type behind req.Body, so metadata
// extraction must use the original body passed to NewRequest.
req, err := http.NewRequest(http.MethodPost, "https://example.com/v1/chat/completions", body)
require.NoError(t, err)
assert.Nil(t, req.GetBody)
assert.Zero(t, req.ContentLength)
_, requestBodyIsReplayable := req.Body.(common.ReplayableBody)
assert.False(t, requestBodyIsReplayable)
ApplyUpstreamBodyMetadata(req, req.Body)
assert.Nil(t, req.GetBody)
assert.Zero(t, req.ContentLength)
ApplyUpstreamBodyMetadata(req, body)
assert.EqualValues(t, len(payload), req.ContentLength)
require.NotNil(t, req.GetBody)
// Drain the primary body as the transport does on the first attempt, then
// make sure GetBody can replay the complete payload repeatedly.
sent, err := io.ReadAll(req.Body)
require.NoError(t, err)
assert.Equal(t, payload, sent)
for i := range 2 {
rc, err := req.GetBody()
require.NoError(t, err)
replay, err := io.ReadAll(rc)
require.NoError(t, err)
require.NoError(t, rc.Close())
assert.Equal(t, payload, replay, "replay %d must equal the original payload", i+1)
}
}
func TestApplyUpstreamBodyMetadataHidesRawBodyStorageCloser(t *testing.T) {
t.Parallel()
payload := []byte(`{"model":"test-model","input":"raw storage"}`)
storage, err := common.CreateBodyStorage(payload)
require.NoError(t, err)
defer storage.Close()
req, err := http.NewRequest(http.MethodPost, "https://example.com/v1/chat/completions", storage)
require.NoError(t, err)
_, exposesStorageBeforeApply := req.Body.(common.BodyStorage)
require.True(t, exposesStorageBeforeApply)
ApplyUpstreamBodyMetadata(req, storage)
_, exposesStorageAfterApply := req.Body.(common.BodyStorage)
assert.False(t, exposesStorageAfterApply)
assert.EqualValues(t, len(payload), req.ContentLength)
require.NotNil(t, req.GetBody)
sent, err := io.ReadAll(req.Body)
require.NoError(t, err)
assert.Equal(t, payload, sent)
require.NoError(t, req.Body.Close())
replayBody, err := req.GetBody()
require.NoError(t, err, "closing the HTTP request body must not close the shared storage")
replay, err := io.ReadAll(replayBody)
require.NoError(t, err)
require.NoError(t, replayBody.Close())
assert.Equal(t, payload, replay)
}
func TestApplyUpstreamBodyMetadataKeepsNativeMetadataForNonReplayableBody(t *testing.T) {
tests := []struct {
name string
body func() io.Reader
}{
{name: "bytes reader", body: func() io.Reader { return bytes.NewReader([]byte("original")) }},
{name: "bytes buffer", body: func() io.Reader { return bytes.NewBufferString("original") }},
{name: "strings reader", body: func() io.Reader { return strings.NewReader("original") }},
}
for _, test := range tests {
t.Run(test.name, func(t *testing.T) {
t.Parallel()
body := test.body()
req, err := http.NewRequest(http.MethodPost, "https://example.com/v1/chat/completions", body)
require.NoError(t, err)
require.NotNil(t, req.GetBody, "net/http must derive GetBody for the concrete reader")
ApplyUpstreamBodyMetadata(req, body)
rc, err := req.GetBody()
require.NoError(t, err)
got, err := io.ReadAll(rc)
require.NoError(t, err)
require.NoError(t, rc.Close())
assert.Equal(t, "original", string(got), "an already correct GetBody must not be overwritten")
assert.EqualValues(t, len("original"), req.ContentLength, "native content length must not be overwritten")
})
}
}
func TestApplyUpstreamBodyMetadataKeepsExistingGetBody(t *testing.T) {
t.Parallel()
payload := []byte(`{"model":"test-model"}`)
body, closer, err := relaycommon.NewOutboundJSONBody(payload)
require.NoError(t, err)
defer closer.Close()
req, err := http.NewRequest(http.MethodPost, "https://example.com/v1/chat/completions", body)
require.NoError(t, err)
req.ContentLength = 99
req.GetBody = func() (io.ReadCloser, error) {
return io.NopCloser(bytes.NewReader([]byte("existing"))), nil
}
ApplyUpstreamBodyMetadata(req, body)
assert.EqualValues(t, len(payload), req.ContentLength)
rc, err := req.GetBody()
require.NoError(t, err)
got, err := io.ReadAll(rc)
require.NoError(t, err)
require.NoError(t, rc.Close())
assert.Equal(t, "existing", string(got))
}
func TestApplyUpstreamBodyMetadataEmptyStorageRemainsReplayable(t *testing.T) {
t.Parallel()
storage, err := common.CreateBodyStorage(nil)
require.NoError(t, err)
defer storage.Close()
body := common.NewReplayableBodyReader(storage)
req, err := http.NewRequest(http.MethodPost, "https://example.com/v1/chat/completions", body)
require.NoError(t, err)
ApplyUpstreamBodyMetadata(req, body)
assert.Zero(t, req.ContentLength)
require.NotNil(t, req.GetBody)
rc, err := req.GetBody()
require.NoError(t, err)
replay, err := io.ReadAll(rc)
require.NoError(t, err)
require.NoError(t, rc.Close())
assert.Empty(t, replay)
}
// stubTaskAdaptor implements just enough of TaskAdaptor for DoTaskApiRequest.
type stubTaskAdaptor struct {
TaskAdaptor
baseURL string
capturedReq *http.Request
}
func (s *stubTaskAdaptor) BuildRequestURL(info *relaycommon.RelayInfo) (string, error) {
return s.baseURL + "/v1/video/generations", nil
}
func (s *stubTaskAdaptor) BuildRequestHeader(c *gin.Context, req *http.Request, info *relaycommon.RelayInfo) error {
s.capturedReq = req
return nil
}
// TestDoTaskApiRequest_KeepsReplayableGetBody guards against reintroducing the
// hand-rolled GetBody override that wrapped the already consumed request
// reader: any transport-level retry would then have silently replayed an empty
// body. net/http derives a correct snapshot-based GetBody from the
// *bytes.Reader bodies the task adaptors pass in, and it must be left intact.
func TestDoTaskApiRequest_KeepsReplayableGetBody(t *testing.T) {
service.InitHttpClient()
payload := []byte(`{"model":"test-model","prompt":"hello"}`)
type receivedBody struct {
body []byte
err error
}
receivedCh := make(chan receivedBody, 1)
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
body, err := io.ReadAll(r.Body)
receivedCh <- receivedBody{body: body, err: err}
w.WriteHeader(http.StatusOK)
}))
defer server.Close()
gin.SetMode(gin.TestMode)
recorder := httptest.NewRecorder()
ctx, _ := gin.CreateTestContext(recorder)
ctx.Request = httptest.NewRequest(http.MethodPost, "/v1/video/generations", bytes.NewReader(payload))
info := &relaycommon.RelayInfo{
ChannelMeta: &relaycommon.ChannelMeta{},
}
adaptor := &stubTaskAdaptor{baseURL: server.URL}
resp, err := DoTaskApiRequest(adaptor, ctx, info, bytes.NewReader(payload))
require.NoError(t, err)
defer resp.Body.Close()
assert.Equal(t, http.StatusOK, resp.StatusCode)
received := <-receivedCh
require.NoError(t, received.err)
assert.Equal(t, payload, received.body)
req := adaptor.capturedReq
require.NotNil(t, req)
require.NotNil(t, req.GetBody)
// Even after the request body has been fully written, GetBody must still
// return the complete payload, repeatedly.
for i := range 2 {
rc, err := req.GetBody()
require.NoError(t, err)
replay, err := io.ReadAll(rc)
require.NoError(t, err)
require.NoError(t, rc.Close())
assert.Equal(t, payload, replay, "replay %d must equal the original payload", i+1)
}
}
type h2ServerResult struct {
err error
streamCount int
attemptBodies [][]byte
}
func acceptH2TestConnection(ln net.Listener) (net.Conn, *http2.Framer, error) {
conn, err := ln.Accept()
if err != nil {
return nil, nil, err
}
_ = conn.SetDeadline(time.Now().Add(15 * time.Second))
preface := make([]byte, len(http2.ClientPreface))
if _, err := io.ReadFull(conn, preface); err != nil {
conn.Close()
return nil, nil, fmt.Errorf("read client preface: %w", err)
}
if !bytes.Equal(preface, []byte(http2.ClientPreface)) {
conn.Close()
return nil, nil, fmt.Errorf("unexpected client preface")
}
framer := http2.NewFramer(conn, conn)
framer.ReadMetaHeaders = hpack.NewDecoder(4096, nil)
if err := framer.WriteSettings(); err != nil {
conn.Close()
return nil, nil, err
}
return conn, framer, nil
}
func readH2TestRequest(framer *http2.Framer) (uint32, []byte, error) {
var streamID uint32
var body []byte
for {
frame, err := framer.ReadFrame()
if err != nil {
return 0, nil, fmt.Errorf("read frame: %w", err)
}
switch f := frame.(type) {
case *http2.SettingsFrame:
if !f.IsAck() {
if err := framer.WriteSettingsAck(); err != nil {
return 0, nil, err
}
}
case *http2.MetaHeadersFrame:
streamID = f.Header().StreamID
if f.StreamEnded() {
return streamID, body, nil
}
case *http2.DataFrame:
if streamID == 0 {
streamID = f.Header().StreamID
}
if f.Header().StreamID != streamID {
continue
}
body = append(body, f.Data()...)
if f.StreamEnded() {
return streamID, body, nil
}
}
}
}
func writeH2TestResponse(framer *http2.Framer, streamID uint32) error {
var hpackBuf bytes.Buffer
henc := hpack.NewEncoder(&hpackBuf)
if err := henc.WriteField(hpack.HeaderField{Name: ":status", Value: "200"}); err != nil {
return err
}
if err := framer.WriteHeaders(http2.HeadersFrameParam{
StreamID: streamID,
BlockFragment: hpackBuf.Bytes(),
EndHeaders: true,
}); err != nil {
return err
}
return framer.WriteData(streamID, true, []byte(`{}`))
}
func awaitH2ServerResult(t *testing.T, resultCh <-chan h2ServerResult) h2ServerResult {
t.Helper()
select {
case result := <-resultCh:
return result
case <-time.After(20 * time.Second):
t.Fatal("timed out waiting for HTTP/2 test server")
return h2ServerResult{}
}
}
// runResetOnFirstStreamServer speaks just enough raw HTTP/2 to emulate an
// upstream that accepts the first request, waits until the request body has
// been fully written, and then resets the stream with REFUSED_STREAM (the
// retry-safe reset some proxy/CDN-fronted upstreams send under load or during
// graceful shutdown, see RFC 9113 section 8.7). When expectRetry is true it
// serves the retried stream a 200 response; otherwise it stops after the reset.
func runResetOnFirstStreamServer(ln net.Listener, expectRetry bool) <-chan h2ServerResult {
resCh := make(chan h2ServerResult, 1)
go func() {
res := h2ServerResult{}
defer func() { resCh <- res }()
conn, framer, err := acceptH2TestConnection(ln)
if err != nil {
res.err = err
return
}
defer conn.Close()
attempts:
for attempt := 0; ; attempt++ {
streamID, body, err := readH2TestRequest(framer)
if err != nil {
res.err = err
return
}
res.streamCount++
res.attemptBodies = append(res.attemptBodies, body)
if attempt == 0 {
if err := framer.WriteRSTStream(streamID, http2.ErrCodeRefusedStream); err != nil {
res.err = err
return
}
if !expectRetry {
break attempts
}
continue
}
if err := writeH2TestResponse(framer, streamID); err != nil {
res.err = err
}
return
}
}()
return resCh
}
func runGoAwayAfterFirstRequestServer(ln net.Listener) <-chan h2ServerResult {
resCh := make(chan h2ServerResult, 1)
go func() {
res := h2ServerResult{}
defer func() { resCh <- res }()
for attempt := range 2 {
conn, framer, err := acceptH2TestConnection(ln)
if err != nil {
res.err = err
return
}
streamID, body, err := readH2TestRequest(framer)
if err != nil {
conn.Close()
res.err = err
return
}
res.streamCount++
res.attemptBodies = append(res.attemptBodies, body)
if attempt == 0 {
err = framer.WriteGoAway(0, http2.ErrCodeNo, nil)
conn.Close()
if err != nil {
res.err = err
return
}
continue
}
err = writeH2TestResponse(framer, streamID)
conn.Close()
if err != nil {
res.err = err
}
return
}
}()
return resCh
}
func newH2PriorKnowledgeClient(ln net.Listener) (*http.Client, *http2.Transport) {
transport := &http2.Transport{
AllowHTTP: true,
DialTLSContext: func(ctx context.Context, network, _ string, _ *tls.Config) (net.Conn, error) {
var dialer net.Dialer
return dialer.DialContext(ctx, network, ln.Addr().String())
},
}
return &http.Client{Transport: transport, Timeout: 15 * time.Second}, transport
}
func newPassThroughBody(t *testing.T, payload []byte) (common.ReplayableBody, common.BodyStorage) {
t.Helper()
storage, err := common.CreateBodyStorage(payload)
require.NoError(t, err)
return common.NewReplayableBodyReader(storage), storage
}
// TestUpstreamGetBody_HTTP2RetryAfterUpstreamStreamReset exercises the actual
// failure this change fixes: an HTTP/2 upstream resets the stream with a
// retryable error after the request body has been written. With GetBody wired
// up the transport must transparently retry, and the retried request must
// carry the complete body.
func TestUpstreamGetBody_HTTP2RetryAfterUpstreamStreamReset(t *testing.T) {
payload := []byte(`{"model":"test-model","messages":[{"role":"user","content":"retry me"}]}`)
ln, err := net.Listen("tcp", "127.0.0.1:0")
require.NoError(t, err)
defer ln.Close()
resCh := runResetOnFirstStreamServer(ln, true)
client, transport := newH2PriorKnowledgeClient(ln)
defer transport.CloseIdleConnections()
// Build the upstream request exactly the way DoApiRequest does: pass the
// original replayable body to the metadata helper after NewRequest.
body, closer, err := relaycommon.NewOutboundJSONBody(payload)
require.NoError(t, err)
defer closer.Close()
req, err := http.NewRequest(http.MethodPost, "http://upstream.test/v1/chat/completions", body)
require.NoError(t, err)
ApplyUpstreamBodyMetadata(req, body)
require.NotNil(t, req.GetBody)
resp, err := client.Do(req)
require.NoError(t, err, "the transport must transparently retry after RST_STREAM")
defer resp.Body.Close()
assert.Equal(t, http.StatusOK, resp.StatusCode)
srv := awaitH2ServerResult(t, resCh)
require.NoError(t, srv.err)
assert.Equal(t, 2, srv.streamCount, "the request must have been attempted twice")
require.Len(t, srv.attemptBodies, 2)
assert.Equal(t, payload, srv.attemptBodies[0], "first attempt must carry the full body")
assert.Equal(t, payload, srv.attemptBodies[1], "the retried request must carry the complete body")
}
func TestUpstreamGetBody_HTTP2RetryAfterUpstreamStreamReset_PassThrough(t *testing.T) {
payload := []byte(`{"model":"test-model","messages":[{"role":"user","content":"pass through"}]}`)
ln, err := net.Listen("tcp", "127.0.0.1:0")
require.NoError(t, err)
defer ln.Close()
resCh := runResetOnFirstStreamServer(ln, true)
client, transport := newH2PriorKnowledgeClient(ln)
defer transport.CloseIdleConnections()
body, storage := newPassThroughBody(t, payload)
defer storage.Close()
req, err := http.NewRequest(http.MethodPost, "http://upstream.test/v1/chat/completions", body)
require.NoError(t, err)
ApplyUpstreamBodyMetadata(req, body)
require.NotNil(t, req.GetBody)
assert.EqualValues(t, len(payload), req.ContentLength)
resp, err := client.Do(req)
require.NoError(t, err, "the transport must transparently retry a pass-through body after RST_STREAM")
defer resp.Body.Close()
assert.Equal(t, http.StatusOK, resp.StatusCode)
srv := awaitH2ServerResult(t, resCh)
require.NoError(t, srv.err)
assert.Equal(t, 2, srv.streamCount)
require.Len(t, srv.attemptBodies, 2)
assert.Equal(t, payload, srv.attemptBodies[0])
assert.Equal(t, payload, srv.attemptBodies[1])
}
func TestUpstreamGetBody_HTTP2RetryAfterGracefulGoAway_PassThrough(t *testing.T) {
payload := []byte(`{"model":"test-model","messages":[{"role":"user","content":"go away"}]}`)
ln, err := net.Listen("tcp", "127.0.0.1:0")
require.NoError(t, err)
defer ln.Close()
resCh := runGoAwayAfterFirstRequestServer(ln)
client, transport := newH2PriorKnowledgeClient(ln)
defer transport.CloseIdleConnections()
body, storage := newPassThroughBody(t, payload)
defer storage.Close()
req, err := http.NewRequest(http.MethodPost, "http://upstream.test/v1/chat/completions", body)
require.NoError(t, err)
ApplyUpstreamBodyMetadata(req, body)
require.NotNil(t, req.GetBody)
resp, err := client.Do(req)
require.NoError(t, err, "the transport must retry on a new connection after graceful GOAWAY")
defer resp.Body.Close()
assert.Equal(t, http.StatusOK, resp.StatusCode)
srv := awaitH2ServerResult(t, resCh)
require.NoError(t, srv.err)
assert.Equal(t, 2, srv.streamCount)
require.Len(t, srv.attemptBodies, 2)
assert.Equal(t, payload, srv.attemptBodies[0])
assert.Equal(t, payload, srv.attemptBodies[1])
}
// TestUpstreamGetBody_HTTP2CannotRetryWithoutGetBody documents the pre-fix
// behavior: without GetBody the transport cannot safely retry once the body
// has been written, and the whole relay request fails.
func TestUpstreamGetBody_HTTP2CannotRetryWithoutGetBody(t *testing.T) {
payload := []byte(`{"model":"test-model","messages":[{"role":"user","content":"retry me"}]}`)
ln, err := net.Listen("tcp", "127.0.0.1:0")
require.NoError(t, err)
defer ln.Close()
resCh := runResetOnFirstStreamServer(ln, false)
client, transport := newH2PriorKnowledgeClient(ln)
defer transport.CloseIdleConnections()
body, closer, err := relaycommon.NewOutboundJSONBody(payload)
require.NoError(t, err)
defer closer.Close()
req, err := http.NewRequest(http.MethodPost, "http://upstream.test/v1/chat/completions", body)
require.NoError(t, err)
req.ContentLength = body.Size()
assert.Nil(t, req.GetBody)
resp, err := client.Do(req) //nolint:bodyclose // Do fails, no body to close
require.Error(t, err)
assert.Nil(t, resp)
require.ErrorContains(t, err, "cannot retry err")
require.ErrorContains(t, err, "Request.Body was written")
srv := awaitH2ServerResult(t, resCh)
require.NoError(t, srv.err)
assert.Equal(t, 1, srv.streamCount)
require.Len(t, srv.attemptBodies, 1)
assert.Equal(t, payload, srv.attemptBodies[0])
}
@@ -0,0 +1,88 @@
package channel
import (
"bytes"
"io"
"net/http"
"net/http/httptest"
"reflect"
"sync/atomic"
"testing"
relaycommon "github.com/QuantumNous/new-api/relay/common"
"github.com/QuantumNous/new-api/service"
"github.com/gin-gonic/gin"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
func TestDoRequestReturnsUpstreamRedirectWithoutFollowing(t *testing.T) {
service.InitHttpClient()
gin.SetMode(gin.TestMode)
sharedClient := service.GetHttpClient()
require.NotNil(t, sharedClient)
require.NotNil(t, sharedClient.CheckRedirect)
originalRedirectPolicy := reflect.ValueOf(sharedClient.CheckRedirect).Pointer()
var targetRequests atomic.Int32
target := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
targetRequests.Add(1)
w.WriteHeader(http.StatusTeapot)
}))
defer target.Close()
const responseBody = "redirect response"
tests := []int{
http.StatusMovedPermanently,
http.StatusFound,
http.StatusSeeOther,
http.StatusTemporaryRedirect,
http.StatusPermanentRedirect,
}
for _, statusCode := range tests {
t.Run(http.StatusText(statusCode), func(t *testing.T) {
targetRequests.Store(0)
var sourceRequests atomic.Int32
type sourceResult struct {
body []byte
err error
}
sourceResultCh := make(chan sourceResult, 1)
source := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
sourceRequests.Add(1)
body, err := io.ReadAll(r.Body)
sourceResultCh <- sourceResult{body: body, err: err}
w.Header().Set("Location", target.URL+"/redirect-target")
w.WriteHeader(statusCode)
_, _ = io.WriteString(w, responseBody)
}))
defer source.Close()
recorder := httptest.NewRecorder()
ctx, _ := gin.CreateTestContext(recorder)
ctx.Request = httptest.NewRequest(http.MethodPost, "/relay", nil)
req, err := http.NewRequest(http.MethodPost, source.URL, bytes.NewReader([]byte("request body")))
require.NoError(t, err)
info := &relaycommon.RelayInfo{ChannelMeta: &relaycommon.ChannelMeta{}}
resp, err := doRequest(ctx, req, info)
require.NoError(t, err)
defer resp.Body.Close()
body, err := io.ReadAll(resp.Body)
require.NoError(t, err)
gotSource := <-sourceResultCh
require.NoError(t, gotSource.err)
assert.Equal(t, statusCode, resp.StatusCode)
assert.Equal(t, target.URL+"/redirect-target", resp.Header.Get("Location"))
assert.Equal(t, responseBody, string(body))
assert.Equal(t, []byte("request body"), gotSource.body)
assert.EqualValues(t, 1, sourceRequests.Load())
assert.Zero(t, targetRequests.Load())
})
}
assert.Equal(t, originalRedirectPolicy, reflect.ValueOf(sharedClient.CheckRedirect).Pointer(), "the cached client must not be mutated")
}
+207
View File
@@ -0,0 +1,207 @@
package channel
import (
"context"
"net/http"
"net/http/httptest"
"testing"
relaycommon "github.com/QuantumNous/new-api/relay/common"
"github.com/gin-gonic/gin"
"github.com/stretchr/testify/require"
)
func TestNewTaskAPIRequestInheritsClientCancellation(t *testing.T) {
recorder := httptest.NewRecorder()
c, _ := gin.CreateTestContext(recorder)
requestContext, cancel := context.WithCancel(context.Background())
c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", nil).WithContext(requestContext)
upstream, err := newTaskAPIRequest(c, "https://provider.example/tasks", nil)
require.NoError(t, err)
cancel()
require.ErrorIs(t, upstream.Context().Err(), context.Canceled)
}
func TestProcessHeaderOverride_ChannelTestSkipsPassthroughRules(t *testing.T) {
t.Parallel()
gin.SetMode(gin.TestMode)
recorder := httptest.NewRecorder()
ctx, _ := gin.CreateTestContext(recorder)
ctx.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", nil)
ctx.Request.Header.Set("X-Trace-Id", "trace-123")
info := &relaycommon.RelayInfo{
IsChannelTest: true,
ChannelMeta: &relaycommon.ChannelMeta{
HeadersOverride: map[string]any{
"*": "",
},
},
}
headers, err := processHeaderOverride(info, ctx)
require.NoError(t, err)
require.Empty(t, headers)
}
func TestProcessHeaderOverride_ChannelTestSkipsClientHeaderPlaceholder(t *testing.T) {
t.Parallel()
gin.SetMode(gin.TestMode)
recorder := httptest.NewRecorder()
ctx, _ := gin.CreateTestContext(recorder)
ctx.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", nil)
ctx.Request.Header.Set("X-Trace-Id", "trace-123")
info := &relaycommon.RelayInfo{
IsChannelTest: true,
ChannelMeta: &relaycommon.ChannelMeta{
HeadersOverride: map[string]any{
"X-Upstream-Trace": "{client_header:X-Trace-Id}",
},
},
}
headers, err := processHeaderOverride(info, ctx)
require.NoError(t, err)
_, ok := headers["x-upstream-trace"]
require.False(t, ok)
}
func TestProcessHeaderOverride_NonTestKeepsClientHeaderPlaceholder(t *testing.T) {
t.Parallel()
gin.SetMode(gin.TestMode)
recorder := httptest.NewRecorder()
ctx, _ := gin.CreateTestContext(recorder)
ctx.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", nil)
ctx.Request.Header.Set("X-Trace-Id", "trace-123")
info := &relaycommon.RelayInfo{
IsChannelTest: false,
ChannelMeta: &relaycommon.ChannelMeta{
HeadersOverride: map[string]any{
"X-Upstream-Trace": "{client_header:X-Trace-Id}",
},
},
}
headers, err := processHeaderOverride(info, ctx)
require.NoError(t, err)
require.Equal(t, "trace-123", headers["x-upstream-trace"])
}
func TestProcessHeaderOverride_RuntimeOverrideIsFinalHeaderMap(t *testing.T) {
t.Parallel()
gin.SetMode(gin.TestMode)
recorder := httptest.NewRecorder()
ctx, _ := gin.CreateTestContext(recorder)
ctx.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", nil)
info := &relaycommon.RelayInfo{
IsChannelTest: false,
UseRuntimeHeadersOverride: true,
RuntimeHeadersOverride: map[string]any{
"x-static": "runtime-value",
"x-runtime": "runtime-only",
},
ChannelMeta: &relaycommon.ChannelMeta{
HeadersOverride: map[string]any{
"X-Static": "legacy-value",
"X-Legacy": "legacy-only",
},
},
}
headers, err := processHeaderOverride(info, ctx)
require.NoError(t, err)
require.Equal(t, "runtime-value", headers["x-static"])
require.Equal(t, "runtime-only", headers["x-runtime"])
_, exists := headers["x-legacy"]
require.False(t, exists)
}
func TestProcessHeaderOverride_PassthroughSkipsAcceptEncoding(t *testing.T) {
t.Parallel()
gin.SetMode(gin.TestMode)
recorder := httptest.NewRecorder()
ctx, _ := gin.CreateTestContext(recorder)
ctx.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", nil)
ctx.Request.Header.Set("X-Trace-Id", "trace-123")
ctx.Request.Header.Set("Accept-Encoding", "gzip")
info := &relaycommon.RelayInfo{
IsChannelTest: false,
ChannelMeta: &relaycommon.ChannelMeta{
HeadersOverride: map[string]any{
"*": "",
},
},
}
headers, err := processHeaderOverride(info, ctx)
require.NoError(t, err)
require.Equal(t, "trace-123", headers["x-trace-id"])
_, hasAcceptEncoding := headers["accept-encoding"]
require.False(t, hasAcceptEncoding)
}
func TestProcessHeaderOverride_PassHeadersTemplateSetsRuntimeHeaders(t *testing.T) {
t.Parallel()
gin.SetMode(gin.TestMode)
recorder := httptest.NewRecorder()
ctx, _ := gin.CreateTestContext(recorder)
ctx.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", nil)
ctx.Request.Header.Set("Originator", "Codex CLI")
ctx.Request.Header.Set("Session_id", "sess-123")
info := &relaycommon.RelayInfo{
IsChannelTest: false,
RequestHeaders: map[string]string{
"Originator": "Codex CLI",
"Session_id": "sess-123",
},
ChannelMeta: &relaycommon.ChannelMeta{
ParamOverride: map[string]any{
"operations": []any{
map[string]any{
"mode": "pass_headers",
"value": []any{"Originator", "Session_id", "X-Codex-Beta-Features"},
},
},
},
HeadersOverride: map[string]any{
"X-Static": "legacy-value",
},
},
}
_, err := relaycommon.ApplyParamOverrideWithRelayInfo([]byte(`{"model":"gpt-4.1"}`), info)
require.NoError(t, err)
require.True(t, info.UseRuntimeHeadersOverride)
require.Equal(t, "Codex CLI", info.RuntimeHeadersOverride["originator"])
require.Equal(t, "sess-123", info.RuntimeHeadersOverride["session_id"])
_, exists := info.RuntimeHeadersOverride["x-codex-beta-features"]
require.False(t, exists)
require.Equal(t, "legacy-value", info.RuntimeHeadersOverride["x-static"])
headers, err := processHeaderOverride(info, ctx)
require.NoError(t, err)
require.Equal(t, "Codex CLI", headers["originator"])
require.Equal(t, "sess-123", headers["session_id"])
_, exists = headers["x-codex-beta-features"]
require.False(t, exists)
upstreamReq := httptest.NewRequest(http.MethodPost, "https://example.com/v1/responses", nil)
applyHeaderOverrideToRequest(upstreamReq, headers)
require.Equal(t, "Codex CLI", upstreamReq.Header.Get("Originator"))
require.Equal(t, "sess-123", upstreamReq.Header.Get("Session_id"))
require.Empty(t, upstreamReq.Header.Get("X-Codex-Beta-Features"))
}
+192
View File
@@ -0,0 +1,192 @@
package aws
import (
"fmt"
"io"
"net/http"
"strings"
"github.com/QuantumNous/new-api/relay/channel"
"github.com/QuantumNous/new-api/relay/channel/claude"
relaycommon "github.com/QuantumNous/new-api/relay/common"
"github.com/QuantumNous/new-api/relaykit/dto"
"github.com/QuantumNous/new-api/relaykit/types"
"github.com/QuantumNous/new-api/service"
"github.com/aws/aws-sdk-go-v2/service/bedrockruntime"
"github.com/pkg/errors"
"github.com/gin-gonic/gin"
)
type ClientMode int
const (
ClientModeApiKey ClientMode = iota + 1
ClientModeAKSK
)
type Adaptor struct {
ClientMode ClientMode
AwsClient *bedrockruntime.Client
AwsModelId string
AwsReq any
IsNova bool
}
func (a *Adaptor) ConvertGeminiRequest(*gin.Context, *relaycommon.RelayInfo, *dto.GeminiChatRequest) (any, error) {
//TODO implement me
return nil, errors.New("not implemented")
}
func (a *Adaptor) ConvertClaudeRequest(c *gin.Context, info *relaycommon.RelayInfo, request *dto.ClaudeRequest) (any, error) {
claudeAdaptor := claude.Adaptor{}
if _, err := claudeAdaptor.ConvertClaudeRequest(c, info, request); err != nil {
return nil, err
}
for i, message := range request.Messages {
updated := false
if !message.IsStringContent() {
content, err := message.ParseContent()
if err != nil {
return nil, errors.Wrap(err, "failed to parse message content")
}
for i2, mediaMessage := range content {
if mediaMessage.Source != nil {
if mediaMessage.Source.Type == "url" {
// 使用统一的文件服务获取图片数据
source := types.NewURLFileSource(mediaMessage.Source.Url)
base64Data, mimeType, err := service.GetBase64Data(c, source, "formatting image for Claude")
if err != nil {
return nil, fmt.Errorf("get file base64 from url failed: %s", err.Error())
}
mediaMessage.Source.MediaType = mimeType
mediaMessage.Source.Data = base64Data
mediaMessage.Source.Url = ""
mediaMessage.Source.Type = "base64"
content[i2] = mediaMessage
updated = true
}
}
}
if updated {
message.SetContent(content)
}
}
if updated {
request.Messages[i] = message
}
}
return request, nil
}
func (a *Adaptor) ConvertAudioRequest(c *gin.Context, info *relaycommon.RelayInfo, request dto.AudioRequest) (io.Reader, error) {
//TODO implement me
return nil, errors.New("not implemented")
}
func (a *Adaptor) ConvertImageRequest(c *gin.Context, info *relaycommon.RelayInfo, request dto.ImageRequest) (any, error) {
//TODO implement me
return nil, errors.New("not implemented")
}
func (a *Adaptor) Init(info *relaycommon.RelayInfo) {
}
func (a *Adaptor) GetRequestURL(info *relaycommon.RelayInfo) (string, error) {
if info.ChannelOtherSettings.AwsKeyType == dto.AwsKeyTypeApiKey {
awsModelId := getAwsModelID(info.UpstreamModelName)
a.ClientMode = ClientModeApiKey
awsSecret := strings.Split(info.ApiKey, "|")
if len(awsSecret) != 2 {
return "", errors.New("invalid aws api key, should be in format of <api-key>|<region>")
}
return fmt.Sprintf("https://bedrock-runtime.%s.amazonaws.com/model/%s/converse", awsModelId, awsSecret[1]), nil
} else {
a.ClientMode = ClientModeAKSK
return "", nil
}
}
func (a *Adaptor) SetupRequestHeader(c *gin.Context, req *http.Header, info *relaycommon.RelayInfo) error {
claude.CommonClaudeHeadersOperation(c, req, info)
if a.ClientMode == ClientModeApiKey {
req.Set("Authorization", "Bearer "+info.ApiKey)
}
return nil
}
func (a *Adaptor) ConvertOpenAIRequest(c *gin.Context, info *relaycommon.RelayInfo, request *dto.GeneralOpenAIRequest) (any, error) {
if request == nil {
return nil, errors.New("request is nil")
}
// 检查是否为Nova模型
if isNovaModel(request.Model) {
novaReq := convertToNovaRequest(request)
a.IsNova = true
return novaReq, nil
}
// 原有的Claude模型处理逻辑
result, err := service.ConvertRequest(c, info, types.RelayFormatClaude, request)
if err != nil {
return nil, errors.Wrap(err, "failed to convert openai request to claude request")
}
claudeReq, ok := result.Value.(*dto.ClaudeRequest)
if !ok {
return nil, fmt.Errorf("expected Anthropic Messages request, got %T", result.Value)
}
info.UpstreamModelName = claudeReq.Model
return claudeReq, err
}
func (a *Adaptor) ConvertRerankRequest(c *gin.Context, relayMode int, request dto.RerankRequest) (any, error) {
return nil, nil
}
func (a *Adaptor) ConvertEmbeddingRequest(c *gin.Context, info *relaycommon.RelayInfo, request dto.EmbeddingRequest) (any, error) {
//TODO implement me
return nil, errors.New("not implemented")
}
func (a *Adaptor) ConvertOpenAIResponsesRequest(c *gin.Context, info *relaycommon.RelayInfo, request dto.OpenAIResponsesRequest) (any, error) {
// TODO implement me
return nil, errors.New("not implemented")
}
func (a *Adaptor) DoRequest(c *gin.Context, info *relaycommon.RelayInfo, requestBody io.Reader) (any, error) {
if a.ClientMode == ClientModeApiKey {
return channel.DoApiRequest(a, c, info, requestBody)
} else {
return doAwsClientRequest(c, info, a, requestBody)
}
}
func (a *Adaptor) DoResponse(c *gin.Context, resp *http.Response, info *relaycommon.RelayInfo) (usage any, err *types.NewAPIError) {
if a.ClientMode == ClientModeApiKey {
claudeAdaptor := claude.Adaptor{}
usage, err = claudeAdaptor.DoResponse(c, resp, info)
} else {
if a.IsNova {
err, usage = handleNovaRequest(c, info, a)
} else {
if info.IsStream {
err, usage = awsStreamHandler(c, info, a)
} else {
err, usage = awsHandler(c, info, a)
}
}
}
return
}
func (a *Adaptor) GetModelList() (models []string) {
for n := range awsModelIDMap {
models = append(models, n)
}
return
}
func (a *Adaptor) GetChannelName() string {
return ChannelName
}
+161
View File
@@ -0,0 +1,161 @@
package aws
import "strings"
var awsModelIDMap = map[string]string{
"claude-3-sonnet-20240229": "anthropic.claude-3-sonnet-20240229-v1:0",
"claude-3-opus-20240229": "anthropic.claude-3-opus-20240229-v1:0",
"claude-3-haiku-20240307": "anthropic.claude-3-haiku-20240307-v1:0",
"claude-3-5-sonnet-20240620": "anthropic.claude-3-5-sonnet-20240620-v1:0",
"claude-3-5-sonnet-20241022": "anthropic.claude-3-5-sonnet-20241022-v2:0",
"claude-3-5-haiku-20241022": "anthropic.claude-3-5-haiku-20241022-v1:0",
"claude-3-7-sonnet-20250219": "anthropic.claude-3-7-sonnet-20250219-v1:0",
"claude-sonnet-4-20250514": "anthropic.claude-sonnet-4-20250514-v1:0",
"claude-opus-4-20250514": "anthropic.claude-opus-4-20250514-v1:0",
"claude-opus-4-1-20250805": "anthropic.claude-opus-4-1-20250805-v1:0",
"claude-sonnet-4-5-20250929": "anthropic.claude-sonnet-4-5-20250929-v1:0",
"claude-sonnet-4-6": "anthropic.claude-sonnet-4-6",
"claude-haiku-4-5-20251001": "anthropic.claude-haiku-4-5-20251001-v1:0",
"claude-opus-4-5-20251101": "anthropic.claude-opus-4-5-20251101-v1:0",
"claude-opus-4-6": "anthropic.claude-opus-4-6-v1",
"claude-opus-4-7": "anthropic.claude-opus-4-7",
"claude-opus-4-8": "anthropic.claude-opus-4-8",
// Nova models
"nova-micro-v1:0": "amazon.nova-micro-v1:0",
"nova-lite-v1:0": "amazon.nova-lite-v1:0",
"nova-pro-v1:0": "amazon.nova-pro-v1:0",
"nova-premier-v1:0": "amazon.nova-premier-v1:0",
"nova-canvas-v1:0": "amazon.nova-canvas-v1:0",
"nova-reel-v1:0": "amazon.nova-reel-v1:0",
"nova-reel-v1:1": "amazon.nova-reel-v1:1",
"nova-sonic-v1:0": "amazon.nova-sonic-v1:0",
}
var awsModelCanCrossRegionMap = map[string]map[string]bool{
"anthropic.claude-3-sonnet-20240229-v1:0": {
"us": true,
"eu": true,
"ap": true,
},
"anthropic.claude-3-opus-20240229-v1:0": {
"us": true,
},
"anthropic.claude-3-haiku-20240307-v1:0": {
"us": true,
"eu": true,
"ap": true,
},
"anthropic.claude-3-5-sonnet-20240620-v1:0": {
"us": true,
"eu": true,
"ap": true,
},
"anthropic.claude-3-5-sonnet-20241022-v2:0": {
"us": true,
"ap": true,
},
"anthropic.claude-3-5-haiku-20241022-v1:0": {
"us": true,
},
"anthropic.claude-3-7-sonnet-20250219-v1:0": {
"us": true,
"ap": true,
"eu": true,
},
"anthropic.claude-sonnet-4-20250514-v1:0": {
"us": true,
"ap": true,
"eu": true,
},
"anthropic.claude-opus-4-20250514-v1:0": {
"us": true,
},
"anthropic.claude-opus-4-1-20250805-v1:0": {
"us": true,
},
"anthropic.claude-sonnet-4-5-20250929-v1:0": {
"us": true,
"ap": true,
"eu": true,
},
"anthropic.claude-sonnet-4-6": {
"us": true,
"ap": true,
"eu": true,
},
"anthropic.claude-opus-4-5-20251101-v1:0": {
"us": true,
"ap": true,
"eu": true,
},
"anthropic.claude-opus-4-6-v1": {
"us": true,
"ap": true,
"eu": true,
},
"anthropic.claude-opus-4-7": {
"us": true,
"ap": true,
"eu": true,
},
"anthropic.claude-opus-4-8": {
"us": true,
"ap": true,
"eu": true,
},
"anthropic.claude-haiku-4-5-20251001-v1:0": {
"us": true,
"ap": true,
"eu": true,
},
// Nova models - all support three major regions
"amazon.nova-micro-v1:0": {
"us": true,
"eu": true,
"apac": true,
},
"amazon.nova-lite-v1:0": {
"us": true,
"eu": true,
"apac": true,
},
"amazon.nova-pro-v1:0": {
"us": true,
"eu": true,
"apac": true,
},
"amazon.nova-premier-v1:0": {
"us": true,
},
"amazon.nova-canvas-v1:0": {
"us": true,
"eu": true,
"apac": true,
},
"amazon.nova-reel-v1:0": {
"us": true,
"eu": true,
"apac": true,
},
"amazon.nova-reel-v1:1": {
"us": true,
},
"amazon.nova-sonic-v1:0": {
"us": true,
"eu": true,
"apac": true,
},
}
var awsRegionCrossModelPrefixMap = map[string]string{
"us": "us",
"eu": "eu",
"ap": "apac",
}
var ChannelName = "aws"
// 判断是否为Nova模型
func isNovaModel(modelId string) bool {
return strings.Contains(modelId, "nova-")
}
+146
View File
@@ -0,0 +1,146 @@
package aws
import (
"context"
"encoding/json"
"io"
"net/http"
"strings"
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/logger"
"github.com/QuantumNous/new-api/relaykit/dto"
)
type AwsClaudeRequest struct {
// AnthropicVersion should be "bedrock-2023-05-31"
AnthropicVersion string `json:"anthropic_version"`
AnthropicBeta json.RawMessage `json:"anthropic_beta,omitempty"`
System any `json:"system,omitempty"`
Messages []dto.ClaudeMessage `json:"messages"`
MaxTokens *uint `json:"max_tokens,omitempty"`
Temperature *float64 `json:"temperature,omitempty"`
TopP *float64 `json:"top_p,omitempty"`
TopK *int `json:"top_k,omitempty"`
StopSequences []string `json:"stop_sequences,omitempty"`
Tools any `json:"tools,omitempty"`
ToolChoice any `json:"tool_choice,omitempty"`
ContextManagement json.RawMessage `json:"context_management,omitempty"`
Thinking *dto.Thinking `json:"thinking,omitempty"`
OutputConfig json.RawMessage `json:"output_config,omitempty"`
//Metadata json.RawMessage `json:"metadata,omitempty"`
}
func formatRequest(requestBody io.Reader, requestHeader http.Header) (*AwsClaudeRequest, error) {
var awsClaudeRequest AwsClaudeRequest
err := common.DecodeJson(requestBody, &awsClaudeRequest)
if err != nil {
return nil, err
}
awsClaudeRequest.AnthropicVersion = "bedrock-2023-05-31"
// check header anthropic-beta
anthropicBetaValues := requestHeader.Get("anthropic-beta")
if len(anthropicBetaValues) > 0 {
var tempArray []string
tempArray = strings.Split(anthropicBetaValues, ",")
if len(tempArray) > 0 {
betaJson, err := json.Marshal(tempArray)
if err != nil {
return nil, err
}
awsClaudeRequest.AnthropicBeta = betaJson
}
}
logger.LogJson(context.Background(), "json", awsClaudeRequest)
return &awsClaudeRequest, nil
}
// NovaMessage Nova模型使用messages-v1格式
type NovaMessage struct {
Role string `json:"role"`
Content []NovaContent `json:"content"`
}
type NovaContent struct {
Text string `json:"text"`
}
type NovaRequest struct {
SchemaVersion string `json:"schemaVersion"` // 请求版本,例如 "1.0"
Messages []NovaMessage `json:"messages"` // 对话消息列表
InferenceConfig *NovaInferenceConfig `json:"inferenceConfig,omitempty"` // 推理配置,可选
}
type NovaInferenceConfig struct {
MaxTokens int `json:"maxTokens,omitempty"` // 最大生成的 token 数
Temperature float64 `json:"temperature,omitempty"` // 随机性 (默认 0.7, 范围 0-1)
TopP float64 `json:"topP,omitempty"` // nucleus sampling (默认 0.9, 范围 0-1)
TopK int `json:"topK,omitempty"` // 限制候选 token 数 (默认 50, 范围 0-128)
StopSequences []string `json:"stopSequences,omitempty"` // 停止生成的序列
}
// 转换OpenAI请求为Nova格式
func convertToNovaRequest(req *dto.GeneralOpenAIRequest) *NovaRequest {
novaMessages := make([]NovaMessage, len(req.Messages))
for i, msg := range req.Messages {
novaMessages[i] = NovaMessage{
Role: msg.Role,
Content: []NovaContent{{Text: msg.StringContent()}},
}
}
novaReq := &NovaRequest{
SchemaVersion: "messages-v1",
Messages: novaMessages,
}
// 设置推理配置
if (req.MaxTokens != nil && *req.MaxTokens != 0) || (req.Temperature != nil && *req.Temperature != 0) || (req.TopP != nil && *req.TopP != 0) || (req.TopK != nil && *req.TopK != 0) || req.Stop != nil {
novaReq.InferenceConfig = &NovaInferenceConfig{}
if req.MaxTokens != nil && *req.MaxTokens != 0 {
novaReq.InferenceConfig.MaxTokens = int(*req.MaxTokens)
}
if req.Temperature != nil && *req.Temperature != 0 {
novaReq.InferenceConfig.Temperature = *req.Temperature
}
if req.TopP != nil && *req.TopP != 0 {
novaReq.InferenceConfig.TopP = *req.TopP
}
if req.TopK != nil && *req.TopK != 0 {
novaReq.InferenceConfig.TopK = *req.TopK
}
if req.Stop != nil {
if stopSequences := parseStopSequences(req.Stop); len(stopSequences) > 0 {
novaReq.InferenceConfig.StopSequences = stopSequences
}
}
}
return novaReq
}
// parseStopSequences 解析停止序列,支持字符串或字符串数组
func parseStopSequences(stop any) []string {
if stop == nil {
return nil
}
switch v := stop.(type) {
case string:
if v != "" {
return []string{v}
}
case []string:
return v
case []any:
var sequences []string
for _, item := range v {
if str, ok := item.(string); ok && str != "" {
sequences = append(sequences, str)
}
}
return sequences
}
return nil
}
+371
View File
@@ -0,0 +1,371 @@
package aws
import (
"context"
"encoding/json"
"fmt"
"io"
"net/http"
"strings"
"time"
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/relay/channel"
"github.com/QuantumNous/new-api/relay/channel/claude"
relaycommon "github.com/QuantumNous/new-api/relay/common"
"github.com/QuantumNous/new-api/relay/helper"
"github.com/QuantumNous/new-api/relaykit/dto"
"github.com/QuantumNous/new-api/relaykit/types"
"github.com/QuantumNous/new-api/service"
"github.com/gin-gonic/gin"
"github.com/pkg/errors"
"github.com/QuantumNous/new-api/setting/model_setting"
"github.com/aws/aws-sdk-go-v2/aws"
"github.com/aws/aws-sdk-go-v2/credentials"
"github.com/aws/aws-sdk-go-v2/service/bedrockruntime"
bedrockruntimeTypes "github.com/aws/aws-sdk-go-v2/service/bedrockruntime/types"
"github.com/aws/smithy-go/auth/bearer"
)
// getAwsErrorStatusCode extracts HTTP status code from AWS SDK error
func getAwsErrorStatusCode(err error) int {
// Check for HTTP response error which contains status code
var httpErr interface{ HTTPStatusCode() int }
if errors.As(err, &httpErr) {
return httpErr.HTTPStatusCode()
}
// Default to 500 if we can't determine the status code
return http.StatusInternalServerError
}
func newAwsInvokeContext(parent context.Context) (context.Context, context.CancelFunc) {
if common.RelayTimeout <= 0 {
return context.WithCancel(parent)
}
return context.WithTimeout(parent, time.Duration(common.RelayTimeout)*time.Second)
}
func newAwsInvokeError(requestContext context.Context, err error, operation string) *types.NewAPIError {
options := make([]types.NewAPIErrorOptions, 0, 1)
if requestContext.Err() != nil {
options = append(options, types.ErrOptionWithSkipRetry())
}
return types.NewOpenAIError(
errors.Wrap(err, operation),
types.ErrorCodeAwsInvokeError,
getAwsErrorStatusCode(err),
options...,
)
}
func newAwsClient(c *gin.Context, info *relaycommon.RelayInfo) (*bedrockruntime.Client, error) {
httpClient, err := service.GetHttpClientWithProxySettings(info.ChannelSetting.Proxy, info.ChannelSetting)
if err != nil {
return nil, fmt.Errorf("new proxy http client failed: %w", err)
}
awsSecret := strings.Split(info.ApiKey, "|")
var client *bedrockruntime.Client
switch len(awsSecret) {
case 2:
apiKey := awsSecret[0]
region := awsSecret[1]
client = bedrockruntime.New(bedrockruntime.Options{
Region: region,
BearerAuthTokenProvider: bearer.StaticTokenProvider{Token: bearer.Token{Value: apiKey}},
HTTPClient: httpClient,
})
case 3:
ak := awsSecret[0]
sk := awsSecret[1]
region := awsSecret[2]
client = bedrockruntime.New(bedrockruntime.Options{
Region: region,
Credentials: aws.NewCredentialsCache(credentials.NewStaticCredentialsProvider(ak, sk, "")),
HTTPClient: httpClient,
})
default:
return nil, errors.New("invalid aws secret key")
}
return client, nil
}
func doAwsClientRequest(c *gin.Context, info *relaycommon.RelayInfo, a *Adaptor, requestBody io.Reader) (any, error) {
awsCli, err := newAwsClient(c, info)
if err != nil {
return nil, types.NewError(err, types.ErrorCodeChannelAwsClientError)
}
a.AwsClient = awsCli
// 获取对应的AWS模型ID
awsModelId := getAwsModelID(info.UpstreamModelName)
awsRegionPrefix := getAwsRegionPrefix(awsCli.Options().Region)
canCrossRegion := awsModelCanCrossRegion(awsModelId, awsRegionPrefix)
if canCrossRegion {
awsModelId = awsModelCrossRegion(awsModelId, awsRegionPrefix)
}
// init empty request.header
requestHeader := http.Header{}
a.SetupRequestHeader(c, &requestHeader, info)
headerOverride, err := channel.ResolveHeaderOverride(info, c)
if err != nil {
return nil, err
}
for key, value := range headerOverride {
requestHeader.Set(key, value)
}
if isNovaModel(awsModelId) {
var novaReq *NovaRequest
err = common.DecodeJson(requestBody, &novaReq)
if err != nil {
return nil, types.NewError(errors.Wrap(err, "decode nova request fail"), types.ErrorCodeBadRequestBody)
}
// 使用InvokeModel API,但使用Nova格式的请求体
awsReq := &bedrockruntime.InvokeModelInput{
ModelId: aws.String(awsModelId),
Accept: aws.String("application/json"),
ContentType: aws.String("application/json"),
}
reqBody, err := common.Marshal(novaReq)
if err != nil {
return nil, types.NewError(errors.Wrap(err, "marshal nova request"), types.ErrorCodeBadResponseBody)
}
awsReq.Body = reqBody
a.AwsReq = awsReq
return nil, nil
} else {
awsClaudeReq, err := formatRequest(requestBody, requestHeader)
if err != nil {
return nil, types.NewError(errors.Wrap(err, "format aws request fail"), types.ErrorCodeBadRequestBody)
}
if info.IsStream {
awsReq := &bedrockruntime.InvokeModelWithResponseStreamInput{
ModelId: aws.String(awsModelId),
Accept: aws.String("application/json"),
ContentType: aws.String("application/json"),
}
awsReq.Body, err = buildAwsRequestBody(c, info, awsClaudeReq)
if err != nil {
return nil, types.NewError(errors.Wrap(err, "marshal aws request fail"), types.ErrorCodeBadRequestBody)
}
a.AwsReq = awsReq
return nil, nil
} else {
awsReq := &bedrockruntime.InvokeModelInput{
ModelId: aws.String(awsModelId),
Accept: aws.String("application/json"),
ContentType: aws.String("application/json"),
}
awsReq.Body, err = buildAwsRequestBody(c, info, awsClaudeReq)
if err != nil {
return nil, types.NewError(errors.Wrap(err, "marshal aws request fail"), types.ErrorCodeBadRequestBody)
}
a.AwsReq = awsReq
return nil, nil
}
}
}
// buildAwsRequestBody prepares the payload for AWS requests, applying passthrough rules when enabled.
func buildAwsRequestBody(c *gin.Context, info *relaycommon.RelayInfo, awsClaudeReq any) ([]byte, error) {
if model_setting.GetGlobalSettings().PassThroughRequestEnabled || info.ChannelSetting.PassThroughBodyEnabled {
storage, err := common.GetBodyStorage(c)
if err != nil {
return nil, errors.Wrap(err, "get request body for pass-through fail")
}
body, err := storage.Bytes()
if err != nil {
return nil, errors.Wrap(err, "get request body bytes fail")
}
var data map[string]any
if err := common.Unmarshal(body, &data); err != nil {
return nil, errors.Wrap(err, "pass-through unmarshal request body fail")
}
delete(data, "model")
delete(data, "stream")
return common.Marshal(data)
}
return common.Marshal(awsClaudeReq)
}
func getAwsRegionPrefix(awsRegionId string) string {
parts := strings.Split(awsRegionId, "-")
regionPrefix := ""
if len(parts) > 0 {
regionPrefix = parts[0]
}
return regionPrefix
}
func awsModelCanCrossRegion(awsModelId, awsRegionPrefix string) bool {
regionSet, exists := awsModelCanCrossRegionMap[awsModelId]
return exists && regionSet[awsRegionPrefix]
}
func awsModelCrossRegion(awsModelId, awsRegionPrefix string) string {
modelPrefix, find := awsRegionCrossModelPrefixMap[awsRegionPrefix]
if !find {
return awsModelId
}
return modelPrefix + "." + awsModelId
}
func getAwsModelID(requestModel string) string {
if awsModelIDName, ok := awsModelIDMap[requestModel]; ok {
return awsModelIDName
}
return requestModel
}
func awsHandler(c *gin.Context, info *relaycommon.RelayInfo, a *Adaptor) (*types.NewAPIError, *dto.Usage) {
requestContext := c.Request.Context()
ctx, cancel := newAwsInvokeContext(requestContext)
defer cancel()
awsResp, err := a.AwsClient.InvokeModel(ctx, a.AwsReq.(*bedrockruntime.InvokeModelInput))
if err != nil {
return newAwsInvokeError(requestContext, err, "InvokeModel"), nil
}
claudeInfo := &claude.ClaudeResponseInfo{
ResponseId: helper.GetResponseID(c),
Created: common.GetTimestamp(),
Model: info.UpstreamModelName,
ResponseText: strings.Builder{},
Usage: &dto.Usage{},
}
// 复制上游 Content-Type 到客户端响应头
if awsResp.ContentType != nil && *awsResp.ContentType != "" {
c.Writer.Header().Set("Content-Type", *awsResp.ContentType)
}
handlerErr := claude.HandleClaudeResponseData(c, info, claudeInfo, nil, awsResp.Body)
if handlerErr != nil {
return handlerErr, nil
}
return nil, claudeInfo.Usage
}
func awsStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, a *Adaptor) (*types.NewAPIError, *dto.Usage) {
requestContext := c.Request.Context()
ctx, cancel := newAwsInvokeContext(requestContext)
defer cancel()
awsResp, err := a.AwsClient.InvokeModelWithResponseStream(ctx, a.AwsReq.(*bedrockruntime.InvokeModelWithResponseStreamInput))
if err != nil {
return newAwsInvokeError(requestContext, err, "InvokeModelWithResponseStream"), nil
}
stream := awsResp.GetStream()
defer stream.Close()
claudeInfo := &claude.ClaudeResponseInfo{
ResponseId: helper.GetResponseID(c),
Created: common.GetTimestamp(),
Model: info.UpstreamModelName,
ResponseText: strings.Builder{},
Usage: &dto.Usage{},
}
events := stream.Events()
streamLoop:
for {
select {
case <-ctx.Done():
break streamLoop
case event, ok := <-events:
if !ok {
break streamLoop
}
if ctx.Err() != nil {
break streamLoop
}
switch v := event.(type) {
case *bedrockruntimeTypes.ResponseStreamMemberChunk:
info.SetFirstResponseTime()
respErr := claude.HandleStreamResponseData(c, info, claudeInfo, string(v.Value.Bytes))
if respErr != nil {
return respErr, nil
}
case *bedrockruntimeTypes.UnknownUnionMember:
fmt.Println("unknown tag:", v.Tag)
return types.NewError(errors.New("unknown response type"), types.ErrorCodeInvalidRequest), nil
default:
fmt.Println("union is nil or unknown type")
return types.NewError(errors.New("nil or unknown response type"), types.ErrorCodeInvalidRequest), nil
}
}
}
_ = stream.Close()
claude.HandleStreamFinalResponse(c, info, claudeInfo)
return nil, claudeInfo.Usage
}
// Nova模型处理函数
func handleNovaRequest(c *gin.Context, info *relaycommon.RelayInfo, a *Adaptor) (*types.NewAPIError, *dto.Usage) {
requestContext := c.Request.Context()
ctx, cancel := newAwsInvokeContext(requestContext)
defer cancel()
awsResp, err := a.AwsClient.InvokeModel(ctx, a.AwsReq.(*bedrockruntime.InvokeModelInput))
if err != nil {
return newAwsInvokeError(requestContext, err, "InvokeModel"), nil
}
// 解析Nova响应
var novaResp struct {
Output struct {
Message struct {
Content []struct {
Text string `json:"text"`
} `json:"content"`
} `json:"message"`
} `json:"output"`
Usage struct {
InputTokens int `json:"inputTokens"`
OutputTokens int `json:"outputTokens"`
TotalTokens int `json:"totalTokens"`
} `json:"usage"`
}
if err := json.Unmarshal(awsResp.Body, &novaResp); err != nil {
return types.NewError(errors.Wrap(err, "unmarshal nova response"), types.ErrorCodeBadResponseBody), nil
}
// 构造OpenAI格式响应
response := dto.OpenAITextResponse{
Id: helper.GetResponseID(c),
Object: "chat.completion",
Created: common.GetTimestamp(),
Model: info.UpstreamModelName,
Choices: []dto.OpenAITextResponseChoice{{
Index: 0,
Message: dto.Message{
Role: "assistant",
Content: novaResp.Output.Message.Content[0].Text,
},
FinishReason: "stop",
}},
Usage: dto.Usage{
PromptTokens: novaResp.Usage.InputTokens,
CompletionTokens: novaResp.Usage.OutputTokens,
TotalTokens: novaResp.Usage.TotalTokens,
},
}
c.JSON(http.StatusOK, response)
return nil, &response.Usage
}
+452
View File
@@ -0,0 +1,452 @@
package aws
import (
"bytes"
"context"
"errors"
"io"
"net/http"
"net/http/httptest"
"sync"
"testing"
"time"
"github.com/QuantumNous/new-api/common"
relaycommon "github.com/QuantumNous/new-api/relay/common"
"github.com/QuantumNous/new-api/relaykit/dto"
relaytypes "github.com/QuantumNous/new-api/relaykit/types"
"github.com/aws/aws-sdk-go-v2/aws"
"github.com/aws/aws-sdk-go-v2/aws/protocol/eventstream"
"github.com/aws/aws-sdk-go-v2/aws/protocol/eventstream/eventstreamapi"
"github.com/aws/aws-sdk-go-v2/credentials"
"github.com/aws/aws-sdk-go-v2/service/bedrockruntime"
"github.com/gin-gonic/gin"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
const awsTestModel = "anthropic.claude-3-5-sonnet-20240620-v1:0"
type awsHTTPClientFunc func(*http.Request) (*http.Response, error)
func (f awsHTTPClientFunc) Do(request *http.Request) (*http.Response, error) {
return f(request)
}
type awsNotifyingResponseWriter struct {
*httptest.ResponseRecorder
notifyOn []byte
notified chan int
once sync.Once
}
func newAwsNotifyingResponseWriter(notifyOn string) *awsNotifyingResponseWriter {
return &awsNotifyingResponseWriter{
ResponseRecorder: httptest.NewRecorder(),
notifyOn: []byte(notifyOn),
notified: make(chan int, 1),
}
}
func (w *awsNotifyingResponseWriter) Write(data []byte) (int, error) {
return w.ResponseRecorder.Write(data)
}
func (w *awsNotifyingResponseWriter) Flush() {
w.ResponseRecorder.Flush()
if bytes.Contains(w.Body.Bytes(), w.notifyOn) {
w.once.Do(func() {
w.notified <- w.Body.Len()
})
}
}
func newAwsTestClient(httpClient bedrockruntime.HTTPClient) *bedrockruntime.Client {
return bedrockruntime.New(bedrockruntime.Options{
Region: "us-east-1",
BaseEndpoint: aws.String("https://bedrock.test"),
Credentials: aws.NewCredentialsCache(credentials.NewStaticCredentialsProvider(
"access-key", "secret-key", "",
)),
HTTPClient: httpClient,
Retryer: aws.NopRetryer{},
})
}
func newAwsTestContext(writer http.ResponseWriter, requestContext context.Context) *gin.Context {
c, _ := gin.CreateTestContext(writer)
c.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", nil).WithContext(requestContext)
return c
}
func newAwsTestRelayInfo() *relaycommon.RelayInfo {
return &relaycommon.RelayInfo{
StartTime: time.Now(),
IsStream: true,
OriginModelName: awsTestModel,
RelayFormat: relaytypes.RelayFormatOpenAI,
ShouldIncludeUsage: true,
ChannelMeta: &relaycommon.ChannelMeta{
UpstreamModelName: awsTestModel,
},
}
}
func newAwsInvokeModelInput() *bedrockruntime.InvokeModelInput {
return &bedrockruntime.InvokeModelInput{
ModelId: aws.String(awsTestModel),
Body: []byte(`{}`),
Accept: aws.String("application/json"),
ContentType: aws.String("application/json"),
}
}
func newAwsStreamInput() *bedrockruntime.InvokeModelWithResponseStreamInput {
return &bedrockruntime.InvokeModelWithResponseStreamInput{
ModelId: aws.String(awsTestModel),
Body: []byte(`{}`),
Accept: aws.String("application/json"),
ContentType: aws.String("application/json"),
}
}
func writeAwsStreamEvent(writer io.Writer, data string) error {
payload, err := common.Marshal(struct {
Bytes []byte `json:"bytes"`
}{Bytes: []byte(data)})
if err != nil {
return err
}
return eventstream.NewEncoder().Encode(writer, eventstream.Message{
Headers: eventstream.Headers{
{Name: eventstreamapi.MessageTypeHeader, Value: eventstream.StringValue(eventstreamapi.EventMessageType)},
{Name: eventstreamapi.EventTypeHeader, Value: eventstream.StringValue("chunk")},
{Name: eventstreamapi.ContentTypeHeader, Value: eventstream.StringValue("application/json")},
},
Payload: payload,
})
}
func newAwsStreamResponse(request *http.Request, body io.ReadCloser) *http.Response {
return &http.Response{
StatusCode: http.StatusOK,
Status: "200 OK",
Header: http.Header{
"Content-Type": []string{"application/vnd.amazon.eventstream"},
"X-Amzn-Bedrock-Content-Type": []string{"application/json"},
},
Body: body,
Request: request,
}
}
func TestDoAwsClientRequest_AppliesRuntimeHeaderOverrideToAnthropicBeta(t *testing.T) {
t.Parallel()
gin.SetMode(gin.TestMode)
recorder := httptest.NewRecorder()
ctx, _ := gin.CreateTestContext(recorder)
ctx.Request = httptest.NewRequest(http.MethodPost, "/v1/messages", nil)
info := &relaycommon.RelayInfo{
OriginModelName: "claude-3-5-sonnet-20240620",
IsStream: false,
UseRuntimeHeadersOverride: true,
RuntimeHeadersOverride: map[string]any{
"anthropic-beta": "computer-use-2025-01-24",
},
ChannelMeta: &relaycommon.ChannelMeta{
ApiKey: "access-key|secret-key|us-east-1",
UpstreamModelName: "claude-3-5-sonnet-20240620",
},
}
requestBody := bytes.NewBufferString(`{"messages":[{"role":"user","content":"hello"}],"max_tokens":128}`)
adaptor := &Adaptor{}
_, err := doAwsClientRequest(ctx, info, adaptor, requestBody)
require.NoError(t, err)
awsReq, ok := adaptor.AwsReq.(*bedrockruntime.InvokeModelInput)
require.True(t, ok)
var payload map[string]any
require.NoError(t, common.Unmarshal(awsReq.Body, &payload))
anthropicBeta, exists := payload["anthropic_beta"]
require.True(t, exists)
values, ok := anthropicBeta.([]any)
require.True(t, ok)
require.Equal(t, []any{"computer-use-2025-01-24"}, values)
}
func TestNewAwsInvokeContextInheritsParent(t *testing.T) {
originalRelayTimeout := common.RelayTimeout
t.Cleanup(func() {
common.RelayTimeout = originalRelayTimeout
})
tests := []struct {
name string
relayTimeout int
wantDeadline bool
}{
{name: "without relay timeout", relayTimeout: 0, wantDeadline: false},
{name: "with relay timeout", relayTimeout: 30, wantDeadline: true},
}
for _, test := range tests {
t.Run(test.name, func(t *testing.T) {
common.RelayTimeout = test.relayTimeout
parent, cancelParent := context.WithCancel(context.Background())
invokeContext, cancelInvoke := newAwsInvokeContext(parent)
defer cancelInvoke()
_, hasDeadline := invokeContext.Deadline()
assert.Equal(t, test.wantDeadline, hasDeadline)
cancelParent()
require.ErrorIs(t, invokeContext.Err(), context.Canceled)
})
}
}
func TestNewAwsInvokeErrorSkipsRetryOnlyForClientCancellation(t *testing.T) {
canceledContext, cancel := context.WithCancel(context.Background())
cancel()
tests := []struct {
name string
requestContext context.Context
err error
wantSkipRetry bool
}{
{
name: "client context canceled",
requestContext: canceledContext,
err: context.Canceled,
wantSkipRetry: true,
},
{
name: "relay timeout with live client context",
requestContext: context.Background(),
err: context.DeadlineExceeded,
wantSkipRetry: false,
},
{
name: "upstream error with live client context",
requestContext: context.Background(),
err: errors.New("upstream failed"),
wantSkipRetry: false,
},
}
for _, test := range tests {
t.Run(test.name, func(t *testing.T) {
err := newAwsInvokeError(test.requestContext, test.err, "InvokeModel")
assert.Equal(t, test.wantSkipRetry, relaytypes.IsSkipRetryError(err))
})
}
}
func TestAwsHandlersCancelSdkRequestAndSkipRetry(t *testing.T) {
originalRelayTimeout := common.RelayTimeout
common.RelayTimeout = 0
t.Cleanup(func() {
common.RelayTimeout = originalRelayTimeout
})
tests := []struct {
name string
request any
handle func(*gin.Context, *relaycommon.RelayInfo, *Adaptor) (*relaytypes.NewAPIError, *dto.Usage)
}{
{name: "non-stream", request: newAwsInvokeModelInput(), handle: awsHandler},
{name: "stream", request: newAwsStreamInput(), handle: awsStreamHandler},
{name: "nova", request: newAwsInvokeModelInput(), handle: handleNovaRequest},
}
for _, test := range tests {
t.Run(test.name, func(t *testing.T) {
requestContext, cancelRequest := context.WithCancel(context.Background())
t.Cleanup(cancelRequest)
upstreamContexts := make(chan context.Context, 1)
client := newAwsTestClient(awsHTTPClientFunc(func(request *http.Request) (*http.Response, error) {
upstreamContexts <- request.Context()
<-request.Context().Done()
return nil, request.Context().Err()
}))
adaptor := &Adaptor{AwsClient: client, AwsReq: test.request}
c := newAwsTestContext(httptest.NewRecorder(), requestContext)
info := newAwsTestRelayInfo()
type handlerResult struct {
err *relaytypes.NewAPIError
usage *dto.Usage
}
results := make(chan handlerResult, 1)
go func() {
err, usage := test.handle(c, info, adaptor)
results <- handlerResult{err: err, usage: usage}
}()
var upstreamContext context.Context
select {
case upstreamContext = <-upstreamContexts:
case result := <-results:
t.Fatalf("handler returned before issuing AWS request: %v", result.err)
case <-time.After(5 * time.Second):
t.Fatal("AWS request did not start")
}
cancelRequest()
var result handlerResult
select {
case result = <-results:
case <-time.After(5 * time.Second):
t.Fatal("handler did not stop after client cancellation")
}
require.ErrorIs(t, upstreamContext.Err(), context.Canceled)
require.NotNil(t, result.err)
assert.True(t, relaytypes.IsSkipRetryError(result.err))
assert.Nil(t, result.usage)
})
}
}
func TestAwsStreamHandlerUsesFinalUpstreamUsage(t *testing.T) {
originalRelayTimeout := common.RelayTimeout
common.RelayTimeout = 0
t.Cleanup(func() {
common.RelayTimeout = originalRelayTimeout
})
events := []string{
`{"type":"message_start","message":{"id":"msg_test","type":"message","role":"assistant","model":"claude-test","content":[],"usage":{"input_tokens":100,"output_tokens":1}}}`,
`{"type":"content_block_start","index":0,"content_block":{"type":"text","text":""}}`,
`{"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"partial"}}`,
`{"type":"message_delta","delta":{"stop_reason":"end_turn"},"usage":{"output_tokens":423}}`,
`{"type":"message_stop"}`,
}
client := newAwsTestClient(awsHTTPClientFunc(func(request *http.Request) (*http.Response, error) {
var body bytes.Buffer
for _, event := range events {
if err := writeAwsStreamEvent(&body, event); err != nil {
return nil, err
}
}
return newAwsStreamResponse(request, io.NopCloser(bytes.NewReader(body.Bytes()))), nil
}))
adaptor := &Adaptor{AwsClient: client, AwsReq: newAwsStreamInput()}
recorder := httptest.NewRecorder()
c := newAwsTestContext(recorder, context.Background())
handlerErr, usage := awsStreamHandler(c, newAwsTestRelayInfo(), adaptor)
require.Nil(t, handlerErr)
require.NotNil(t, usage)
require.NotNil(t, usage.BillingUsage)
require.NotNil(t, usage.BillingUsage.ClaudeUsage)
assert.Equal(t, 100, usage.BillingUsage.ClaudeUsage.InputTokens)
assert.Equal(t, 423, usage.BillingUsage.ClaudeUsage.OutputTokens)
assert.Contains(t, recorder.Body.String(), "[DONE]")
}
func TestAwsStreamHandlerStopsAtClientCancellation(t *testing.T) {
originalRelayTimeout := common.RelayTimeout
common.RelayTimeout = 0
t.Cleanup(func() {
common.RelayTimeout = originalRelayTimeout
})
requestContext, cancelRequest := context.WithCancel(context.Background())
t.Cleanup(cancelRequest)
releaseFinal := make(chan struct{})
var releaseFinalOnce sync.Once
release := func() {
releaseFinalOnce.Do(func() {
close(releaseFinal)
})
}
t.Cleanup(release)
producerResults := make(chan error, 1)
upstreamContexts := make(chan context.Context, 1)
client := newAwsTestClient(awsHTTPClientFunc(func(request *http.Request) (*http.Response, error) {
upstreamContexts <- request.Context()
reader, writer := io.Pipe()
go func() {
defer writer.Close()
initialEvents := []string{
`{"type":"message_start","message":{"id":"msg_test","type":"message","role":"assistant","model":"claude-test","content":[],"usage":{"input_tokens":100,"output_tokens":1}}}`,
`{"type":"content_block_start","index":0,"content_block":{"type":"text","text":""}}`,
`{"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"partial"}}`,
}
for _, event := range initialEvents {
if err := writeAwsStreamEvent(writer, event); err != nil {
producerResults <- err
return
}
}
<-releaseFinal
producerResults <- writeAwsStreamEvent(writer, `{"type":"message_delta","delta":{"stop_reason":"end_turn"},"usage":{"output_tokens":423}}`)
}()
return newAwsStreamResponse(request, reader), nil
}))
responseWriter := newAwsNotifyingResponseWriter("partial")
c := newAwsTestContext(responseWriter, requestContext)
adaptor := &Adaptor{AwsClient: client, AwsReq: newAwsStreamInput()}
type handlerResult struct {
err *relaytypes.NewAPIError
usage *dto.Usage
}
results := make(chan handlerResult, 1)
go func() {
err, usage := awsStreamHandler(c, newAwsTestRelayInfo(), adaptor)
results <- handlerResult{err: err, usage: usage}
}()
var upstreamContext context.Context
select {
case upstreamContext = <-upstreamContexts:
case <-time.After(5 * time.Second):
t.Fatal("AWS stream request did not start")
}
var bodyLengthBeforeCancel int
select {
case bodyLengthBeforeCancel = <-responseWriter.notified:
case <-time.After(5 * time.Second):
t.Fatal("partial response was not written")
}
cancelRequest()
var result handlerResult
select {
case result = <-results:
case <-time.After(5 * time.Second):
t.Fatal("stream handler did not stop after client cancellation")
}
require.ErrorIs(t, upstreamContext.Err(), context.Canceled)
require.Nil(t, result.err)
require.NotNil(t, result.usage)
assert.Equal(t, bodyLengthBeforeCancel, responseWriter.Body.Len())
assert.NotContains(t, responseWriter.Body.String(), "[DONE]")
release()
select {
case producerErr := <-producerResults:
require.Error(t, producerErr)
case <-time.After(5 * time.Second):
t.Fatal("upstream producer did not observe the closed stream")
}
}
+169
View File
@@ -0,0 +1,169 @@
package baidu
import (
"errors"
"fmt"
"io"
"net/http"
"strings"
"github.com/QuantumNous/new-api/relay/channel"
relaycommon "github.com/QuantumNous/new-api/relay/common"
"github.com/QuantumNous/new-api/relay/constant"
"github.com/QuantumNous/new-api/relaykit/dto"
"github.com/QuantumNous/new-api/relaykit/types"
"github.com/gin-gonic/gin"
)
type Adaptor struct {
}
func (a *Adaptor) ConvertGeminiRequest(*gin.Context, *relaycommon.RelayInfo, *dto.GeminiChatRequest) (any, error) {
//TODO implement me
return nil, errors.New("not implemented")
}
func (a *Adaptor) ConvertClaudeRequest(*gin.Context, *relaycommon.RelayInfo, *dto.ClaudeRequest) (any, error) {
//TODO implement me
panic("implement me")
}
func (a *Adaptor) ConvertAudioRequest(c *gin.Context, info *relaycommon.RelayInfo, request dto.AudioRequest) (io.Reader, error) {
//TODO implement me
return nil, errors.New("not implemented")
}
func (a *Adaptor) ConvertImageRequest(c *gin.Context, info *relaycommon.RelayInfo, request dto.ImageRequest) (any, error) {
//TODO implement me
return nil, errors.New("not implemented")
}
func (a *Adaptor) Init(info *relaycommon.RelayInfo) {
}
func (a *Adaptor) GetRequestURL(info *relaycommon.RelayInfo) (string, error) {
// https://cloud.baidu.com/doc/WENXINWORKSHOP/s/clntwmv7t
suffix := "chat/"
if strings.HasPrefix(info.UpstreamModelName, "Embedding") {
suffix = "embeddings/"
}
if strings.HasPrefix(info.UpstreamModelName, "bge-large") {
suffix = "embeddings/"
}
if strings.HasPrefix(info.UpstreamModelName, "tao-8k") {
suffix = "embeddings/"
}
switch info.UpstreamModelName {
case "ERNIE-4.0":
suffix += "completions_pro"
case "ERNIE-Bot-4":
suffix += "completions_pro"
case "ERNIE-Bot":
suffix += "completions"
case "ERNIE-Bot-turbo":
suffix += "eb-instant"
case "ERNIE-Speed":
suffix += "ernie_speed"
case "ERNIE-4.0-8K":
suffix += "completions_pro"
case "ERNIE-3.5-8K":
suffix += "completions"
case "ERNIE-3.5-8K-0205":
suffix += "ernie-3.5-8k-0205"
case "ERNIE-3.5-8K-1222":
suffix += "ernie-3.5-8k-1222"
case "ERNIE-Bot-8K":
suffix += "ernie_bot_8k"
case "ERNIE-3.5-4K-0205":
suffix += "ernie-3.5-4k-0205"
case "ERNIE-Speed-8K":
suffix += "ernie_speed"
case "ERNIE-Speed-128K":
suffix += "ernie-speed-128k"
case "ERNIE-Lite-8K-0922":
suffix += "eb-instant"
case "ERNIE-Lite-8K-0308":
suffix += "ernie-lite-8k"
case "ERNIE-Tiny-8K":
suffix += "ernie-tiny-8k"
case "BLOOMZ-7B":
suffix += "bloomz_7b1"
case "Embedding-V1":
suffix += "embedding-v1"
case "bge-large-zh":
suffix += "bge_large_zh"
case "bge-large-en":
suffix += "bge_large_en"
case "tao-8k":
suffix += "tao_8k"
default:
suffix += strings.ToLower(info.UpstreamModelName)
}
fullRequestURL := fmt.Sprintf("%s/rpc/2.0/ai_custom/v1/wenxinworkshop/%s", info.ChannelBaseUrl, suffix)
var accessToken string
var err error
if accessToken, err = getBaiduAccessToken(info.ApiKey); err != nil {
return "", err
}
fullRequestURL += "?access_token=" + accessToken
return fullRequestURL, nil
}
func (a *Adaptor) SetupRequestHeader(c *gin.Context, req *http.Header, info *relaycommon.RelayInfo) error {
channel.SetupApiRequestHeader(info, c, req)
req.Set("Authorization", "Bearer "+info.ApiKey)
return nil
}
func (a *Adaptor) ConvertOpenAIRequest(c *gin.Context, info *relaycommon.RelayInfo, request *dto.GeneralOpenAIRequest) (any, error) {
if request == nil {
return nil, errors.New("request is nil")
}
switch info.RelayMode {
default:
baiduRequest := requestOpenAI2Baidu(*request)
return baiduRequest, nil
}
}
func (a *Adaptor) ConvertRerankRequest(c *gin.Context, relayMode int, request dto.RerankRequest) (any, error) {
return nil, nil
}
func (a *Adaptor) ConvertEmbeddingRequest(c *gin.Context, info *relaycommon.RelayInfo, request dto.EmbeddingRequest) (any, error) {
baiduEmbeddingRequest := embeddingRequestOpenAI2Baidu(request)
return baiduEmbeddingRequest, nil
}
func (a *Adaptor) ConvertOpenAIResponsesRequest(c *gin.Context, info *relaycommon.RelayInfo, request dto.OpenAIResponsesRequest) (any, error) {
// TODO implement me
return nil, errors.New("not implemented")
}
func (a *Adaptor) DoRequest(c *gin.Context, info *relaycommon.RelayInfo, requestBody io.Reader) (any, error) {
return channel.DoApiRequest(a, c, info, requestBody)
}
func (a *Adaptor) DoResponse(c *gin.Context, resp *http.Response, info *relaycommon.RelayInfo) (usage any, err *types.NewAPIError) {
if info.IsStream {
err, usage = baiduStreamHandler(c, info, resp)
} else {
switch info.RelayMode {
case constant.RelayModeEmbeddings:
err, usage = baiduEmbeddingHandler(c, info, resp)
default:
err, usage = baiduHandler(c, info, resp)
}
}
return
}
func (a *Adaptor) GetModelList() []string {
return ModelList
}
func (a *Adaptor) GetChannelName() string {
return ChannelName
}
+22
View File
@@ -0,0 +1,22 @@
package baidu
var ModelList = []string{
"ERNIE-4.0-8K",
"ERNIE-3.5-8K",
"ERNIE-3.5-8K-0205",
"ERNIE-3.5-8K-1222",
"ERNIE-Bot-8K",
"ERNIE-3.5-4K-0205",
"ERNIE-Speed-8K",
"ERNIE-Speed-128K",
"ERNIE-Lite-8K-0922",
"ERNIE-Lite-8K-0308",
"ERNIE-Tiny-8K",
"BLOOMZ-7B",
"Embedding-V1",
"bge-large-zh",
"bge-large-en",
"tao-8k",
}
var ChannelName = "baidu"
+80
View File
@@ -0,0 +1,80 @@
package baidu
import (
"encoding/json"
"time"
"github.com/QuantumNous/new-api/relaykit/dto"
)
type BaiduMessage struct {
Role string `json:"role"`
Content string `json:"content"`
}
type BaiduChatRequest struct {
Messages []BaiduMessage `json:"messages"`
Temperature *float64 `json:"temperature,omitempty"`
TopP float64 `json:"top_p,omitempty"`
PenaltyScore float64 `json:"penalty_score,omitempty"`
Stream bool `json:"stream,omitempty"`
System string `json:"system,omitempty"`
DisableSearch bool `json:"disable_search,omitempty"`
EnableCitation bool `json:"enable_citation,omitempty"`
MaxOutputTokens *int `json:"max_output_tokens,omitempty"`
UserId json.RawMessage `json:"user_id,omitempty"`
}
type Error struct {
ErrorCode int `json:"error_code"`
ErrorMsg string `json:"error_msg"`
}
type BaiduChatResponse struct {
Id string `json:"id"`
Object string `json:"object"`
Created int64 `json:"created"`
Result string `json:"result"`
IsTruncated bool `json:"is_truncated"`
NeedClearHistory bool `json:"need_clear_history"`
Usage dto.Usage `json:"usage"`
Error
}
type BaiduChatStreamResponse struct {
BaiduChatResponse
SentenceId int `json:"sentence_id"`
IsEnd bool `json:"is_end"`
}
type BaiduEmbeddingRequest struct {
Input []string `json:"input"`
}
type BaiduEmbeddingData struct {
Object string `json:"object"`
Embedding []float64 `json:"embedding"`
Index int `json:"index"`
}
type BaiduEmbeddingResponse struct {
Id string `json:"id"`
Object string `json:"object"`
Created int64 `json:"created"`
Data []BaiduEmbeddingData `json:"data"`
Usage dto.Usage `json:"usage"`
Error
}
type BaiduAccessToken struct {
AccessToken string `json:"access_token"`
Error string `json:"error,omitempty"`
ErrorDescription string `json:"error_description,omitempty"`
ExpiresIn int64 `json:"expires_in,omitempty"`
ExpiresAt time.Time `json:"-"`
}
type BaiduTokenResponse struct {
ExpiresIn int `json:"expires_in"`
AccessToken string `json:"access_token"`
}
+246
View File
@@ -0,0 +1,246 @@
package baidu
import (
"encoding/json"
"errors"
"fmt"
"io"
"net/http"
"strings"
"sync"
"time"
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/constant"
relaycommon "github.com/QuantumNous/new-api/relay/common"
"github.com/QuantumNous/new-api/relay/helper"
"github.com/QuantumNous/new-api/relaykit/dto"
"github.com/QuantumNous/new-api/relaykit/types"
"github.com/QuantumNous/new-api/service"
"github.com/samber/lo"
"github.com/gin-gonic/gin"
)
// https://cloud.baidu.com/doc/WENXINWORKSHOP/s/flfmc9do2
var baiduTokenStore sync.Map
func requestOpenAI2Baidu(request dto.GeneralOpenAIRequest) *BaiduChatRequest {
baiduRequest := BaiduChatRequest{
Temperature: request.Temperature,
TopP: lo.FromPtrOr(request.TopP, 0),
PenaltyScore: lo.FromPtrOr(request.FrequencyPenalty, 0),
Stream: lo.FromPtrOr(request.Stream, false),
DisableSearch: false,
EnableCitation: false,
UserId: request.User,
}
if request.GetMaxTokens() != 0 {
maxTokens := int(request.GetMaxTokens())
if request.GetMaxTokens() == 1 {
maxTokens = 2
}
baiduRequest.MaxOutputTokens = &maxTokens
}
for _, message := range request.Messages {
if message.Role == "system" {
baiduRequest.System = message.StringContent()
} else {
baiduRequest.Messages = append(baiduRequest.Messages, BaiduMessage{
Role: message.Role,
Content: message.StringContent(),
})
}
}
return &baiduRequest
}
func responseBaidu2OpenAI(response *BaiduChatResponse) *dto.OpenAITextResponse {
choice := dto.OpenAITextResponseChoice{
Index: 0,
Message: dto.Message{
Role: "assistant",
Content: response.Result,
},
FinishReason: "stop",
}
fullTextResponse := dto.OpenAITextResponse{
Id: response.Id,
Object: "chat.completion",
Created: response.Created,
Choices: []dto.OpenAITextResponseChoice{choice},
Usage: response.Usage,
}
return &fullTextResponse
}
func streamResponseBaidu2OpenAI(baiduResponse *BaiduChatStreamResponse) *dto.ChatCompletionsStreamResponse {
var choice dto.ChatCompletionsStreamResponseChoice
choice.Delta.SetContentString(baiduResponse.Result)
if baiduResponse.IsEnd {
choice.FinishReason = &constant.FinishReasonStop
}
response := dto.ChatCompletionsStreamResponse{
Id: baiduResponse.Id,
Object: "chat.completion.chunk",
Created: baiduResponse.Created,
Model: "ernie-bot",
Choices: []dto.ChatCompletionsStreamResponseChoice{choice},
}
return &response
}
func embeddingRequestOpenAI2Baidu(request dto.EmbeddingRequest) *BaiduEmbeddingRequest {
return &BaiduEmbeddingRequest{
Input: request.ParseInput(),
}
}
func embeddingResponseBaidu2OpenAI(response *BaiduEmbeddingResponse) *dto.OpenAIEmbeddingResponse {
openAIEmbeddingResponse := dto.OpenAIEmbeddingResponse{
Object: "list",
Data: make([]dto.OpenAIEmbeddingResponseItem, 0, len(response.Data)),
Model: "baidu-embedding",
Usage: response.Usage,
}
for _, item := range response.Data {
openAIEmbeddingResponse.Data = append(openAIEmbeddingResponse.Data, dto.OpenAIEmbeddingResponseItem{
Object: item.Object,
Index: item.Index,
Embedding: item.Embedding,
})
}
return &openAIEmbeddingResponse
}
func baiduStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http.Response) (*types.NewAPIError, *dto.Usage) {
usage := &dto.Usage{}
helper.StreamScannerHandler(c, resp, info, func(data string, sr *helper.StreamResult) {
var baiduResponse BaiduChatStreamResponse
if err := common.Unmarshal([]byte(data), &baiduResponse); err != nil {
common.SysLog("error unmarshalling stream response: " + err.Error())
sr.Error(err)
return
}
if baiduResponse.Usage.TotalTokens != 0 {
usage.TotalTokens = baiduResponse.Usage.TotalTokens
usage.PromptTokens = baiduResponse.Usage.PromptTokens
usage.CompletionTokens = baiduResponse.Usage.TotalTokens - baiduResponse.Usage.PromptTokens
}
response := streamResponseBaidu2OpenAI(&baiduResponse)
if err := helper.ObjectData(c, response); err != nil {
common.SysLog("error sending stream response: " + err.Error())
sr.Error(err)
}
})
service.CloseResponseBodyGracefully(resp)
return nil, usage
}
func baiduHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http.Response) (*types.NewAPIError, *dto.Usage) {
var baiduResponse BaiduChatResponse
responseBody, err := io.ReadAll(resp.Body)
if err != nil {
return types.NewError(err, types.ErrorCodeBadResponseBody), nil
}
service.CloseResponseBodyGracefully(resp)
err = json.Unmarshal(responseBody, &baiduResponse)
if err != nil {
return types.NewError(err, types.ErrorCodeBadResponseBody), nil
}
if baiduResponse.ErrorMsg != "" {
return types.NewError(fmt.Errorf("%s", baiduResponse.ErrorMsg), types.ErrorCodeBadResponseBody), nil
}
fullTextResponse := responseBaidu2OpenAI(&baiduResponse)
jsonResponse, err := json.Marshal(fullTextResponse)
if err != nil {
return types.NewError(err, types.ErrorCodeBadResponseBody), nil
}
c.Writer.Header().Set("Content-Type", "application/json")
c.Writer.WriteHeader(resp.StatusCode)
_, err = c.Writer.Write(jsonResponse)
return nil, &fullTextResponse.Usage
}
func baiduEmbeddingHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http.Response) (*types.NewAPIError, *dto.Usage) {
var baiduResponse BaiduEmbeddingResponse
responseBody, err := io.ReadAll(resp.Body)
if err != nil {
return types.NewError(err, types.ErrorCodeBadResponseBody), nil
}
service.CloseResponseBodyGracefully(resp)
err = json.Unmarshal(responseBody, &baiduResponse)
if err != nil {
return types.NewError(err, types.ErrorCodeBadResponseBody), nil
}
if baiduResponse.ErrorMsg != "" {
return types.NewError(fmt.Errorf("%s", baiduResponse.ErrorMsg), types.ErrorCodeBadResponseBody), nil
}
fullTextResponse := embeddingResponseBaidu2OpenAI(&baiduResponse)
jsonResponse, err := json.Marshal(fullTextResponse)
if err != nil {
return types.NewError(err, types.ErrorCodeBadResponseBody), nil
}
c.Writer.Header().Set("Content-Type", "application/json")
c.Writer.WriteHeader(resp.StatusCode)
_, err = c.Writer.Write(jsonResponse)
return nil, &fullTextResponse.Usage
}
func getBaiduAccessToken(apiKey string) (string, error) {
if val, ok := baiduTokenStore.Load(apiKey); ok {
var accessToken BaiduAccessToken
if accessToken, ok = val.(BaiduAccessToken); ok {
// soon this will expire
if time.Now().Add(time.Hour).After(accessToken.ExpiresAt) {
go func() {
_, _ = getBaiduAccessTokenHelper(apiKey)
}()
}
return accessToken.AccessToken, nil
}
}
accessToken, err := getBaiduAccessTokenHelper(apiKey)
if err != nil {
return "", err
}
if accessToken == nil {
return "", errors.New("getBaiduAccessToken return a nil token")
}
return (*accessToken).AccessToken, nil
}
func getBaiduAccessTokenHelper(apiKey string) (*BaiduAccessToken, error) {
parts := strings.Split(apiKey, "|")
if len(parts) != 2 {
return nil, errors.New("invalid baidu apikey")
}
req, err := http.NewRequest("POST", fmt.Sprintf("https://aip.baidubce.com/oauth/2.0/token?grant_type=client_credentials&client_id=%s&client_secret=%s",
parts[0], parts[1]), nil)
if err != nil {
return nil, err
}
req.Header.Add("Content-Type", "application/json")
req.Header.Add("Accept", "application/json")
res, err := service.GetHttpClient().Do(req)
if err != nil {
return nil, err
}
defer res.Body.Close()
var accessToken BaiduAccessToken
err = json.NewDecoder(res.Body).Decode(&accessToken)
if err != nil {
return nil, err
}
if accessToken.Error != "" {
return nil, errors.New(accessToken.Error + ": " + accessToken.ErrorDescription)
}
if accessToken.AccessToken == "" {
return nil, errors.New("getBaiduAccessTokenHelper get empty access token")
}
accessToken.ExpiresAt = time.Now().Add(time.Duration(accessToken.ExpiresIn) * time.Second)
baiduTokenStore.Store(apiKey, accessToken)
return &accessToken, nil
}
+130
View File
@@ -0,0 +1,130 @@
package baidu_v2
import (
"errors"
"fmt"
"io"
"net/http"
"strings"
"github.com/QuantumNous/new-api/relay/channel"
"github.com/QuantumNous/new-api/relay/channel/openai"
relaycommon "github.com/QuantumNous/new-api/relay/common"
"github.com/QuantumNous/new-api/relay/constant"
"github.com/QuantumNous/new-api/relaykit/dto"
"github.com/QuantumNous/new-api/relaykit/types"
"github.com/gin-gonic/gin"
)
type Adaptor struct {
}
func (a *Adaptor) ConvertGeminiRequest(*gin.Context, *relaycommon.RelayInfo, *dto.GeminiChatRequest) (any, error) {
//TODO implement me
return nil, errors.New("not implemented")
}
func (a *Adaptor) ConvertClaudeRequest(c *gin.Context, info *relaycommon.RelayInfo, req *dto.ClaudeRequest) (any, error) {
adaptor := openai.Adaptor{}
return adaptor.ConvertClaudeRequest(c, info, req)
}
func (a *Adaptor) ConvertAudioRequest(c *gin.Context, info *relaycommon.RelayInfo, request dto.AudioRequest) (io.Reader, error) {
//TODO implement me
return nil, errors.New("not implemented")
}
func (a *Adaptor) ConvertImageRequest(c *gin.Context, info *relaycommon.RelayInfo, request dto.ImageRequest) (any, error) {
//TODO implement me
return nil, errors.New("not implemented")
}
func (a *Adaptor) Init(info *relaycommon.RelayInfo) {
}
func (a *Adaptor) GetRequestURL(info *relaycommon.RelayInfo) (string, error) {
switch info.RelayMode {
case constant.RelayModeChatCompletions:
return fmt.Sprintf("%s/v2/chat/completions", info.ChannelBaseUrl), nil
case constant.RelayModeEmbeddings:
return fmt.Sprintf("%s/v2/embeddings", info.ChannelBaseUrl), nil
case constant.RelayModeImagesGenerations:
return fmt.Sprintf("%s/v2/images/generations", info.ChannelBaseUrl), nil
case constant.RelayModeImagesEdits:
return fmt.Sprintf("%s/v2/images/edits", info.ChannelBaseUrl), nil
case constant.RelayModeRerank:
return fmt.Sprintf("%s/v2/rerank", info.ChannelBaseUrl), nil
default:
}
return "", fmt.Errorf("unsupported relay mode: %d", info.RelayMode)
}
func (a *Adaptor) SetupRequestHeader(c *gin.Context, req *http.Header, info *relaycommon.RelayInfo) error {
channel.SetupApiRequestHeader(info, c, req)
keyParts := strings.Split(info.ApiKey, "|")
if len(keyParts) == 0 || keyParts[0] == "" {
return errors.New("invalid API key: authorization token is required")
}
if len(keyParts) > 1 {
if keyParts[1] != "" {
req.Set("appid", keyParts[1])
}
}
req.Set("Authorization", "Bearer "+keyParts[0])
return nil
}
func (a *Adaptor) ConvertOpenAIRequest(c *gin.Context, info *relaycommon.RelayInfo, request *dto.GeneralOpenAIRequest) (any, error) {
if request == nil {
return nil, errors.New("request is nil")
}
if strings.HasSuffix(info.UpstreamModelName, "-search") {
info.UpstreamModelName = strings.TrimSuffix(info.UpstreamModelName, "-search")
request.Model = info.UpstreamModelName
if len(request.WebSearch) == 0 {
toMap := request.ToMap()
toMap["web_search"] = map[string]any{
"enable": true,
"enable_citation": true,
"enable_trace": true,
"enable_status": false,
}
return toMap, nil
}
return request, nil
}
return request, nil
}
func (a *Adaptor) ConvertRerankRequest(c *gin.Context, relayMode int, request dto.RerankRequest) (any, error) {
return nil, errors.New("not implemented")
}
func (a *Adaptor) ConvertEmbeddingRequest(c *gin.Context, info *relaycommon.RelayInfo, request dto.EmbeddingRequest) (any, error) {
//TODO implement me
return nil, errors.New("not implemented")
}
func (a *Adaptor) ConvertOpenAIResponsesRequest(c *gin.Context, info *relaycommon.RelayInfo, request dto.OpenAIResponsesRequest) (any, error) {
// TODO implement me
return nil, errors.New("not implemented")
}
func (a *Adaptor) DoRequest(c *gin.Context, info *relaycommon.RelayInfo, requestBody io.Reader) (any, error) {
return channel.DoApiRequest(a, c, info, requestBody)
}
func (a *Adaptor) DoResponse(c *gin.Context, resp *http.Response, info *relaycommon.RelayInfo) (usage any, err *types.NewAPIError) {
adaptor := openai.Adaptor{}
usage, err = adaptor.DoResponse(c, resp, info)
return
}
func (a *Adaptor) GetModelList() []string {
return ModelList
}
func (a *Adaptor) GetChannelName() string {
return ChannelName
}
+29
View File
@@ -0,0 +1,29 @@
package baidu_v2
var ModelList = []string{
"ernie-4.0-8k-latest",
"ernie-4.0-8k-preview",
"ernie-4.0-8k",
"ernie-4.0-turbo-8k-latest",
"ernie-4.0-turbo-8k-preview",
"ernie-4.0-turbo-8k",
"ernie-4.0-turbo-128k",
"ernie-3.5-8k-preview",
"ernie-3.5-8k",
"ernie-3.5-128k",
"ernie-speed-8k",
"ernie-speed-128k",
"ernie-speed-pro-128k",
"ernie-lite-8k",
"ernie-lite-pro-128k",
"ernie-tiny-8k",
"ernie-char-8k",
"ernie-char-fiction-8k",
"ernie-novel-8k",
"deepseek-v3",
"deepseek-r1",
"deepseek-r1-distill-qwen-32b",
"deepseek-r1-distill-qwen-14b",
}
var ChannelName = "volcengine"
+172
View File
@@ -0,0 +1,172 @@
package claude
import (
"errors"
"fmt"
"io"
"net/http"
"net/url"
"github.com/QuantumNous/new-api/relay/channel"
relaycommon "github.com/QuantumNous/new-api/relay/common"
"github.com/QuantumNous/new-api/relaykit/dto"
"github.com/QuantumNous/new-api/relaykit/relayconvert"
"github.com/QuantumNous/new-api/relaykit/types"
"github.com/QuantumNous/new-api/service"
"github.com/QuantumNous/new-api/setting/model_setting"
"github.com/gin-gonic/gin"
)
type Adaptor struct {
}
func (a *Adaptor) ConvertGeminiRequest(c *gin.Context, info *relaycommon.RelayInfo, request *dto.GeminiChatRequest) (any, error) {
if request == nil {
return nil, errors.New("request is nil")
}
result, err := service.ConvertRequest(c, info, types.RelayFormatClaude, request)
if err != nil {
return nil, err
}
return result.Value, nil
}
func (a *Adaptor) ConvertClaudeRequest(c *gin.Context, info *relaycommon.RelayInfo, request *dto.ClaudeRequest) (any, error) {
if request.MaxTokens != nil && *request.MaxTokens == 0 {
request.MaxTokens = nil
}
if err := relayconvert.ApplyClaudeThinkingModel(request, info); err != nil {
return nil, err
}
if request.MaxTokens == nil {
defaultMaxTokens := uint(model_setting.GetClaudeSettings().GetDefaultMaxTokens(request.Model))
request.MaxTokens = &defaultMaxTokens
}
// ApplyClaudeThinkingModel no longer rewrites request.Model. Do not write
// a still-suffixed name back over the entry-normalized UpstreamModelName
// (AWS/Vertex look up getAwsModelID / claudeModelMap from that field).
if info.UpstreamModelName == "" {
info.UpstreamModelName = request.Model
}
return request, nil
}
func (a *Adaptor) ConvertAudioRequest(c *gin.Context, info *relaycommon.RelayInfo, request dto.AudioRequest) (io.Reader, error) {
//TODO implement me
return nil, errors.New("not implemented")
}
func (a *Adaptor) ConvertImageRequest(c *gin.Context, info *relaycommon.RelayInfo, request dto.ImageRequest) (any, error) {
//TODO implement me
return nil, errors.New("not implemented")
}
func (a *Adaptor) Init(info *relaycommon.RelayInfo) {
}
func (a *Adaptor) GetRequestURL(info *relaycommon.RelayInfo) (string, error) {
requestURL := fmt.Sprintf("%s/v1/messages", info.ChannelBaseUrl)
if !shouldAppendClaudeBetaQuery(info) {
return requestURL, nil
}
parsedURL, err := url.Parse(requestURL)
if err != nil {
return "", err
}
query := parsedURL.Query()
query.Set("beta", "true")
parsedURL.RawQuery = query.Encode()
return parsedURL.String(), nil
}
func shouldAppendClaudeBetaQuery(info *relaycommon.RelayInfo) bool {
if info == nil {
return false
}
if info.IsClaudeBetaQuery {
return true
}
if info.ChannelOtherSettings.ClaudeBetaQuery {
return true
}
return false
}
func CommonClaudeHeadersOperation(c *gin.Context, req *http.Header, info *relaycommon.RelayInfo) {
// common headers operation
anthropicBeta := c.Request.Header.Get("anthropic-beta")
if anthropicBeta != "" {
req.Set("anthropic-beta", anthropicBeta)
}
model_setting.GetClaudeSettings().WriteHeaders(info.OriginModelName, req)
}
func (a *Adaptor) SetupRequestHeader(c *gin.Context, req *http.Header, info *relaycommon.RelayInfo) error {
channel.SetupApiRequestHeader(info, c, req)
req.Set("x-api-key", info.ApiKey)
anthropicVersion := c.Request.Header.Get("anthropic-version")
if anthropicVersion == "" {
anthropicVersion = "2023-06-01"
}
req.Set("anthropic-version", anthropicVersion)
CommonClaudeHeadersOperation(c, req, info)
return nil
}
func (a *Adaptor) ConvertOpenAIRequest(c *gin.Context, info *relaycommon.RelayInfo, request *dto.GeneralOpenAIRequest) (any, error) {
if request == nil {
return nil, errors.New("request is nil")
}
result, err := service.ConvertRequest(c, info, types.RelayFormatClaude, request)
if err != nil {
return nil, err
}
return result.Value, nil
}
func (a *Adaptor) ConvertRerankRequest(c *gin.Context, relayMode int, request dto.RerankRequest) (any, error) {
return nil, nil
}
func (a *Adaptor) ConvertEmbeddingRequest(c *gin.Context, info *relaycommon.RelayInfo, request dto.EmbeddingRequest) (any, error) {
//TODO implement me
return nil, errors.New("not implemented")
}
func (a *Adaptor) ConvertOpenAIResponsesRequest(c *gin.Context, info *relaycommon.RelayInfo, request dto.OpenAIResponsesRequest) (any, error) {
result, err := service.ConvertRequest(c, info, types.RelayFormatClaude, &request)
if err != nil {
return nil, err
}
claudeRequest, ok := result.Value.(*dto.ClaudeRequest)
if !ok {
return nil, fmt.Errorf("expected Anthropic Messages request, got %T", result.Value)
}
return claudeRequest, nil
}
func (a *Adaptor) DoRequest(c *gin.Context, info *relaycommon.RelayInfo, requestBody io.Reader) (any, error) {
return channel.DoApiRequest(a, c, info, requestBody)
}
func (a *Adaptor) DoResponse(c *gin.Context, resp *http.Response, info *relaycommon.RelayInfo) (usage any, err *types.NewAPIError) {
info.FinalRequestRelayFormat = types.RelayFormatClaude
if info.RelayFormat == types.RelayFormatOpenAIResponses && info.IsStream {
return ClaudeResponsesStreamHandler(c, resp, info)
}
if info.IsStream {
return ClaudeStreamHandler(c, resp, info)
} else {
return ClaudeHandler(c, resp, info)
}
}
func (a *Adaptor) GetModelList() []string {
return ModelList
}
func (a *Adaptor) GetChannelName() string {
return ChannelName
}
+225
View File
@@ -0,0 +1,225 @@
package claude
import (
"net/http/httptest"
"testing"
"github.com/QuantumNous/new-api/common"
relaycommon "github.com/QuantumNous/new-api/relay/common"
"github.com/QuantumNous/new-api/relay/helper"
"github.com/QuantumNous/new-api/relaykit/dto"
"github.com/QuantumNous/new-api/setting/model_setting"
"github.com/gin-gonic/gin"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
func TestConvertClaudeRequestTreatsZeroMaxTokensAsUnset(t *testing.T) {
zero := uint(0)
req := &dto.ClaudeRequest{
Model: "claude-sonnet-4-5",
MaxTokens: &zero,
Messages: []dto.ClaudeMessage{
{Role: "user", Content: "hello"},
},
}
info := &relaycommon.RelayInfo{
ChannelMeta: &relaycommon.ChannelMeta{
UpstreamModelName: "claude-sonnet-4-5",
},
}
out, err := (&Adaptor{}).ConvertClaudeRequest(nil, info, req)
require.NoError(t, err)
converted, ok := out.(*dto.ClaudeRequest)
require.True(t, ok)
require.NotNil(t, converted.MaxTokens)
assert.Equal(t, uint(model_setting.GetClaudeSettings().GetDefaultMaxTokens(req.Model)), *converted.MaxTokens)
}
func TestConvertClaudeRequestPreservesNativeClaudeCodeThinking(t *testing.T) {
budget := 10000
maxTokens := uint(20000)
temperature := 0.7
topP := 0.9
req := &dto.ClaudeRequest{
Model: "claude-opus-4-8",
MaxTokens: &maxTokens,
Temperature: &temperature,
TopP: &topP,
Thinking: &dto.Thinking{Type: "enabled", BudgetTokens: &budget},
OutputConfig: []byte(`{"effort":"high"}`),
Messages: []dto.ClaudeMessage{
{Role: "user", Content: "hello"},
},
}
info := &relaycommon.RelayInfo{
OriginModelName: req.Model,
ChannelMeta: &relaycommon.ChannelMeta{
UpstreamModelName: req.Model,
},
}
out, err := (&Adaptor{}).ConvertClaudeRequest(nil, info, req)
require.NoError(t, err)
converted, ok := out.(*dto.ClaudeRequest)
require.True(t, ok)
require.NotNil(t, converted.Thinking)
assert.Equal(t, "enabled", converted.Thinking.Type)
require.NotNil(t, converted.Thinking.BudgetTokens)
assert.Equal(t, budget, *converted.Thinking.BudgetTokens)
assert.Equal(t, temperature, *converted.Temperature)
assert.Equal(t, topP, *converted.TopP)
assert.JSONEq(t, `{"effort":"high"}`, string(converted.OutputConfig))
assert.Empty(t, info.ConversionDiagnostics())
}
func TestConvertClaudeRequestZeroMaxTokensStillRaisesThinkingBudget(t *testing.T) {
gin.SetMode(gin.TestMode)
c, _ := gin.CreateTestContext(httptest.NewRecorder())
zero := uint(0)
original := &dto.ClaudeRequest{
Model: "claude-3-7-sonnet-thinking",
MaxTokens: &zero,
Messages: []dto.ClaudeMessage{
{Role: "user", Content: "hello"},
},
}
info := &relaycommon.RelayInfo{
OriginModelName: "claude-3-7-sonnet-thinking",
Request: original,
ChannelMeta: &relaycommon.ChannelMeta{
UpstreamModelName: "claude-3-7-sonnet-thinking",
},
}
outbound, err := common.DeepCopy(original)
require.NoError(t, err)
require.NoError(t, helper.ModelMappedHelper(c, info, outbound))
err = helper.ApplyReasoningModelSuffix(nil, info, outbound)
require.NoError(t, err)
out, err := (&Adaptor{}).ConvertClaudeRequest(nil, info, outbound)
require.NoError(t, err)
converted, ok := out.(*dto.ClaudeRequest)
require.True(t, ok)
assert.Equal(t, "claude-3-7-sonnet", converted.Model)
require.NotNil(t, converted.Thinking)
require.NotNil(t, converted.MaxTokens)
assert.Greater(t, *converted.MaxTokens, uint(1024))
}
func TestConvertClaudeRequestDoesNotOverwriteTrimmedUpstreamModelName(t *testing.T) {
req := &dto.ClaudeRequest{
Model: "claude-3-7-sonnet-thinking",
Messages: []dto.ClaudeMessage{
{Role: "user", Content: "hello"},
},
}
info := &relaycommon.RelayInfo{
ChannelMeta: &relaycommon.ChannelMeta{
UpstreamModelName: "claude-3-7-sonnet",
},
}
_, err := (&Adaptor{}).ConvertClaudeRequest(nil, info, req)
require.NoError(t, err)
assert.Equal(t, "claude-3-7-sonnet", info.UpstreamModelName)
}
func geminiToClaudeInfo() *relaycommon.RelayInfo {
return &relaycommon.RelayInfo{
OriginModelName: "claude-3-7-sonnet",
ChannelMeta: &relaycommon.ChannelMeta{
UpstreamModelName: "claude-3-7-sonnet",
},
}
}
func TestConvertGeminiRequestMapsSystemInstructionToolsAndMultimodal(t *testing.T) {
req := &dto.GeminiChatRequest{
Contents: []dto.GeminiChatContent{
{
Role: "user",
Parts: []dto.GeminiPart{
{Text: "What is in this image?"},
{InlineData: &dto.GeminiInlineData{MimeType: "image/png", Data: "aGVsbG8="}},
},
},
},
SystemInstructions: &dto.GeminiChatContent{
Parts: []dto.GeminiPart{{Text: "You are a helpful assistant."}},
},
}
req.SetTools([]dto.GeminiChatTool{
{
FunctionDeclarations: []dto.FunctionRequest{
{
Name: "lookup",
Description: "Lookup data",
Parameters: map[string]any{
"type": "object",
"properties": map[string]any{"q": map[string]any{"type": "string"}},
},
},
},
},
})
out, err := (&Adaptor{}).ConvertGeminiRequest(nil, geminiToClaudeInfo(), req)
require.NoError(t, err)
converted, ok := out.(*dto.ClaudeRequest)
require.True(t, ok)
system := converted.ParseSystem()
require.NotEmpty(t, system)
assert.Contains(t, system[0].GetText(), "You are a helpful assistant.")
require.NotEmpty(t, converted.Messages)
assert.Equal(t, "user", converted.Messages[0].Role)
blocks, parseErr := converted.Messages[0].ParseContent()
require.NoError(t, parseErr)
var foundImage bool
for _, block := range blocks {
if block.Type == "image" || (block.Source != nil && block.Source.Type == "base64") {
foundImage = true
break
}
}
assert.True(t, foundImage)
require.NotNil(t, converted.Tools)
tools, err := common.Marshal(converted.Tools)
require.NoError(t, err)
assert.Contains(t, string(tools), `"lookup"`)
require.NotNil(t, converted.MaxTokens)
assert.Greater(t, *converted.MaxTokens, uint(0))
}
func TestConvertGeminiRequestThinkingConfigUsesReasoningIntent(t *testing.T) {
budget := 1024
maxTokens := uint(4096)
req := &dto.GeminiChatRequest{
Contents: []dto.GeminiChatContent{
{Role: "user", Parts: []dto.GeminiPart{{Text: "think"}}},
},
GenerationConfig: dto.GeminiChatGenerationConfig{
MaxOutputTokens: &maxTokens,
ThinkingConfig: &dto.GeminiThinkingConfig{ThinkingBudget: &budget},
},
}
out, err := (&Adaptor{}).ConvertGeminiRequest(nil, geminiToClaudeInfo(), req)
require.NoError(t, err)
converted, ok := out.(*dto.ClaudeRequest)
require.True(t, ok)
require.NotNil(t, converted.Thinking)
assert.Equal(t, "enabled", converted.Thinking.Type)
require.NotNil(t, converted.Thinking.BudgetTokens)
assert.Equal(t, 1024, *converted.Thinking.BudgetTokens)
}
func TestConvertGeminiRequestNilRequest(t *testing.T) {
_, err := (&Adaptor{}).ConvertGeminiRequest(nil, geminiToClaudeInfo(), nil)
require.Error(t, err)
}
+45
View File
@@ -0,0 +1,45 @@
package claude
var ModelList = []string{
"claude-3-sonnet-20240229",
"claude-3-opus-20240229",
"claude-3-haiku-20240307",
"claude-3-5-haiku-20241022",
"claude-haiku-4-5-20251001",
"claude-3-5-sonnet-20240620",
"claude-3-5-sonnet-20241022",
"claude-3-7-sonnet-20250219",
"claude-3-7-sonnet-20250219-thinking",
"claude-sonnet-4-20250514",
"claude-sonnet-4-20250514-thinking",
"claude-opus-4-20250514",
"claude-opus-4-20250514-thinking",
"claude-opus-4-1-20250805",
"claude-opus-4-1-20250805-thinking",
"claude-sonnet-4-5-20250929",
"claude-sonnet-4-5-20250929-thinking",
"claude-opus-4-5-20251101",
"claude-opus-4-5-20251101-thinking",
"claude-opus-4-6",
"claude-opus-4-6-max",
"claude-opus-4-6-high",
"claude-opus-4-6-medium",
"claude-opus-4-6-low",
"claude-sonnet-4-6",
"claude-opus-4-7",
"claude-opus-4-7-max",
"claude-opus-4-7-xhigh",
"claude-opus-4-7-high",
"claude-opus-4-7-medium",
"claude-opus-4-7-low",
"claude-opus-4-7-thinking",
"claude-opus-4-8",
"claude-opus-4-8-max",
"claude-opus-4-8-xhigh",
"claude-opus-4-8-high",
"claude-opus-4-8-medium",
"claude-opus-4-8-low",
"claude-opus-4-8-thinking",
}
var ChannelName = "claude"
+95
View File
@@ -0,0 +1,95 @@
package claude
//
//type ClaudeMetadata struct {
// UserId string `json:"user_id"`
//}
//
//type ClaudeMediaMessage struct {
// Type string `json:"type"`
// Text string `json:"text,omitempty"`
// Source *ClaudeMessageSource `json:"source,omitempty"`
// Usage *ClaudeUsage `json:"usage,omitempty"`
// StopReason *string `json:"stop_reason,omitempty"`
// PartialJson string `json:"partial_json,omitempty"`
// Thinking string `json:"thinking,omitempty"`
// Signature string `json:"signature,omitempty"`
// Delta string `json:"delta,omitempty"`
// // tool_calls
// Id string `json:"id,omitempty"`
// Name string `json:"name,omitempty"`
// Input any `json:"input,omitempty"`
// Content string `json:"content,omitempty"`
// ToolUseId string `json:"tool_use_id,omitempty"`
//}
//
//type ClaudeMessageSource struct {
// Type string `json:"type"`
// MediaType string `json:"media_type"`
// Data string `json:"data"`
//}
//
//type ClaudeMessage struct {
// Role string `json:"role"`
// Content any `json:"content"`
//}
//
//type Tool struct {
// Name string `json:"name"`
// Description string `json:"description,omitempty"`
// InputSchema map[string]interface{} `json:"input_schema"`
//}
//
//type InputSchema struct {
// Type string `json:"type"`
// Properties any `json:"properties,omitempty"`
// Required any `json:"required,omitempty"`
//}
//
//type ClaudeRequest struct {
// Model string `json:"model"`
// Prompt string `json:"prompt,omitempty"`
// System string `json:"system,omitempty"`
// Messages []ClaudeMessage `json:"messages,omitempty"`
// MaxTokens uint `json:"max_tokens,omitempty"`
// MaxTokensToSample uint `json:"max_tokens_to_sample,omitempty"`
// StopSequences []string `json:"stop_sequences,omitempty"`
// Temperature *float64 `json:"temperature,omitempty"`
// TopP float64 `json:"top_p,omitempty"`
// TopK int `json:"top_k,omitempty"`
// //ClaudeMetadata `json:"metadata,omitempty"`
// Stream bool `json:"stream,omitempty"`
// Tools any `json:"tools,omitempty"`
// ToolChoice any `json:"tool_choice,omitempty"`
// Thinking *Thinking `json:"thinking,omitempty"`
//}
//
//type Thinking struct {
// Type string `json:"type"`
// BudgetTokens int `json:"budget_tokens"`
//}
//
//type ClaudeError struct {
// Type string `json:"type"`
// Message string `json:"message"`
//}
//
//type ClaudeResponse struct {
// Id string `json:"id"`
// Type string `json:"type"`
// Content []ClaudeMediaMessage `json:"content"`
// Completion string `json:"completion"`
// StopReason string `json:"stop_reason"`
// Model string `json:"model"`
// Error ClaudeError `json:"error"`
// Usage ClaudeUsage `json:"usage"`
// Index int `json:"index"` // stream only
// ContentBlock *ClaudeMediaMessage `json:"content_block"`
// Delta *ClaudeMediaMessage `json:"delta"` // stream only
// Message *ClaudeResponse `json:"message"` // stream only: message_start
//}
//
//type ClaudeUsage struct {
// InputTokens int `json:"input_tokens"`
// OutputTokens int `json:"output_tokens"`
//}
@@ -0,0 +1,129 @@
package claude
import (
"testing"
relaycommon "github.com/QuantumNous/new-api/relay/common"
"github.com/QuantumNous/new-api/relaykit/dto"
"github.com/QuantumNous/new-api/setting/model_setting"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"github.com/tidwall/gjson"
)
func TestPatchClaudeMessageDeltaUsageDataPreserveUnknownFields(t *testing.T) {
originalData := `{"type":"message_delta","delta":{"stop_reason":"end_turn","stop_sequence":null},"usage":{"output_tokens":53},"vendor_meta":{"trace_id":"trace_001"}}`
usage := &dto.ClaudeUsage{
InputTokens: 100,
CacheReadInputTokens: 30,
CacheCreationInputTokens: 50,
}
patchedData := patchClaudeMessageDeltaUsageData(originalData, usage)
require.Equal(t, "message_delta", gjson.Get(patchedData, "type").String())
require.Equal(t, "end_turn", gjson.Get(patchedData, "delta.stop_reason").String())
require.Equal(t, "trace_001", gjson.Get(patchedData, "vendor_meta.trace_id").String())
require.EqualValues(t, 53, gjson.Get(patchedData, "usage.output_tokens").Int())
require.EqualValues(t, 100, gjson.Get(patchedData, "usage.input_tokens").Int())
require.EqualValues(t, 30, gjson.Get(patchedData, "usage.cache_read_input_tokens").Int())
require.EqualValues(t, 50, gjson.Get(patchedData, "usage.cache_creation_input_tokens").Int())
}
func TestPatchClaudeMessageDeltaUsageDataZeroValueChecks(t *testing.T) {
originalData := `{"type":"message_delta","usage":{"output_tokens":53,"input_tokens":9,"cache_read_input_tokens":0}}`
usage := &dto.ClaudeUsage{
InputTokens: 100,
CacheReadInputTokens: 30,
CacheCreationInputTokens: 0,
}
patchedData := patchClaudeMessageDeltaUsageData(originalData, usage)
require.EqualValues(t, 9, gjson.Get(patchedData, "usage.input_tokens").Int())
require.EqualValues(t, 30, gjson.Get(patchedData, "usage.cache_read_input_tokens").Int())
assert.False(t, gjson.Get(patchedData, "usage.cache_creation_input_tokens").Exists())
}
func TestShouldSkipClaudeMessageDeltaUsagePatch(t *testing.T) {
originGlobalPassThrough := model_setting.GetGlobalSettings().PassThroughRequestEnabled
t.Cleanup(func() {
model_setting.GetGlobalSettings().PassThroughRequestEnabled = originGlobalPassThrough
})
model_setting.GetGlobalSettings().PassThroughRequestEnabled = true
assert.True(t, shouldSkipClaudeMessageDeltaUsagePatch(&relaycommon.RelayInfo{}))
model_setting.GetGlobalSettings().PassThroughRequestEnabled = false
assert.True(t, shouldSkipClaudeMessageDeltaUsagePatch(&relaycommon.RelayInfo{
ChannelMeta: &relaycommon.ChannelMeta{ChannelSetting: dto.ChannelSettings{PassThroughBodyEnabled: true}},
}))
assert.False(t, shouldSkipClaudeMessageDeltaUsagePatch(&relaycommon.RelayInfo{
ChannelMeta: &relaycommon.ChannelMeta{ChannelSetting: dto.ChannelSettings{PassThroughBodyEnabled: false}},
}))
}
func TestBuildMessageDeltaPatchUsage(t *testing.T) {
t.Run("merge missing fields from claudeInfo", func(t *testing.T) {
claudeResponse := &dto.ClaudeResponse{Usage: &dto.ClaudeUsage{OutputTokens: 53}}
claudeInfo := &ClaudeResponseInfo{
Usage: &dto.Usage{
PromptTokens: 100,
PromptTokensDetails: dto.InputTokenDetails{
CachedTokens: 30,
CachedCreationTokens: 50,
},
ClaudeCacheCreation5mTokens: 10,
ClaudeCacheCreation1hTokens: 20,
},
}
usage := buildMessageDeltaPatchUsage(claudeResponse, claudeInfo)
require.NotNil(t, usage)
require.EqualValues(t, 100, usage.InputTokens)
require.EqualValues(t, 30, usage.CacheReadInputTokens)
require.EqualValues(t, 50, usage.CacheCreationInputTokens)
require.EqualValues(t, 53, usage.OutputTokens)
require.NotNil(t, usage.CacheCreation)
require.EqualValues(t, 30, usage.CacheCreation.Ephemeral5mInputTokens)
require.EqualValues(t, 20, usage.CacheCreation.Ephemeral1hInputTokens)
})
t.Run("keep upstream non-zero values", func(t *testing.T) {
claudeResponse := &dto.ClaudeResponse{Usage: &dto.ClaudeUsage{
InputTokens: 9,
CacheReadInputTokens: 7,
CacheCreationInputTokens: 6,
}}
claudeInfo := &ClaudeResponseInfo{Usage: &dto.Usage{
PromptTokens: 100,
PromptTokensDetails: dto.InputTokenDetails{
CachedTokens: 30,
CachedCreationTokens: 50,
},
}}
usage := buildMessageDeltaPatchUsage(claudeResponse, claudeInfo)
require.EqualValues(t, 9, usage.InputTokens)
require.EqualValues(t, 7, usage.CacheReadInputTokens)
require.EqualValues(t, 6, usage.CacheCreationInputTokens)
})
t.Run("default aggregate cache creation to 5m when split missing", func(t *testing.T) {
claudeResponse := &dto.ClaudeResponse{Usage: &dto.ClaudeUsage{
OutputTokens: 53,
CacheCreationInputTokens: 50,
}}
claudeInfo := &ClaudeResponseInfo{Usage: &dto.Usage{
PromptTokensDetails: dto.InputTokenDetails{
CachedCreationTokens: 50,
},
}}
usage := buildMessageDeltaPatchUsage(claudeResponse, claudeInfo)
require.NotNil(t, usage)
require.NotNil(t, usage.CacheCreation)
require.EqualValues(t, 50, usage.CacheCreation.Ephemeral5mInputTokens)
require.EqualValues(t, 0, usage.CacheCreation.Ephemeral1hInputTokens)
})
}
+415
View File
@@ -0,0 +1,415 @@
package claude
import (
"fmt"
"io"
"net/http"
"strings"
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/constant"
"github.com/QuantumNous/new-api/logger"
relaycommon "github.com/QuantumNous/new-api/relay/common"
"github.com/QuantumNous/new-api/relay/helper"
"github.com/QuantumNous/new-api/relaykit/dto"
"github.com/QuantumNous/new-api/relaykit/relayconvert"
"github.com/QuantumNous/new-api/relaykit/types"
"github.com/QuantumNous/new-api/service"
"github.com/QuantumNous/new-api/setting/model_setting"
"github.com/gin-gonic/gin"
)
func stopReasonClaude2OpenAI(reason string) string {
return relayconvert.StopReasonClaudeToOpenAI(reason)
}
func maybeMarkClaudeRefusal(c *gin.Context, stopReason string) {
if c == nil {
return
}
if strings.EqualFold(stopReason, "refusal") {
common.SetContextKey(c, constant.ContextKeyAdminRejectReason, "claude_stop_reason=refusal")
}
}
func StreamResponseClaude2OpenAI(claudeResponse *dto.ClaudeResponse) *dto.ChatCompletionsStreamResponse {
return relayconvert.StreamResponseClaude2OpenAI(claudeResponse)
}
func ResponseClaude2OpenAI(claudeResponse *dto.ClaudeResponse) *dto.OpenAITextResponse {
return relayconvert.ResponseClaude2OpenAI(claudeResponse)
}
type ClaudeResponseInfo = relayconvert.ClaudeResponseInfo
func cacheCreationTokensForOpenAIUsage(usage *dto.Usage) int {
if usage == nil {
return 0
}
openAIUsage := relayconvert.UsageFromClaudeUsage(usage)
if openAIUsage == nil {
return 0
}
return openAIUsage.PromptTokens - usage.PromptTokens - usage.PromptTokensDetails.CachedTokens
}
func buildOpenAIStyleUsageFromClaudeUsage(usage *dto.Usage) dto.Usage {
mapped := relayconvert.UsageFromClaudeUsage(usage)
if mapped == nil {
return dto.Usage{}
}
return *mapped
}
func buildMessageDeltaPatchUsage(claudeResponse *dto.ClaudeResponse, claudeInfo *ClaudeResponseInfo) *dto.ClaudeUsage {
return relayconvert.BuildMessageDeltaPatchUsage(claudeResponse, claudeInfo)
}
func shouldSkipClaudeMessageDeltaUsagePatch(info *relaycommon.RelayInfo) bool {
if model_setting.GetGlobalSettings().PassThroughRequestEnabled {
return true
}
if info == nil {
return false
}
return info.ChannelSetting.PassThroughBodyEnabled
}
func patchClaudeMessageDeltaUsageData(data string, usage *dto.ClaudeUsage) string {
return relayconvert.PatchClaudeMessageDeltaUsageData(data, usage)
}
func FormatClaudeResponseInfo(claudeResponse *dto.ClaudeResponse, oaiResponse *dto.ChatCompletionsStreamResponse, claudeInfo *ClaudeResponseInfo) bool {
return relayconvert.FormatClaudeResponseInfo(claudeResponse, oaiResponse, claudeInfo)
}
func HandleStreamResponseData(c *gin.Context, info *relaycommon.RelayInfo, claudeInfo *ClaudeResponseInfo, data string) *types.NewAPIError {
var claudeResponse dto.ClaudeResponse
err := common.UnmarshalJsonStr(data, &claudeResponse)
if err != nil {
common.SysLog("error unmarshalling stream response: " + err.Error())
return types.NewError(err, types.ErrorCodeBadResponseBody)
}
if claudeError := claudeResponse.GetClaudeError(); claudeError != nil && claudeError.Type != "" {
return types.WithClaudeError(*claudeError, http.StatusInternalServerError)
}
if claudeResponse.StopReason != "" {
maybeMarkClaudeRefusal(c, claudeResponse.StopReason)
}
if claudeResponse.Delta != nil && claudeResponse.Delta.StopReason != nil {
maybeMarkClaudeRefusal(c, *claudeResponse.Delta.StopReason)
}
if info.RelayFormat == types.RelayFormatClaude {
FormatClaudeResponseInfo(&claudeResponse, nil, claudeInfo)
if claudeResponse.Type == "message_start" {
// message_start, 获取usage
if claudeResponse.Message != nil {
info.UpstreamModelName = claudeResponse.Message.Model
}
} else if claudeResponse.Type == "message_delta" {
// 确保 message_delta 的 usage 包含完整的 input_tokens 和 cache 相关字段
// 解决 AWS Bedrock 等上游返回的 message_delta 缺少这些字段的问题
if !shouldSkipClaudeMessageDeltaUsagePatch(info) {
data = patchClaudeMessageDeltaUsageData(data, buildMessageDeltaPatchUsage(&claudeResponse, claudeInfo))
}
}
countClaudeStreamBillableTools(c, info, &claudeResponse)
helper.ClaudeChunkData(c, claudeResponse, data)
} else if info.RelayFormat == types.RelayFormatOpenAI {
state, err := claudeToChatStreamState(info)
if err != nil {
return types.NewError(err, types.ErrorCodeBadResponseBody)
}
response, err := state.ConvertChunk(&claudeResponse)
if err != nil {
return types.NewError(err, types.ErrorCodeBadResponseBody)
}
if !FormatClaudeResponseInfo(&claudeResponse, response, claudeInfo) {
return nil
}
countClaudeStreamBillableTools(c, info, &claudeResponse)
if response == nil {
return nil
}
err = helper.ObjectData(c, response)
if err != nil {
logger.LogError(c, "send_stream_response_failed: "+err.Error())
}
} else if info.RelayFormat == types.RelayFormatGemini {
state, err := claudeToGeminiStreamState(info)
if err != nil {
return types.NewError(err, types.ErrorCodeBadResponseBody)
}
results, err := service.ConvertStreamResponseChunk(c, info, state, &claudeResponse)
if err != nil {
return types.NewError(err, types.ErrorCodeBadResponseBody)
}
if !FormatClaudeResponseInfo(&claudeResponse, nil, claudeInfo) {
return nil
}
countClaudeStreamBillableTools(c, info, &claudeResponse)
if sendErr := sendGeminiStreamResults(c, results); sendErr != nil {
return sendErr
}
}
return nil
}
func claudeToChatStreamState(info *relaycommon.RelayInfo) (*relayconvert.ClaudeToChatStreamState, error) {
if info != nil && info.ClaudeToChatStreamState != nil {
state, ok := info.ClaudeToChatStreamState.(*relayconvert.ClaudeToChatStreamState)
if !ok || state == nil {
return nil, fmt.Errorf("invalid Claude-to-Chat stream state %T", info.ClaudeToChatStreamState)
}
return state, nil
}
state := relayconvert.NewClaudeToChatStreamState()
if info != nil {
info.ClaudeToChatStreamState = state
}
return state, nil
}
func claudeToGeminiStreamState(info *relaycommon.RelayInfo) (*relayconvert.ResponseStreamState, error) {
if info != nil && info.ChatToGeminiStreamState != nil {
state, ok := info.ChatToGeminiStreamState.(*relayconvert.ResponseStreamState)
if !ok || state == nil {
return nil, fmt.Errorf("invalid Claude-to-Gemini stream state %T", info.ChatToGeminiStreamState)
}
return state, nil
}
state, err := relayconvert.NewResponseStreamState(types.RelayFormatClaude, types.RelayFormatGemini, relayconvert.ResponseStreamOptions{})
if err != nil {
return nil, err
}
if info != nil {
info.ChatToGeminiStreamState = state
}
return state, nil
}
func sendGeminiStreamResults(c *gin.Context, results []relayconvert.ResponseResult) *types.NewAPIError {
for _, result := range results {
geminiResponse, ok := result.Value.(*dto.GeminiChatResponse)
if !ok {
return types.NewError(fmt.Errorf("expected Gemini stream response, got %T", result.Value), types.ErrorCodeBadResponseBody)
}
if geminiResponse == nil {
continue
}
data, err := common.Marshal(geminiResponse)
if err != nil {
return types.NewError(err, types.ErrorCodeBadResponseBody)
}
c.Render(-1, common.CustomEvent{Data: "data: " + string(data)})
_ = helper.FlushWriter(c)
}
return nil
}
func countClaudeStreamBillableTools(c *gin.Context, info *relaycommon.RelayInfo, claudeResponse *dto.ClaudeResponse) {
if claudeResponse == nil {
return
}
if claudeResponse.Type == "content_block_start" &&
claudeResponse.ContentBlock != nil &&
claudeResponse.ContentBlock.Type == "tool_use" {
info.CountBillableToolCall(dto.BuildInCallToolUse, claudeResponse.ContentBlock.Name)
}
if claudeResponse.Type == "message_delta" &&
claudeResponse.Usage != nil &&
claudeResponse.Usage.ServerToolUse != nil &&
claudeResponse.Usage.ServerToolUse.WebSearchRequests > 0 {
c.Set("claude_web_search_requests", claudeResponse.Usage.ServerToolUse.WebSearchRequests)
}
}
func HandleStreamFinalResponse(c *gin.Context, info *relaycommon.RelayInfo, claudeInfo *ClaudeResponseInfo) {
if claudeInfo.Usage.PromptTokens == 0 {
//上游出错
}
if claudeInfo.Usage.CompletionTokens == 0 || !claudeInfo.Done {
if common.DebugEnabled {
common.SysLog("claude response usage is not complete, maybe upstream error")
}
// 只补缺失字段,不整份覆盖——保留 message_start 已拿到的 cache 字段
fallback := service.ResponseText2Usage(c, claudeInfo.ResponseText.String(), info.UpstreamModelName, info.GetEstimatePromptTokens())
if claudeInfo.Usage.CompletionTokens == 0 ||
(!claudeInfo.Done && fallback.CompletionTokens > claudeInfo.Usage.CompletionTokens) {
claudeInfo.Usage.CompletionTokens = fallback.CompletionTokens
}
if claudeInfo.Usage.PromptTokens == 0 {
claudeInfo.Usage.PromptTokens = fallback.PromptTokens
}
claudeInfo.Usage.TotalTokens = claudeInfo.Usage.PromptTokens + claudeInfo.Usage.CompletionTokens
}
if claudeInfo.Usage != nil {
claudeInfo.Usage.UsageSemantic = "anthropic"
}
relayconvert.FinalizeClaudeStreamBillingUsage(claudeInfo)
if info.RelayFormat == types.RelayFormatClaude {
//
} else if info.RelayFormat == types.RelayFormatOpenAI {
if info.ShouldIncludeUsage {
openAIUsage := buildOpenAIStyleUsageFromClaudeUsage(claudeInfo.Usage)
response := helper.GenerateFinalUsageResponse(claudeInfo.ResponseId, claudeInfo.Created, info.UpstreamModelName, openAIUsage)
err := helper.ObjectData(c, response)
if err != nil {
common.SysLog("send final response failed: " + err.Error())
}
}
helper.Done(c)
} else if info.RelayFormat == types.RelayFormatGemini {
state, err := claudeToGeminiStreamState(info)
if err != nil {
common.SysLog("error creating Gemini stream state: " + err.Error())
return
}
results, err := service.FinalizeStreamResponse(c, info, state)
if err != nil {
common.SysLog("error finalizing Gemini stream response: " + err.Error())
return
}
if sendErr := sendGeminiStreamResults(c, results); sendErr != nil {
common.SysLog("send final Gemini stream response failed: " + sendErr.Error())
}
}
}
func ClaudeStreamHandler(c *gin.Context, resp *http.Response, info *relaycommon.RelayInfo) (*dto.Usage, *types.NewAPIError) {
claudeInfo := &ClaudeResponseInfo{
ResponseId: helper.GetResponseID(c),
Created: common.GetTimestamp(),
Model: info.UpstreamModelName,
ResponseText: strings.Builder{},
Usage: &dto.Usage{},
}
var err *types.NewAPIError
helper.StreamScannerHandler(c, resp, info, func(data string, sr *helper.StreamResult) {
err = HandleStreamResponseData(c, info, claudeInfo, data)
if err != nil {
sr.Stop(err)
}
})
if err != nil {
return nil, err
}
HandleStreamFinalResponse(c, info, claudeInfo)
return claudeInfo.Usage, nil
}
func HandleClaudeResponseData(c *gin.Context, info *relaycommon.RelayInfo, claudeInfo *ClaudeResponseInfo, httpResp *http.Response, data []byte) *types.NewAPIError {
var claudeResponse dto.ClaudeResponse
err := common.Unmarshal(data, &claudeResponse)
if err != nil {
return types.NewError(err, types.ErrorCodeBadResponseBody)
}
if claudeError := claudeResponse.GetClaudeError(); claudeError != nil && claudeError.Type != "" {
return types.WithClaudeError(*claudeError, http.StatusInternalServerError)
}
maybeMarkClaudeRefusal(c, claudeResponse.StopReason)
if claudeInfo.Usage == nil {
claudeInfo.Usage = &dto.Usage{}
}
if claudeResponse.Usage != nil {
claudeInfo.Usage.PromptTokens = claudeResponse.Usage.InputTokens
claudeInfo.Usage.CompletionTokens = claudeResponse.Usage.OutputTokens
claudeInfo.Usage.TotalTokens = claudeResponse.Usage.InputTokens + claudeResponse.Usage.OutputTokens
claudeInfo.Usage.UsageSemantic = "anthropic"
claudeInfo.Usage.BillingUsage = dto.CloneBillingUsage(claudeResponse.Usage.BillingUsage)
if claudeInfo.Usage.BillingUsage == nil {
claudeInfo.Usage.BillingUsage = dto.NewClaudeMessagesBillingUsage(claudeResponse.Usage)
}
claudeInfo.Usage.PromptTokensDetails.CachedTokens = claudeResponse.Usage.CacheReadInputTokens
claudeInfo.Usage.PromptTokensDetails.CachedCreationTokens = claudeResponse.Usage.CacheCreationInputTokens
claudeInfo.Usage.ClaudeCacheCreation5mTokens = claudeResponse.Usage.GetCacheCreation5mTokens()
claudeInfo.Usage.ClaudeCacheCreation1hTokens = claudeResponse.Usage.GetCacheCreation1hTokens()
}
var responseData []byte
switch info.RelayFormat {
case types.RelayFormatOpenAI:
openaiResponse := ResponseClaude2OpenAI(&claudeResponse)
openaiResponse.Usage = buildOpenAIStyleUsageFromClaudeUsage(claudeInfo.Usage)
responseData, err = common.Marshal(openaiResponse)
if err != nil {
return types.NewError(err, types.ErrorCodeBadResponseBody)
}
case types.RelayFormatOpenAIResponses:
convertResult, err := service.ConvertResponse(c, info, types.RelayFormatOpenAIResponses, &claudeResponse)
if err != nil {
return types.NewError(err, types.ErrorCodeBadResponseBody)
}
responsesResponse, ok := convertResult.Value.(*dto.OpenAIResponsesResponse)
if !ok {
return types.NewError(fmt.Errorf("expected OpenAI Responses response, got %T", convertResult.Value), types.ErrorCodeBadResponseBody)
}
if responseID := helper.GetResponseID(c); responseID != "" {
responsesResponse.ID = responseID
}
responseData, err = common.Marshal(responsesResponse)
if err != nil {
return types.NewError(err, types.ErrorCodeBadResponseBody)
}
case types.RelayFormatClaude:
responseData = data
case types.RelayFormatGemini:
{
convertResult, convertErr := service.ConvertResponse(c, info, types.RelayFormatGemini, &claudeResponse)
if convertErr != nil {
return types.NewError(convertErr, types.ErrorCodeBadResponseBody)
}
geminiResponse, ok := convertResult.Value.(*dto.GeminiChatResponse)
if !ok {
return types.NewError(fmt.Errorf("expected Gemini generateContent response, got %T", convertResult.Value), types.ErrorCodeBadResponseBody)
}
responseData, err = common.Marshal(geminiResponse)
if err != nil {
return types.NewError(err, types.ErrorCodeBadResponseBody)
}
}
}
if claudeResponse.Usage != nil && claudeResponse.Usage.ServerToolUse != nil && claudeResponse.Usage.ServerToolUse.WebSearchRequests > 0 {
c.Set("claude_web_search_requests", claudeResponse.Usage.ServerToolUse.WebSearchRequests)
}
for _, block := range claudeResponse.Content {
if block.Type == "tool_use" {
info.CountBillableToolCall(dto.BuildInCallToolUse, block.Name)
}
}
service.IOCopyBytesGracefully(c, httpResp, responseData)
return nil
}
func ClaudeHandler(c *gin.Context, resp *http.Response, info *relaycommon.RelayInfo) (*dto.Usage, *types.NewAPIError) {
defer service.CloseResponseBodyGracefully(resp)
claudeInfo := &ClaudeResponseInfo{
ResponseId: helper.GetResponseID(c),
Created: common.GetTimestamp(),
Model: info.UpstreamModelName,
ResponseText: strings.Builder{},
Usage: &dto.Usage{},
}
responseBody, err := io.ReadAll(resp.Body)
if err != nil {
return nil, types.NewError(err, types.ErrorCodeBadResponseBody)
}
logger.LogDebug(c, "responseBody: %s", responseBody)
handleErr := HandleClaudeResponseData(c, info, claudeInfo, resp, responseBody)
if handleErr != nil {
return nil, handleErr
}
return claudeInfo.Usage, nil
}
+402
View File
@@ -0,0 +1,402 @@
package claude
import (
"net/http/httptest"
"strings"
"testing"
"github.com/QuantumNous/new-api/common"
relaycommon "github.com/QuantumNous/new-api/relay/common"
"github.com/QuantumNous/new-api/relay/helper"
"github.com/QuantumNous/new-api/relaykit/dto"
"github.com/QuantumNous/new-api/relaykit/relayconvert"
"github.com/gin-gonic/gin"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
func commonPointer[T any](value T) *T {
return &value
}
func TestResponseOpenAI2ClaudeToolUseInputIsObject(t *testing.T) {
tests := []struct {
name string
args string
want map[string]any
}{
{name: "object", args: `{"q":"x"}`, want: map[string]any{"q": "x"}},
{name: "empty", args: "", want: map[string]any{}},
{name: "invalid", args: "{", want: map[string]any{}},
{name: "null", args: "null", want: map[string]any{}},
{name: "array", args: `["x"]`, want: map[string]any{}},
{name: "string", args: `"x"`, want: map[string]any{}},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
msg := dto.Message{Role: "assistant"}
msg.SetToolCalls([]dto.ToolCallRequest{
{
ID: "call_1",
Type: "function",
Function: dto.FunctionRequest{
Name: "lookup",
Arguments: tt.args,
},
},
})
resp := relayconvert.ResponseOpenAI2Claude(&dto.OpenAITextResponse{
Id: "chatcmpl_1",
Model: "gpt-test",
Choices: []dto.OpenAITextResponseChoice{
{Message: msg, FinishReason: "tool_calls"},
},
}, nil)
require.Len(t, resp.Content, 1)
assert.Equal(t, "tool_use", resp.Content[0].Type)
assert.Equal(t, tt.want, resp.Content[0].Input)
})
}
}
func TestFormatClaudeResponseInfo_MessageStart(t *testing.T) {
claudeInfo := &ClaudeResponseInfo{
Usage: &dto.Usage{},
}
claudeResponse := &dto.ClaudeResponse{
Type: "message_start",
Message: &dto.ClaudeMediaMessage{
Id: "msg_123",
Model: "claude-3-5-sonnet",
Usage: &dto.ClaudeUsage{
InputTokens: 100,
OutputTokens: 1,
CacheCreationInputTokens: 50,
CacheReadInputTokens: 30,
},
},
}
ok := FormatClaudeResponseInfo(claudeResponse, nil, claudeInfo)
if !ok {
t.Fatal("expected true")
}
if claudeInfo.Usage.PromptTokens != 100 {
t.Errorf("PromptTokens = %d, want 100", claudeInfo.Usage.PromptTokens)
}
if claudeInfo.Usage.PromptTokensDetails.CachedTokens != 30 {
t.Errorf("CachedTokens = %d, want 30", claudeInfo.Usage.PromptTokensDetails.CachedTokens)
}
if claudeInfo.Usage.PromptTokensDetails.CachedCreationTokens != 50 {
t.Errorf("CachedCreationTokens = %d, want 50", claudeInfo.Usage.PromptTokensDetails.CachedCreationTokens)
}
if claudeInfo.ResponseId != "msg_123" {
t.Errorf("ResponseId = %s, want msg_123", claudeInfo.ResponseId)
}
if claudeInfo.Model != "claude-3-5-sonnet" {
t.Errorf("Model = %s, want claude-3-5-sonnet", claudeInfo.Model)
}
}
func TestFormatClaudeResponseInfo_MessageDelta_FullUsage(t *testing.T) {
// message_start 先积累 usage
claudeInfo := &ClaudeResponseInfo{
Usage: &dto.Usage{
PromptTokens: 100,
PromptTokensDetails: dto.InputTokenDetails{
CachedTokens: 30,
CachedCreationTokens: 50,
},
CompletionTokens: 1,
},
}
// message_delta 带完整 usage(原生 Anthropic 场景)
claudeResponse := &dto.ClaudeResponse{
Type: "message_delta",
Usage: &dto.ClaudeUsage{
InputTokens: 100,
OutputTokens: 200,
CacheCreationInputTokens: 50,
CacheReadInputTokens: 30,
},
}
ok := FormatClaudeResponseInfo(claudeResponse, nil, claudeInfo)
if !ok {
t.Fatal("expected true")
}
if claudeInfo.Usage.PromptTokens != 100 {
t.Errorf("PromptTokens = %d, want 100", claudeInfo.Usage.PromptTokens)
}
if claudeInfo.Usage.CompletionTokens != 200 {
t.Errorf("CompletionTokens = %d, want 200", claudeInfo.Usage.CompletionTokens)
}
if claudeInfo.Usage.TotalTokens != 300 {
t.Errorf("TotalTokens = %d, want 300", claudeInfo.Usage.TotalTokens)
}
if !claudeInfo.Done {
t.Error("expected Done = true")
}
}
func TestFormatClaudeResponseInfo_MessageDelta_OnlyOutputTokens(t *testing.T) {
// 模拟 Bedrock: message_start 已积累 usage
claudeInfo := &ClaudeResponseInfo{
Usage: &dto.Usage{
PromptTokens: 100,
PromptTokensDetails: dto.InputTokenDetails{
CachedTokens: 30,
CachedCreationTokens: 50,
},
CompletionTokens: 1,
ClaudeCacheCreation5mTokens: 10,
ClaudeCacheCreation1hTokens: 20,
},
}
// Bedrock 的 message_delta 只有 output_tokens,缺少 input_tokens 和 cache 字段
claudeResponse := &dto.ClaudeResponse{
Type: "message_delta",
Usage: &dto.ClaudeUsage{
OutputTokens: 200,
// InputTokens, CacheCreationInputTokens, CacheReadInputTokens 都是 0
},
}
ok := FormatClaudeResponseInfo(claudeResponse, nil, claudeInfo)
if !ok {
t.Fatal("expected true")
}
// PromptTokens 应保持 message_start 的值(因为 message_delta 的 InputTokens=0,不更新)
if claudeInfo.Usage.PromptTokens != 100 {
t.Errorf("PromptTokens = %d, want 100", claudeInfo.Usage.PromptTokens)
}
if claudeInfo.Usage.CompletionTokens != 200 {
t.Errorf("CompletionTokens = %d, want 200", claudeInfo.Usage.CompletionTokens)
}
if claudeInfo.Usage.TotalTokens != 300 {
t.Errorf("TotalTokens = %d, want 300", claudeInfo.Usage.TotalTokens)
}
// cache 字段应保持 message_start 的值
if claudeInfo.Usage.PromptTokensDetails.CachedTokens != 30 {
t.Errorf("CachedTokens = %d, want 30", claudeInfo.Usage.PromptTokensDetails.CachedTokens)
}
if claudeInfo.Usage.PromptTokensDetails.CachedCreationTokens != 50 {
t.Errorf("CachedCreationTokens = %d, want 50", claudeInfo.Usage.PromptTokensDetails.CachedCreationTokens)
}
if claudeInfo.Usage.ClaudeCacheCreation5mTokens != 10 {
t.Errorf("ClaudeCacheCreation5mTokens = %d, want 10", claudeInfo.Usage.ClaudeCacheCreation5mTokens)
}
if claudeInfo.Usage.ClaudeCacheCreation1hTokens != 20 {
t.Errorf("ClaudeCacheCreation1hTokens = %d, want 20", claudeInfo.Usage.ClaudeCacheCreation1hTokens)
}
if !claudeInfo.Done {
t.Error("expected Done = true")
}
}
func TestFormatClaudeResponseInfo_NilClaudeInfo(t *testing.T) {
claudeResponse := &dto.ClaudeResponse{Type: "message_start"}
ok := FormatClaudeResponseInfo(claudeResponse, nil, nil)
if ok {
t.Error("expected false for nil claudeInfo")
}
}
func TestFormatClaudeResponseInfo_ContentBlockDelta(t *testing.T) {
text := "hello"
claudeInfo := &ClaudeResponseInfo{
Usage: &dto.Usage{},
ResponseText: strings.Builder{},
}
claudeResponse := &dto.ClaudeResponse{
Type: "content_block_delta",
Delta: &dto.ClaudeMediaMessage{
Text: &text,
},
}
ok := FormatClaudeResponseInfo(claudeResponse, nil, claudeInfo)
if !ok {
t.Fatal("expected true")
}
if claudeInfo.ResponseText.String() != "hello" {
t.Errorf("ResponseText = %q, want %q", claudeInfo.ResponseText.String(), "hello")
}
}
func TestBuildOpenAIStyleUsageFromClaudeUsage(t *testing.T) {
usage := &dto.Usage{
PromptTokens: 100,
CompletionTokens: 20,
PromptTokensDetails: dto.InputTokenDetails{
CachedTokens: 30,
CachedCreationTokens: 50,
},
ClaudeCacheCreation5mTokens: 10,
ClaudeCacheCreation1hTokens: 20,
UsageSemantic: "anthropic",
}
openAIUsage := buildOpenAIStyleUsageFromClaudeUsage(usage)
if openAIUsage.PromptTokens != 180 {
t.Fatalf("PromptTokens = %d, want 180", openAIUsage.PromptTokens)
}
if openAIUsage.InputTokens != 180 {
t.Fatalf("InputTokens = %d, want 180", openAIUsage.InputTokens)
}
if openAIUsage.TotalTokens != 200 {
t.Fatalf("TotalTokens = %d, want 200", openAIUsage.TotalTokens)
}
if openAIUsage.UsageSemantic != "openai" {
t.Fatalf("UsageSemantic = %s, want openai", openAIUsage.UsageSemantic)
}
if openAIUsage.UsageSource != "anthropic" {
t.Fatalf("UsageSource = %s, want anthropic", openAIUsage.UsageSource)
}
}
func TestBuildOpenAIStyleUsageFromClaudeUsagePreservesCacheCreationRemainder(t *testing.T) {
tests := []struct {
name string
cachedCreationTokens int
cacheCreationTokens5m int
cacheCreationTokens1h int
expectedTotalInputToken int
}{
{
name: "prefers aggregate when it includes remainder",
cachedCreationTokens: 50,
cacheCreationTokens5m: 10,
cacheCreationTokens1h: 20,
expectedTotalInputToken: 180,
},
{
name: "falls back to split tokens when aggregate missing",
cachedCreationTokens: 0,
cacheCreationTokens5m: 10,
cacheCreationTokens1h: 20,
expectedTotalInputToken: 160,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
usage := &dto.Usage{
PromptTokens: 100,
CompletionTokens: 20,
PromptTokensDetails: dto.InputTokenDetails{
CachedTokens: 30,
CachedCreationTokens: tt.cachedCreationTokens,
},
ClaudeCacheCreation5mTokens: tt.cacheCreationTokens5m,
ClaudeCacheCreation1hTokens: tt.cacheCreationTokens1h,
UsageSemantic: "anthropic",
}
openAIUsage := buildOpenAIStyleUsageFromClaudeUsage(usage)
if openAIUsage.PromptTokens != tt.expectedTotalInputToken {
t.Fatalf("PromptTokens = %d, want %d", openAIUsage.PromptTokens, tt.expectedTotalInputToken)
}
if openAIUsage.InputTokens != tt.expectedTotalInputToken {
t.Fatalf("InputTokens = %d, want %d", openAIUsage.InputTokens, tt.expectedTotalInputToken)
}
})
}
}
func TestBuildOpenAIStyleUsageFromClaudeUsageDefaultsAggregateCacheCreationTo5m(t *testing.T) {
usage := &dto.Usage{
PromptTokens: 100,
CompletionTokens: 20,
PromptTokensDetails: dto.InputTokenDetails{
CachedTokens: 30,
CachedCreationTokens: 50,
},
UsageSemantic: "anthropic",
}
openAIUsage := buildOpenAIStyleUsageFromClaudeUsage(usage)
require.Equal(t, 50, openAIUsage.ClaudeCacheCreation5mTokens)
require.Equal(t, 0, openAIUsage.ClaudeCacheCreation1hTokens)
}
func applyOpenAIChatReasoningThroughHandlerOrder(t *testing.T, original dto.GeneralOpenAIRequest) (*dto.GeneralOpenAIRequest, *relaycommon.RelayInfo) {
t.Helper()
gin.SetMode(gin.TestMode)
c, _ := gin.CreateTestContext(httptest.NewRecorder())
info := &relaycommon.RelayInfo{
OriginModelName: original.Model,
Request: &original,
ChannelMeta: &relaycommon.ChannelMeta{
UpstreamModelName: original.Model,
},
}
outbound, err := common.DeepCopy(&original)
require.NoError(t, err)
require.NoError(t, helper.ModelMappedHelper(c, info, outbound))
err = helper.ApplyReasoningModelSuffix(nil, info, outbound)
require.NoError(t, err)
return outbound, info
}
func TestOpenAIChatRequestToClaudeMessages_ClaudeOpus48HighUsesAdaptiveThinking(t *testing.T) {
original := dto.GeneralOpenAIRequest{
Model: "claude-opus-4-8-high",
Temperature: commonPointer(0.7),
TopP: commonPointer(0.9),
TopK: commonPointer(40),
Messages: []dto.Message{
{
Role: "user",
Content: "hello",
},
},
}
outbound, info := applyOpenAIChatReasoningThroughHandlerOrder(t, original)
claudeRequest, err := relayconvert.OpenAIChatRequestToClaudeMessages(nil, info, *outbound)
require.NoError(t, err)
require.Equal(t, "claude-opus-4-8", claudeRequest.Model)
require.NotNil(t, claudeRequest.Thinking)
require.Equal(t, "adaptive", claudeRequest.Thinking.Type)
require.Equal(t, "summarized", claudeRequest.Thinking.Display)
require.JSONEq(t, `{"effort":"high"}`, string(claudeRequest.OutputConfig))
require.Nil(t, claudeRequest.Temperature)
require.Nil(t, claudeRequest.TopP)
require.Nil(t, claudeRequest.TopK)
}
func TestOpenAIChatRequestToClaudeMessages_ClaudeOpus48ThinkingUsesAdaptiveHighEffort(t *testing.T) {
original := dto.GeneralOpenAIRequest{
Model: "claude-opus-4-8-thinking",
Temperature: commonPointer(0.7),
TopP: commonPointer(0.9),
TopK: commonPointer(40),
Messages: []dto.Message{
{
Role: "user",
Content: "hello",
},
},
}
outbound, info := applyOpenAIChatReasoningThroughHandlerOrder(t, original)
claudeRequest, err := relayconvert.OpenAIChatRequestToClaudeMessages(nil, info, *outbound)
require.NoError(t, err)
require.Equal(t, "claude-opus-4-8", claudeRequest.Model)
require.NotNil(t, claudeRequest.Thinking)
require.Equal(t, "adaptive", claudeRequest.Thinking.Type)
require.Equal(t, "summarized", claudeRequest.Thinking.Display)
require.JSONEq(t, `{"effort":"high"}`, string(claudeRequest.OutputConfig))
require.Nil(t, claudeRequest.Temperature)
require.Nil(t, claudeRequest.TopP)
require.Nil(t, claudeRequest.TopK)
}
+180
View File
@@ -0,0 +1,180 @@
package claude
import (
"fmt"
"net/http"
"strings"
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/logger"
relaycommon "github.com/QuantumNous/new-api/relay/common"
"github.com/QuantumNous/new-api/relay/helper"
"github.com/QuantumNous/new-api/relaykit/dto"
"github.com/QuantumNous/new-api/relaykit/relayconvert"
"github.com/QuantumNous/new-api/relaykit/types"
"github.com/QuantumNous/new-api/service"
"github.com/gin-gonic/gin"
)
func ClaudeResponsesStreamHandler(c *gin.Context, resp *http.Response, info *relaycommon.RelayInfo) (*dto.Usage, *types.NewAPIError) {
responseID := helper.GetResponseID(c)
created := common.GetTimestamp()
state, err := relayconvert.NewResponseStreamState(types.RelayFormatClaude, types.RelayFormatOpenAIResponses, relayconvert.ResponseStreamOptions{
ID: responseID,
Model: info.UpstreamModelName,
Created: created,
EmitSequenceNumber: true,
})
if err != nil {
return nil, types.NewOpenAIError(err, types.ErrorCodeBadResponse, http.StatusInternalServerError)
}
hostedBridge := relayconvert.NewClaudeHostedStreamBridge()
claudeInfo := &ClaudeResponseInfo{
ResponseId: responseID,
Created: created,
Model: info.UpstreamModelName,
ResponseText: strings.Builder{},
Usage: &dto.Usage{},
}
var streamErr *types.NewAPIError
// streamFailed means a Responses-native terminal error was sent successfully.
// In that case the scanner stops without a transport error and the partial
// upstream usage remains billable.
streamFailed := false
sendResponsesEvent := func(eventType string, payload dto.ResponsesStreamResponse) bool {
payload.Type = eventType
data, err := common.Marshal(payload)
if err != nil {
streamErr = types.NewOpenAIError(err, types.ErrorCodeJsonMarshalFailed, http.StatusInternalServerError)
return false
}
if err := helper.ResponseChunkData(c, dto.ResponsesStreamResponse{Type: eventType}, string(data)); err != nil {
streamErr = types.NewOpenAIError(err, types.ErrorCodeBadResponse, http.StatusInternalServerError)
return false
}
return true
}
sendResult := func(result relayconvert.ResponseResult) bool {
event, ok := result.Value.(relayconvert.ChatToResponsesStreamEvent)
if !ok {
streamErr = types.NewOpenAIError(
fmt.Errorf("expected OpenAI Responses stream event, got %T", result.Value),
types.ErrorCodeBadResponse,
http.StatusInternalServerError,
)
return false
}
return sendResponsesEvent(event.Type, event.Payload)
}
failResponsesStream := func(err error) bool {
failureResults, handled := state.FailResponsesStream("server_error", err.Error(), "")
if !handled {
return false
}
for _, result := range failureResults {
if !sendResult(result) {
return true
}
}
streamFailed = true
return true
}
helper.StreamScannerHandler(c, resp, info, func(data string, sr *helper.StreamResult) {
var claudeResponse dto.ClaudeResponse
if err := common.UnmarshalJsonStr(data, &claudeResponse); err != nil {
logger.LogError(c, "failed to unmarshal Claude stream event: "+err.Error())
if failResponsesStream(err) {
// A nil streamErr here is intentional: the protocol-level failure
// event was delivered, so only the scanner needs to stop.
sr.Stop(streamErr)
return
}
streamErr = types.NewOpenAIError(err, types.ErrorCodeBadResponseBody, http.StatusInternalServerError)
sr.Stop(streamErr)
return
}
if claudeError := claudeResponse.GetClaudeError(); claudeError != nil && claudeError.Type != "" {
if failResponsesStream(fmt.Errorf("%s", claudeError.Message)) {
sr.Stop(streamErr)
return
}
streamErr = types.WithClaudeError(*claudeError, http.StatusInternalServerError)
sr.Stop(streamErr)
return
}
if claudeResponse.StopReason != "" {
maybeMarkClaudeRefusal(c, claudeResponse.StopReason)
}
if claudeResponse.Delta != nil && claudeResponse.Delta.StopReason != nil {
maybeMarkClaudeRefusal(c, *claudeResponse.Delta.StopReason)
}
if claudeResponse.Type == "message_start" && claudeResponse.Message != nil {
info.UpstreamModelName = claudeResponse.Message.Model
}
FormatClaudeResponseInfo(&claudeResponse, nil, claudeInfo)
countClaudeStreamBillableTools(c, info, &claudeResponse)
hostedEvents, consumed, err := hostedBridge.Convert(&claudeResponse, state)
if err != nil {
if failResponsesStream(err) {
sr.Stop(streamErr)
return
}
streamErr = types.NewOpenAIError(err, types.ErrorCodeBadResponse, http.StatusInternalServerError)
sr.Stop(streamErr)
return
}
for _, event := range hostedEvents {
if !sendResponsesEvent(event.Type, event.Payload) {
sr.Stop(streamErr)
return
}
}
if consumed {
return
}
results, err := service.ConvertStreamResponseChunk(c, info, state, &claudeResponse)
if err != nil {
if failResponsesStream(err) {
sr.Stop(streamErr)
return
}
streamErr = types.NewOpenAIError(err, types.ErrorCodeBadResponse, http.StatusInternalServerError)
sr.Stop(streamErr)
return
}
for _, result := range results {
if !sendResult(result) {
sr.Stop(streamErr)
return
}
}
})
if streamErr != nil {
return nil, streamErr
}
if streamFailed {
return claudeInfo.Usage, nil
}
HandleStreamFinalResponse(c, info, claudeInfo)
openAIUsage := buildOpenAIStyleUsageFromClaudeUsage(claudeInfo.Usage)
state.SetUsage(&openAIUsage)
finalResults, err := service.FinalizeStreamResponse(c, info, state)
if err != nil {
if failResponsesStream(err) {
return claudeInfo.Usage, streamErr
}
return nil, types.NewOpenAIError(err, types.ErrorCodeBadResponse, http.StatusInternalServerError)
}
for _, result := range finalResults {
if !sendResult(result) {
return nil, streamErr
}
}
return claudeInfo.Usage, nil
}
+76
View File
@@ -0,0 +1,76 @@
package claude
import (
"net/http/httptest"
"testing"
relaycommon "github.com/QuantumNous/new-api/relay/common"
"github.com/QuantumNous/new-api/relaykit/dto"
"github.com/QuantumNous/new-api/relaykit/types"
"github.com/QuantumNous/new-api/setting/operation_setting"
"github.com/gin-gonic/gin"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
func TestHandleClaudeResponseDataCountsToolUse(t *testing.T) {
gin.SetMode(gin.TestMode)
operation_setting.SetToolPriceForTest("lookup_fn", 3.0)
t.Cleanup(func() {
operation_setting.DeleteToolPriceForTest("lookup_fn")
})
w := httptest.NewRecorder()
c, _ := gin.CreateTestContext(w)
info := &relaycommon.RelayInfo{
OriginModelName: "claude-3-7-sonnet",
RelayFormat: types.RelayFormatClaude,
}
claudeInfo := &ClaudeResponseInfo{Usage: &dto.Usage{}}
data := []byte(`{
"type":"message",
"content":[
{"type":"text","text":"hi"},
{"type":"tool_use","id":"tu1","name":"lookup_fn","input":{}},
{"type":"server_tool_use","id":"stu1","name":"web_search","input":{}}
],
"usage":{"input_tokens":1,"output_tokens":1}
}`)
err := HandleClaudeResponseData(c, info, claudeInfo, nil, data)
require.Nil(t, err)
require.NotNil(t, info.ResponsesUsageInfo)
require.Contains(t, info.ResponsesUsageInfo.BuiltInTools, "lookup_fn")
assert.Equal(t, 1, info.ResponsesUsageInfo.BuiltInTools["lookup_fn"].CallCount)
assert.NotContains(t, info.ResponsesUsageInfo.BuiltInTools, "web_search")
}
func TestCountClaudeStreamBillableToolsSetsWebSearchRequests(t *testing.T) {
gin.SetMode(gin.TestMode)
w := httptest.NewRecorder()
c, _ := gin.CreateTestContext(w)
info := &relaycommon.RelayInfo{OriginModelName: "claude-3-7-sonnet"}
countClaudeStreamBillableTools(c, info, &dto.ClaudeResponse{
Type: "message_delta",
Usage: &dto.ClaudeUsage{
ServerToolUse: &dto.ClaudeServerToolUse{WebSearchRequests: 3},
},
})
assert.Equal(t, 3, c.GetInt("claude_web_search_requests"))
operation_setting.SetToolPriceForTest("stream_fn", 2.0)
t.Cleanup(func() {
operation_setting.DeleteToolPriceForTest("stream_fn")
})
countClaudeStreamBillableTools(c, info, &dto.ClaudeResponse{
Type: "content_block_start",
ContentBlock: &dto.ClaudeMediaMessage{
Type: "tool_use",
Name: "stream_fn",
},
})
require.Contains(t, info.ResponsesUsageInfo.BuiltInTools, "stream_fn")
assert.Equal(t, 1, info.ResponsesUsageInfo.BuiltInTools["stream_fn"].CallCount)
}
+135
View File
@@ -0,0 +1,135 @@
package cloudflare
import (
"bytes"
"errors"
"fmt"
"io"
"net/http"
"github.com/QuantumNous/new-api/relay/channel"
"github.com/QuantumNous/new-api/relay/channel/openai"
relaycommon "github.com/QuantumNous/new-api/relay/common"
"github.com/QuantumNous/new-api/relay/constant"
"github.com/QuantumNous/new-api/relaykit/dto"
"github.com/QuantumNous/new-api/relaykit/types"
"github.com/gin-gonic/gin"
)
type Adaptor struct {
}
func (a *Adaptor) ConvertGeminiRequest(*gin.Context, *relaycommon.RelayInfo, *dto.GeminiChatRequest) (any, error) {
//TODO implement me
return nil, errors.New("not implemented")
}
func (a *Adaptor) ConvertClaudeRequest(*gin.Context, *relaycommon.RelayInfo, *dto.ClaudeRequest) (any, error) {
//TODO implement me
panic("implement me")
}
func (a *Adaptor) Init(info *relaycommon.RelayInfo) {
}
func (a *Adaptor) GetRequestURL(info *relaycommon.RelayInfo) (string, error) {
switch info.RelayMode {
case constant.RelayModeChatCompletions:
return fmt.Sprintf("%s/client/v4/accounts/%s/ai/v1/chat/completions", info.ChannelBaseUrl, info.ApiVersion), nil
case constant.RelayModeEmbeddings:
return fmt.Sprintf("%s/client/v4/accounts/%s/ai/v1/embeddings", info.ChannelBaseUrl, info.ApiVersion), nil
case constant.RelayModeResponses:
return fmt.Sprintf("%s/client/v4/accounts/%s/ai/v1/responses", info.ChannelBaseUrl, info.ApiVersion), nil
default:
return fmt.Sprintf("%s/client/v4/accounts/%s/ai/run/%s", info.ChannelBaseUrl, info.ApiVersion, info.UpstreamModelName), nil
}
}
func (a *Adaptor) SetupRequestHeader(c *gin.Context, req *http.Header, info *relaycommon.RelayInfo) error {
channel.SetupApiRequestHeader(info, c, req)
req.Set("Authorization", fmt.Sprintf("Bearer %s", info.ApiKey))
return nil
}
func (a *Adaptor) ConvertOpenAIRequest(c *gin.Context, info *relaycommon.RelayInfo, request *dto.GeneralOpenAIRequest) (any, error) {
if request == nil {
return nil, errors.New("request is nil")
}
switch info.RelayMode {
case constant.RelayModeCompletions:
return convertCf2CompletionsRequest(*request), nil
default:
return request, nil
}
}
func (a *Adaptor) ConvertOpenAIResponsesRequest(c *gin.Context, info *relaycommon.RelayInfo, request dto.OpenAIResponsesRequest) (any, error) {
return request, nil
}
func (a *Adaptor) DoRequest(c *gin.Context, info *relaycommon.RelayInfo, requestBody io.Reader) (any, error) {
return channel.DoApiRequest(a, c, info, requestBody)
}
func (a *Adaptor) ConvertRerankRequest(c *gin.Context, relayMode int, request dto.RerankRequest) (any, error) {
return request, nil
}
func (a *Adaptor) ConvertEmbeddingRequest(c *gin.Context, info *relaycommon.RelayInfo, request dto.EmbeddingRequest) (any, error) {
return request, nil
}
func (a *Adaptor) ConvertAudioRequest(c *gin.Context, info *relaycommon.RelayInfo, request dto.AudioRequest) (io.Reader, error) {
// 添加文件字段
file, _, err := c.Request.FormFile("file")
if err != nil {
return nil, errors.New("file is required")
}
defer file.Close()
// 打开临时文件用于保存上传的文件内容
requestBody := &bytes.Buffer{}
// 将上传的文件内容复制到临时文件
if _, err := io.Copy(requestBody, file); err != nil {
return nil, err
}
return requestBody, nil
}
func (a *Adaptor) ConvertImageRequest(c *gin.Context, info *relaycommon.RelayInfo, request dto.ImageRequest) (any, error) {
//TODO implement me
return nil, errors.New("not implemented")
}
func (a *Adaptor) DoResponse(c *gin.Context, resp *http.Response, info *relaycommon.RelayInfo) (usage any, err *types.NewAPIError) {
switch info.RelayMode {
case constant.RelayModeEmbeddings:
fallthrough
case constant.RelayModeChatCompletions:
if info.IsStream {
err, usage = cfStreamHandler(c, info, resp)
} else {
err, usage = cfHandler(c, info, resp)
}
case constant.RelayModeResponses:
if info.IsStream {
usage, err = openai.OaiResponsesStreamHandler(c, info, resp)
} else {
usage, err = openai.OaiResponsesHandler(c, info, resp)
}
case constant.RelayModeAudioTranslation:
fallthrough
case constant.RelayModeAudioTranscription:
err, usage = cfSTTHandler(c, info, resp)
}
return
}
func (a *Adaptor) GetModelList() []string {
return ModelList
}
func (a *Adaptor) GetChannelName() string {
return ChannelName
}
+39
View File
@@ -0,0 +1,39 @@
package cloudflare
var ModelList = []string{
"@cf/meta/llama-3.1-8b-instruct",
"@cf/meta/llama-2-7b-chat-fp16",
"@cf/meta/llama-2-7b-chat-int8",
"@cf/mistral/mistral-7b-instruct-v0.1",
"@hf/thebloke/deepseek-coder-6.7b-base-awq",
"@hf/thebloke/deepseek-coder-6.7b-instruct-awq",
"@cf/deepseek-ai/deepseek-math-7b-base",
"@cf/deepseek-ai/deepseek-math-7b-instruct",
"@cf/thebloke/discolm-german-7b-v1-awq",
"@cf/tiiuae/falcon-7b-instruct",
"@cf/google/gemma-2b-it-lora",
"@hf/google/gemma-7b-it",
"@cf/google/gemma-7b-it-lora",
"@hf/nousresearch/hermes-2-pro-mistral-7b",
"@hf/thebloke/llama-2-13b-chat-awq",
"@cf/meta-llama/llama-2-7b-chat-hf-lora",
"@cf/meta/llama-3-8b-instruct",
"@hf/thebloke/llamaguard-7b-awq",
"@hf/thebloke/mistral-7b-instruct-v0.1-awq",
"@hf/mistralai/mistral-7b-instruct-v0.2",
"@cf/mistral/mistral-7b-instruct-v0.2-lora",
"@hf/thebloke/neural-chat-7b-v3-1-awq",
"@cf/openchat/openchat-3.5-0106",
"@hf/thebloke/openhermes-2.5-mistral-7b-awq",
"@cf/microsoft/phi-2",
"@cf/qwen/qwen1.5-0.5b-chat",
"@cf/qwen/qwen1.5-1.8b-chat",
"@cf/qwen/qwen1.5-14b-chat-awq",
"@cf/qwen/qwen1.5-7b-chat-awq",
"@cf/defog/sqlcoder-7b-2",
"@hf/nexusflow/starling-lm-7b-beta",
"@cf/tinyllama/tinyllama-1.1b-chat-v1.0",
"@hf/thebloke/zephyr-7b-beta-awq",
}
var ChannelName = "cloudflare"
+21
View File
@@ -0,0 +1,21 @@
package cloudflare
import "github.com/QuantumNous/new-api/relaykit/dto"
type CfRequest struct {
Messages []dto.Message `json:"messages,omitempty"`
Lora string `json:"lora,omitempty"`
MaxTokens uint `json:"max_tokens,omitempty"`
Prompt string `json:"prompt,omitempty"`
Raw bool `json:"raw,omitempty"`
Stream bool `json:"stream,omitempty"`
Temperature *float64 `json:"temperature,omitempty"`
}
type CfAudioResponse struct {
Result CfSTTResult `json:"result"`
}
type CfSTTResult struct {
Text string `json:"text"`
}
@@ -0,0 +1,148 @@
package cloudflare
import (
"bufio"
"encoding/json"
"io"
"net/http"
"strings"
"time"
"github.com/QuantumNous/new-api/logger"
relaycommon "github.com/QuantumNous/new-api/relay/common"
"github.com/QuantumNous/new-api/relay/helper"
"github.com/QuantumNous/new-api/relaykit/dto"
"github.com/QuantumNous/new-api/relaykit/types"
"github.com/QuantumNous/new-api/service"
"github.com/samber/lo"
"github.com/gin-gonic/gin"
)
func convertCf2CompletionsRequest(textRequest dto.GeneralOpenAIRequest) *CfRequest {
p, _ := textRequest.Prompt.(string)
return &CfRequest{
Prompt: p,
MaxTokens: textRequest.GetMaxTokens(),
Stream: lo.FromPtrOr(textRequest.Stream, false),
Temperature: textRequest.Temperature,
}
}
func cfStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http.Response) (*types.NewAPIError, *dto.Usage) {
scanner := helper.NewStreamScanner(resp.Body)
scanner.Split(bufio.ScanLines)
helper.SetEventStreamHeaders(c)
id := helper.GetResponseID(c)
var responseText string
isFirst := true
for scanner.Scan() {
data := scanner.Text()
if len(data) < len("data: ") {
continue
}
data = strings.TrimPrefix(data, "data: ")
data = strings.TrimSuffix(data, "\r")
if data == "[DONE]" {
break
}
var response dto.ChatCompletionsStreamResponse
err := json.Unmarshal([]byte(data), &response)
if err != nil {
logger.LogError(c, "error_unmarshalling_stream_response: "+err.Error())
continue
}
for _, choice := range response.Choices {
choice.Delta.Role = "assistant"
responseText += choice.Delta.GetContentString()
}
response.Id = id
response.Model = info.UpstreamModelName
err = helper.ObjectData(c, response)
if isFirst {
isFirst = false
info.FirstResponseTime = time.Now()
}
if err != nil {
logger.LogError(c, "error_rendering_stream_response: "+err.Error())
}
}
if err := scanner.Err(); err != nil {
logger.LogError(c, "error_scanning_stream_response: "+err.Error())
}
usage := service.ResponseText2Usage(c, responseText, info.UpstreamModelName, info.GetEstimatePromptTokens())
if info.ShouldIncludeUsage {
response := helper.GenerateFinalUsageResponse(id, info.StartTime.Unix(), info.UpstreamModelName, *usage)
err := helper.ObjectData(c, response)
if err != nil {
logger.LogError(c, "error_rendering_final_usage_response: "+err.Error())
}
}
helper.Done(c)
service.CloseResponseBodyGracefully(resp)
return nil, usage
}
func cfHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http.Response) (*types.NewAPIError, *dto.Usage) {
responseBody, err := io.ReadAll(resp.Body)
if err != nil {
return types.NewError(err, types.ErrorCodeBadResponseBody), nil
}
service.CloseResponseBodyGracefully(resp)
var response dto.TextResponse
err = json.Unmarshal(responseBody, &response)
if err != nil {
return types.NewError(err, types.ErrorCodeBadResponseBody), nil
}
response.Model = info.UpstreamModelName
var responseText string
for _, choice := range response.Choices {
responseText += choice.Message.StringContent()
}
usage := service.ResponseText2Usage(c, responseText, info.UpstreamModelName, info.GetEstimatePromptTokens())
response.Usage = *usage
response.Id = helper.GetResponseID(c)
jsonResponse, err := json.Marshal(response)
if err != nil {
return types.NewError(err, types.ErrorCodeBadResponseBody), nil
}
c.Writer.Header().Set("Content-Type", "application/json")
c.Writer.WriteHeader(resp.StatusCode)
_, _ = c.Writer.Write(jsonResponse)
return nil, usage
}
func cfSTTHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http.Response) (*types.NewAPIError, *dto.Usage) {
var cfResp CfAudioResponse
responseBody, err := io.ReadAll(resp.Body)
if err != nil {
return types.NewError(err, types.ErrorCodeBadResponseBody), nil
}
service.CloseResponseBodyGracefully(resp)
err = json.Unmarshal(responseBody, &cfResp)
if err != nil {
return types.NewError(err, types.ErrorCodeBadResponseBody), nil
}
audioResp := &dto.AudioResponse{
Text: cfResp.Result.Text,
}
jsonResponse, err := json.Marshal(audioResp)
if err != nil {
return types.NewError(err, types.ErrorCodeBadResponseBody), nil
}
c.Writer.Header().Set("Content-Type", "application/json")
c.Writer.WriteHeader(resp.StatusCode)
_, _ = c.Writer.Write(jsonResponse)
usage := service.ResponseText2Usage(c, cfResp.Result.Text, info.UpstreamModelName, info.GetEstimatePromptTokens())
return nil, usage
}
+200
View File
@@ -0,0 +1,200 @@
package codex
import (
"encoding/json"
"errors"
"io"
"net/http"
"strings"
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/relay/channel"
"github.com/QuantumNous/new-api/relay/channel/openai"
relaycommon "github.com/QuantumNous/new-api/relay/common"
relayconstant "github.com/QuantumNous/new-api/relay/constant"
"github.com/QuantumNous/new-api/relaykit/dto"
"github.com/QuantumNous/new-api/relaykit/types"
"github.com/gin-gonic/gin"
)
type Adaptor struct {
}
func (a *Adaptor) ConvertGeminiRequest(c *gin.Context, info *relaycommon.RelayInfo, request *dto.GeminiChatRequest) (any, error) {
return nil, errors.New("codex channel: endpoint not supported")
}
func (a *Adaptor) ConvertClaudeRequest(*gin.Context, *relaycommon.RelayInfo, *dto.ClaudeRequest) (any, error) {
return nil, errors.New("codex channel: /v1/messages endpoint not supported")
}
func (a *Adaptor) ConvertAudioRequest(c *gin.Context, info *relaycommon.RelayInfo, request dto.AudioRequest) (io.Reader, error) {
return nil, errors.New("codex channel: endpoint not supported")
}
func (a *Adaptor) ConvertImageRequest(c *gin.Context, info *relaycommon.RelayInfo, request dto.ImageRequest) (any, error) {
return nil, errors.New("codex channel: endpoint not supported")
}
func (a *Adaptor) Init(info *relaycommon.RelayInfo) {
}
func (a *Adaptor) ConvertOpenAIRequest(c *gin.Context, info *relaycommon.RelayInfo, request *dto.GeneralOpenAIRequest) (any, error) {
return nil, errors.New("codex channel: /v1/chat/completions endpoint not supported")
}
func (a *Adaptor) ConvertRerankRequest(c *gin.Context, relayMode int, request dto.RerankRequest) (any, error) {
return nil, errors.New("codex channel: /v1/rerank endpoint not supported")
}
func (a *Adaptor) ConvertEmbeddingRequest(c *gin.Context, info *relaycommon.RelayInfo, request dto.EmbeddingRequest) (any, error) {
return nil, errors.New("codex channel: /v1/embeddings endpoint not supported")
}
func (a *Adaptor) ConvertOpenAIResponsesRequest(c *gin.Context, info *relaycommon.RelayInfo, request dto.OpenAIResponsesRequest) (any, error) {
isCompact := info != nil && info.RelayMode == relayconstant.RelayModeResponsesCompact
if info != nil && info.ChannelSetting.SystemPrompt != "" {
systemPrompt := info.ChannelSetting.SystemPrompt
if len(request.Instructions) == 0 {
if b, err := common.Marshal(systemPrompt); err == nil {
request.Instructions = b
} else {
return nil, err
}
} else if info.ChannelSetting.SystemPromptOverride {
var existing string
if err := common.Unmarshal(request.Instructions, &existing); err == nil {
existing = strings.TrimSpace(existing)
if existing == "" {
if b, err := common.Marshal(systemPrompt); err == nil {
request.Instructions = b
} else {
return nil, err
}
} else {
if b, err := common.Marshal(systemPrompt + "\n" + existing); err == nil {
request.Instructions = b
} else {
return nil, err
}
}
} else {
if b, err := common.Marshal(systemPrompt); err == nil {
request.Instructions = b
} else {
return nil, err
}
}
}
}
// Codex backend requires the `instructions` field to be present.
// Keep it consistent with Codex CLI behavior by defaulting to an empty string.
if len(request.Instructions) == 0 {
request.Instructions = json.RawMessage(`""`)
}
if isCompact {
return request, nil
}
// codex: store must be false
request.Store = json.RawMessage("false")
// rm max_output_tokens
request.MaxOutputTokens = nil
request.Temperature = nil
request.FrequencyPenalty = nil
request.PresencePenalty = nil
return request, nil
}
func (a *Adaptor) DoRequest(c *gin.Context, info *relaycommon.RelayInfo, requestBody io.Reader) (any, error) {
return channel.DoApiRequest(a, c, info, requestBody)
}
func (a *Adaptor) DoResponse(c *gin.Context, resp *http.Response, info *relaycommon.RelayInfo) (usage any, err *types.NewAPIError) {
switch info.RelayMode {
case relayconstant.RelayModeAlphaSearch:
// Alpha search responses are handled by relay.AlphaSearchHelper.
return nil, types.NewError(errors.New("codex channel: alpha search response should be handled by AlphaSearchHelper"), types.ErrorCodeInvalidRequest)
case relayconstant.RelayModeResponsesCompact:
return openai.OaiResponsesCompactionHandler(c, resp)
case relayconstant.RelayModeResponses:
if info.IsStream {
return openai.OaiResponsesStreamHandler(c, info, resp)
}
return openai.OaiResponsesHandler(c, info, resp)
default:
return nil, types.NewError(errors.New("codex channel: endpoint not supported"), types.ErrorCodeInvalidRequest)
}
}
func (a *Adaptor) GetModelList() []string {
return ModelList
}
func (a *Adaptor) GetChannelName() string {
return ChannelName
}
func (a *Adaptor) GetRequestURL(info *relaycommon.RelayInfo) (string, error) {
var path string
switch info.RelayMode {
case relayconstant.RelayModeResponses:
path = "/backend-api/codex/responses"
case relayconstant.RelayModeResponsesCompact:
path = "/backend-api/codex/responses/compact"
case relayconstant.RelayModeAlphaSearch:
path = "/backend-api/codex/alpha/search"
default:
return "", errors.New("codex channel: only /v1/responses, /v1/responses/compact and /v1/alpha/search are supported")
}
return relaycommon.GetFullRequestURL(info.ChannelBaseUrl, path, info.ChannelType), nil
}
func (a *Adaptor) SetupRequestHeader(c *gin.Context, req *http.Header, info *relaycommon.RelayInfo) error {
channel.SetupApiRequestHeader(info, c, req)
key := strings.TrimSpace(info.ApiKey)
if !strings.HasPrefix(key, "{") {
return errors.New("codex channel: key must be a JSON object")
}
oauthKey, err := ParseOAuthKey(key)
if err != nil {
return err
}
accessToken := strings.TrimSpace(oauthKey.AccessToken)
accountID := strings.TrimSpace(oauthKey.AccountID)
if accessToken == "" {
return errors.New("codex channel: access_token is required")
}
if accountID == "" {
return errors.New("codex channel: account_id is required")
}
req.Set("Authorization", "Bearer "+accessToken)
req.Set("chatgpt-account-id", accountID)
if req.Get("OpenAI-Beta") == "" {
req.Set("OpenAI-Beta", "responses=experimental")
}
if req.Get("originator") == "" {
req.Set("originator", "codex_cli_rs")
}
// chatgpt.com/backend-api/codex/responses is strict about Content-Type.
// Clients may omit it or include parameters like `application/json; charset=utf-8`,
// which can be rejected by the upstream. Force the exact media type.
req.Set("Content-Type", "application/json")
if info.IsStream {
req.Set("Accept", "text/event-stream")
} else if req.Get("Accept") == "" {
req.Set("Accept", "application/json")
}
return nil
}
+56
View File
@@ -0,0 +1,56 @@
package codex
import (
"encoding/json"
"testing"
"github.com/QuantumNous/new-api/constant"
relaycommon "github.com/QuantumNous/new-api/relay/common"
relayconstant "github.com/QuantumNous/new-api/relay/constant"
"github.com/QuantumNous/new-api/relaykit/dto"
"github.com/samber/lo"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
func TestGetRequestURLAlphaSearch(t *testing.T) {
adaptor := &Adaptor{}
info := &relaycommon.RelayInfo{
ChannelMeta: &relaycommon.ChannelMeta{
ChannelType: constant.ChannelTypeCodex,
ChannelBaseUrl: "https://chatgpt.com",
},
RelayMode: relayconstant.RelayModeAlphaSearch,
}
url, err := adaptor.GetRequestURL(info)
require.NoError(t, err)
assert.Equal(t, "https://chatgpt.com/backend-api/codex/alpha/search", url)
}
// The Codex backend rejects these fields, so the adaptor clears them rather
// than forwarding what the client sent.
func TestConvertOpenAIResponsesRequestDropsPenalties(t *testing.T) {
adaptor := &Adaptor{}
info := &relaycommon.RelayInfo{
ChannelMeta: &relaycommon.ChannelMeta{ChannelType: constant.ChannelTypeCodex},
RelayMode: relayconstant.RelayModeResponses,
}
converted, err := adaptor.ConvertOpenAIResponsesRequest(nil, info, dto.OpenAIResponsesRequest{
Model: "gpt-5-codex",
Input: json.RawMessage(`"hello"`),
MaxOutputTokens: lo.ToPtr(uint(128)),
Temperature: lo.ToPtr(1.0),
FrequencyPenalty: json.RawMessage(`1.5`),
PresencePenalty: json.RawMessage(`1.5`),
})
require.NoError(t, err)
request, ok := converted.(dto.OpenAIResponsesRequest)
require.True(t, ok)
assert.Nil(t, request.MaxOutputTokens)
assert.Nil(t, request.Temperature)
assert.Nil(t, request.FrequencyPenalty)
assert.Nil(t, request.PresencePenalty)
}
+14
View File
@@ -0,0 +1,14 @@
package codex
var ModelList = []string{
"gpt-5.6-sol",
"gpt-5.6-terra",
"gpt-5.6-luna",
"gpt-5.5",
"gpt-5.4",
"gpt-5.4-mini",
"gpt-5.3-codex-spark",
"codex-auto-review",
}
const ChannelName = "codex"
+30
View File
@@ -0,0 +1,30 @@
package codex
import (
"errors"
"github.com/QuantumNous/new-api/common"
)
type OAuthKey struct {
IDToken string `json:"id_token,omitempty"`
AccessToken string `json:"access_token,omitempty"`
RefreshToken string `json:"refresh_token,omitempty"`
AccountID string `json:"account_id,omitempty"`
LastRefresh string `json:"last_refresh,omitempty"`
Email string `json:"email,omitempty"`
Type string `json:"type,omitempty"`
Expired string `json:"expired,omitempty"`
}
func ParseOAuthKey(raw string) (*OAuthKey, error) {
if raw == "" {
return nil, errors.New("codex channel: empty oauth key")
}
var key OAuthKey
if err := common.Unmarshal([]byte(raw), &key); err != nil {
return nil, errors.New("codex channel: invalid oauth key json")
}
return &key, nil
}
+99
View File
@@ -0,0 +1,99 @@
package cohere
import (
"errors"
"fmt"
"io"
"net/http"
"github.com/QuantumNous/new-api/relay/channel"
relaycommon "github.com/QuantumNous/new-api/relay/common"
"github.com/QuantumNous/new-api/relay/constant"
"github.com/QuantumNous/new-api/relaykit/dto"
"github.com/QuantumNous/new-api/relaykit/types"
"github.com/gin-gonic/gin"
)
type Adaptor struct {
}
func (a *Adaptor) ConvertGeminiRequest(*gin.Context, *relaycommon.RelayInfo, *dto.GeminiChatRequest) (any, error) {
//TODO implement me
return nil, errors.New("not implemented")
}
func (a *Adaptor) ConvertClaudeRequest(*gin.Context, *relaycommon.RelayInfo, *dto.ClaudeRequest) (any, error) {
//TODO implement me
panic("implement me")
}
func (a *Adaptor) ConvertAudioRequest(c *gin.Context, info *relaycommon.RelayInfo, request dto.AudioRequest) (io.Reader, error) {
//TODO implement me
return nil, errors.New("not implemented")
}
func (a *Adaptor) ConvertImageRequest(c *gin.Context, info *relaycommon.RelayInfo, request dto.ImageRequest) (any, error) {
//TODO implement me
return nil, errors.New("not implemented")
}
func (a *Adaptor) Init(info *relaycommon.RelayInfo) {
}
func (a *Adaptor) GetRequestURL(info *relaycommon.RelayInfo) (string, error) {
if info.RelayMode == constant.RelayModeRerank {
return fmt.Sprintf("%s/v1/rerank", info.ChannelBaseUrl), nil
} else {
return fmt.Sprintf("%s/v1/chat", info.ChannelBaseUrl), nil
}
}
func (a *Adaptor) SetupRequestHeader(c *gin.Context, req *http.Header, info *relaycommon.RelayInfo) error {
channel.SetupApiRequestHeader(info, c, req)
req.Set("Authorization", fmt.Sprintf("Bearer %s", info.ApiKey))
return nil
}
func (a *Adaptor) ConvertOpenAIRequest(c *gin.Context, info *relaycommon.RelayInfo, request *dto.GeneralOpenAIRequest) (any, error) {
return requestOpenAI2Cohere(*request), nil
}
func (a *Adaptor) ConvertOpenAIResponsesRequest(c *gin.Context, info *relaycommon.RelayInfo, request dto.OpenAIResponsesRequest) (any, error) {
// TODO implement me
return nil, errors.New("not implemented")
}
func (a *Adaptor) DoRequest(c *gin.Context, info *relaycommon.RelayInfo, requestBody io.Reader) (any, error) {
return channel.DoApiRequest(a, c, info, requestBody)
}
func (a *Adaptor) ConvertRerankRequest(c *gin.Context, relayMode int, request dto.RerankRequest) (any, error) {
return requestConvertRerank2Cohere(request), nil
}
func (a *Adaptor) ConvertEmbeddingRequest(c *gin.Context, info *relaycommon.RelayInfo, request dto.EmbeddingRequest) (any, error) {
//TODO implement me
return nil, errors.New("not implemented")
}
func (a *Adaptor) DoResponse(c *gin.Context, resp *http.Response, info *relaycommon.RelayInfo) (usage any, err *types.NewAPIError) {
if info.RelayMode == constant.RelayModeRerank {
usage, err = cohereRerankHandler(c, resp, info)
} else {
if info.IsStream {
usage, err = cohereStreamHandler(c, info, resp) // TODO: fix this
} else {
usage, err = cohereHandler(c, info, resp)
}
}
return
}
func (a *Adaptor) GetModelList() []string {
return ModelList
}
func (a *Adaptor) GetChannelName() string {
return ChannelName
}
+12
View File
@@ -0,0 +1,12 @@
package cohere
var ModelList = []string{
"command-a-03-2025",
"command-r", "command-r-plus",
"command-r-08-2024", "command-r-plus-08-2024",
"c4ai-aya-23-35b", "c4ai-aya-23-8b",
"command-light", "command-light-nightly", "command", "command-nightly",
"rerank-english-v3.0", "rerank-multilingual-v3.0", "rerank-english-v2.0", "rerank-multilingual-v2.0",
}
var ChannelName = "cohere"
+60
View File
@@ -0,0 +1,60 @@
package cohere
import "github.com/QuantumNous/new-api/relaykit/dto"
type CohereRequest struct {
Model string `json:"model"`
ChatHistory []ChatHistory `json:"chat_history"`
Message string `json:"message"`
Stream bool `json:"stream"`
MaxTokens uint `json:"max_tokens"`
SafetyMode string `json:"safety_mode,omitempty"`
}
type ChatHistory struct {
Role string `json:"role"`
Message string `json:"message"`
}
type CohereResponse struct {
IsFinished bool `json:"is_finished"`
EventType string `json:"event_type"`
Text string `json:"text,omitempty"`
FinishReason string `json:"finish_reason,omitempty"`
Response *CohereResponseResult `json:"response"`
}
type CohereResponseResult struct {
ResponseId string `json:"response_id"`
FinishReason string `json:"finish_reason,omitempty"`
Text string `json:"text"`
Meta CohereMeta `json:"meta"`
}
type CohereRerankRequest struct {
Documents []any `json:"documents"`
Query string `json:"query"`
Model string `json:"model"`
TopN int `json:"top_n"`
ReturnDocuments bool `json:"return_documents"`
}
type CohereRerankResponseResult struct {
Results []dto.RerankResponseResult `json:"results"`
Meta CohereMeta `json:"meta"`
}
type CohereMeta struct {
//Tokens CohereTokens `json:"tokens"`
BilledUnits CohereBilledUnits `json:"billed_units"`
}
type CohereBilledUnits struct {
InputTokens int `json:"input_tokens"`
OutputTokens int `json:"output_tokens"`
}
type CohereTokens struct {
InputTokens int `json:"input_tokens"`
OutputTokens int `json:"output_tokens"`
}
+253
View File
@@ -0,0 +1,253 @@
package cohere
import (
"encoding/json"
"io"
"net/http"
"strings"
"time"
"github.com/QuantumNous/new-api/common"
relaycommon "github.com/QuantumNous/new-api/relay/common"
"github.com/QuantumNous/new-api/relay/helper"
"github.com/QuantumNous/new-api/relaykit/dto"
"github.com/QuantumNous/new-api/relaykit/types"
"github.com/QuantumNous/new-api/service"
"github.com/gin-gonic/gin"
"github.com/samber/lo"
)
func requestOpenAI2Cohere(textRequest dto.GeneralOpenAIRequest) *CohereRequest {
cohereReq := CohereRequest{
Model: textRequest.Model,
ChatHistory: []ChatHistory{},
Message: "",
Stream: lo.FromPtrOr(textRequest.Stream, false),
MaxTokens: textRequest.GetMaxTokens(),
}
if common.CohereSafetySetting != "NONE" {
cohereReq.SafetyMode = common.CohereSafetySetting
}
if cohereReq.MaxTokens == 0 {
cohereReq.MaxTokens = 4000
}
for _, msg := range textRequest.Messages {
if msg.Role == "user" {
cohereReq.Message = msg.StringContent()
} else {
var role string
if msg.Role == "assistant" {
role = "CHATBOT"
} else if msg.Role == "system" {
role = "SYSTEM"
} else {
role = "USER"
}
cohereReq.ChatHistory = append(cohereReq.ChatHistory, ChatHistory{
Role: role,
Message: msg.StringContent(),
})
}
}
return &cohereReq
}
func requestConvertRerank2Cohere(rerankRequest dto.RerankRequest) *CohereRerankRequest {
topN := lo.FromPtrOr(rerankRequest.TopN, 1)
if topN <= 0 {
topN = 1
}
cohereReq := CohereRerankRequest{
Query: rerankRequest.Query,
Documents: rerankRequest.Documents,
Model: rerankRequest.Model,
TopN: topN,
ReturnDocuments: true,
}
return &cohereReq
}
func stopReasonCohere2OpenAI(reason string) string {
switch reason {
case "COMPLETE":
return "stop"
case "MAX_TOKENS":
return "max_tokens"
default:
return reason
}
}
func cohereStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http.Response) (*dto.Usage, *types.NewAPIError) {
responseId := helper.GetResponseID(c)
createdTime := common.GetTimestamp()
usage := &dto.Usage{}
responseText := ""
scanner := helper.NewStreamScanner(resp.Body)
scanner.Split(func(data []byte, atEOF bool) (advance int, token []byte, err error) {
if atEOF && len(data) == 0 {
return 0, nil, nil
}
if i := strings.Index(string(data), "\n"); i >= 0 {
return i + 1, data[0:i], nil
}
if atEOF {
return len(data), data, nil
}
return 0, nil, nil
})
dataChan := make(chan string)
stopChan := make(chan bool)
go func() {
for scanner.Scan() {
data := scanner.Text()
dataChan <- data
}
if err := scanner.Err(); err != nil {
common.SysLog("error reading stream: " + err.Error())
}
stopChan <- true
}()
helper.SetEventStreamHeaders(c)
isFirst := true
c.Stream(func(w io.Writer) bool {
select {
case data := <-dataChan:
if isFirst {
isFirst = false
info.FirstResponseTime = time.Now()
}
data = strings.TrimSuffix(data, "\r")
var cohereResp CohereResponse
err := json.Unmarshal([]byte(data), &cohereResp)
if err != nil {
common.SysLog("error unmarshalling stream response: " + err.Error())
return true
}
var openaiResp dto.ChatCompletionsStreamResponse
openaiResp.Id = responseId
openaiResp.Created = createdTime
openaiResp.Object = "chat.completion.chunk"
openaiResp.Model = info.UpstreamModelName
if cohereResp.IsFinished {
finishReason := stopReasonCohere2OpenAI(cohereResp.FinishReason)
openaiResp.Choices = []dto.ChatCompletionsStreamResponseChoice{
{
Delta: dto.ChatCompletionsStreamResponseChoiceDelta{},
Index: 0,
FinishReason: &finishReason,
},
}
if cohereResp.Response != nil {
usage.PromptTokens = cohereResp.Response.Meta.BilledUnits.InputTokens
usage.CompletionTokens = cohereResp.Response.Meta.BilledUnits.OutputTokens
}
} else {
openaiResp.Choices = []dto.ChatCompletionsStreamResponseChoice{
{
Delta: dto.ChatCompletionsStreamResponseChoiceDelta{
Role: "assistant",
Content: &cohereResp.Text,
},
Index: 0,
},
}
responseText += cohereResp.Text
}
jsonStr, err := json.Marshal(openaiResp)
if err != nil {
common.SysLog("error marshalling stream response: " + err.Error())
return true
}
c.Render(-1, common.CustomEvent{Data: "data: " + string(jsonStr)})
return true
case <-stopChan:
c.Render(-1, common.CustomEvent{Data: "data: [DONE]"})
return false
}
})
if usage.PromptTokens == 0 {
usage = service.ResponseText2Usage(c, responseText, info.UpstreamModelName, info.GetEstimatePromptTokens())
}
return usage, nil
}
func cohereHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http.Response) (*dto.Usage, *types.NewAPIError) {
createdTime := common.GetTimestamp()
responseBody, err := io.ReadAll(resp.Body)
if err != nil {
return nil, types.NewError(err, types.ErrorCodeBadResponseBody)
}
service.CloseResponseBodyGracefully(resp)
var cohereResp CohereResponseResult
err = json.Unmarshal(responseBody, &cohereResp)
if err != nil {
return nil, types.NewError(err, types.ErrorCodeBadResponseBody)
}
usage := dto.Usage{}
usage.PromptTokens = cohereResp.Meta.BilledUnits.InputTokens
usage.CompletionTokens = cohereResp.Meta.BilledUnits.OutputTokens
usage.TotalTokens = cohereResp.Meta.BilledUnits.InputTokens + cohereResp.Meta.BilledUnits.OutputTokens
var openaiResp dto.TextResponse
openaiResp.Id = cohereResp.ResponseId
openaiResp.Created = createdTime
openaiResp.Object = "chat.completion"
openaiResp.Model = info.UpstreamModelName
openaiResp.Usage = usage
openaiResp.Choices = []dto.OpenAITextResponseChoice{
{
Index: 0,
Message: dto.Message{Content: cohereResp.Text, Role: "assistant"},
FinishReason: stopReasonCohere2OpenAI(cohereResp.FinishReason),
},
}
jsonResponse, err := json.Marshal(openaiResp)
if err != nil {
return nil, types.NewError(err, types.ErrorCodeBadResponseBody)
}
c.Writer.Header().Set("Content-Type", "application/json")
c.Writer.WriteHeader(resp.StatusCode)
_, _ = c.Writer.Write(jsonResponse)
return &usage, nil
}
func cohereRerankHandler(c *gin.Context, resp *http.Response, info *relaycommon.RelayInfo) (*dto.Usage, *types.NewAPIError) {
responseBody, err := io.ReadAll(resp.Body)
if err != nil {
return nil, types.NewError(err, types.ErrorCodeBadResponseBody)
}
service.CloseResponseBodyGracefully(resp)
var cohereResp CohereRerankResponseResult
err = json.Unmarshal(responseBody, &cohereResp)
if err != nil {
return nil, types.NewError(err, types.ErrorCodeBadResponseBody)
}
usage := dto.Usage{}
if cohereResp.Meta.BilledUnits.InputTokens == 0 {
usage.PromptTokens = info.GetEstimatePromptTokens()
usage.CompletionTokens = 0
usage.TotalTokens = info.GetEstimatePromptTokens()
} else {
usage.PromptTokens = cohereResp.Meta.BilledUnits.InputTokens
usage.CompletionTokens = cohereResp.Meta.BilledUnits.OutputTokens
usage.TotalTokens = cohereResp.Meta.BilledUnits.InputTokens + cohereResp.Meta.BilledUnits.OutputTokens
}
var rerankResp dto.RerankResponse
rerankResp.Results = cohereResp.Results
rerankResp.Usage = usage
jsonResponse, err := json.Marshal(rerankResp)
if err != nil {
return nil, types.NewError(err, types.ErrorCodeBadResponseBody)
}
c.Writer.Header().Set("Content-Type", "application/json")
c.Writer.WriteHeader(resp.StatusCode)
_, err = c.Writer.Write(jsonResponse)
return &usage, nil
}
+139
View File
@@ -0,0 +1,139 @@
package coze
import (
"encoding/json"
"errors"
"fmt"
"io"
"net/http"
"time"
"github.com/QuantumNous/new-api/relay/channel"
"github.com/QuantumNous/new-api/relay/common"
"github.com/QuantumNous/new-api/relaykit/dto"
"github.com/QuantumNous/new-api/relaykit/types"
"github.com/gin-gonic/gin"
)
type Adaptor struct {
}
func (a *Adaptor) ConvertGeminiRequest(*gin.Context, *common.RelayInfo, *dto.GeminiChatRequest) (any, error) {
//TODO implement me
return nil, errors.New("not implemented")
}
// ConvertAudioRequest implements channel.Adaptor.
func (a *Adaptor) ConvertAudioRequest(c *gin.Context, info *common.RelayInfo, request dto.AudioRequest) (io.Reader, error) {
return nil, errors.New("not implemented")
}
// ConvertClaudeRequest implements channel.Adaptor.
func (a *Adaptor) ConvertClaudeRequest(c *gin.Context, info *common.RelayInfo, request *dto.ClaudeRequest) (any, error) {
return nil, errors.New("not implemented")
}
// ConvertEmbeddingRequest implements channel.Adaptor.
func (a *Adaptor) ConvertEmbeddingRequest(c *gin.Context, info *common.RelayInfo, request dto.EmbeddingRequest) (any, error) {
return nil, errors.New("not implemented")
}
// ConvertImageRequest implements channel.Adaptor.
func (a *Adaptor) ConvertImageRequest(c *gin.Context, info *common.RelayInfo, request dto.ImageRequest) (any, error) {
return nil, errors.New("not implemented")
}
// ConvertOpenAIRequest implements channel.Adaptor.
func (a *Adaptor) ConvertOpenAIRequest(c *gin.Context, info *common.RelayInfo, request *dto.GeneralOpenAIRequest) (any, error) {
if request == nil {
return nil, errors.New("request is nil")
}
return convertCozeChatRequest(c, *request), nil
}
// ConvertOpenAIResponsesRequest implements channel.Adaptor.
func (a *Adaptor) ConvertOpenAIResponsesRequest(c *gin.Context, info *common.RelayInfo, request dto.OpenAIResponsesRequest) (any, error) {
return nil, errors.New("not implemented")
}
// ConvertRerankRequest implements channel.Adaptor.
func (a *Adaptor) ConvertRerankRequest(c *gin.Context, relayMode int, request dto.RerankRequest) (any, error) {
return nil, errors.New("not implemented")
}
// DoRequest implements channel.Adaptor.
func (a *Adaptor) DoRequest(c *gin.Context, info *common.RelayInfo, requestBody io.Reader) (any, error) {
if info.IsStream {
return channel.DoApiRequest(a, c, info, requestBody)
}
// 首先发送创建消息请求,成功后再发送获取消息请求
// 发送创建消息请求
resp, err := channel.DoApiRequest(a, c, info, requestBody)
if err != nil {
return nil, err
}
// 解析 resp
var cozeResponse CozeChatResponse
respBody, err := io.ReadAll(resp.Body)
if err != nil {
return nil, err
}
err = json.Unmarshal(respBody, &cozeResponse)
if cozeResponse.Code != 0 {
return nil, errors.New(cozeResponse.Msg)
}
c.Set("coze_conversation_id", cozeResponse.Data.ConversationId)
c.Set("coze_chat_id", cozeResponse.Data.Id)
// 轮询检查消息是否完成
for {
err, isComplete := checkIfChatComplete(a, c, info)
if err != nil {
return nil, err
} else {
if isComplete {
break
}
}
time.Sleep(time.Second * 1)
}
// 发送获取消息请求
return getChatDetail(a, c, info)
}
// DoResponse implements channel.Adaptor.
func (a *Adaptor) DoResponse(c *gin.Context, resp *http.Response, info *common.RelayInfo) (usage any, err *types.NewAPIError) {
if info.IsStream {
usage, err = cozeChatStreamHandler(c, info, resp)
} else {
usage, err = cozeChatHandler(c, info, resp)
}
return
}
// GetChannelName implements channel.Adaptor.
func (a *Adaptor) GetChannelName() string {
return ChannelName
}
// GetModelList implements channel.Adaptor.
func (a *Adaptor) GetModelList() []string {
return ModelList
}
// GetRequestURL implements channel.Adaptor.
func (a *Adaptor) GetRequestURL(info *common.RelayInfo) (string, error) {
return fmt.Sprintf("%s/v3/chat", info.ChannelBaseUrl), nil
}
// Init implements channel.Adaptor.
func (a *Adaptor) Init(info *common.RelayInfo) {
}
// SetupRequestHeader implements channel.Adaptor.
func (a *Adaptor) SetupRequestHeader(c *gin.Context, req *http.Header, info *common.RelayInfo) error {
channel.SetupApiRequestHeader(info, c, req)
req.Set("Authorization", "Bearer "+info.ApiKey)
return nil
}
+30
View File
@@ -0,0 +1,30 @@
package coze
var ModelList = []string{
"moonshot-v1-8k",
"moonshot-v1-32k",
"moonshot-v1-128k",
"Baichuan4",
"abab6.5s-chat-pro",
"glm-4-0520",
"qwen-max",
"deepseek-r1",
"deepseek-v3",
"deepseek-r1-distill-qwen-32b",
"deepseek-r1-distill-qwen-7b",
"step-1v-8k",
"step-1.5v-mini",
"Doubao-pro-32k",
"Doubao-pro-256k",
"Doubao-lite-128k",
"Doubao-lite-32k",
"Doubao-vision-lite-32k",
"Doubao-vision-pro-32k",
"Doubao-1.5-pro-vision-32k",
"Doubao-1.5-lite-32k",
"Doubao-1.5-pro-32k",
"Doubao-1.5-thinking-pro",
"Doubao-1.5-pro-256k",
}
var ChannelName = "coze"
+78
View File
@@ -0,0 +1,78 @@
package coze
import "encoding/json"
type CozeError struct {
Code int `json:"code"`
Message string `json:"message"`
}
type CozeEnterMessage struct {
Role string `json:"role"`
Type string `json:"type,omitempty"`
Content any `json:"content,omitempty"`
MetaData json.RawMessage `json:"meta_data,omitempty"`
ContentType string `json:"content_type,omitempty"`
}
type CozeChatRequest struct {
BotId string `json:"bot_id"`
UserId json.RawMessage `json:"user_id"`
AdditionalMessages []CozeEnterMessage `json:"additional_messages,omitempty"`
Stream bool `json:"stream,omitempty"`
CustomVariables json.RawMessage `json:"custom_variables,omitempty"`
AutoSaveHistory bool `json:"auto_save_history,omitempty"`
MetaData json.RawMessage `json:"meta_data,omitempty"`
ExtraParams json.RawMessage `json:"extra_params,omitempty"`
ShortcutCommand json.RawMessage `json:"shortcut_command,omitempty"`
Parameters json.RawMessage `json:"parameters,omitempty"`
}
type CozeChatResponse struct {
Code int `json:"code"`
Msg string `json:"msg"`
Data CozeChatResponseData `json:"data"`
}
type CozeChatResponseData struct {
Id string `json:"id"`
ConversationId string `json:"conversation_id"`
BotId string `json:"bot_id"`
CreatedAt int64 `json:"created_at"`
LastError CozeError `json:"last_error"`
Status string `json:"status"`
Usage CozeChatUsage `json:"usage"`
}
type CozeChatUsage struct {
TokenCount int `json:"token_count"`
OutputCount int `json:"output_count"`
InputCount int `json:"input_count"`
}
type CozeChatDetailResponse struct {
Data []CozeChatV3MessageDetail `json:"data"`
Code int `json:"code"`
Msg string `json:"msg"`
Detail CozeResponseDetail `json:"detail"`
}
type CozeChatV3MessageDetail struct {
Id string `json:"id"`
Role string `json:"role"`
Type string `json:"type"`
BotId string `json:"bot_id"`
ChatId string `json:"chat_id"`
Content json.RawMessage `json:"content"`
MetaData json.RawMessage `json:"meta_data"`
CreatedAt int64 `json:"created_at"`
SectionId string `json:"section_id"`
UpdatedAt int64 `json:"updated_at"`
ContentType string `json:"content_type"`
ConversationId string `json:"conversation_id"`
ReasoningContent string `json:"reasoning_content"`
}
type CozeResponseDetail struct {
Logid string `json:"logid"`
}
+292
View File
@@ -0,0 +1,292 @@
package coze
import (
"bufio"
"encoding/json"
"errors"
"fmt"
"io"
"net/http"
"strings"
"github.com/QuantumNous/new-api/common"
relaycommon "github.com/QuantumNous/new-api/relay/common"
"github.com/QuantumNous/new-api/relay/helper"
"github.com/QuantumNous/new-api/relaykit/dto"
"github.com/QuantumNous/new-api/relaykit/types"
"github.com/QuantumNous/new-api/service"
"github.com/samber/lo"
"github.com/gin-gonic/gin"
)
func convertCozeChatRequest(c *gin.Context, request dto.GeneralOpenAIRequest) *CozeChatRequest {
var messages []CozeEnterMessage
// 将 request的messages的role为user的content转换为CozeMessage
for _, message := range request.Messages {
if message.Role == "user" {
messages = append(messages, CozeEnterMessage{
Role: "user",
Content: message.Content,
// TODO: support more content type
ContentType: "text",
})
}
}
user := request.User
if len(user) == 0 {
user = json.RawMessage(helper.GetResponseID(c))
}
cozeRequest := &CozeChatRequest{
BotId: c.GetString("bot_id"),
UserId: user,
AdditionalMessages: messages,
Stream: lo.FromPtrOr(request.Stream, false),
}
return cozeRequest
}
func cozeChatHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http.Response) (*dto.Usage, *types.NewAPIError) {
responseBody, err := io.ReadAll(resp.Body)
if err != nil {
return nil, types.NewError(err, types.ErrorCodeBadResponseBody)
}
service.CloseResponseBodyGracefully(resp)
// convert coze response to openai response
var response dto.TextResponse
var cozeResponse CozeChatDetailResponse
response.Model = info.UpstreamModelName
err = json.Unmarshal(responseBody, &cozeResponse)
if err != nil {
return nil, types.NewError(err, types.ErrorCodeBadResponseBody)
}
if cozeResponse.Code != 0 {
return nil, types.NewError(errors.New(cozeResponse.Msg), types.ErrorCodeBadResponseBody)
}
// 从上下文获取 usage
var usage dto.Usage
usage.PromptTokens = c.GetInt("coze_input_count")
usage.CompletionTokens = c.GetInt("coze_output_count")
usage.TotalTokens = c.GetInt("coze_token_count")
response.Usage = usage
response.Id = helper.GetResponseID(c)
var responseContent json.RawMessage
for _, data := range cozeResponse.Data {
if data.Type == "answer" {
responseContent = data.Content
response.Created = data.CreatedAt
}
}
// 添加 response.Choices
response.Choices = []dto.OpenAITextResponseChoice{
{
Index: 0,
Message: dto.Message{Role: "assistant", Content: responseContent},
FinishReason: "stop",
},
}
jsonResponse, err := json.Marshal(response)
if err != nil {
return nil, types.NewError(err, types.ErrorCodeBadResponseBody)
}
c.Writer.Header().Set("Content-Type", "application/json")
c.Writer.WriteHeader(resp.StatusCode)
_, _ = c.Writer.Write(jsonResponse)
return &usage, nil
}
func cozeChatStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http.Response) (*dto.Usage, *types.NewAPIError) {
scanner := helper.NewStreamScanner(resp.Body)
scanner.Split(bufio.ScanLines)
helper.SetEventStreamHeaders(c)
id := helper.GetResponseID(c)
var responseText string
var currentEvent string
var currentData string
var usage = &dto.Usage{}
for scanner.Scan() {
line := scanner.Text()
if line == "" {
if currentEvent != "" && currentData != "" {
// handle last event
handleCozeEvent(c, currentEvent, currentData, &responseText, usage, id, info)
currentEvent = ""
currentData = ""
}
continue
}
if strings.HasPrefix(line, "event:") {
currentEvent = strings.TrimSpace(line[6:])
continue
}
if strings.HasPrefix(line, "data:") {
currentData = strings.TrimSpace(line[5:])
continue
}
}
// Last event
if currentEvent != "" && currentData != "" {
handleCozeEvent(c, currentEvent, currentData, &responseText, usage, id, info)
}
if err := scanner.Err(); err != nil {
return nil, types.NewError(err, types.ErrorCodeBadResponseBody)
}
helper.Done(c)
if usage.TotalTokens == 0 {
usage = service.ResponseText2Usage(c, responseText, info.UpstreamModelName, c.GetInt("coze_input_count"))
}
return usage, nil
}
func handleCozeEvent(c *gin.Context, event string, data string, responseText *string, usage *dto.Usage, id string, info *relaycommon.RelayInfo) {
switch event {
case "conversation.chat.completed":
// 将 data 解析为 CozeChatResponseData
var chatData CozeChatResponseData
err := json.Unmarshal([]byte(data), &chatData)
if err != nil {
common.SysLog("error_unmarshalling_stream_response: " + err.Error())
return
}
usage.PromptTokens = chatData.Usage.InputCount
usage.CompletionTokens = chatData.Usage.OutputCount
usage.TotalTokens = chatData.Usage.TokenCount
finishReason := "stop"
stopResponse := helper.GenerateStopResponse(id, common.GetTimestamp(), info.UpstreamModelName, finishReason)
helper.ObjectData(c, stopResponse)
case "conversation.message.delta":
// 将 data 解析为 CozeChatV3MessageDetail
var messageData CozeChatV3MessageDetail
err := json.Unmarshal([]byte(data), &messageData)
if err != nil {
common.SysLog("error_unmarshalling_stream_response: " + err.Error())
return
}
var content string
err = json.Unmarshal(messageData.Content, &content)
if err != nil {
common.SysLog("error_unmarshalling_stream_response: " + err.Error())
return
}
*responseText += content
openaiResponse := dto.ChatCompletionsStreamResponse{
Id: id,
Object: "chat.completion.chunk",
Created: common.GetTimestamp(),
Model: info.UpstreamModelName,
}
choice := dto.ChatCompletionsStreamResponseChoice{
Index: 0,
}
choice.Delta.SetContentString(content)
openaiResponse.Choices = append(openaiResponse.Choices, choice)
helper.ObjectData(c, openaiResponse)
case "error":
var errorData CozeError
err := json.Unmarshal([]byte(data), &errorData)
if err != nil {
common.SysLog("error_unmarshalling_stream_response: " + err.Error())
return
}
common.SysLog(fmt.Sprintf("stream event error: %v %v", errorData.Code, errorData.Message))
}
}
func checkIfChatComplete(a *Adaptor, c *gin.Context, info *relaycommon.RelayInfo) (error, bool) {
requestURL := fmt.Sprintf("%s/v3/chat/retrieve", info.ChannelBaseUrl)
requestURL = requestURL + "?conversation_id=" + c.GetString("coze_conversation_id") + "&chat_id=" + c.GetString("coze_chat_id")
// 将 conversationId和chatId作为参数发送get请求
req, err := http.NewRequest("GET", requestURL, nil)
if err != nil {
return err, false
}
err = a.SetupRequestHeader(c, &req.Header, info)
if err != nil {
return err, false
}
resp, err := doRequest(req, info) // 调用 doRequest
if err != nil {
return err, false
}
if resp == nil { // 确保在 doRequest 失败时 resp 不为 nil 导致 panic
return fmt.Errorf("resp is nil"), false
}
defer resp.Body.Close() // 确保响应体被关闭
// 解析 resp 到 CozeChatResponse
var cozeResponse CozeChatResponse
responseBody, err := io.ReadAll(resp.Body)
if err != nil {
return fmt.Errorf("read response body failed: %w", err), false
}
err = json.Unmarshal(responseBody, &cozeResponse)
if err != nil {
return fmt.Errorf("unmarshal response body failed: %w", err), false
}
if cozeResponse.Data.Status == "completed" {
// 在上下文设置 usage
c.Set("coze_token_count", cozeResponse.Data.Usage.TokenCount)
c.Set("coze_output_count", cozeResponse.Data.Usage.OutputCount)
c.Set("coze_input_count", cozeResponse.Data.Usage.InputCount)
return nil, true
} else if cozeResponse.Data.Status == "failed" || cozeResponse.Data.Status == "canceled" || cozeResponse.Data.Status == "requires_action" {
return fmt.Errorf("chat status: %s", cozeResponse.Data.Status), false
} else {
return nil, false
}
}
func getChatDetail(a *Adaptor, c *gin.Context, info *relaycommon.RelayInfo) (*http.Response, error) {
requestURL := fmt.Sprintf("%s/v3/chat/message/list", info.ChannelBaseUrl)
requestURL = requestURL + "?conversation_id=" + c.GetString("coze_conversation_id") + "&chat_id=" + c.GetString("coze_chat_id")
req, err := http.NewRequest("GET", requestURL, nil)
if err != nil {
return nil, fmt.Errorf("new request failed: %w", err)
}
err = a.SetupRequestHeader(c, &req.Header, info)
if err != nil {
return nil, fmt.Errorf("setup request header failed: %w", err)
}
resp, err := doRequest(req, info)
if err != nil {
return nil, fmt.Errorf("do request failed: %w", err)
}
return resp, nil
}
func doRequest(req *http.Request, info *relaycommon.RelayInfo) (*http.Response, error) {
client, err := service.GetHttpClientWithProxySettings(info.ChannelSetting.Proxy, info.ChannelSetting)
if err != nil {
return nil, fmt.Errorf("new proxy http client failed: %w", err)
}
resp, err := client.Do(req)
if err != nil { // 增加对 client.Do(req) 返回错误的检查
return nil, fmt.Errorf("client.Do failed: %w", err)
}
// _ = resp.Body.Close()
return resp, nil
}
+223
View File
@@ -0,0 +1,223 @@
package deepseek
import (
"errors"
"fmt"
"io"
"net/http"
"strings"
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/relay/channel"
"github.com/QuantumNous/new-api/relay/channel/claude"
"github.com/QuantumNous/new-api/relay/channel/openai"
relaycommon "github.com/QuantumNous/new-api/relay/common"
"github.com/QuantumNous/new-api/relay/constant"
"github.com/QuantumNous/new-api/relaykit/dto"
"github.com/QuantumNous/new-api/relaykit/types"
"github.com/QuantumNous/new-api/setting/model_setting"
"github.com/QuantumNous/new-api/setting/reasoning"
"github.com/gin-gonic/gin"
)
type Adaptor struct {
}
func (a *Adaptor) ConvertGeminiRequest(*gin.Context, *relaycommon.RelayInfo, *dto.GeminiChatRequest) (any, error) {
//TODO implement me
return nil, errors.New("not implemented")
}
func (a *Adaptor) ConvertClaudeRequest(c *gin.Context, info *relaycommon.RelayInfo, req *dto.ClaudeRequest) (any, error) {
adaptor := claude.Adaptor{}
convertedRequest, err := adaptor.ConvertClaudeRequest(c, info, req)
if err != nil {
return nil, err
}
claudeRequest, ok := convertedRequest.(*dto.ClaudeRequest)
if !ok {
return convertedRequest, nil
}
if err := applyDeepSeekV4ClaudeThinkingSuffix(info, claudeRequest); err != nil {
return nil, err
}
return claudeRequest, nil
}
func (a *Adaptor) ConvertAudioRequest(c *gin.Context, info *relaycommon.RelayInfo, request dto.AudioRequest) (io.Reader, error) {
//TODO implement me
return nil, errors.New("not implemented")
}
func (a *Adaptor) ConvertImageRequest(c *gin.Context, info *relaycommon.RelayInfo, request dto.ImageRequest) (any, error) {
//TODO implement me
return nil, errors.New("not implemented")
}
func (a *Adaptor) Init(info *relaycommon.RelayInfo) {
}
func (a *Adaptor) GetRequestURL(info *relaycommon.RelayInfo) (string, error) {
fimBaseUrl := info.ChannelBaseUrl
switch info.RelayFormat {
case types.RelayFormatClaude:
return fmt.Sprintf("%s/anthropic/v1/messages", info.ChannelBaseUrl), nil
default:
if !strings.HasSuffix(info.ChannelBaseUrl, "/beta") {
fimBaseUrl += "/beta"
}
switch info.RelayMode {
case constant.RelayModeCompletions:
return fmt.Sprintf("%s/completions", fimBaseUrl), nil
case constant.RelayModeResponses:
return fmt.Sprintf("%s/responses", info.ChannelBaseUrl), nil
default:
return fmt.Sprintf("%s/v1/chat/completions", info.ChannelBaseUrl), nil
}
}
}
func (a *Adaptor) SetupRequestHeader(c *gin.Context, req *http.Header, info *relaycommon.RelayInfo) error {
channel.SetupApiRequestHeader(info, c, req)
req.Set("Authorization", "Bearer "+info.ApiKey)
return nil
}
func (a *Adaptor) ConvertOpenAIRequest(c *gin.Context, info *relaycommon.RelayInfo, request *dto.GeneralOpenAIRequest) (any, error) {
if request == nil {
return nil, errors.New("request is nil")
}
if err := applyDeepSeekV4OpenAIThinkingSuffix(info, request); err != nil {
return nil, err
}
return request, nil
}
func applyDeepSeekV4OpenAIThinkingSuffix(info *relaycommon.RelayInfo, request *dto.GeneralOpenAIRequest) error {
modelName := request.Model
if info != nil && info.ChannelMeta != nil && info.UpstreamModelName != "" {
modelName = info.UpstreamModelName
}
if model_setting.ShouldPreserveThinkingSuffix(modelName) || info != nil && model_setting.ShouldPreserveThinkingSuffix(info.OriginModelName) {
return nil
}
baseModel, thinkingType, effort, ok := reasoning.ParseDeepSeekV4ThinkingSuffix(modelName)
if !ok {
return nil
}
thinking, err := common.Marshal(map[string]string{
"type": thinkingType,
})
if err != nil {
return fmt.Errorf("error marshalling thinking: %w", err)
}
request.Model = baseModel
request.THINKING = thinking
request.ReasoningEffort = effort
if info != nil {
if info.ChannelMeta != nil {
info.UpstreamModelName = baseModel
}
info.SetReasoningEffort(effort)
}
return nil
}
func applyDeepSeekV4ClaudeThinkingSuffix(info *relaycommon.RelayInfo, request *dto.ClaudeRequest) error {
modelName := request.Model
if info != nil && info.ChannelMeta != nil && info.UpstreamModelName != "" {
modelName = info.UpstreamModelName
}
if model_setting.ShouldPreserveThinkingSuffix(modelName) || info != nil && model_setting.ShouldPreserveThinkingSuffix(info.OriginModelName) {
return nil
}
baseModel, thinkingType, effort, ok := reasoning.ParseDeepSeekV4ThinkingSuffix(modelName)
if !ok {
return nil
}
request.Model = baseModel
request.Thinking = &dto.Thinking{Type: thinkingType}
if effort == "" {
request.OutputConfig = nil
} else {
outputConfig, err := common.Marshal(map[string]string{
"effort": effort,
})
if err != nil {
return fmt.Errorf("error marshalling output_config: %w", err)
}
request.OutputConfig = outputConfig
}
if info != nil {
if info.ChannelMeta != nil {
info.UpstreamModelName = baseModel
}
info.SetReasoningEffort(effort)
}
return nil
}
func (a *Adaptor) ConvertRerankRequest(c *gin.Context, relayMode int, request dto.RerankRequest) (any, error) {
return nil, nil
}
func (a *Adaptor) ConvertEmbeddingRequest(c *gin.Context, info *relaycommon.RelayInfo, request dto.EmbeddingRequest) (any, error) {
//TODO implement me
return nil, errors.New("not implemented")
}
func (a *Adaptor) ConvertOpenAIResponsesRequest(_ *gin.Context, info *relaycommon.RelayInfo, request dto.OpenAIResponsesRequest) (any, error) {
applyDeepSeekV4ResponsesThinkingSuffix(info, &request)
return request, nil
}
func applyDeepSeekV4ResponsesThinkingSuffix(info *relaycommon.RelayInfo, request *dto.OpenAIResponsesRequest) {
modelName := request.Model
if info != nil && info.ChannelMeta != nil && info.UpstreamModelName != "" {
modelName = info.UpstreamModelName
}
if model_setting.ShouldPreserveThinkingSuffix(modelName) || info != nil && model_setting.ShouldPreserveThinkingSuffix(info.OriginModelName) {
return
}
baseModel, thinkingType, effort, ok := reasoning.ParseDeepSeekV4ThinkingSuffix(modelName)
if ok {
if thinkingType == "disabled" {
effort = "none"
}
request.Model = baseModel
if request.Reasoning == nil {
request.Reasoning = &dto.Reasoning{}
}
request.Reasoning.Effort = effort
if info != nil && info.ChannelMeta != nil {
info.UpstreamModelName = baseModel
}
}
if info != nil && request.Reasoning != nil {
info.SetReasoningEffort(request.Reasoning.Effort)
}
}
func (a *Adaptor) DoRequest(c *gin.Context, info *relaycommon.RelayInfo, requestBody io.Reader) (any, error) {
return channel.DoApiRequest(a, c, info, requestBody)
}
func (a *Adaptor) DoResponse(c *gin.Context, resp *http.Response, info *relaycommon.RelayInfo) (usage any, err *types.NewAPIError) {
switch info.RelayFormat {
case types.RelayFormatClaude:
adaptor := claude.Adaptor{}
return adaptor.DoResponse(c, resp, info)
default:
adaptor := openai.Adaptor{}
return adaptor.DoResponse(c, resp, info)
}
}
func (a *Adaptor) GetModelList() []string {
return ModelList
}
func (a *Adaptor) GetChannelName() string {
return ChannelName
}
+9
View File
@@ -0,0 +1,9 @@
package deepseek
var ModelList = []string{
"deepseek-chat", "deepseek-reasoner",
"deepseek-v4-flash", "deepseek-v4-flash-none", "deepseek-v4-flash-max",
"deepseek-v4-pro", "deepseek-v4-pro-none", "deepseek-v4-pro-max",
}
var ChannelName = "deepseek"
+119
View File
@@ -0,0 +1,119 @@
package dify
import (
"errors"
"fmt"
"io"
"net/http"
"github.com/QuantumNous/new-api/relay/channel"
relaycommon "github.com/QuantumNous/new-api/relay/common"
"github.com/QuantumNous/new-api/relaykit/dto"
"github.com/QuantumNous/new-api/relaykit/types"
"github.com/gin-gonic/gin"
)
const (
BotTypeChatFlow = 1 // chatflow default
BotTypeAgent = 2
BotTypeWorkFlow = 3
BotTypeCompletion = 4
)
type Adaptor struct {
BotType int
}
func (a *Adaptor) ConvertGeminiRequest(*gin.Context, *relaycommon.RelayInfo, *dto.GeminiChatRequest) (any, error) {
//TODO implement me
return nil, errors.New("not implemented")
}
func (a *Adaptor) ConvertClaudeRequest(*gin.Context, *relaycommon.RelayInfo, *dto.ClaudeRequest) (any, error) {
//TODO implement me
panic("implement me")
}
func (a *Adaptor) ConvertAudioRequest(c *gin.Context, info *relaycommon.RelayInfo, request dto.AudioRequest) (io.Reader, error) {
//TODO implement me
return nil, errors.New("not implemented")
}
func (a *Adaptor) ConvertImageRequest(c *gin.Context, info *relaycommon.RelayInfo, request dto.ImageRequest) (any, error) {
//TODO implement me
return nil, errors.New("not implemented")
}
func (a *Adaptor) Init(info *relaycommon.RelayInfo) {
//if strings.HasPrefix(info.UpstreamModelName, "agent") {
// a.BotType = BotTypeAgent
//} else if strings.HasPrefix(info.UpstreamModelName, "workflow") {
// a.BotType = BotTypeWorkFlow
//} else if strings.HasPrefix(info.UpstreamModelName, "chat") {
// a.BotType = BotTypeCompletion
//} else {
//}
a.BotType = BotTypeChatFlow
}
func (a *Adaptor) GetRequestURL(info *relaycommon.RelayInfo) (string, error) {
switch a.BotType {
case BotTypeWorkFlow:
return fmt.Sprintf("%s/v1/workflows/run", info.ChannelBaseUrl), nil
case BotTypeCompletion:
return fmt.Sprintf("%s/v1/completion-messages", info.ChannelBaseUrl), nil
case BotTypeAgent:
fallthrough
default:
return fmt.Sprintf("%s/v1/chat-messages", info.ChannelBaseUrl), nil
}
}
func (a *Adaptor) SetupRequestHeader(c *gin.Context, req *http.Header, info *relaycommon.RelayInfo) error {
channel.SetupApiRequestHeader(info, c, req)
req.Set("Authorization", "Bearer "+info.ApiKey)
return nil
}
func (a *Adaptor) ConvertOpenAIRequest(c *gin.Context, info *relaycommon.RelayInfo, request *dto.GeneralOpenAIRequest) (any, error) {
if request == nil {
return nil, errors.New("request is nil")
}
return requestOpenAI2Dify(c, info, *request), nil
}
func (a *Adaptor) ConvertRerankRequest(c *gin.Context, relayMode int, request dto.RerankRequest) (any, error) {
return nil, nil
}
func (a *Adaptor) ConvertEmbeddingRequest(c *gin.Context, info *relaycommon.RelayInfo, request dto.EmbeddingRequest) (any, error) {
//TODO implement me
return nil, errors.New("not implemented")
}
func (a *Adaptor) ConvertOpenAIResponsesRequest(c *gin.Context, info *relaycommon.RelayInfo, request dto.OpenAIResponsesRequest) (any, error) {
// TODO implement me
return nil, errors.New("not implemented")
}
func (a *Adaptor) DoRequest(c *gin.Context, info *relaycommon.RelayInfo, requestBody io.Reader) (any, error) {
return channel.DoApiRequest(a, c, info, requestBody)
}
func (a *Adaptor) DoResponse(c *gin.Context, resp *http.Response, info *relaycommon.RelayInfo) (usage any, err *types.NewAPIError) {
if info.IsStream {
return difyStreamHandler(c, info, resp)
} else {
return difyHandler(c, info, resp)
}
}
func (a *Adaptor) GetModelList() []string {
return ModelList
}
func (a *Adaptor) GetChannelName() string {
return ChannelName
}
+5
View File
@@ -0,0 +1,5 @@
package dify
var ModelList []string
var ChannelName = "dify"
+47
View File
@@ -0,0 +1,47 @@
package dify
import (
"github.com/QuantumNous/new-api/relaykit/dto"
)
type DifyChatRequest struct {
Inputs map[string]interface{} `json:"inputs"`
Query string `json:"query"`
ResponseMode string `json:"response_mode"`
User string `json:"user"`
AutoGenerateName bool `json:"auto_generate_name"`
Files []DifyFile `json:"files"`
}
type DifyFile struct {
Type string `json:"type"`
TransferMode string `json:"transfer_mode"`
URL string `json:"url,omitempty"`
UploadFileId string `json:"upload_file_id,omitempty"`
}
type DifyMetaData struct {
Usage dto.Usage `json:"usage"`
}
type DifyData struct {
WorkflowId string `json:"workflow_id"`
NodeId string `json:"node_id"`
NodeType string `json:"node_type"`
Status string `json:"status"`
}
type DifyChatCompletionResponse struct {
ConversationId string `json:"conversation_id"`
Answer string `json:"answer"`
CreateAt int64 `json:"create_at"`
MetaData DifyMetaData `json:"metadata"`
}
type DifyChunkChatCompletionResponse struct {
Event string `json:"event"`
ConversationId string `json:"conversation_id"`
Answer string `json:"answer"`
Data DifyData `json:"data"`
MetaData DifyMetaData `json:"metadata"`
}
+301
View File
@@ -0,0 +1,301 @@
package dify
import (
"bytes"
"encoding/base64"
"encoding/json"
"fmt"
"io"
"mime/multipart"
"net/http"
"os"
"strings"
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/constant"
relaycommon "github.com/QuantumNous/new-api/relay/common"
"github.com/QuantumNous/new-api/relay/helper"
"github.com/QuantumNous/new-api/relaykit/dto"
"github.com/QuantumNous/new-api/relaykit/types"
"github.com/QuantumNous/new-api/service"
"github.com/samber/lo"
"github.com/gin-gonic/gin"
)
func uploadDifyFile(c *gin.Context, info *relaycommon.RelayInfo, user string, media dto.MediaContent) *DifyFile {
uploadUrl := fmt.Sprintf("%s/v1/files/upload", info.ChannelBaseUrl)
switch media.Type {
case dto.ContentTypeImageURL:
// Decode base64 data
imageMedia := media.GetImageMedia()
base64Data := imageMedia.Url
// Remove base64 prefix if exists (e.g., "data:image/jpeg;base64,")
if idx := strings.Index(base64Data, ","); idx != -1 {
base64Data = base64Data[idx+1:]
}
// Decode base64 string
decodedData, err := base64.StdEncoding.DecodeString(base64Data)
if err != nil {
common.SysLog("failed to decode base64: " + err.Error())
return nil
}
// Create temporary file
tempFile, err := os.CreateTemp("", "dify-upload-*")
if err != nil {
common.SysLog("failed to create temp file: " + err.Error())
return nil
}
defer tempFile.Close()
defer os.Remove(tempFile.Name())
// Write decoded data to temp file
if _, err := tempFile.Write(decodedData); err != nil {
common.SysLog("failed to write to temp file: " + err.Error())
return nil
}
// Create multipart form
body := &bytes.Buffer{}
writer := multipart.NewWriter(body)
// Add user field
if err := writer.WriteField("user", user); err != nil {
common.SysLog("failed to add user field: " + err.Error())
return nil
}
// Create form file with proper mime type
mimeType := imageMedia.MimeType
if mimeType == "" {
mimeType = "image/jpeg" // default mime type
}
// Create form file
part, err := writer.CreateFormFile("file", fmt.Sprintf("image.%s", strings.TrimPrefix(mimeType, "image/")))
if err != nil {
common.SysLog("failed to create form file: " + err.Error())
return nil
}
// Copy file content to form
if _, err = io.Copy(part, bytes.NewReader(decodedData)); err != nil {
common.SysLog("failed to copy file content: " + err.Error())
return nil
}
writer.Close()
// Create HTTP request
req, err := http.NewRequest("POST", uploadUrl, body)
if err != nil {
common.SysLog("failed to create request: " + err.Error())
return nil
}
req.Header.Set("Content-Type", writer.FormDataContentType())
req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", info.ApiKey))
// Send request
client := service.GetHttpClient()
resp, err := client.Do(req)
if err != nil {
common.SysLog("failed to send request: " + err.Error())
return nil
}
defer resp.Body.Close()
// Parse response
var result struct {
Id string `json:"id"`
}
if err := json.NewDecoder(resp.Body).Decode(&result); err != nil {
common.SysLog("failed to decode response: " + err.Error())
return nil
}
return &DifyFile{
UploadFileId: result.Id,
Type: "image",
TransferMode: "local_file",
}
}
return nil
}
func requestOpenAI2Dify(c *gin.Context, info *relaycommon.RelayInfo, request dto.GeneralOpenAIRequest) *DifyChatRequest {
difyReq := DifyChatRequest{
Inputs: make(map[string]interface{}),
AutoGenerateName: false,
}
user := request.User
if len(user) == 0 {
user = json.RawMessage(helper.GetResponseID(c))
}
var stringUser string
err := json.Unmarshal(user, &stringUser)
if err != nil {
common.SysLog("failed to unmarshal user: " + err.Error())
stringUser = helper.GetResponseID(c)
}
difyReq.User = stringUser
files := make([]DifyFile, 0)
var content strings.Builder
for _, message := range request.Messages {
if message.Role == "system" {
content.WriteString("SYSTEM: \n" + message.StringContent() + "\n")
} else if message.Role == "assistant" {
content.WriteString("ASSISTANT: \n" + message.StringContent() + "\n")
} else {
parseContent := message.ParseContent()
for _, mediaContent := range parseContent {
switch mediaContent.Type {
case dto.ContentTypeText:
content.WriteString("USER: \n" + mediaContent.Text + "\n")
case dto.ContentTypeImageURL:
media := mediaContent.GetImageMedia()
var file *DifyFile
if media.IsRemoteImage() {
// 修复 #2083: 远程图片分支此前未初始化 file,
// 导致 file.Type = ... 触发 nil pointer dereference
// 而 panic(500: "invalid memory address or nil pointer dereference")。
file = &DifyFile{
Type: media.MimeType,
TransferMode: "remote_url",
URL: media.Url,
}
} else {
file = uploadDifyFile(c, info, difyReq.User, mediaContent)
}
if file != nil {
files = append(files, *file)
}
}
}
}
}
difyReq.Query = content.String()
difyReq.Files = files
mode := "blocking"
if lo.FromPtrOr(request.Stream, false) {
mode = "streaming"
}
difyReq.ResponseMode = mode
return &difyReq
}
func streamResponseDify2OpenAI(difyResponse DifyChunkChatCompletionResponse) *dto.ChatCompletionsStreamResponse {
response := dto.ChatCompletionsStreamResponse{
Object: "chat.completion.chunk",
Created: common.GetTimestamp(),
Model: "dify",
}
var choice dto.ChatCompletionsStreamResponseChoice
if strings.HasPrefix(difyResponse.Event, "workflow_") {
if constant.DifyDebug {
text := "Workflow: " + difyResponse.Data.WorkflowId
if difyResponse.Event == "workflow_finished" {
text += " " + difyResponse.Data.Status
}
choice.Delta.SetReasoningContent(text + "\n")
}
} else if strings.HasPrefix(difyResponse.Event, "node_") {
if constant.DifyDebug {
text := "Node: " + difyResponse.Data.NodeType
if difyResponse.Event == "node_finished" {
text += " " + difyResponse.Data.Status
}
choice.Delta.SetReasoningContent(text + "\n")
}
} else if difyResponse.Event == "message" || difyResponse.Event == "agent_message" {
if difyResponse.Answer == "<details style=\"color:gray;background-color: #f8f8f8;padding: 8px;border-radius: 4px;\" open> <summary> Thinking... </summary>\n" {
difyResponse.Answer = "<think>"
} else if difyResponse.Answer == "</details>" {
difyResponse.Answer = "</think>"
}
choice.Delta.SetContentString(difyResponse.Answer)
}
response.Choices = append(response.Choices, choice)
return &response
}
func difyStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http.Response) (*dto.Usage, *types.NewAPIError) {
var responseText string
usage := &dto.Usage{}
var nodeToken int
helper.SetEventStreamHeaders(c)
helper.StreamScannerHandler(c, resp, info, func(data string, sr *helper.StreamResult) {
var difyResponse DifyChunkChatCompletionResponse
if err := json.Unmarshal([]byte(data), &difyResponse); err != nil {
common.SysLog("error unmarshalling stream response: " + err.Error())
sr.Error(err)
return
}
if difyResponse.Event == "message_end" {
usage = &difyResponse.MetaData.Usage
sr.Done()
return
} else if difyResponse.Event == "error" {
sr.Stop(fmt.Errorf("dify error event"))
return
}
openaiResponse := *streamResponseDify2OpenAI(difyResponse)
if len(openaiResponse.Choices) != 0 {
responseText += openaiResponse.Choices[0].Delta.GetContentString()
if openaiResponse.Choices[0].Delta.ReasoningContent != nil {
nodeToken += 1
}
}
if err := helper.ObjectData(c, openaiResponse); err != nil {
common.SysLog(err.Error())
sr.Error(err)
}
})
helper.Done(c)
if usage.TotalTokens == 0 {
usage = service.ResponseText2Usage(c, responseText, info.UpstreamModelName, info.GetEstimatePromptTokens())
}
usage.CompletionTokens += nodeToken
return usage, nil
}
func difyHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http.Response) (*dto.Usage, *types.NewAPIError) {
var difyResponse DifyChatCompletionResponse
responseBody, err := io.ReadAll(resp.Body)
if err != nil {
return nil, types.NewError(err, types.ErrorCodeBadResponseBody)
}
service.CloseResponseBodyGracefully(resp)
err = json.Unmarshal(responseBody, &difyResponse)
if err != nil {
return nil, types.NewError(err, types.ErrorCodeBadResponseBody)
}
fullTextResponse := dto.OpenAITextResponse{
Id: difyResponse.ConversationId,
Object: "chat.completion",
Created: common.GetTimestamp(),
Usage: difyResponse.MetaData.Usage,
}
choice := dto.OpenAITextResponseChoice{
Index: 0,
Message: dto.Message{
Role: "assistant",
Content: difyResponse.Answer,
},
FinishReason: "stop",
}
fullTextResponse.Choices = append(fullTextResponse.Choices, choice)
jsonResponse, err := json.Marshal(fullTextResponse)
if err != nil {
return nil, types.NewError(err, types.ErrorCodeBadResponseBody)
}
c.Writer.Header().Set("Content-Type", "application/json")
c.Writer.WriteHeader(resp.StatusCode)
c.Writer.Write(jsonResponse)
return &difyResponse.MetaData.Usage, nil
}
+290
View File
@@ -0,0 +1,290 @@
package gemini
import (
"errors"
"fmt"
"io"
"net/http"
"strings"
"github.com/QuantumNous/new-api/relay/channel"
relaycommon "github.com/QuantumNous/new-api/relay/common"
"github.com/QuantumNous/new-api/relay/constant"
"github.com/QuantumNous/new-api/relaykit/dto"
"github.com/QuantumNous/new-api/relaykit/relayconvert"
"github.com/QuantumNous/new-api/relaykit/types"
"github.com/QuantumNous/new-api/service"
"github.com/QuantumNous/new-api/setting/model_setting"
"github.com/gin-gonic/gin"
"github.com/samber/lo"
)
type Adaptor struct {
}
func (a *Adaptor) ConvertGeminiRequest(c *gin.Context, info *relaycommon.RelayInfo, request *dto.GeminiChatRequest) (any, error) {
if err := relayconvert.ApplyGeminiThinkingConfigChecked(request, info); err != nil {
return nil, err
}
if len(request.Contents) > 0 {
for i, content := range request.Contents {
if i == 0 {
if request.Contents[0].Role == "" {
request.Contents[0].Role = "user"
}
}
for _, part := range content.Parts {
if part.FileData != nil {
if part.FileData.MimeType == "" && strings.Contains(part.FileData.FileUri, "www.youtube.com") {
part.FileData.MimeType = "video/webm"
}
}
}
}
}
return request, nil
}
func (a *Adaptor) ConvertClaudeRequest(c *gin.Context, info *relaycommon.RelayInfo, req *dto.ClaudeRequest) (any, error) {
result, err := service.ConvertRequest(c, info, types.RelayFormatGemini, req)
if err != nil {
return nil, err
}
geminiRequest, ok := result.Value.(*dto.GeminiChatRequest)
if !ok {
return nil, fmt.Errorf("expected Gemini generateContent request, got %T", result.Value)
}
return geminiRequest, nil
}
func (a *Adaptor) ConvertAudioRequest(c *gin.Context, info *relaycommon.RelayInfo, request dto.AudioRequest) (io.Reader, error) {
//TODO implement me
return nil, errors.New("not implemented")
}
func (a *Adaptor) ConvertImageRequest(c *gin.Context, info *relaycommon.RelayInfo, request dto.ImageRequest) (any, error) {
if !strings.HasPrefix(info.UpstreamModelName, "imagen") {
return nil, errors.New("not supported model for image generation, only imagen models are supported")
}
// convert size to aspect ratio but allow user to specify aspect ratio
aspectRatio := "1:1" // default aspect ratio
size := strings.TrimSpace(request.Size)
if size != "" {
if strings.Contains(size, ":") {
aspectRatio = size
} else {
switch size {
case "256x256", "512x512", "1024x1024":
aspectRatio = "1:1"
case "1536x1024":
aspectRatio = "3:2"
case "1024x1536":
aspectRatio = "2:3"
case "1024x1792":
aspectRatio = "9:16"
case "1792x1024":
aspectRatio = "16:9"
}
}
}
// build gemini imagen request
geminiRequest := dto.GeminiImageRequest{
Instances: []dto.GeminiImageInstance{
{
Prompt: request.Prompt,
},
},
Parameters: dto.GeminiImageParameters{
SampleCount: int(lo.FromPtrOr(request.N, uint(1))),
AspectRatio: aspectRatio,
PersonGeneration: "allow_adult", // default allow adult
},
}
// Set imageSize when quality parameter is specified
// Map quality parameter to imageSize (only supported by Standard and Ultra models)
// quality values: auto, high, medium, low (for gpt-image-1), hd, standard (for dall-e-3)
// imageSize values: 1K (default), 2K
// https://ai.google.dev/gemini-api/docs/imagen
// https://platform.openai.com/docs/api-reference/images/create
if request.Quality != "" {
imageSize := "1K" // default
switch request.Quality {
case "hd", "high":
imageSize = "2K"
case "2K":
imageSize = "2K"
case "standard", "medium", "low", "auto", "1K":
imageSize = "1K"
default:
// unknown quality value, default to 1K
imageSize = "1K"
}
geminiRequest.Parameters.ImageSize = imageSize
}
return geminiRequest, nil
}
func (a *Adaptor) Init(info *relaycommon.RelayInfo) {
}
func (a *Adaptor) GetRequestURL(info *relaycommon.RelayInfo) (string, error) {
version := model_setting.GetGeminiVersionSetting(info.UpstreamModelName)
if strings.HasPrefix(info.UpstreamModelName, "imagen") {
return fmt.Sprintf("%s/%s/models/%s:predict", info.ChannelBaseUrl, version, info.UpstreamModelName), nil
}
if strings.HasPrefix(info.UpstreamModelName, "text-embedding") ||
strings.HasPrefix(info.UpstreamModelName, "embedding") ||
strings.HasPrefix(info.UpstreamModelName, "gemini-embedding") {
action := "embedContent"
if info.IsGeminiBatchEmbedding {
action = "batchEmbedContents"
}
return fmt.Sprintf("%s/%s/models/%s:%s", info.ChannelBaseUrl, version, info.UpstreamModelName, action), nil
}
action := "generateContent"
if info.IsStream {
action = "streamGenerateContent?alt=sse"
if info.RelayMode == constant.RelayModeGemini {
info.DisablePing = true
}
}
return fmt.Sprintf("%s/%s/models/%s:%s", info.ChannelBaseUrl, version, info.UpstreamModelName, action), nil
}
func (a *Adaptor) SetupRequestHeader(c *gin.Context, req *http.Header, info *relaycommon.RelayInfo) error {
channel.SetupApiRequestHeader(info, c, req)
req.Set("x-goog-api-key", info.ApiKey)
return nil
}
func (a *Adaptor) ConvertOpenAIRequest(c *gin.Context, info *relaycommon.RelayInfo, request *dto.GeneralOpenAIRequest) (any, error) {
if request == nil {
return nil, errors.New("request is nil")
}
result, err := service.ConvertRequest(c, info, types.RelayFormatGemini, request)
if err != nil {
return nil, err
}
return result.Value, nil
}
func (a *Adaptor) ConvertRerankRequest(c *gin.Context, relayMode int, request dto.RerankRequest) (any, error) {
return nil, nil
}
func (a *Adaptor) ConvertEmbeddingRequest(c *gin.Context, info *relaycommon.RelayInfo, request dto.EmbeddingRequest) (any, error) {
if request.Input == nil {
return nil, errors.New("input is required")
}
inputs := request.ParseInput()
if len(inputs) == 0 {
return nil, errors.New("input is empty")
}
// We always build a batch-style payload with `requests`, so ensure we call the
// batch endpoint upstream to avoid payload/endpoint mismatches.
info.IsGeminiBatchEmbedding = true
// process all inputs
geminiRequests := make([]map[string]any, 0, len(inputs))
for _, input := range inputs {
geminiRequest := map[string]any{
"model": fmt.Sprintf("models/%s", info.UpstreamModelName),
"content": dto.GeminiChatContent{
Parts: []dto.GeminiPart{
{
Text: input,
},
},
},
}
// set specific parameters for different models
// https://ai.google.dev/api/embeddings?hl=zh-cn#method:-models.embedcontent
switch info.UpstreamModelName {
case "text-embedding-004", "gemini-embedding-exp-03-07", "gemini-embedding-001":
// Only newer models introduced after 2024 support OutputDimensionality
dimensions := lo.FromPtrOr(request.Dimensions, 0)
if dimensions > 0 {
geminiRequest["outputDimensionality"] = dimensions
}
}
geminiRequests = append(geminiRequests, geminiRequest)
}
return map[string]any{
"requests": geminiRequests,
}, nil
}
func (a *Adaptor) ConvertOpenAIResponsesRequest(c *gin.Context, info *relaycommon.RelayInfo, request dto.OpenAIResponsesRequest) (any, error) {
result, err := service.ConvertRequest(c, info, types.RelayFormatGemini, &request)
if err != nil {
return nil, err
}
geminiRequest, ok := result.Value.(*dto.GeminiChatRequest)
if !ok {
return nil, fmt.Errorf("expected Gemini generateContent request, got %T", result.Value)
}
return geminiRequest, nil
}
func (a *Adaptor) DoRequest(c *gin.Context, info *relaycommon.RelayInfo, requestBody io.Reader) (any, error) {
return channel.DoApiRequest(a, c, info, requestBody)
}
func (a *Adaptor) DoResponse(c *gin.Context, resp *http.Response, info *relaycommon.RelayInfo) (usage any, err *types.NewAPIError) {
if info.RelayMode == constant.RelayModeResponses {
if info.IsStream {
return GeminiResponsesStreamHandler(c, info, resp)
}
return GeminiResponsesHandler(c, info, resp)
}
if info.RelayMode == constant.RelayModeGemini {
if strings.Contains(info.RequestURLPath, ":embedContent") ||
strings.Contains(info.RequestURLPath, ":batchEmbedContents") {
return NativeGeminiEmbeddingHandler(c, resp, info)
}
if info.IsStream {
return GeminiTextGenerationStreamHandler(c, info, resp)
} else {
return GeminiTextGenerationHandler(c, info, resp)
}
}
if strings.HasPrefix(info.UpstreamModelName, "imagen") {
return GeminiImageHandler(c, info, resp)
}
// check if the model is an embedding model
if strings.HasPrefix(info.UpstreamModelName, "text-embedding") ||
strings.HasPrefix(info.UpstreamModelName, "embedding") ||
strings.HasPrefix(info.UpstreamModelName, "gemini-embedding") {
return GeminiEmbeddingHandler(c, info, resp)
}
if info.IsStream {
return GeminiChatStreamHandler(c, info, resp)
} else {
return GeminiChatHandler(c, info, resp)
}
}
func (a *Adaptor) GetModelList() []string {
return ModelList
}
func (a *Adaptor) GetChannelName() string {
return ChannelName
}
@@ -0,0 +1,176 @@
package gemini
import (
"testing"
"github.com/QuantumNous/new-api/common"
relaycommon "github.com/QuantumNous/new-api/relay/common"
"github.com/QuantumNous/new-api/relaykit/dto"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"github.com/tidwall/gjson"
)
func TestConvertOpenAIResponsesRequestToGeminiInstructionsAndInput(t *testing.T) {
got := mustConvertResponsesToGemini(t, dto.OpenAIResponsesRequest{
Model: "gemini-test",
Instructions: mustGeminiRawMessage(t, "system rules"),
Input: mustGeminiRawMessage(t, "hello"),
})
require.NotNil(t, got.SystemInstructions)
require.Len(t, got.SystemInstructions.Parts, 1)
assert.Equal(t, "system rules", got.SystemInstructions.Parts[0].Text)
require.Len(t, got.Contents, 1)
assert.Equal(t, "user", got.Contents[0].Role)
require.Len(t, got.Contents[0].Parts, 1)
assert.Equal(t, "hello", got.Contents[0].Parts[0].Text)
}
func TestConvertOpenAIResponsesRequestToGeminiFunctionToolAndChoice(t *testing.T) {
got := mustConvertResponsesToGemini(t, dto.OpenAIResponsesRequest{
Model: "gemini-test",
Input: mustGeminiRawMessage(t, "lookup weather"),
Tools: mustGeminiRawMessage(t, []map[string]any{
{
"type": "function",
"name": "lookup",
"description": "Lookup data",
"parameters": map[string]any{
"type": "object",
"properties": map[string]any{
"q": map[string]any{"type": "string"},
},
},
},
{"type": "custom", "name": "freeform"},
}),
ToolChoice: mustGeminiRawMessage(t, map[string]any{
"type": "function",
"name": "lookup",
}),
})
tools := got.GetTools()
require.Len(t, tools, 1)
assert.Equal(t, "lookup", gjson.GetBytes(got.Tools, "0.functionDeclarations.0.name").String())
assert.Equal(t, "Lookup data", gjson.GetBytes(got.Tools, "0.functionDeclarations.0.description").String())
require.NotNil(t, got.ToolConfig)
require.NotNil(t, got.ToolConfig.FunctionCallingConfig)
assert.Equal(t, dto.FunctionCallingConfigMode("ANY"), got.ToolConfig.FunctionCallingConfig.Mode)
assert.Equal(t, []string{"lookup"}, got.ToolConfig.FunctionCallingConfig.AllowedFunctionNames)
}
func TestConvertOpenAIResponsesRequestToGeminiFunctionCallConversation(t *testing.T) {
got := mustConvertResponsesToGemini(t, dto.OpenAIResponsesRequest{
Model: "gemini-test",
Input: mustGeminiRawMessage(t, []map[string]any{
{
"role": "assistant",
"content": []map[string]any{
{"type": "output_text", "text": "I will call."},
},
},
{
"type": "function_call",
"call_id": "call_1",
"name": "lookup",
"arguments": map[string]any{"q": "x"},
},
{
"type": "function_call_output",
"call_id": "call_1",
"output": map[string]any{"ok": true},
},
}),
Tools: mustGeminiRawMessage(t, []map[string]any{
{"type": "function", "name": "lookup", "parameters": map[string]any{"type": "object"}},
}),
})
require.Len(t, got.Contents, 2)
assert.Equal(t, "model", got.Contents[0].Role)
require.Len(t, got.Contents[0].Parts, 2)
require.NotNil(t, got.Contents[0].Parts[0].FunctionCall)
assert.Equal(t, "lookup", got.Contents[0].Parts[0].FunctionCall.FunctionName)
assert.Equal(t, map[string]any{"q": "x"}, got.Contents[0].Parts[0].FunctionCall.Arguments)
assert.Equal(t, "I will call.", got.Contents[0].Parts[1].Text)
assert.Equal(t, "user", got.Contents[1].Role)
require.Len(t, got.Contents[1].Parts, 1)
require.NotNil(t, got.Contents[1].Parts[0].FunctionResponse)
assert.Equal(t, "lookup", got.Contents[1].Parts[0].FunctionResponse.Name)
assert.Equal(t, map[string]any{"ok": true}, got.Contents[1].Parts[0].FunctionResponse.Response)
}
func TestConvertOpenAIResponsesRequestToGeminiSkipsCustomToolCalls(t *testing.T) {
got := mustConvertResponsesToGemini(t, dto.OpenAIResponsesRequest{
Model: "gemini-test",
Input: mustGeminiRawMessage(t, []map[string]any{
{
"role": "assistant",
"content": []map[string]any{
{"type": "output_text", "text": "before custom"},
},
},
{
"type": "custom_tool_call",
"call_id": "call_custom",
"name": "apply_patch",
"input": "patch body",
},
{
"type": "custom_tool_call_output",
"call_id": "call_custom",
"output": "ok",
},
{
"type": "function_call_output",
"call_id": "call_custom",
"output": "legacy custom output",
},
{
"role": "user",
"content": "next turn",
},
}),
Tools: mustGeminiRawMessage(t, []map[string]any{
{"type": "custom", "name": "apply_patch"},
{"type": "unknown", "name": "unknown"},
}),
})
assert.Empty(t, got.GetTools())
require.Len(t, got.Contents, 2)
assert.Equal(t, "model", got.Contents[0].Role)
require.Len(t, got.Contents[0].Parts, 1)
assert.Equal(t, "before custom", got.Contents[0].Parts[0].Text)
assert.Nil(t, got.Contents[0].Parts[0].FunctionCall)
assert.Equal(t, "user", got.Contents[1].Role)
require.Len(t, got.Contents[1].Parts, 1)
assert.Equal(t, "next turn", got.Contents[1].Parts[0].Text)
assert.Nil(t, got.Contents[1].Parts[0].FunctionResponse)
}
func mustConvertResponsesToGemini(t *testing.T, req dto.OpenAIResponsesRequest) *dto.GeminiChatRequest {
t.Helper()
info := &relaycommon.RelayInfo{
OriginModelName: req.Model,
ChannelMeta: &relaycommon.ChannelMeta{
UpstreamModelName: req.Model,
},
}
got, err := (&Adaptor{}).ConvertOpenAIResponsesRequest(nil, info, req)
require.NoError(t, err)
geminiReq, ok := got.(*dto.GeminiChatRequest)
require.True(t, ok)
return geminiReq
}
func mustGeminiRawMessage(t *testing.T, value any) []byte {
t.Helper()
raw, err := common.Marshal(value)
require.NoError(t, err)
return raw
}
+44
View File
@@ -0,0 +1,44 @@
package gemini
var ModelList = []string{
// stable version
"gemini-2.5-flash", "gemini-2.5-pro", "gemini-2.0-flash",
"gemini-2.0-flash-001", "gemini-2.0-flash-lite-001", "gemini-2.0-flash-lite",
"gemini-2.5-flash-lite",
"gemini-3-pro-image", "gemini-3.1-flash-image",
// latest version
"gemini-flash-latest", "gemini-flash-lite-latest", "gemini-pro-latest",
"gemini-2.5-flash-native-audio-latest",
// preview version
"gemini-2.5-flash-preview-tts", "gemini-2.5-pro-preview-tts",
"gemini-2.5-flash-image", "gemini-2.5-flash-lite-preview-09-2025",
"gemini-3-pro-preview", "gemini-3-flash-preview", "gemini-3.1-pro-preview",
"gemini-3.1-pro-preview-customtools", "gemini-3.1-flash-lite-preview",
"gemini-3-pro-image-preview", "nano-banana-pro-preview",
"gemini-3.1-flash-image-preview", "gemini-robotics-er-1.5-preview",
"gemini-2.5-computer-use-preview-10-2025", "deep-research-pro-preview-12-2025",
"gemini-2.5-flash-native-audio-preview-09-2025", "gemini-2.5-flash-native-audio-preview-12-2025",
// gemma models
"gemma-3-1b-it", "gemma-3-4b-it", "gemma-3-12b-it",
"gemma-3-27b-it", "gemma-3n-e4b-it", "gemma-3n-e2b-it",
// embedding models
"gemini-embedding-001", "gemini-embedding-2-preview",
// imagen models
"imagen-4.0-generate-001", "imagen-4.0-ultra-generate-001",
"imagen-4.0-fast-generate-001",
// veo models
"veo-2.0-generate-001", "veo-3.0-generate-001", "veo-3.0-fast-generate-001",
"veo-3.1-generate-preview", "veo-3.1-fast-generate-preview",
// other models
"aqa",
}
var SafetySettingList = []string{
"HARM_CATEGORY_HARASSMENT",
"HARM_CATEGORY_HATE_SPEECH",
"HARM_CATEGORY_SEXUALLY_EXPLICIT",
"HARM_CATEGORY_DANGEROUS_CONTENT",
//"HARM_CATEGORY_CIVIC_INTEGRITY", This item is deprecated!
}
var ChannelName = "google gemini"
@@ -0,0 +1,94 @@
package gemini
import (
"fmt"
"io"
"net/http"
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/constant"
"github.com/QuantumNous/new-api/logger"
relaycommon "github.com/QuantumNous/new-api/relay/common"
"github.com/QuantumNous/new-api/relay/helper"
"github.com/QuantumNous/new-api/relaykit/dto"
"github.com/QuantumNous/new-api/relaykit/types"
"github.com/QuantumNous/new-api/service"
"github.com/gin-gonic/gin"
)
func GeminiTextGenerationHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http.Response) (*dto.Usage, *types.NewAPIError) {
defer service.CloseResponseBodyGracefully(resp)
// 读取响应体
responseBody, err := io.ReadAll(resp.Body)
if err != nil {
return nil, types.NewOpenAIError(err, types.ErrorCodeBadResponseBody, http.StatusInternalServerError)
}
logger.LogDebug(c, "Gemini native response body: %s", responseBody)
// 解析为 Gemini 原生响应格式
var geminiResponse dto.GeminiChatResponse
err = common.Unmarshal(responseBody, &geminiResponse)
if err != nil {
return nil, types.NewOpenAIError(err, types.ErrorCodeBadResponseBody, http.StatusInternalServerError)
}
countGeminiBillableFunctionCalls(info, &geminiResponse)
if len(geminiResponse.Candidates) == 0 && geminiResponse.PromptFeedback != nil && geminiResponse.PromptFeedback.BlockReason != nil {
common.SetContextKey(c, constant.ContextKeyAdminRejectReason, fmt.Sprintf("gemini_block_reason=%s", *geminiResponse.PromptFeedback.BlockReason))
}
// 计算使用量(优先上游 UsageMetadata,缺失时本地估算并保留 Gemini 计费语义)
usage := buildUsageFromGeminiResponse(c, info, &geminiResponse)
service.IOCopyBytesGracefully(c, resp, responseBody)
return &usage, nil
}
func NativeGeminiEmbeddingHandler(c *gin.Context, resp *http.Response, info *relaycommon.RelayInfo) (*dto.Usage, *types.NewAPIError) {
defer service.CloseResponseBodyGracefully(resp)
responseBody, err := io.ReadAll(resp.Body)
if err != nil {
return nil, types.NewOpenAIError(err, types.ErrorCodeBadResponseBody, http.StatusInternalServerError)
}
logger.LogDebug(c, "Gemini native embedding response body: %s", responseBody)
usage := service.ResponseText2Usage(c, "", info.UpstreamModelName, info.GetEstimatePromptTokens())
if info.IsGeminiBatchEmbedding {
var geminiResponse dto.GeminiBatchEmbeddingResponse
err = common.Unmarshal(responseBody, &geminiResponse)
if err != nil {
return nil, types.NewOpenAIError(err, types.ErrorCodeBadResponseBody, http.StatusInternalServerError)
}
} else {
var geminiResponse dto.GeminiEmbeddingResponse
err = common.Unmarshal(responseBody, &geminiResponse)
if err != nil {
return nil, types.NewOpenAIError(err, types.ErrorCodeBadResponseBody, http.StatusInternalServerError)
}
}
service.IOCopyBytesGracefully(c, resp, responseBody)
return usage, nil
}
func GeminiTextGenerationStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http.Response) (*dto.Usage, *types.NewAPIError) {
helper.SetEventStreamHeaders(c)
return geminiStreamHandler(c, info, resp, func(data string, geminiResponse *dto.GeminiChatResponse) bool {
err := helper.StringData(c, data)
if err != nil {
logger.LogError(c, "failed to write stream data: "+err.Error())
return false
}
info.SendResponseCount++
return true
})
}
+611
View File
@@ -0,0 +1,611 @@
package gemini
import (
"context"
"errors"
"fmt"
"io"
"net/http"
"strings"
"time"
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/constant"
"github.com/QuantumNous/new-api/logger"
"github.com/QuantumNous/new-api/relay/channel/openai"
relaycommon "github.com/QuantumNous/new-api/relay/common"
"github.com/QuantumNous/new-api/relay/helper"
"github.com/QuantumNous/new-api/relaykit/dto"
"github.com/QuantumNous/new-api/relaykit/relayconvert"
"github.com/QuantumNous/new-api/relaykit/types"
"github.com/QuantumNous/new-api/service"
"github.com/gin-gonic/gin"
)
func buildUsageFromGeminiMetadata(metadata *dto.GeminiUsageMetadata, fallbackPromptTokens int) dto.Usage {
usage := relayconvert.UsageFromGeminiMetadata(metadata, fallbackPromptTokens)
if usage == nil {
return dto.Usage{}
}
return *usage
}
func attachEstimatedGeminiBillingUsage(usage *dto.Usage) *dto.Usage {
if usage != nil && usage.BillingUsage == nil {
usage.BillingUsage = dto.NewEstimatedGeminiChatBillingUsage(usage)
}
return usage
}
// patchGeminiZeroCompletionUsage estimates completion tokens locally when upstream
// usageMetadata was billable but reported zero completion tokens even though output
// content was actually received. Typical case: the client aborts a stream before the
// final chunk that carries candidatesTokenCount, leaving prompt-only metadata; without
// this patch the output side would settle at zero quota.
func patchGeminiZeroCompletionUsage(c *gin.Context, info *relaycommon.RelayInfo, usage *dto.Usage, responseText string, imageCount int) {
if usage == nil || usage.CompletionTokens > 0 {
return
}
if responseText == "" && imageCount == 0 {
return
}
estimated := service.ResponseText2Usage(c, responseText, info.UpstreamModelName, usage.PromptTokens)
usage.CompletionTokens = estimated.CompletionTokens
if imageCount != 0 && usage.CompletionTokens == 0 {
usage.CompletionTokens = imageCount * 1400
}
usage.TotalTokens = usage.PromptTokens + usage.CompletionTokens
// Settlement prefers BillingUsage, so fill the missing completion in the
// original upstream dialect without discarding cache or modality details.
if usage.BillingUsage != nil {
usage.BillingUsage = dto.CloneBillingUsageWithEstimatedCompletion(usage.BillingUsage, usage.CompletionTokens)
} else {
usage.BillingUsage = dto.NewEstimatedGeminiChatBillingUsage(usage)
}
}
func geminiResponseUsageText(response *dto.GeminiChatResponse) string {
if response == nil {
return ""
}
var text strings.Builder
for _, candidate := range response.Candidates {
for _, part := range candidate.Content.Parts {
if part.Text != "" {
text.WriteString(part.Text)
}
}
}
return text.String()
}
func markGeminiGoogleSearchCall(c *gin.Context, response *dto.GeminiChatResponse) {
if c == nil || response == nil {
return
}
for _, candidate := range response.Candidates {
if candidate.GroundingMetadata != nil && len(candidate.GroundingMetadata.WebSearchQueries) > 0 {
c.Set("gemini_google_search_call", true)
return
}
}
}
func countGeminiBillableFunctionCalls(info *relaycommon.RelayInfo, response *dto.GeminiChatResponse) {
if info == nil || response == nil {
return
}
for _, candidate := range response.Candidates {
for _, part := range candidate.Content.Parts {
if part.FunctionCall == nil {
continue
}
if part.FunctionCall.WillContinue != nil && *part.FunctionCall.WillContinue {
continue
}
info.CountBillableToolCall(dto.BuildInCallFunctionCall, part.FunctionCall.FunctionName)
}
}
}
func buildUsageFromGeminiResponse(c *gin.Context, info *relaycommon.RelayInfo, response *dto.GeminiChatResponse) dto.Usage {
metadata := response.GetUsageMetadata()
if dto.HasGeminiUsageMetadataTokens(metadata) {
usage := buildUsageFromGeminiMetadata(metadata, info.GetEstimatePromptTokens())
patchGeminiZeroCompletionUsage(c, info, &usage, geminiResponseUsageText(response), geminiResponseInlineImageCount(response))
return usage
}
usage := service.ResponseText2Usage(c, geminiResponseUsageText(response), info.UpstreamModelName, info.GetEstimatePromptTokens())
attachEstimatedGeminiBillingUsage(usage)
return *usage
}
func geminiResponseInlineImageCount(response *dto.GeminiChatResponse) int {
if response == nil {
return 0
}
count := 0
for _, candidate := range response.Candidates {
for _, part := range candidate.Content.Parts {
if part.InlineData != nil && part.InlineData.MimeType != "" {
count++
}
}
}
return count
}
func responseGeminiChat2OpenAI(c *gin.Context, response *dto.GeminiChatResponse) *dto.OpenAITextResponse {
return relayconvert.ResponseGeminiChat2OpenAI(helper.GetResponseID(c), common.GetTimestamp(), response)
}
func streamResponseGeminiChat2OpenAI(geminiResponse *dto.GeminiChatResponse) (*dto.ChatCompletionsStreamResponse, bool) {
return relayconvert.StreamResponseGeminiChat2OpenAI(geminiResponse)
}
func handleStream(c *gin.Context, info *relaycommon.RelayInfo, resp *dto.ChatCompletionsStreamResponse) error {
streamData, err := common.Marshal(resp)
if err != nil {
return fmt.Errorf("failed to marshal stream response: %w", err)
}
err = openai.HandleStreamFormat(c, info, string(streamData), info.ChannelSetting.ForceFormat, info.ChannelSetting.ThinkingToContent)
if err != nil {
return fmt.Errorf("failed to handle stream format: %w", err)
}
return nil
}
func handleFinalStream(c *gin.Context, info *relaycommon.RelayInfo, resp *dto.ChatCompletionsStreamResponse) error {
streamData, err := common.Marshal(resp)
if err != nil {
return fmt.Errorf("failed to marshal stream response: %w", err)
}
openai.HandleFinalResponse(c, info, string(streamData), resp.Id, resp.Created, resp.Model, resp.GetSystemFingerprint(), resp.Usage, false)
return nil
}
func geminiStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http.Response, callback func(data string, geminiResponse *dto.GeminiChatResponse) bool) (*dto.Usage, *types.NewAPIError) {
var usage = &dto.Usage{}
var imageCount int
var hasBillableUsageMetadata bool
var streamErr error
var accumulatedUsageMetadata *dto.GeminiUsageMetadata
responseText := strings.Builder{}
helper.StreamScannerHandler(c, resp, info, func(data string, sr *helper.StreamResult) {
var geminiResponse dto.GeminiChatResponse
if err := common.UnmarshalJsonStr(data, &geminiResponse); err != nil {
streamErr = fmt.Errorf("unmarshal Gemini stream response: %w", err)
sr.Stop(streamErr)
return
}
if len(geminiResponse.Candidates) == 0 && geminiResponse.PromptFeedback != nil && geminiResponse.PromptFeedback.BlockReason != nil {
common.SetContextKey(c, constant.ContextKeyAdminRejectReason, fmt.Sprintf("gemini_block_reason=%s", *geminiResponse.PromptFeedback.BlockReason))
}
markGeminiGoogleSearchCall(c, &geminiResponse)
countGeminiBillableFunctionCalls(info, &geminiResponse)
// 统计图片数量
for _, candidate := range geminiResponse.Candidates {
for _, part := range candidate.Content.Parts {
if part.InlineData != nil && part.InlineData.MimeType != "" {
imageCount++
}
if part.Text != "" {
responseText.WriteString(part.Text)
}
}
}
// 更新使用量统计
if metadata := geminiResponse.GetUsageMetadata(); dto.HasGeminiUsageMetadataTokens(metadata) {
accumulatedUsageMetadata = dto.MergeGeminiUsageMetadataNonZero(accumulatedUsageMetadata, metadata)
mappedUsage := buildUsageFromGeminiMetadata(accumulatedUsageMetadata, info.GetEstimatePromptTokens())
*usage = mappedUsage
hasBillableUsageMetadata = true
}
if !callback(data, &geminiResponse) {
if isGeminiDownstreamStop(c, info) {
sr.Stop(nil)
return
}
streamErr = errors.New("Gemini stream callback stopped")
sr.Stop(streamErr)
}
})
if !hasBillableUsageMetadata {
if info.ReceivedResponseCount > 0 {
usage = service.ResponseText2Usage(c, responseText.String(), info.UpstreamModelName, info.GetEstimatePromptTokens())
} else {
usage = &dto.Usage{}
}
if imageCount != 0 && usage.CompletionTokens == 0 {
usage.CompletionTokens = imageCount * 1400
usage.TotalTokens = usage.PromptTokens + usage.CompletionTokens
common.SetContextKey(c, constant.ContextKeyLocalCountTokens, true)
}
attachEstimatedGeminiBillingUsage(usage)
} else {
patchGeminiZeroCompletionUsage(c, info, usage, responseText.String(), imageCount)
}
if streamErr != nil {
return usage, types.NewOpenAIError(streamErr, types.ErrorCodeBadResponseBody, http.StatusInternalServerError)
}
if info.StreamStatus != nil && !info.StreamStatus.IsNormalEnd() {
logger.LogWarn(c, fmt.Sprintf("Gemini stream ended unexpectedly: %s", info.StreamStatus.Summary()))
}
return usage, nil
}
func isGeminiDownstreamStop(c *gin.Context, info *relaycommon.RelayInfo) bool {
if c != nil && c.Request != nil && c.Request.Context().Err() != nil {
return true
}
return info != nil && info.StreamStatus != nil &&
info.StreamStatus.EndReason == relaycommon.StreamEndReasonClientGone
}
func GeminiChatStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http.Response) (*dto.Usage, *types.NewAPIError) {
id := helper.GetResponseID(c)
createAt := common.GetTimestamp()
finishReason := constant.FinishReasonStop
toolCallIndexByChoice := make(map[int]map[string]int)
nextToolCallIndexByChoice := make(map[int]int)
usage, err := geminiStreamHandler(c, info, resp, func(data string, geminiResponse *dto.GeminiChatResponse) bool {
response, isStop := streamResponseGeminiChat2OpenAI(geminiResponse)
response.Id = id
response.Created = createAt
response.Model = info.UpstreamModelName
if response.IsToolCall() {
finishReason = constant.FinishReasonToolCalls
if info.RelayFormat == types.RelayFormatClaude {
for choiceIdx := range response.Choices {
response.Choices[choiceIdx].FinishReason = nil
}
}
}
for choiceIdx := range response.Choices {
choiceKey := response.Choices[choiceIdx].Index
for toolIdx := range response.Choices[choiceIdx].Delta.ToolCalls {
tool := &response.Choices[choiceIdx].Delta.ToolCalls[toolIdx]
if tool.ID == "" {
continue
}
m := toolCallIndexByChoice[choiceKey]
if m == nil {
m = make(map[string]int)
toolCallIndexByChoice[choiceKey] = m
}
if idx, ok := m[tool.ID]; ok {
tool.SetIndex(idx)
continue
}
idx := nextToolCallIndexByChoice[choiceKey]
nextToolCallIndexByChoice[choiceKey] = idx + 1
m[tool.ID] = idx
tool.SetIndex(idx)
}
}
logger.LogDebug(c, "info.SendResponseCount = %d", info.SendResponseCount)
if info.SendResponseCount == 0 {
// send first response
emptyResponse := helper.GenerateStartEmptyResponse(id, createAt, info.UpstreamModelName, nil)
// Claude message_start is emitted from this first OpenAI chunk.
// Carry upstream usage when the current Gemini frame provided it.
emptyResponse.Usage = response.Usage
if response.IsToolCall() {
if len(emptyResponse.Choices) > 0 && len(response.Choices) > 0 {
toolCalls := response.Choices[0].Delta.ToolCalls
copiedToolCalls := make([]dto.ToolCallResponse, len(toolCalls))
for idx := range toolCalls {
copiedToolCalls[idx] = toolCalls[idx]
copiedToolCalls[idx].Function.Arguments = ""
}
emptyResponse.Choices[0].Delta.ToolCalls = copiedToolCalls
}
finishReason = constant.FinishReasonToolCalls
err := handleStream(c, info, emptyResponse)
if err != nil {
logger.LogError(c, err.Error())
}
response.ClearToolCalls()
if response.IsFinished() {
response.Choices[0].FinishReason = nil
}
} else {
err := handleStream(c, info, emptyResponse)
if err != nil {
logger.LogError(c, err.Error())
}
}
}
err := handleStream(c, info, response)
if err != nil {
logger.LogError(c, err.Error())
}
if isStop {
if info.RelayFormat != types.RelayFormatClaude {
_ = handleStream(c, info, helper.GenerateStopResponse(id, createAt, info.UpstreamModelName, finishReason))
}
}
return true
})
if err != nil {
return usage, err
}
response := helper.GenerateFinalUsageResponse(id, createAt, info.UpstreamModelName, *usage)
if info.RelayFormat == types.RelayFormatClaude && info.ClaudeConvertInfo != nil && !info.ClaudeConvertInfo.Done {
response = helper.GenerateStopResponse(id, createAt, info.UpstreamModelName, finishReason)
response.Usage = usage
}
handleErr := handleFinalStream(c, info, response)
if handleErr != nil {
common.SysLog("send final response failed: " + handleErr.Error())
}
return usage, nil
}
func GeminiChatHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http.Response) (*dto.Usage, *types.NewAPIError) {
responseBody, err := io.ReadAll(resp.Body)
if err != nil {
return nil, types.NewOpenAIError(err, types.ErrorCodeBadResponseBody, http.StatusInternalServerError)
}
service.CloseResponseBodyGracefully(resp)
logger.LogDebug(c, "Gemini response body: %s", responseBody)
var geminiResponse dto.GeminiChatResponse
err = common.Unmarshal(responseBody, &geminiResponse)
if err != nil {
return nil, types.NewOpenAIError(err, types.ErrorCodeBadResponseBody, http.StatusInternalServerError)
}
markGeminiGoogleSearchCall(c, &geminiResponse)
countGeminiBillableFunctionCalls(info, &geminiResponse)
if len(geminiResponse.Candidates) == 0 {
usage := buildUsageFromGeminiResponse(c, info, &geminiResponse)
var newAPIError *types.NewAPIError
if geminiResponse.PromptFeedback != nil && geminiResponse.PromptFeedback.BlockReason != nil {
common.SetContextKey(c, constant.ContextKeyAdminRejectReason, fmt.Sprintf("gemini_block_reason=%s", *geminiResponse.PromptFeedback.BlockReason))
newAPIError = types.NewOpenAIError(
errors.New("request blocked by Gemini API: "+*geminiResponse.PromptFeedback.BlockReason),
types.ErrorCodePromptBlocked,
http.StatusBadRequest,
)
} else {
common.SetContextKey(c, constant.ContextKeyAdminRejectReason, "gemini_empty_candidates")
newAPIError = types.NewOpenAIError(
errors.New("empty response from Gemini API"),
types.ErrorCodeEmptyResponse,
http.StatusInternalServerError,
)
}
service.ResetStatusCode(newAPIError, c.GetString("status_code_mapping"))
switch info.RelayFormat {
case types.RelayFormatClaude:
c.JSON(newAPIError.StatusCode, gin.H{
"type": "error",
"error": newAPIError.ToClaudeError(),
})
default:
c.JSON(newAPIError.StatusCode, gin.H{
"error": newAPIError.ToOpenAIError(),
})
}
return &usage, nil
}
fullTextResponse := responseGeminiChat2OpenAI(c, &geminiResponse)
fullTextResponse.Model = info.UpstreamModelName
usage := buildUsageFromGeminiResponse(c, info, &geminiResponse)
fullTextResponse.Usage = usage
switch info.RelayFormat {
case types.RelayFormatOpenAI:
responseBody, err = common.Marshal(fullTextResponse)
if err != nil {
return nil, types.NewError(err, types.ErrorCodeBadResponseBody)
}
case types.RelayFormatClaude:
convertResult, err := service.ConvertResponse(c, info, types.RelayFormatClaude, fullTextResponse)
if err != nil {
return nil, types.NewError(err, types.ErrorCodeBadResponseBody)
}
claudeRespStr, err := common.Marshal(convertResult.Value)
if err != nil {
return nil, types.NewError(err, types.ErrorCodeBadResponseBody)
}
responseBody = claudeRespStr
case types.RelayFormatGemini:
break
}
service.IOCopyBytesGracefully(c, resp, responseBody)
return &usage, nil
}
func GeminiEmbeddingHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http.Response) (*dto.Usage, *types.NewAPIError) {
defer service.CloseResponseBodyGracefully(resp)
responseBody, readErr := io.ReadAll(resp.Body)
if readErr != nil {
return nil, types.NewOpenAIError(readErr, types.ErrorCodeBadResponseBody, http.StatusInternalServerError)
}
var geminiResponse dto.GeminiBatchEmbeddingResponse
if jsonErr := common.Unmarshal(responseBody, &geminiResponse); jsonErr != nil {
return nil, types.NewOpenAIError(jsonErr, types.ErrorCodeBadResponseBody, http.StatusInternalServerError)
}
// convert to openai format response
openAIResponse := dto.OpenAIEmbeddingResponse{
Object: "list",
Data: make([]dto.OpenAIEmbeddingResponseItem, 0, len(geminiResponse.Embeddings)),
Model: info.UpstreamModelName,
}
for i, embedding := range geminiResponse.Embeddings {
openAIResponse.Data = append(openAIResponse.Data, dto.OpenAIEmbeddingResponseItem{
Object: "embedding",
Embedding: embedding.Values,
Index: i,
})
}
// calculate usage
// https://ai.google.dev/gemini-api/docs/pricing?hl=zh-cn#text-embedding-004
// Google has not yet clarified how embedding models will be billed
// refer to openai billing method to use input tokens billing
// https://platform.openai.com/docs/guides/embeddings#what-are-embeddings
usage := service.ResponseText2Usage(c, "", info.UpstreamModelName, info.GetEstimatePromptTokens())
openAIResponse.Usage = *usage
jsonResponse, jsonErr := common.Marshal(openAIResponse)
if jsonErr != nil {
return nil, types.NewOpenAIError(jsonErr, types.ErrorCodeBadResponseBody, http.StatusInternalServerError)
}
service.IOCopyBytesGracefully(c, resp, jsonResponse)
return usage, nil
}
func GeminiImageHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http.Response) (*dto.Usage, *types.NewAPIError) {
responseBody, readErr := io.ReadAll(resp.Body)
if readErr != nil {
return nil, types.NewOpenAIError(readErr, types.ErrorCodeBadResponseBody, http.StatusInternalServerError)
}
_ = resp.Body.Close()
var geminiResponse dto.GeminiImageResponse
if jsonErr := common.Unmarshal(responseBody, &geminiResponse); jsonErr != nil {
return nil, types.NewOpenAIError(jsonErr, types.ErrorCodeBadResponseBody, http.StatusInternalServerError)
}
if len(geminiResponse.Predictions) == 0 {
return nil, types.NewOpenAIError(errors.New("no images generated"), types.ErrorCodeBadResponseBody, http.StatusInternalServerError)
}
// convert to openai format response
openAIResponse := dto.ImageResponse{
Created: common.GetTimestamp(),
Data: make([]dto.ImageData, 0, len(geminiResponse.Predictions)),
}
for _, prediction := range geminiResponse.Predictions {
if prediction.RaiFilteredReason != "" {
continue // skip filtered image
}
openAIResponse.Data = append(openAIResponse.Data, dto.ImageData{
B64Json: prediction.BytesBase64Encoded,
})
}
jsonResponse, jsonErr := common.Marshal(openAIResponse)
if jsonErr != nil {
return nil, types.NewError(jsonErr, types.ErrorCodeBadResponseBody)
}
c.Writer.Header().Set("Content-Type", "application/json")
c.Writer.WriteHeader(resp.StatusCode)
_, _ = c.Writer.Write(jsonResponse)
// https://github.com/google-gemini/cookbook/blob/719a27d752aac33f39de18a8d3cb42a70874917e/quickstarts/Counting_Tokens.ipynb
// each image has fixed 258 tokens
const imageTokens = 258
generatedImages := len(openAIResponse.Data)
usage := &dto.Usage{
PromptTokens: imageTokens * generatedImages, // each generated image has fixed 258 tokens
CompletionTokens: 0, // image generation does not calculate completion tokens
TotalTokens: imageTokens * generatedImages,
}
return usage, nil
}
type GeminiModelsResponse struct {
Models []dto.GeminiModel `json:"models"`
NextPageToken string `json:"nextPageToken"`
}
func FetchGeminiModels(baseURL, apiKey, proxyURL string) ([]string, error) {
client, err := service.GetHttpClientWithProxy(proxyURL)
if err != nil {
return nil, fmt.Errorf("创建HTTP客户端失败: %v", err)
}
allModels := make([]string, 0)
nextPageToken := ""
maxPages := 100 // Safety limit to prevent infinite loops
for range maxPages {
url := fmt.Sprintf("%s/v1beta/models", baseURL)
if nextPageToken != "" {
url = fmt.Sprintf("%s?pageToken=%s", url, nextPageToken)
}
ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
request, err := http.NewRequestWithContext(ctx, "GET", url, nil)
if err != nil {
cancel()
return nil, fmt.Errorf("创建请求失败: %v", err)
}
request.Header.Set("x-goog-api-key", apiKey)
response, err := client.Do(request)
if err != nil {
cancel()
return nil, fmt.Errorf("请求失败: %v", err)
}
if response.StatusCode != http.StatusOK {
body, _ := io.ReadAll(response.Body)
response.Body.Close()
cancel()
return nil, fmt.Errorf("服务器返回错误 %d: %s", response.StatusCode, string(body))
}
body, err := io.ReadAll(response.Body)
response.Body.Close()
cancel()
if err != nil {
return nil, fmt.Errorf("读取响应失败: %v", err)
}
var modelsResponse GeminiModelsResponse
if err = common.Unmarshal(body, &modelsResponse); err != nil {
return nil, fmt.Errorf("解析响应失败: %v", err)
}
for _, model := range modelsResponse.Models {
modelNameValue, ok := model.Name.(string)
if !ok {
continue
}
modelName := strings.TrimPrefix(modelNameValue, "models/")
allModels = append(allModels, modelName)
}
nextPageToken = modelsResponse.NextPageToken
if nextPageToken == "" {
break
}
}
return allModels, nil
}
@@ -0,0 +1,641 @@
package gemini
import (
"bytes"
"io"
"net/http"
"net/http/httptest"
"strings"
"testing"
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/constant"
relaycommon "github.com/QuantumNous/new-api/relay/common"
"github.com/QuantumNous/new-api/relaykit/dto"
"github.com/QuantumNous/new-api/relaykit/types"
"github.com/gin-gonic/gin"
"github.com/stretchr/testify/require"
)
func TestStreamResponseGeminiChat2OpenAIAttachesUsageMetadata(t *testing.T) {
t.Parallel()
withUsage, isStop := streamResponseGeminiChat2OpenAI(&dto.GeminiChatResponse{
Candidates: []dto.GeminiChatCandidate{{
Content: dto.GeminiChatContent{
Role: "model",
Parts: []dto.GeminiPart{{Text: "hello"}},
},
}},
UsageMetadata: dto.GeminiUsageMetadata{
PromptTokenCount: 3868,
CandidatesTokenCount: 0,
TotalTokenCount: 3868,
},
})
require.False(t, isStop)
require.NotNil(t, withUsage)
require.NotNil(t, withUsage.Usage)
require.Equal(t, 3868, withUsage.Usage.PromptTokens)
require.Equal(t, 3868, withUsage.Usage.TotalTokens)
require.NotNil(t, withUsage.Usage.BillingUsage)
require.Equal(t, dto.BillingUsageSourceGeminiChat, withUsage.Usage.BillingUsage.Source)
require.Equal(t, dto.BillingUsageSemanticGemini, withUsage.Usage.BillingUsage.Semantic)
require.NotNil(t, withUsage.Usage.BillingUsage.GeminiUsageMetadata)
require.Equal(t, 3868, withUsage.Usage.BillingUsage.GeminiUsageMetadata.PromptTokenCount)
require.False(t, withUsage.Usage.BillingUsage.Estimated)
withoutUsage, _ := streamResponseGeminiChat2OpenAI(&dto.GeminiChatResponse{
Candidates: []dto.GeminiChatCandidate{{
Content: dto.GeminiChatContent{
Role: "model",
Parts: []dto.GeminiPart{{Text: "hello"}},
},
}},
})
require.NotNil(t, withoutUsage)
require.Nil(t, withoutUsage.Usage)
}
func TestGeminiChatStreamHandlerClaudeFirstFrameUsesUpstreamUsage(t *testing.T) {
gin.SetMode(gin.TestMode)
recorder := httptest.NewRecorder()
c, _ := gin.CreateTestContext(recorder)
c.Request = httptest.NewRequest(http.MethodPost, "/v1/messages", nil)
oldStreamingTimeout := constant.StreamingTimeout
constant.StreamingTimeout = 300
t.Cleanup(func() {
constant.StreamingTimeout = oldStreamingTimeout
})
info := &relaycommon.RelayInfo{
RelayFormat: types.RelayFormatClaude,
OriginModelName: "gemini-2.5-flash",
ChannelMeta: &relaycommon.ChannelMeta{
UpstreamModelName: "gemini-2.5-flash",
},
ClaudeConvertInfo: &relaycommon.ClaudeConvertInfo{
LastMessagesType: relaycommon.LastMessageTypeNone,
},
}
info.SetEstimatePromptTokens(4994)
chunkData, err := common.Marshal(dto.GeminiChatResponse{
Candidates: []dto.GeminiChatCandidate{{
Content: dto.GeminiChatContent{
Role: "model",
Parts: []dto.GeminiPart{{Text: "hello"}},
},
}},
UsageMetadata: dto.GeminiUsageMetadata{
PromptTokenCount: 3868,
TotalTokenCount: 3868,
},
})
require.NoError(t, err)
resp := &http.Response{
Body: io.NopCloser(bytes.NewReader([]byte("data: " + string(chunkData) + "\n" + "data: [DONE]\n"))),
}
usage, newAPIError := GeminiChatStreamHandler(c, info, resp)
require.Nil(t, newAPIError)
require.NotNil(t, usage)
require.Equal(t, 3868, usage.PromptTokens)
var startUsage, deltaUsage *dto.ClaudeUsage
for line := range strings.SplitSeq(recorder.Body.String(), "\n") {
payload, ok := strings.CutPrefix(strings.TrimSpace(line), "data: ")
if !ok {
continue
}
var event dto.ClaudeResponse
if err := common.UnmarshalJsonStr(payload, &event); err != nil {
continue
}
switch event.Type {
case "message_start":
if event.Message != nil {
startUsage = event.Message.Usage
}
case "message_delta":
deltaUsage = event.Usage
}
}
require.NotNil(t, startUsage)
require.Equal(t, 3868, startUsage.InputTokens)
require.NotNil(t, startUsage.BillingUsage)
require.Equal(t, dto.BillingUsageSourceGeminiChat, startUsage.BillingUsage.Source)
require.Equal(t, dto.BillingUsageSemanticGemini, startUsage.BillingUsage.Semantic)
require.NotNil(t, startUsage.BillingUsage.GeminiUsageMetadata)
require.Equal(t, 3868, startUsage.BillingUsage.GeminiUsageMetadata.PromptTokenCount)
require.False(t, startUsage.BillingUsage.Estimated)
require.NotNil(t, deltaUsage)
require.Equal(t, 3868, deltaUsage.InputTokens)
require.NotNil(t, deltaUsage.BillingUsage)
require.Equal(t, dto.BillingUsageSourceGeminiChat, deltaUsage.BillingUsage.Source)
require.Equal(t, dto.BillingUsageSemanticGemini, deltaUsage.BillingUsage.Semantic)
require.NotNil(t, deltaUsage.BillingUsage.GeminiUsageMetadata)
require.Equal(t, 3868, deltaUsage.BillingUsage.GeminiUsageMetadata.PromptTokenCount)
}
func TestGeminiChatHandlerCompletionTokensExcludeToolUsePromptTokens(t *testing.T) {
t.Parallel()
gin.SetMode(gin.TestMode)
c, _ := gin.CreateTestContext(httptest.NewRecorder())
c.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", nil)
info := &relaycommon.RelayInfo{
RelayFormat: types.RelayFormatGemini,
OriginModelName: "gemini-3-flash-preview",
ChannelMeta: &relaycommon.ChannelMeta{
UpstreamModelName: "gemini-3-flash-preview",
},
}
payload := dto.GeminiChatResponse{
Candidates: []dto.GeminiChatCandidate{
{
Content: dto.GeminiChatContent{
Role: "model",
Parts: []dto.GeminiPart{
{Text: "ok"},
},
},
},
},
UsageMetadata: dto.GeminiUsageMetadata{
PromptTokenCount: 151,
ToolUsePromptTokenCount: 18329,
CandidatesTokenCount: 1089,
ThoughtsTokenCount: 1120,
TotalTokenCount: 20689,
},
}
body, err := common.Marshal(payload)
require.NoError(t, err)
resp := &http.Response{
Body: io.NopCloser(bytes.NewReader(body)),
}
usage, newAPIError := GeminiChatHandler(c, info, resp)
require.Nil(t, newAPIError)
require.NotNil(t, usage)
require.Equal(t, 18480, usage.PromptTokens)
require.Equal(t, 2209, usage.CompletionTokens)
require.Equal(t, 20689, usage.TotalTokens)
require.Equal(t, 1120, usage.CompletionTokenDetails.ReasoningTokens)
}
func TestGeminiStreamHandlerCompletionTokensExcludeToolUsePromptTokens(t *testing.T) {
gin.SetMode(gin.TestMode)
c, _ := gin.CreateTestContext(httptest.NewRecorder())
c.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", nil)
oldStreamingTimeout := constant.StreamingTimeout
constant.StreamingTimeout = 300
t.Cleanup(func() {
constant.StreamingTimeout = oldStreamingTimeout
})
info := &relaycommon.RelayInfo{
OriginModelName: "gemini-3-flash-preview",
ChannelMeta: &relaycommon.ChannelMeta{
UpstreamModelName: "gemini-3-flash-preview",
},
}
chunk := dto.GeminiChatResponse{
Candidates: []dto.GeminiChatCandidate{
{
Content: dto.GeminiChatContent{
Role: "model",
Parts: []dto.GeminiPart{
{Text: "partial"},
},
},
},
},
UsageMetadata: dto.GeminiUsageMetadata{
PromptTokenCount: 151,
ToolUsePromptTokenCount: 18329,
CandidatesTokenCount: 1089,
ThoughtsTokenCount: 1120,
TotalTokenCount: 20689,
},
}
chunkData, err := common.Marshal(chunk)
require.NoError(t, err)
streamBody := []byte("data: " + string(chunkData) + "\n" + "data: [DONE]\n")
resp := &http.Response{
Body: io.NopCloser(bytes.NewReader(streamBody)),
}
usage, newAPIError := geminiStreamHandler(c, info, resp, func(_ string, _ *dto.GeminiChatResponse) bool {
return true
})
require.Nil(t, newAPIError)
require.NotNil(t, usage)
require.Equal(t, 18480, usage.PromptTokens)
require.Equal(t, 2209, usage.CompletionTokens)
require.Equal(t, 20689, usage.TotalTokens)
require.Equal(t, 1120, usage.CompletionTokenDetails.ReasoningTokens)
}
func TestGeminiTextGenerationHandlerPromptTokensIncludeToolUsePromptTokens(t *testing.T) {
t.Parallel()
gin.SetMode(gin.TestMode)
c, _ := gin.CreateTestContext(httptest.NewRecorder())
c.Request = httptest.NewRequest(http.MethodPost, "/v1beta/models/gemini-3-flash-preview:generateContent", nil)
info := &relaycommon.RelayInfo{
OriginModelName: "gemini-3-flash-preview",
ChannelMeta: &relaycommon.ChannelMeta{
UpstreamModelName: "gemini-3-flash-preview",
},
}
payload := dto.GeminiChatResponse{
Candidates: []dto.GeminiChatCandidate{
{
Content: dto.GeminiChatContent{
Role: "model",
Parts: []dto.GeminiPart{
{Text: "ok"},
},
},
},
},
UsageMetadata: dto.GeminiUsageMetadata{
PromptTokenCount: 151,
ToolUsePromptTokenCount: 18329,
CandidatesTokenCount: 1089,
ThoughtsTokenCount: 1120,
TotalTokenCount: 20689,
},
}
body, err := common.Marshal(payload)
require.NoError(t, err)
resp := &http.Response{
Body: io.NopCloser(bytes.NewReader(body)),
}
usage, newAPIError := GeminiTextGenerationHandler(c, info, resp)
require.Nil(t, newAPIError)
require.NotNil(t, usage)
require.Equal(t, 18480, usage.PromptTokens)
require.Equal(t, 2209, usage.CompletionTokens)
require.Equal(t, 20689, usage.TotalTokens)
require.Equal(t, 1120, usage.CompletionTokenDetails.ReasoningTokens)
}
func TestGeminiChatHandlerUsesEstimatedPromptTokensWhenUsagePromptMissing(t *testing.T) {
t.Parallel()
gin.SetMode(gin.TestMode)
c, _ := gin.CreateTestContext(httptest.NewRecorder())
c.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", nil)
info := &relaycommon.RelayInfo{
RelayFormat: types.RelayFormatGemini,
OriginModelName: "gemini-3-flash-preview",
ChannelMeta: &relaycommon.ChannelMeta{
UpstreamModelName: "gemini-3-flash-preview",
},
}
info.SetEstimatePromptTokens(20)
payload := dto.GeminiChatResponse{
Candidates: []dto.GeminiChatCandidate{
{
Content: dto.GeminiChatContent{
Role: "model",
Parts: []dto.GeminiPart{
{Text: "ok"},
},
},
},
},
UsageMetadata: dto.GeminiUsageMetadata{
PromptTokenCount: 0,
ToolUsePromptTokenCount: 0,
CandidatesTokenCount: 90,
ThoughtsTokenCount: 10,
TotalTokenCount: 110,
},
}
body, err := common.Marshal(payload)
require.NoError(t, err)
resp := &http.Response{
Body: io.NopCloser(bytes.NewReader(body)),
}
usage, newAPIError := GeminiChatHandler(c, info, resp)
require.Nil(t, newAPIError)
require.NotNil(t, usage)
require.Equal(t, 20, usage.PromptTokens)
require.Equal(t, 100, usage.CompletionTokens)
require.Equal(t, 110, usage.TotalTokens)
}
func TestGeminiStreamHandlerUsesEstimatedPromptTokensWhenUsagePromptMissing(t *testing.T) {
gin.SetMode(gin.TestMode)
c, _ := gin.CreateTestContext(httptest.NewRecorder())
c.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", nil)
oldStreamingTimeout := constant.StreamingTimeout
constant.StreamingTimeout = 300
t.Cleanup(func() {
constant.StreamingTimeout = oldStreamingTimeout
})
info := &relaycommon.RelayInfo{
OriginModelName: "gemini-3-flash-preview",
ChannelMeta: &relaycommon.ChannelMeta{
UpstreamModelName: "gemini-3-flash-preview",
},
}
info.SetEstimatePromptTokens(20)
chunk := dto.GeminiChatResponse{
Candidates: []dto.GeminiChatCandidate{
{
Content: dto.GeminiChatContent{
Role: "model",
Parts: []dto.GeminiPart{
{Text: "partial"},
},
},
},
},
UsageMetadata: dto.GeminiUsageMetadata{
PromptTokenCount: 0,
ToolUsePromptTokenCount: 0,
CandidatesTokenCount: 90,
ThoughtsTokenCount: 10,
TotalTokenCount: 110,
},
}
chunkData, err := common.Marshal(chunk)
require.NoError(t, err)
streamBody := []byte("data: " + string(chunkData) + "\n" + "data: [DONE]\n")
resp := &http.Response{
Body: io.NopCloser(bytes.NewReader(streamBody)),
}
usage, newAPIError := geminiStreamHandler(c, info, resp, func(_ string, _ *dto.GeminiChatResponse) bool {
return true
})
require.Nil(t, newAPIError)
require.NotNil(t, usage)
require.Equal(t, 20, usage.PromptTokens)
require.Equal(t, 100, usage.CompletionTokens)
require.Equal(t, 110, usage.TotalTokens)
}
func TestGeminiTextGenerationHandlerUsesEstimatedPromptTokensWhenUsagePromptMissing(t *testing.T) {
t.Parallel()
gin.SetMode(gin.TestMode)
c, _ := gin.CreateTestContext(httptest.NewRecorder())
c.Request = httptest.NewRequest(http.MethodPost, "/v1beta/models/gemini-3-flash-preview:generateContent", nil)
info := &relaycommon.RelayInfo{
OriginModelName: "gemini-3-flash-preview",
ChannelMeta: &relaycommon.ChannelMeta{
UpstreamModelName: "gemini-3-flash-preview",
},
}
info.SetEstimatePromptTokens(20)
payload := dto.GeminiChatResponse{
Candidates: []dto.GeminiChatCandidate{
{
Content: dto.GeminiChatContent{
Role: "model",
Parts: []dto.GeminiPart{
{Text: "ok"},
},
},
},
},
UsageMetadata: dto.GeminiUsageMetadata{
PromptTokenCount: 0,
ToolUsePromptTokenCount: 0,
CandidatesTokenCount: 90,
ThoughtsTokenCount: 10,
TotalTokenCount: 110,
},
}
body, err := common.Marshal(payload)
require.NoError(t, err)
resp := &http.Response{
Body: io.NopCloser(bytes.NewReader(body)),
}
usage, newAPIError := GeminiTextGenerationHandler(c, info, resp)
require.Nil(t, newAPIError)
require.NotNil(t, usage)
require.Equal(t, 20, usage.PromptTokens)
require.Equal(t, 100, usage.CompletionTokens)
require.Equal(t, 110, usage.TotalTokens)
}
func TestGeminiChatHandlerMissingUsageMetadataBuildsEstimatedBillingUsage(t *testing.T) {
t.Parallel()
gin.SetMode(gin.TestMode)
c, _ := gin.CreateTestContext(httptest.NewRecorder())
c.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", nil)
info := &relaycommon.RelayInfo{
RelayFormat: types.RelayFormatGemini,
OriginModelName: "gemini-3-flash-preview",
ChannelMeta: &relaycommon.ChannelMeta{
UpstreamModelName: "gemini-3-flash-preview",
},
}
info.SetEstimatePromptTokens(20)
body := []byte(`{"candidates":[{"content":{"role":"model","parts":[{"text":"ok"}]}}]}`)
resp := &http.Response{
Body: io.NopCloser(bytes.NewReader(body)),
}
usage, newAPIError := GeminiChatHandler(c, info, resp)
require.Nil(t, newAPIError)
require.NotNil(t, usage)
require.Equal(t, 20, usage.PromptTokens)
require.NotNil(t, usage.BillingUsage)
require.True(t, usage.BillingUsage.Estimated)
require.Equal(t, dto.BillingUsageSourceGeminiChat, usage.BillingUsage.Source)
require.Equal(t, dto.BillingUsageSemanticGemini, usage.BillingUsage.Semantic)
require.NotNil(t, usage.BillingUsage.GeminiUsageMetadata)
require.Equal(t, usage.PromptTokens, usage.BillingUsage.GeminiUsageMetadata.PromptTokenCount)
require.Equal(t, usage.CompletionTokens, usage.BillingUsage.GeminiUsageMetadata.CandidatesTokenCount)
require.True(t, common.GetContextKeyBool(c, constant.ContextKeyLocalCountTokens))
}
func TestGeminiStreamHandlerPromptOnlyUsageMetadataEstimatesCompletionTokens(t *testing.T) {
gin.SetMode(gin.TestMode)
c, _ := gin.CreateTestContext(httptest.NewRecorder())
c.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", nil)
oldStreamingTimeout := constant.StreamingTimeout
constant.StreamingTimeout = 300
t.Cleanup(func() {
constant.StreamingTimeout = oldStreamingTimeout
})
info := &relaycommon.RelayInfo{
OriginModelName: "gemini-3-flash-preview",
ChannelMeta: &relaycommon.ChannelMeta{
UpstreamModelName: "gemini-3-flash-preview",
},
}
info.SetEstimatePromptTokens(20)
// Simulates a client aborting the stream before the final chunk: text was
// streamed but the last observed usageMetadata only carries prompt tokens.
chunk := dto.GeminiChatResponse{
Candidates: []dto.GeminiChatCandidate{
{
Content: dto.GeminiChatContent{
Role: "model",
Parts: []dto.GeminiPart{
{Text: "partial streamed answer before disconnect"},
},
},
},
},
UsageMetadata: dto.GeminiUsageMetadata{
PromptTokenCount: 151,
TotalTokenCount: 151,
},
}
chunkData, err := common.Marshal(chunk)
require.NoError(t, err)
streamBody := []byte("data: " + string(chunkData) + "\n" + "data: [DONE]\n")
resp := &http.Response{
Body: io.NopCloser(bytes.NewReader(streamBody)),
}
usage, newAPIError := geminiStreamHandler(c, info, resp, func(_ string, _ *dto.GeminiChatResponse) bool {
return true
})
require.Nil(t, newAPIError)
require.NotNil(t, usage)
require.Equal(t, 151, usage.PromptTokens)
require.Greater(t, usage.CompletionTokens, 0)
require.Equal(t, usage.PromptTokens+usage.CompletionTokens, usage.TotalTokens)
require.NotNil(t, usage.BillingUsage)
require.True(t, usage.BillingUsage.Estimated)
require.NotNil(t, usage.BillingUsage.GeminiUsageMetadata)
require.Equal(t, usage.CompletionTokens, usage.BillingUsage.GeminiUsageMetadata.CandidatesTokenCount)
}
func TestGeminiChatHandlerPromptOnlyUsageMetadataEstimatesCompletionTokens(t *testing.T) {
t.Parallel()
gin.SetMode(gin.TestMode)
c, _ := gin.CreateTestContext(httptest.NewRecorder())
c.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", nil)
info := &relaycommon.RelayInfo{
RelayFormat: types.RelayFormatGemini,
OriginModelName: "gemini-3-flash-preview",
ChannelMeta: &relaycommon.ChannelMeta{
UpstreamModelName: "gemini-3-flash-preview",
},
}
payload := dto.GeminiChatResponse{
Candidates: []dto.GeminiChatCandidate{
{
Content: dto.GeminiChatContent{
Role: "model",
Parts: []dto.GeminiPart{
{Text: "answer text without candidate token count"},
},
},
},
},
UsageMetadata: dto.GeminiUsageMetadata{
PromptTokenCount: 151,
TotalTokenCount: 151,
},
}
body, err := common.Marshal(payload)
require.NoError(t, err)
resp := &http.Response{
Body: io.NopCloser(bytes.NewReader(body)),
}
usage, newAPIError := GeminiChatHandler(c, info, resp)
require.Nil(t, newAPIError)
require.NotNil(t, usage)
require.Equal(t, 151, usage.PromptTokens)
require.Greater(t, usage.CompletionTokens, 0)
require.Equal(t, usage.PromptTokens+usage.CompletionTokens, usage.TotalTokens)
require.NotNil(t, usage.BillingUsage)
require.True(t, usage.BillingUsage.Estimated)
}
func TestGeminiStreamHandlerEmptyUsageMetadataBuildsEstimatedBillingUsage(t *testing.T) {
gin.SetMode(gin.TestMode)
c, _ := gin.CreateTestContext(httptest.NewRecorder())
c.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", nil)
oldStreamingTimeout := constant.StreamingTimeout
constant.StreamingTimeout = 300
t.Cleanup(func() {
constant.StreamingTimeout = oldStreamingTimeout
})
info := &relaycommon.RelayInfo{
OriginModelName: "gemini-3-flash-preview",
ChannelMeta: &relaycommon.ChannelMeta{
UpstreamModelName: "gemini-3-flash-preview",
},
}
info.SetEstimatePromptTokens(20)
streamBody := []byte("data: {\"candidates\":[{\"content\":{\"role\":\"model\",\"parts\":[{\"text\":\"partial\"}]}}],\"usageMetadata\":{}}\n" + "data: [DONE]\n")
resp := &http.Response{
Body: io.NopCloser(bytes.NewReader(streamBody)),
}
usage, newAPIError := geminiStreamHandler(c, info, resp, func(_ string, _ *dto.GeminiChatResponse) bool {
return true
})
require.Nil(t, newAPIError)
require.NotNil(t, usage)
require.Equal(t, 20, usage.PromptTokens)
require.NotNil(t, usage.BillingUsage)
require.True(t, usage.BillingUsage.Estimated)
require.Equal(t, dto.BillingUsageSourceGeminiChat, usage.BillingUsage.Source)
require.NotNil(t, usage.BillingUsage.GeminiUsageMetadata)
require.Equal(t, usage.PromptTokens, usage.BillingUsage.GeminiUsageMetadata.PromptTokenCount)
require.Equal(t, usage.CompletionTokens, usage.BillingUsage.GeminiUsageMetadata.CandidatesTokenCount)
require.True(t, common.GetContextKeyBool(c, constant.ContextKeyLocalCountTokens))
}
+201
View File
@@ -0,0 +1,201 @@
package gemini
import (
"errors"
"fmt"
"io"
"net/http"
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/constant"
"github.com/QuantumNous/new-api/logger"
relaycommon "github.com/QuantumNous/new-api/relay/common"
"github.com/QuantumNous/new-api/relay/helper"
"github.com/QuantumNous/new-api/relaykit/dto"
"github.com/QuantumNous/new-api/relaykit/relayconvert"
"github.com/QuantumNous/new-api/relaykit/types"
"github.com/QuantumNous/new-api/service"
"github.com/gin-gonic/gin"
)
func GeminiResponsesHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http.Response) (*dto.Usage, *types.NewAPIError) {
defer service.CloseResponseBodyGracefully(resp)
responseBody, err := io.ReadAll(resp.Body)
if err != nil {
return nil, types.NewOpenAIError(err, types.ErrorCodeBadResponseBody, http.StatusInternalServerError)
}
logger.LogDebug(c, "Gemini responses response body: %s", responseBody)
var geminiResponse dto.GeminiChatResponse
if err := common.Unmarshal(responseBody, &geminiResponse); err != nil {
return nil, types.NewOpenAIError(err, types.ErrorCodeBadResponseBody, http.StatusInternalServerError)
}
markGeminiGoogleSearchCall(c, &geminiResponse)
countGeminiBillableFunctionCalls(info, &geminiResponse)
if len(geminiResponse.Candidates) == 0 {
usage := buildUsageFromGeminiResponse(c, info, &geminiResponse)
if geminiResponse.PromptFeedback != nil && geminiResponse.PromptFeedback.BlockReason != nil {
common.SetContextKey(c, constant.ContextKeyAdminRejectReason, fmt.Sprintf("gemini_block_reason=%s", *geminiResponse.PromptFeedback.BlockReason))
return &usage, types.NewOpenAIError(
errors.New("request blocked by Gemini API: "+*geminiResponse.PromptFeedback.BlockReason),
types.ErrorCodePromptBlocked,
http.StatusBadRequest,
)
}
common.SetContextKey(c, constant.ContextKeyAdminRejectReason, "gemini_empty_candidates")
return &usage, types.NewOpenAIError(
errors.New("empty response from Gemini API"),
types.ErrorCodeEmptyResponse,
http.StatusInternalServerError,
)
}
usage := buildUsageFromGeminiResponse(c, info, &geminiResponse)
convertResult, err := service.ConvertResponse(c, info, types.RelayFormatOpenAIResponses, &geminiResponse)
if err != nil {
return nil, types.NewOpenAIError(err, types.ErrorCodeBadResponseBody, http.StatusInternalServerError)
}
responsesResp, ok := convertResult.Value.(*dto.OpenAIResponsesResponse)
if !ok {
return nil, types.NewOpenAIError(fmt.Errorf("expected OpenAI responses response, got %T", convertResult.Value), types.ErrorCodeBadResponseBody, http.StatusInternalServerError)
}
if responseID := helper.GetResponseID(c); responseID != "" {
responsesResp.ID = responseID
}
responsesResp.Model = info.UpstreamModelName
responsesResp.Usage = relayconvert.UsageFromChatUsage(&usage)
responseBody, err = common.Marshal(responsesResp)
if err != nil {
return nil, types.NewOpenAIError(err, types.ErrorCodeJsonMarshalFailed, http.StatusInternalServerError)
}
service.IOCopyBytesGracefully(c, resp, responseBody)
return &usage, nil
}
func GeminiResponsesStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http.Response) (*dto.Usage, *types.NewAPIError) {
responseID := helper.GetResponseID(c)
created := common.GetTimestamp()
state, err := relayconvert.NewResponseStreamState(types.RelayFormatGemini, types.RelayFormatOpenAIResponses, relayconvert.ResponseStreamOptions{
ID: responseID,
Model: info.UpstreamModelName,
Created: created,
EmitSequenceNumber: true,
})
if err != nil {
return nil, types.NewOpenAIError(err, types.ErrorCodeBadResponse, http.StatusInternalServerError)
}
hostedBridge := relayconvert.NewGeminiHostedStreamBridge()
var streamErr *types.NewAPIError
sendEvent := func(event relayconvert.ChatToResponsesStreamEvent) bool {
data, err := common.Marshal(event.Payload)
if err != nil {
streamErr = types.NewOpenAIError(err, types.ErrorCodeJsonMarshalFailed, http.StatusInternalServerError)
return false
}
if err := helper.ResponseChunkData(c, dto.ResponsesStreamResponse{Type: event.Type}, string(data)); err != nil {
if info.StreamStatus != nil {
info.StreamStatus.SetEndReason(relaycommon.StreamEndReasonClientGone, err)
}
return false
}
return true
}
failResponsesStream := func(err error) bool {
failureResults, handled := state.FailResponsesStream("server_error", err.Error(), "")
if !handled {
return false
}
for _, result := range failureResults {
event, ok := result.Value.(relayconvert.ChatToResponsesStreamEvent)
if !ok {
streamErr = types.NewOpenAIError(fmt.Errorf("expected OAI responses stream event, got %T", result.Value), types.ErrorCodeBadResponse, http.StatusInternalServerError)
return true
}
if !sendEvent(event) {
return true
}
}
return true
}
sendChunk := func(chunk *dto.GeminiChatResponse) bool {
results, err := service.ConvertStreamResponseChunk(c, info, state, chunk)
if err != nil {
if failResponsesStream(err) {
return false
}
streamErr = types.NewOpenAIError(err, types.ErrorCodeBadResponse, http.StatusInternalServerError)
return false
}
for _, result := range results {
event, ok := result.Value.(relayconvert.ChatToResponsesStreamEvent)
if !ok {
streamErr = types.NewOpenAIError(fmt.Errorf("expected OAI responses stream event, got %T", result.Value), types.ErrorCodeBadResponse, http.StatusInternalServerError)
return false
}
if !sendEvent(event) {
return false
}
}
return true
}
usage, streamAPIError := geminiStreamHandler(c, info, resp, func(_ string, geminiResponse *dto.GeminiChatResponse) bool {
hostedBridge.Observe(geminiResponse)
return sendChunk(geminiResponse)
})
if streamAPIError != nil {
if failResponsesStream(streamAPIError) && streamErr == nil {
return usage, nil
}
return usage, streamAPIError
}
if info.StreamStatus != nil && !info.StreamStatus.IsNormalEnd() {
if info.StreamStatus.EndReason != relaycommon.StreamEndReasonClientGone {
failResponsesStream(fmt.Errorf("gemini stream ended unexpectedly: %s", info.StreamStatus.Summary()))
}
return usage, nil
}
if streamErr != nil {
return nil, streamErr
}
hostedEvents, err := hostedBridge.Finalize(state)
if err != nil {
return nil, types.NewOpenAIError(err, types.ErrorCodeBadResponse, http.StatusInternalServerError)
}
for _, event := range hostedEvents {
if !sendEvent(event) {
if streamErr != nil {
return usage, streamErr
}
return usage, nil
}
}
if usage != nil {
state.SetUsage(usage)
}
finalResults, err := service.FinalizeStreamResponse(c, info, state)
if err != nil {
if failResponsesStream(err) {
return usage, streamErr
}
return nil, types.NewOpenAIError(err, types.ErrorCodeBadResponse, http.StatusInternalServerError)
}
for _, result := range finalResults {
event, ok := result.Value.(relayconvert.ChatToResponsesStreamEvent)
if !ok {
return nil, types.NewOpenAIError(fmt.Errorf("expected OAI responses stream event, got %T", result.Value), types.ErrorCodeBadResponse, http.StatusInternalServerError)
}
if !sendEvent(event) {
if streamErr != nil {
return usage, streamErr
}
return usage, nil
}
}
return usage, nil
}
@@ -0,0 +1,205 @@
package gemini
import (
"bytes"
"errors"
"io"
"net/http"
"net/http/httptest"
"strings"
"testing"
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/constant"
relaycommon "github.com/QuantumNous/new-api/relay/common"
relayconstant "github.com/QuantumNous/new-api/relay/constant"
"github.com/QuantumNous/new-api/relaykit/dto"
"github.com/QuantumNous/new-api/relaykit/types"
"github.com/gin-gonic/gin"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
func TestGeminiResponsesHandlerReturnsOpenAIResponsesJSON(t *testing.T) {
gin.SetMode(gin.TestMode)
recorder := httptest.NewRecorder()
c, _ := gin.CreateTestContext(recorder)
c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", nil)
c.Set(common.RequestIdKey, "gemini-responses-test")
info := newGeminiResponsesRelayInfo(false)
payload := dto.GeminiChatResponse{
Candidates: []dto.GeminiChatCandidate{
{
Content: dto.GeminiChatContent{
Role: "model",
Parts: []dto.GeminiPart{
{Text: "hello"},
},
},
},
},
UsageMetadata: dto.GeminiUsageMetadata{
PromptTokenCount: 2,
CandidatesTokenCount: 3,
TotalTokenCount: 5,
},
}
body, err := common.Marshal(payload)
require.NoError(t, err)
usage, newAPIError := GeminiResponsesHandler(c, info, &http.Response{
Body: io.NopCloser(bytes.NewReader(body)),
})
require.Nil(t, newAPIError)
require.NotNil(t, usage)
assert.Equal(t, 2, usage.PromptTokens)
assert.Equal(t, 3, usage.CompletionTokens)
got := recorder.Body.String()
assert.Contains(t, got, `"object":"response"`)
assert.Contains(t, got, `"status":"completed"`)
assert.Contains(t, got, `"type":"output_text"`)
assert.Contains(t, got, `"text":"hello"`)
assert.Contains(t, got, `"input_tokens":2`)
assert.Contains(t, got, `"output_tokens":3`)
assert.NotContains(t, got, `"choices"`)
assert.NotContains(t, got, `"candidates"`)
}
func TestGeminiResponsesHandlerClosesBodyOnReadError(t *testing.T) {
gin.SetMode(gin.TestMode)
recorder := httptest.NewRecorder()
c, _ := gin.CreateTestContext(recorder)
c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", nil)
c.Set(common.RequestIdKey, "gemini-responses-read-error-test")
body := &failingReadCloser{}
usage, newAPIError := GeminiResponsesHandler(c, newGeminiResponsesRelayInfo(false), &http.Response{Body: body})
require.Nil(t, usage)
require.NotNil(t, newAPIError)
assert.True(t, body.closed)
}
func TestGeminiResponsesStreamHandlerReturnsOpenAIResponsesSSE(t *testing.T) {
gin.SetMode(gin.TestMode)
recorder := httptest.NewRecorder()
c, _ := gin.CreateTestContext(recorder)
c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", nil)
c.Set(common.RequestIdKey, "gemini-responses-stream-test")
oldStreamingTimeout := constant.StreamingTimeout
constant.StreamingTimeout = 300
t.Cleanup(func() { constant.StreamingTimeout = oldStreamingTimeout })
info := newGeminiResponsesRelayInfo(true)
first := dto.GeminiChatResponse{
Candidates: []dto.GeminiChatCandidate{
{
Content: dto.GeminiChatContent{
Role: "model",
Parts: []dto.GeminiPart{
{Text: "hello"},
},
},
},
},
UsageMetadata: dto.GeminiUsageMetadata{
PromptTokenCount: 2,
CandidatesTokenCount: 3,
TotalTokenCount: 5,
},
}
stop := "STOP"
final := dto.GeminiChatResponse{
Candidates: []dto.GeminiChatCandidate{
{
FinishReason: &stop,
Content: dto.GeminiChatContent{
Role: "model",
Parts: []dto.GeminiPart{{Text: ""}},
},
},
},
UsageMetadata: dto.GeminiUsageMetadata{
PromptTokenCount: 2,
CandidatesTokenCount: 3,
TotalTokenCount: 5,
},
}
firstData, err := common.Marshal(first)
require.NoError(t, err)
finalData, err := common.Marshal(final)
require.NoError(t, err)
streamBody := strings.Join([]string{
"data: " + string(firstData),
"",
"data: " + string(finalData),
"",
"data: [DONE]",
"",
}, "\n")
usage, newAPIError := GeminiResponsesStreamHandler(c, info, &http.Response{
Body: io.NopCloser(strings.NewReader(streamBody)),
})
require.Nil(t, newAPIError)
require.NotNil(t, usage)
assert.Equal(t, 5, usage.TotalTokens)
got := recorder.Body.String()
assert.Equal(t, "text/event-stream", recorder.Header().Get("Content-Type"))
assert.Contains(t, got, `event: response.created`)
assert.Contains(t, got, `event: response.output_text.delta`)
assert.Contains(t, got, `"delta":"hello"`)
assert.Contains(t, got, `event: response.completed`)
assert.Contains(t, got, `"input_tokens":2`)
assert.Contains(t, got, `"output_tokens":3`)
assert.NotContains(t, got, `"choices"`)
assert.NotContains(t, got, `"candidates"`)
requireOrderedGeminiResponsesSubstrings(t, got,
`event: response.created`,
`event: response.output_item.added`,
`event: response.output_text.delta`,
`event: response.output_text.done`,
`event: response.completed`,
)
}
func newGeminiResponsesRelayInfo(isStream bool) *relaycommon.RelayInfo {
return &relaycommon.RelayInfo{
IsStream: isStream,
RelayMode: relayconstant.RelayModeResponses,
RelayFormat: types.RelayFormatOpenAIResponses,
RequestURLPath: "/v1/responses",
DisablePing: true,
OriginModelName: "gemini-test",
ChannelMeta: &relaycommon.ChannelMeta{
UpstreamModelName: "gemini-test",
},
}
}
type failingReadCloser struct {
closed bool
}
func (r *failingReadCloser) Read([]byte) (int, error) {
return 0, errors.New("read failed")
}
func (r *failingReadCloser) Close() error {
r.closed = true
return nil
}
func requireOrderedGeminiResponsesSubstrings(t *testing.T, s string, parts ...string) {
t.Helper()
offset := 0
for _, part := range parts {
idx := strings.Index(s[offset:], part)
require.NotEqualf(t, -1, idx, "missing %q after byte offset %d", part, offset)
offset += idx + len(part)
}
}
+144
View File
@@ -0,0 +1,144 @@
package jimeng
import (
"encoding/json"
"errors"
"fmt"
"io"
"net/http"
"github.com/QuantumNous/new-api/relay/channel"
"github.com/QuantumNous/new-api/relay/channel/openai"
relaycommon "github.com/QuantumNous/new-api/relay/common"
relayconstant "github.com/QuantumNous/new-api/relay/constant"
"github.com/QuantumNous/new-api/relaykit/dto"
"github.com/QuantumNous/new-api/relaykit/types"
"github.com/gin-gonic/gin"
)
type Adaptor struct {
}
func (a *Adaptor) ConvertGeminiRequest(*gin.Context, *relaycommon.RelayInfo, *dto.GeminiChatRequest) (any, error) {
//TODO implement me
return nil, errors.New("not implemented")
}
func (a *Adaptor) ConvertClaudeRequest(*gin.Context, *relaycommon.RelayInfo, *dto.ClaudeRequest) (any, error) {
return nil, errors.New("not implemented")
}
func (a *Adaptor) Init(info *relaycommon.RelayInfo) {
}
func (a *Adaptor) GetRequestURL(info *relaycommon.RelayInfo) (string, error) {
return fmt.Sprintf("%s/?Action=CVProcess&Version=2022-08-31", info.ChannelBaseUrl), nil
}
func (a *Adaptor) SetupRequestHeader(c *gin.Context, header *http.Header, info *relaycommon.RelayInfo) error {
return errors.New("not implemented")
}
func (a *Adaptor) ConvertOpenAIRequest(c *gin.Context, info *relaycommon.RelayInfo, request *dto.GeneralOpenAIRequest) (any, error) {
if request == nil {
return nil, errors.New("request is nil")
}
return request, nil
}
type LogoInfo struct {
AddLogo bool `json:"add_logo,omitempty"`
Position int `json:"position,omitempty"`
Language int `json:"language,omitempty"`
Opacity float64 `json:"opacity,omitempty"`
LogoTextContent string `json:"logo_text_content,omitempty"`
}
type imageRequestPayload struct {
ReqKey string `json:"req_key"` // Service identifier, fixed value: jimeng_high_aes_general_v21_L
Prompt string `json:"prompt"` // Prompt for image generation, supports both Chinese and English
Seed int64 `json:"seed,omitempty"` // Random seed, default -1 (random)
Width int `json:"width,omitempty"` // Image width, default 512, range [256, 768]
Height int `json:"height,omitempty"` // Image height, default 512, range [256, 768]
UsePreLLM bool `json:"use_pre_llm,omitempty"` // Enable text expansion, default true
UseSR bool `json:"use_sr,omitempty"` // Enable super resolution, default true
ReturnURL bool `json:"return_url,omitempty"` // Whether to return image URL (valid for 24 hours)
LogoInfo LogoInfo `json:"logo_info,omitempty"` // Watermark information
ImageUrls []string `json:"image_urls,omitempty"` // Image URLs for input
BinaryData []string `json:"binary_data_base64,omitempty"` // Base64 encoded binary data
}
func (a *Adaptor) ConvertImageRequest(c *gin.Context, info *relaycommon.RelayInfo, request dto.ImageRequest) (any, error) {
payload := imageRequestPayload{
ReqKey: request.Model,
Prompt: request.Prompt,
}
if request.ResponseFormat == "" || request.ResponseFormat == "url" {
payload.ReturnURL = true // Default to returning image URLs
}
if len(request.ExtraFields) > 0 {
if err := json.Unmarshal(request.ExtraFields, &payload); err != nil {
return nil, fmt.Errorf("failed to unmarshal extra fields: %w", err)
}
}
return payload, nil
}
func (a *Adaptor) ConvertRerankRequest(c *gin.Context, relayMode int, request dto.RerankRequest) (any, error) {
return nil, errors.New("not implemented")
}
func (a *Adaptor) ConvertEmbeddingRequest(c *gin.Context, info *relaycommon.RelayInfo, request dto.EmbeddingRequest) (any, error) {
return nil, errors.New("not implemented")
}
func (a *Adaptor) ConvertAudioRequest(c *gin.Context, info *relaycommon.RelayInfo, request dto.AudioRequest) (io.Reader, error) {
return nil, errors.New("not implemented")
}
func (a *Adaptor) ConvertOpenAIResponsesRequest(c *gin.Context, info *relaycommon.RelayInfo, request dto.OpenAIResponsesRequest) (any, error) {
return nil, errors.New("not implemented")
}
func (a *Adaptor) DoRequest(c *gin.Context, info *relaycommon.RelayInfo, requestBody io.Reader) (any, error) {
fullRequestURL, err := a.GetRequestURL(info)
if err != nil {
return nil, fmt.Errorf("get request url failed: %w", err)
}
req, err := http.NewRequest(c.Request.Method, fullRequestURL, requestBody)
if err != nil {
return nil, fmt.Errorf("new request failed: %w", err)
}
channel.ApplyUpstreamBodyMetadata(req, requestBody)
err = Sign(c, req, info.ApiKey)
if err != nil {
return nil, fmt.Errorf("setup request header failed: %w", err)
}
resp, err := channel.DoRequest(c, req, info)
if err != nil {
return nil, fmt.Errorf("do request failed: %w", err)
}
return resp, nil
}
func (a *Adaptor) DoResponse(c *gin.Context, resp *http.Response, info *relaycommon.RelayInfo) (usage any, err *types.NewAPIError) {
if info.RelayMode == relayconstant.RelayModeImagesGenerations {
usage, err = jimengImageHandler(c, resp, info)
} else if info.IsStream {
usage, err = openai.OaiStreamHandler(c, info, resp)
} else {
usage, err = openai.OpenaiHandler(c, info, resp)
}
return
}
func (a *Adaptor) GetModelList() []string {
return ModelList
}
func (a *Adaptor) GetChannelName() string {
return ChannelName
}
+9
View File
@@ -0,0 +1,9 @@
package jimeng
const (
ChannelName = "jimeng"
)
var ModelList = []string{
"jimeng_high_aes_general_v21_L",
}
+90
View File
@@ -0,0 +1,90 @@
package jimeng
import (
"encoding/json"
"fmt"
"io"
"net/http"
relaycommon "github.com/QuantumNous/new-api/relay/common"
"github.com/QuantumNous/new-api/relaykit/dto"
"github.com/QuantumNous/new-api/relaykit/types"
"github.com/QuantumNous/new-api/service"
"github.com/gin-gonic/gin"
)
type ImageResponse struct {
Code int `json:"code"`
Message string `json:"message"`
Data struct {
BinaryDataBase64 []string `json:"binary_data_base64"`
ImageUrls []string `json:"image_urls"`
RephraseResult string `json:"rephraser_result"`
RequestID string `json:"request_id"`
// Other fields are omitted for brevity
} `json:"data"`
RequestID string `json:"request_id"`
Status int `json:"status"`
TimeElapsed string `json:"time_elapsed"`
}
func responseJimeng2OpenAIImage(_ *gin.Context, response *ImageResponse, info *relaycommon.RelayInfo) *dto.ImageResponse {
imageResponse := dto.ImageResponse{
Created: info.StartTime.Unix(),
}
for _, base64Data := range response.Data.BinaryDataBase64 {
imageResponse.Data = append(imageResponse.Data, dto.ImageData{
B64Json: base64Data,
})
}
for _, imageUrl := range response.Data.ImageUrls {
imageResponse.Data = append(imageResponse.Data, dto.ImageData{
Url: imageUrl,
})
}
return &imageResponse
}
// jimengImageHandler handles the Jimeng image generation response
func jimengImageHandler(c *gin.Context, resp *http.Response, info *relaycommon.RelayInfo) (*dto.Usage, *types.NewAPIError) {
var jimengResponse ImageResponse
responseBody, err := io.ReadAll(resp.Body)
if err != nil {
return nil, types.NewOpenAIError(err, types.ErrorCodeReadResponseBodyFailed, http.StatusInternalServerError)
}
service.CloseResponseBodyGracefully(resp)
err = json.Unmarshal(responseBody, &jimengResponse)
if err != nil {
return nil, types.NewOpenAIError(err, types.ErrorCodeBadResponseBody, http.StatusInternalServerError)
}
// Check if the response indicates an error
if jimengResponse.Code != 10000 {
return nil, types.WithOpenAIError(types.OpenAIError{
Message: jimengResponse.Message,
Type: "jimeng_error",
Param: "",
Code: fmt.Sprintf("%d", jimengResponse.Code),
}, resp.StatusCode)
}
// Convert Jimeng response to OpenAI format
fullTextResponse := responseJimeng2OpenAIImage(c, &jimengResponse, info)
jsonResponse, err := json.Marshal(fullTextResponse)
if err != nil {
return nil, types.NewError(err, types.ErrorCodeBadResponseBody)
}
c.Writer.Header().Set("Content-Type", "application/json")
c.Writer.WriteHeader(resp.StatusCode)
_, err = c.Writer.Write(jsonResponse)
if err != nil {
return nil, types.NewError(err, types.ErrorCodeBadResponseBody)
}
return &dto.Usage{}, nil
}
+177
View File
@@ -0,0 +1,177 @@
package jimeng
import (
"bytes"
"crypto/hmac"
"crypto/sha256"
"encoding/hex"
"encoding/json"
"errors"
"fmt"
"io"
"net/http"
"net/url"
"sort"
"strings"
"time"
"github.com/QuantumNous/new-api/logger"
"github.com/gin-gonic/gin"
)
// SignRequestForJimeng 对即梦 API 请求进行签名,支持 http.Request 或 header+url+body 方式
//func SignRequestForJimeng(req *http.Request, accessKey, secretKey string) error {
// var bodyBytes []byte
// var err error
//
// if req.Body != nil {
// bodyBytes, err = io.ReadAll(req.Body)
// if err != nil {
// return fmt.Errorf("read request body failed: %w", err)
// }
// _ = req.Body.Close()
// req.Body = io.NopCloser(bytes.NewBuffer(bodyBytes)) // rewind
// } else {
// bodyBytes = []byte{}
// }
//
// return signJimengHeaders(&req.Header, req.Method, req.URL, bodyBytes, accessKey, secretKey)
//}
const HexPayloadHashKey = "HexPayloadHash"
func SetPayloadHash(c *gin.Context, req any) error {
body, err := json.Marshal(req)
if err != nil {
return err
}
logger.LogInfo(c, fmt.Sprintf("SetPayloadHash body: %s", body))
payloadHash := sha256.Sum256(body)
hexPayloadHash := hex.EncodeToString(payloadHash[:])
c.Set(HexPayloadHashKey, hexPayloadHash)
return nil
}
func getPayloadHash(c *gin.Context) string {
return c.GetString(HexPayloadHashKey)
}
func Sign(c *gin.Context, req *http.Request, apiKey string) error {
header := req.Header
var bodyBytes []byte
var err error
if req.Body != nil {
bodyBytes, err = io.ReadAll(req.Body)
if err != nil {
return err
}
_ = req.Body.Close()
req.Body = io.NopCloser(bytes.NewBuffer(bodyBytes)) // Rewind
}
payloadHash := sha256.Sum256(bodyBytes)
hexPayloadHash := hex.EncodeToString(payloadHash[:])
method := c.Request.Method
u := req.URL
keyParts := strings.Split(apiKey, "|")
if len(keyParts) != 2 {
return errors.New("invalid api key format for jimeng: expected 'ak|sk'")
}
accessKey := strings.TrimSpace(keyParts[0])
secretKey := strings.TrimSpace(keyParts[1])
t := time.Now().UTC()
xDate := t.Format("20060102T150405Z")
shortDate := t.Format("20060102")
host := u.Host
header.Set("Host", host)
header.Set("X-Date", xDate)
header.Set("X-Content-Sha256", hexPayloadHash)
// Sort and encode query parameters to create canonical query string
queryParams := u.Query()
sortedKeys := make([]string, 0, len(queryParams))
for k := range queryParams {
sortedKeys = append(sortedKeys, k)
}
sort.Strings(sortedKeys)
var queryParts []string
for _, k := range sortedKeys {
values := queryParams[k]
sort.Strings(values)
for _, v := range values {
queryParts = append(queryParts, fmt.Sprintf("%s=%s", url.QueryEscape(k), url.QueryEscape(v)))
}
}
canonicalQueryString := strings.Join(queryParts, "&")
headersToSign := map[string]string{
"host": host,
"x-date": xDate,
"x-content-sha256": hexPayloadHash,
}
if header.Get("Content-Type") == "" {
header.Set("Content-Type", "application/json")
}
headersToSign["content-type"] = header.Get("Content-Type")
var signedHeaderKeys []string
for k := range headersToSign {
signedHeaderKeys = append(signedHeaderKeys, k)
}
sort.Strings(signedHeaderKeys)
var canonicalHeaders strings.Builder
for _, k := range signedHeaderKeys {
canonicalHeaders.WriteString(k)
canonicalHeaders.WriteString(":")
canonicalHeaders.WriteString(strings.TrimSpace(headersToSign[k]))
canonicalHeaders.WriteString("\n")
}
signedHeaders := strings.Join(signedHeaderKeys, ";")
canonicalRequest := fmt.Sprintf("%s\n%s\n%s\n%s\n%s\n%s",
method,
u.Path,
canonicalQueryString,
canonicalHeaders.String(),
signedHeaders,
hexPayloadHash,
)
hashedCanonicalRequest := sha256.Sum256([]byte(canonicalRequest))
hexHashedCanonicalRequest := hex.EncodeToString(hashedCanonicalRequest[:])
region := "cn-north-1"
serviceName := "cv"
credentialScope := fmt.Sprintf("%s/%s/%s/request", shortDate, region, serviceName)
stringToSign := fmt.Sprintf("HMAC-SHA256\n%s\n%s\n%s",
xDate,
credentialScope,
hexHashedCanonicalRequest,
)
kDate := hmacSHA256([]byte(secretKey), []byte(shortDate))
kRegion := hmacSHA256(kDate, []byte(region))
kService := hmacSHA256(kRegion, []byte(serviceName))
kSigning := hmacSHA256(kService, []byte("request"))
signature := hex.EncodeToString(hmacSHA256(kSigning, []byte(stringToSign)))
authorization := fmt.Sprintf("HMAC-SHA256 Credential=%s/%s, SignedHeaders=%s, Signature=%s",
accessKey,
credentialScope,
signedHeaders,
signature,
)
header.Set("Authorization", authorization)
return nil
}
// hmacSHA256 计算 HMAC-SHA256
func hmacSHA256(key []byte, data []byte) []byte {
h := hmac.New(sha256.New, key)
h.Write(data)
return h.Sum(nil)
}
+98
View File
@@ -0,0 +1,98 @@
package jina
import (
"errors"
"fmt"
"io"
"net/http"
"github.com/QuantumNous/new-api/relay/channel"
"github.com/QuantumNous/new-api/relay/channel/openai"
relaycommon "github.com/QuantumNous/new-api/relay/common"
"github.com/QuantumNous/new-api/relay/common_handler"
"github.com/QuantumNous/new-api/relay/constant"
"github.com/QuantumNous/new-api/relaykit/dto"
"github.com/QuantumNous/new-api/relaykit/types"
"github.com/gin-gonic/gin"
)
type Adaptor struct {
}
func (a *Adaptor) ConvertGeminiRequest(*gin.Context, *relaycommon.RelayInfo, *dto.GeminiChatRequest) (any, error) {
//TODO implement me
return nil, errors.New("not implemented")
}
func (a *Adaptor) ConvertClaudeRequest(*gin.Context, *relaycommon.RelayInfo, *dto.ClaudeRequest) (any, error) {
//TODO implement me
panic("implement me")
}
func (a *Adaptor) ConvertAudioRequest(c *gin.Context, info *relaycommon.RelayInfo, request dto.AudioRequest) (io.Reader, error) {
//TODO implement me
return nil, errors.New("not implemented")
}
func (a *Adaptor) ConvertImageRequest(c *gin.Context, info *relaycommon.RelayInfo, request dto.ImageRequest) (any, error) {
//TODO implement me
return nil, errors.New("not implemented")
}
func (a *Adaptor) Init(info *relaycommon.RelayInfo) {
}
func (a *Adaptor) GetRequestURL(info *relaycommon.RelayInfo) (string, error) {
if info.RelayMode == constant.RelayModeRerank {
return fmt.Sprintf("%s/v1/rerank", info.ChannelBaseUrl), nil
} else if info.RelayMode == constant.RelayModeEmbeddings {
return fmt.Sprintf("%s/v1/embeddings", info.ChannelBaseUrl), nil
}
return "", errors.New("invalid relay mode")
}
func (a *Adaptor) SetupRequestHeader(c *gin.Context, req *http.Header, info *relaycommon.RelayInfo) error {
channel.SetupApiRequestHeader(info, c, req)
req.Set("Authorization", fmt.Sprintf("Bearer %s", info.ApiKey))
return nil
}
func (a *Adaptor) ConvertOpenAIRequest(c *gin.Context, info *relaycommon.RelayInfo, request *dto.GeneralOpenAIRequest) (any, error) {
return request, nil
}
func (a *Adaptor) ConvertOpenAIResponsesRequest(c *gin.Context, info *relaycommon.RelayInfo, request dto.OpenAIResponsesRequest) (any, error) {
// TODO implement me
return nil, errors.New("not implemented")
}
func (a *Adaptor) DoRequest(c *gin.Context, info *relaycommon.RelayInfo, requestBody io.Reader) (any, error) {
return channel.DoApiRequest(a, c, info, requestBody)
}
func (a *Adaptor) ConvertRerankRequest(c *gin.Context, relayMode int, request dto.RerankRequest) (any, error) {
return request, nil
}
func (a *Adaptor) ConvertEmbeddingRequest(c *gin.Context, info *relaycommon.RelayInfo, request dto.EmbeddingRequest) (any, error) {
request.EncodingFormat = ""
return request, nil
}
func (a *Adaptor) DoResponse(c *gin.Context, resp *http.Response, info *relaycommon.RelayInfo) (usage any, err *types.NewAPIError) {
if info.RelayMode == constant.RelayModeRerank {
usage, err = common_handler.RerankHandler(c, info, resp)
} else if info.RelayMode == constant.RelayModeEmbeddings {
usage, err = openai.OpenaiHandler(c, info, resp)
}
return
}
func (a *Adaptor) GetModelList() []string {
return ModelList
}
func (a *Adaptor) GetChannelName() string {
return ChannelName
}
+9
View File
@@ -0,0 +1,9 @@
package jina
var ModelList = []string{
"jina-clip-v1",
"jina-reranker-v2-base-multilingual",
"jina-reranker-m0",
}
var ChannelName = "jina"
+1
View File
@@ -0,0 +1 @@
package jina
+9
View File
@@ -0,0 +1,9 @@
package lingyiwanwu
// https://platform.lingyiwanwu.com/docs
var ModelList = []string{
"yi-large", "yi-medium", "yi-vision", "yi-medium-200k", "yi-spark", "yi-large-rag", "yi-large-turbo", "yi-large-preview", "yi-large-rag-preview",
}
var ChannelName = "lingyiwanwu"
+147
View File
@@ -0,0 +1,147 @@
package minimax
import (
"bytes"
"encoding/json"
"errors"
"fmt"
"io"
"net/http"
"github.com/QuantumNous/new-api/relay/channel"
"github.com/QuantumNous/new-api/relay/channel/claude"
"github.com/QuantumNous/new-api/relay/channel/openai"
relaycommon "github.com/QuantumNous/new-api/relay/common"
"github.com/QuantumNous/new-api/relay/constant"
"github.com/QuantumNous/new-api/relaykit/dto"
"github.com/QuantumNous/new-api/relaykit/types"
"github.com/gin-gonic/gin"
"github.com/samber/lo"
)
type Adaptor struct {
}
func (a *Adaptor) ConvertGeminiRequest(*gin.Context, *relaycommon.RelayInfo, *dto.GeminiChatRequest) (any, error) {
return nil, errors.New("not implemented")
}
func (a *Adaptor) ConvertClaudeRequest(c *gin.Context, info *relaycommon.RelayInfo, req *dto.ClaudeRequest) (any, error) {
adaptor := claude.Adaptor{}
return adaptor.ConvertClaudeRequest(c, info, req)
}
func (a *Adaptor) ConvertAudioRequest(c *gin.Context, info *relaycommon.RelayInfo, request dto.AudioRequest) (io.Reader, error) {
if info.RelayMode != constant.RelayModeAudioSpeech {
return nil, errors.New("unsupported audio relay mode")
}
voiceID := request.Voice
speed := lo.FromPtrOr(request.Speed, 0.0)
outputFormat := request.ResponseFormat
minimaxRequest := MiniMaxTTSRequest{
Model: info.OriginModelName,
Text: request.Input,
VoiceSetting: VoiceSetting{
VoiceID: voiceID,
Speed: speed,
},
AudioSetting: &AudioSetting{
Format: outputFormat,
},
OutputFormat: outputFormat,
}
// 同步扩展字段的厂商自定义metadata
if len(request.Metadata) > 0 {
if err := json.Unmarshal(request.Metadata, &minimaxRequest); err != nil {
return nil, fmt.Errorf("error unmarshalling metadata to minimax request: %w", err)
}
}
jsonData, err := json.Marshal(minimaxRequest)
if err != nil {
return nil, fmt.Errorf("error marshalling minimax request: %w", err)
}
if outputFormat != "hex" {
outputFormat = "url"
}
c.Set("response_format", outputFormat)
// Debug: log the request structure
// fmt.Printf("MiniMax TTS Request: %s\n", string(jsonData))
return bytes.NewReader(jsonData), nil
}
func (a *Adaptor) ConvertImageRequest(c *gin.Context, info *relaycommon.RelayInfo, request dto.ImageRequest) (any, error) {
if info.RelayMode != constant.RelayModeImagesGenerations {
return nil, fmt.Errorf("unsupported image relay mode: %d", info.RelayMode)
}
return oaiImage2MiniMaxImageRequest(request), nil
}
func (a *Adaptor) Init(info *relaycommon.RelayInfo) {
}
func (a *Adaptor) GetRequestURL(info *relaycommon.RelayInfo) (string, error) {
return GetRequestURL(info)
}
func (a *Adaptor) SetupRequestHeader(c *gin.Context, req *http.Header, info *relaycommon.RelayInfo) error {
channel.SetupApiRequestHeader(info, c, req)
req.Set("Authorization", "Bearer "+info.ApiKey)
return nil
}
func (a *Adaptor) ConvertOpenAIRequest(c *gin.Context, info *relaycommon.RelayInfo, request *dto.GeneralOpenAIRequest) (any, error) {
if request == nil {
return nil, errors.New("request is nil")
}
return request, nil
}
func (a *Adaptor) ConvertRerankRequest(c *gin.Context, relayMode int, request dto.RerankRequest) (any, error) {
return nil, nil
}
func (a *Adaptor) ConvertEmbeddingRequest(c *gin.Context, info *relaycommon.RelayInfo, request dto.EmbeddingRequest) (any, error) {
return request, nil
}
func (a *Adaptor) ConvertOpenAIResponsesRequest(c *gin.Context, info *relaycommon.RelayInfo, request dto.OpenAIResponsesRequest) (any, error) {
return nil, errors.New("not implemented")
}
func (a *Adaptor) DoRequest(c *gin.Context, info *relaycommon.RelayInfo, requestBody io.Reader) (any, error) {
return channel.DoApiRequest(a, c, info, requestBody)
}
func (a *Adaptor) DoResponse(c *gin.Context, resp *http.Response, info *relaycommon.RelayInfo) (usage any, err *types.NewAPIError) {
if info.RelayMode == constant.RelayModeAudioSpeech {
return handleTTSResponse(c, resp, info)
}
if info.RelayMode == constant.RelayModeImagesGenerations {
return miniMaxImageHandler(c, resp, info)
}
switch info.RelayFormat {
case types.RelayFormatClaude:
adaptor := claude.Adaptor{}
return adaptor.DoResponse(c, resp, info)
default:
adaptor := openai.Adaptor{}
return adaptor.DoResponse(c, resp, info)
}
}
func (a *Adaptor) GetModelList() []string {
return ModelList
}
func (a *Adaptor) GetChannelName() string {
return ChannelName
}
+137
View File
@@ -0,0 +1,137 @@
package minimax
import (
"encoding/json"
"net/http"
"net/http/httptest"
"strings"
"testing"
"time"
relaycommon "github.com/QuantumNous/new-api/relay/common"
relayconstant "github.com/QuantumNous/new-api/relay/constant"
"github.com/QuantumNous/new-api/relaykit/dto"
"github.com/gin-gonic/gin"
)
func TestGetRequestURLForImageGeneration(t *testing.T) {
t.Parallel()
info := &relaycommon.RelayInfo{
RelayMode: relayconstant.RelayModeImagesGenerations,
ChannelMeta: &relaycommon.ChannelMeta{
ChannelBaseUrl: "https://api.minimax.chat",
},
}
got, err := GetRequestURL(info)
if err != nil {
t.Fatalf("GetRequestURL returned error: %v", err)
}
want := "https://api.minimax.chat/v1/image_generation"
if got != want {
t.Fatalf("GetRequestURL() = %q, want %q", got, want)
}
}
func TestConvertImageRequest(t *testing.T) {
t.Parallel()
adaptor := &Adaptor{}
info := &relaycommon.RelayInfo{
RelayMode: relayconstant.RelayModeImagesGenerations,
OriginModelName: "image-01",
}
request := dto.ImageRequest{
Model: "image-01",
Prompt: "a red fox in snowfall",
Size: "1536x1024",
ResponseFormat: "url",
N: uintPtr(2),
}
got, err := adaptor.ConvertImageRequest(gin.CreateTestContextOnly(httptest.NewRecorder(), gin.New()), info, request)
if err != nil {
t.Fatalf("ConvertImageRequest returned error: %v", err)
}
body, err := json.Marshal(got)
if err != nil {
t.Fatalf("json.Marshal returned error: %v", err)
}
var payload map[string]any
if err := json.Unmarshal(body, &payload); err != nil {
t.Fatalf("json.Unmarshal returned error: %v", err)
}
if payload["model"] != "image-01" {
t.Fatalf("model = %#v, want %q", payload["model"], "image-01")
}
if payload["prompt"] != request.Prompt {
t.Fatalf("prompt = %#v, want %q", payload["prompt"], request.Prompt)
}
if payload["n"] != float64(2) {
t.Fatalf("n = %#v, want 2", payload["n"])
}
if payload["aspect_ratio"] != "3:2" {
t.Fatalf("aspect_ratio = %#v, want %q", payload["aspect_ratio"], "3:2")
}
if payload["response_format"] != "url" {
t.Fatalf("response_format = %#v, want %q", payload["response_format"], "url")
}
}
func TestDoResponseForImageGeneration(t *testing.T) {
t.Parallel()
gin.SetMode(gin.TestMode)
recorder := httptest.NewRecorder()
c, _ := gin.CreateTestContext(recorder)
info := &relaycommon.RelayInfo{
RelayMode: relayconstant.RelayModeImagesGenerations,
StartTime: time.Unix(1700000000, 0),
}
resp := &http.Response{
StatusCode: http.StatusOK,
Header: make(http.Header),
Body: httptest.NewRecorder().Result().Body,
}
resp.Body = ioNopCloser(`{"data":{"image_urls":["https://example.com/minimax.png"]}}`)
adaptor := &Adaptor{}
usage, err := adaptor.DoResponse(c, resp, info)
if err != nil {
t.Fatalf("DoResponse returned error: %v", err)
}
if usage == nil {
t.Fatalf("DoResponse returned nil usage")
}
body := recorder.Body.String()
if !strings.Contains(body, `"url":"https://example.com/minimax.png"`) {
t.Fatalf("response body = %s, want OpenAI image response with image URL", body)
}
if strings.Contains(body, `"image_urls"`) {
t.Fatalf("response body = %s, should not expose raw MiniMax image_urls payload", body)
}
}
type nopReadCloser struct {
*strings.Reader
}
func (n nopReadCloser) Close() error {
return nil
}
func ioNopCloser(body string) nopReadCloser {
return nopReadCloser{Reader: strings.NewReader(body)}
}
func uintPtr(v uint) *uint {
return &v
}
+28
View File
@@ -0,0 +1,28 @@
package minimax
// https://www.minimaxi.com/document/guides/chat-model/V2?id=65e0736ab2845de20908e2dd
var ModelList = []string{
"abab6.5-chat",
"abab6.5s-chat",
"abab6-chat",
"abab5.5-chat",
"abab5.5s-chat",
"MiniMax-M2.7",
"MiniMax-M2.7-highspeed",
"speech-2.5-hd-preview",
"speech-2.5-turbo-preview",
"speech-02-hd",
"speech-02-turbo",
"speech-01-hd",
"speech-01-turbo",
"MiniMax-M2.1",
"MiniMax-M2.1-highspeed",
"MiniMax-M2",
"MiniMax-M2.5",
"MiniMax-M2.5-highspeed",
"image-01",
"image-01-live",
}
var ChannelName = "minimax"
+213
View File
@@ -0,0 +1,213 @@
package minimax
import (
"fmt"
"io"
"net/http"
"strconv"
"strings"
"github.com/QuantumNous/new-api/common"
relaycommon "github.com/QuantumNous/new-api/relay/common"
"github.com/QuantumNous/new-api/relaykit/dto"
"github.com/QuantumNous/new-api/relaykit/types"
"github.com/QuantumNous/new-api/service"
"github.com/gin-gonic/gin"
)
type MiniMaxImageRequest struct {
Model string `json:"model"`
Prompt string `json:"prompt"`
AspectRatio string `json:"aspect_ratio,omitempty"`
ResponseFormat string `json:"response_format,omitempty"`
N int `json:"n,omitempty"`
PromptOptimizer *bool `json:"prompt_optimizer,omitempty"`
AigcWatermark *bool `json:"aigc_watermark,omitempty"`
}
type MiniMaxImageResponse struct {
ID string `json:"id"`
Data struct {
ImageURLs []string `json:"image_urls"`
ImageBase64 []string `json:"image_base64"`
} `json:"data"`
Metadata map[string]any `json:"metadata"`
BaseResp struct {
StatusCode int `json:"status_code"`
StatusMsg string `json:"status_msg"`
} `json:"base_resp"`
}
func oaiImage2MiniMaxImageRequest(request dto.ImageRequest) MiniMaxImageRequest {
responseFormat := normalizeMiniMaxResponseFormat(request.ResponseFormat)
minimaxRequest := MiniMaxImageRequest{
Model: request.Model,
Prompt: request.Prompt,
ResponseFormat: responseFormat,
N: 1,
AigcWatermark: request.Watermark,
}
if request.Model == "" {
minimaxRequest.Model = "image-01"
}
if request.N != nil && *request.N > 0 {
minimaxRequest.N = int(*request.N)
}
if aspectRatio := aspectRatioFromImageRequest(request); aspectRatio != "" {
minimaxRequest.AspectRatio = aspectRatio
}
if raw, ok := request.Extra["prompt_optimizer"]; ok {
var promptOptimizer bool
if err := common.Unmarshal(raw, &promptOptimizer); err == nil {
minimaxRequest.PromptOptimizer = &promptOptimizer
}
}
return minimaxRequest
}
func aspectRatioFromImageRequest(request dto.ImageRequest) string {
if raw, ok := request.Extra["aspect_ratio"]; ok {
var aspectRatio string
if err := common.Unmarshal(raw, &aspectRatio); err == nil && aspectRatio != "" {
return aspectRatio
}
}
switch request.Size {
case "1024x1024":
return "1:1"
case "1792x1024":
return "16:9"
case "1024x1792":
return "9:16"
case "1536x1024", "1248x832":
return "3:2"
case "1024x1536", "832x1248":
return "2:3"
case "1152x864":
return "4:3"
case "864x1152":
return "3:4"
case "1344x576":
return "21:9"
}
width, height, ok := parseImageSize(request.Size)
if !ok {
return ""
}
ratio := reduceAspectRatio(width, height)
switch ratio {
case "1:1", "16:9", "4:3", "3:2", "2:3", "3:4", "9:16", "21:9":
return ratio
default:
return ""
}
}
func parseImageSize(size string) (int, int, bool) {
parts := strings.Split(size, "x")
if len(parts) != 2 {
return 0, 0, false
}
width, err := strconv.Atoi(parts[0])
if err != nil {
return 0, 0, false
}
height, err := strconv.Atoi(parts[1])
if err != nil {
return 0, 0, false
}
if width <= 0 || height <= 0 {
return 0, 0, false
}
return width, height, true
}
func reduceAspectRatio(width, height int) string {
divisor := gcd(width, height)
return fmt.Sprintf("%d:%d", width/divisor, height/divisor)
}
func gcd(a, b int) int {
for b != 0 {
a, b = b, a%b
}
if a == 0 {
return 1
}
return a
}
func normalizeMiniMaxResponseFormat(responseFormat string) string {
switch strings.ToLower(responseFormat) {
case "", "url":
return "url"
case "b64_json", "base64":
return "base64"
default:
return responseFormat
}
}
func responseMiniMax2OpenAIImage(response *MiniMaxImageResponse, info *relaycommon.RelayInfo) (*dto.ImageResponse, error) {
imageResponse := &dto.ImageResponse{
Created: info.StartTime.Unix(),
}
for _, imageURL := range response.Data.ImageURLs {
imageResponse.Data = append(imageResponse.Data, dto.ImageData{Url: imageURL})
}
for _, imageBase64 := range response.Data.ImageBase64 {
imageResponse.Data = append(imageResponse.Data, dto.ImageData{B64Json: imageBase64})
}
if len(response.Metadata) > 0 {
metadata, err := common.Marshal(response.Metadata)
if err != nil {
return nil, err
}
imageResponse.Metadata = metadata
}
return imageResponse, nil
}
func miniMaxImageHandler(c *gin.Context, resp *http.Response, info *relaycommon.RelayInfo) (*dto.Usage, *types.NewAPIError) {
responseBody, err := io.ReadAll(resp.Body)
if err != nil {
return nil, types.NewOpenAIError(err, types.ErrorCodeReadResponseBodyFailed, http.StatusInternalServerError)
}
service.CloseResponseBodyGracefully(resp)
var minimaxResponse MiniMaxImageResponse
if err := common.Unmarshal(responseBody, &minimaxResponse); err != nil {
return nil, types.NewOpenAIError(err, types.ErrorCodeBadResponseBody, http.StatusInternalServerError)
}
if minimaxResponse.BaseResp.StatusCode != 0 {
return nil, types.WithOpenAIError(types.OpenAIError{
Message: minimaxResponse.BaseResp.StatusMsg,
Type: "minimax_image_error",
Code: fmt.Sprintf("%d", minimaxResponse.BaseResp.StatusCode),
}, resp.StatusCode)
}
openAIResponse, err := responseMiniMax2OpenAIImage(&minimaxResponse, info)
if err != nil {
return nil, types.NewError(err, types.ErrorCodeBadResponseBody)
}
jsonResponse, err := common.Marshal(openAIResponse)
if err != nil {
return nil, types.NewError(err, types.ErrorCodeBadResponseBody)
}
c.Writer.Header().Set("Content-Type", "application/json")
c.Writer.WriteHeader(resp.StatusCode)
if _, err := c.Writer.Write(jsonResponse); err != nil {
return nil, types.NewError(err, types.ErrorCodeBadResponseBody)
}
return &dto.Usage{}, nil
}
+32
View File
@@ -0,0 +1,32 @@
package minimax
import (
"fmt"
channelconstant "github.com/QuantumNous/new-api/constant"
relaycommon "github.com/QuantumNous/new-api/relay/common"
"github.com/QuantumNous/new-api/relay/constant"
"github.com/QuantumNous/new-api/relaykit/types"
)
func GetRequestURL(info *relaycommon.RelayInfo) (string, error) {
baseUrl := info.ChannelBaseUrl
if baseUrl == "" {
baseUrl = channelconstant.GetChannelBaseURL(channelconstant.ChannelTypeMiniMax)
}
switch info.RelayFormat {
case types.RelayFormatClaude:
return fmt.Sprintf("%s/anthropic/v1/messages", info.ChannelBaseUrl), nil
default:
switch info.RelayMode {
case constant.RelayModeChatCompletions:
return fmt.Sprintf("%s/v1/text/chatcompletion_v2", baseUrl), nil
case constant.RelayModeImagesGenerations:
return fmt.Sprintf("%s/v1/image_generation", baseUrl), nil
case constant.RelayModeAudioSpeech:
return fmt.Sprintf("%s/v1/t2a_v2", baseUrl), nil
default:
return "", fmt.Errorf("unsupported relay mode: %d", info.RelayMode)
}
}
}
+198
View File
@@ -0,0 +1,198 @@
package minimax
import (
"encoding/hex"
"encoding/json"
"errors"
"fmt"
"io"
"net/http"
"strings"
relaycommon "github.com/QuantumNous/new-api/relay/common"
"github.com/QuantumNous/new-api/relaykit/dto"
"github.com/QuantumNous/new-api/relaykit/types"
"github.com/QuantumNous/new-api/service"
"github.com/gin-gonic/gin"
)
type MiniMaxTTSRequest struct {
Model string `json:"model"`
Text string `json:"text"`
Stream bool `json:"stream,omitempty"`
StreamOptions *StreamOptions `json:"stream_options,omitempty"`
VoiceSetting VoiceSetting `json:"voice_setting"`
PronunciationDict *PronunciationDict `json:"pronunciation_dict,omitempty"`
AudioSetting *AudioSetting `json:"audio_setting,omitempty"`
TimbreWeights []TimbreWeight `json:"timbre_weights,omitempty"`
LanguageBoost string `json:"language_boost,omitempty"`
VoiceModify *VoiceModify `json:"voice_modify,omitempty"`
SubtitleEnable bool `json:"subtitle_enable,omitempty"`
OutputFormat string `json:"output_format,omitempty"`
AigcWatermark bool `json:"aigc_watermark,omitempty"`
}
type StreamOptions struct {
ExcludeAggregatedAudio bool `json:"exclude_aggregated_audio,omitempty"`
}
type VoiceSetting struct {
VoiceID string `json:"voice_id"`
Speed float64 `json:"speed,omitempty"`
Vol float64 `json:"vol,omitempty"`
Pitch int `json:"pitch,omitempty"`
Emotion string `json:"emotion,omitempty"`
TextNormalization bool `json:"text_normalization,omitempty"`
LatexRead bool `json:"latex_read,omitempty"`
}
type PronunciationDict struct {
Tone []string `json:"tone,omitempty"`
}
type AudioSetting struct {
SampleRate int `json:"sample_rate,omitempty"`
Bitrate int `json:"bitrate,omitempty"`
Format string `json:"format,omitempty"`
Channel int `json:"channel,omitempty"`
ForceCbr bool `json:"force_cbr,omitempty"`
}
type TimbreWeight struct {
VoiceID string `json:"voice_id"`
Weight int `json:"weight"`
}
type VoiceModify struct {
Pitch int `json:"pitch,omitempty"`
Intensity int `json:"intensity,omitempty"`
Timbre int `json:"timbre,omitempty"`
SoundEffects string `json:"sound_effects,omitempty"`
}
type MiniMaxTTSResponse struct {
Data MiniMaxTTSData `json:"data"`
ExtraInfo MiniMaxExtraInfo `json:"extra_info"`
TraceID string `json:"trace_id"`
BaseResp MiniMaxBaseResp `json:"base_resp"`
}
type MiniMaxTTSData struct {
Audio string `json:"audio"`
Status int `json:"status"`
}
type MiniMaxExtraInfo struct {
UsageCharacters int64 `json:"usage_characters"`
}
type MiniMaxBaseResp struct {
StatusCode int64 `json:"status_code"`
StatusMsg string `json:"status_msg"`
}
func getContentTypeByFormat(format string) string {
contentTypeMap := map[string]string{
"mp3": "audio/mpeg",
"wav": "audio/wav",
"flac": "audio/flac",
"aac": "audio/aac",
"pcm": "audio/pcm",
}
if ct, ok := contentTypeMap[format]; ok {
return ct
}
return "audio/mpeg" // default to mp3
}
func handleTTSResponse(c *gin.Context, resp *http.Response, info *relaycommon.RelayInfo) (usage any, err *types.NewAPIError) {
body, readErr := io.ReadAll(resp.Body)
if readErr != nil {
return nil, types.NewErrorWithStatusCode(
fmt.Errorf("failed to read minimax response: %w", readErr),
types.ErrorCodeReadResponseBodyFailed,
http.StatusInternalServerError,
)
}
defer resp.Body.Close()
// Parse response
var minimaxResp MiniMaxTTSResponse
if unmarshalErr := json.Unmarshal(body, &minimaxResp); unmarshalErr != nil {
return nil, types.NewErrorWithStatusCode(
fmt.Errorf("failed to unmarshal minimax TTS response: %w", unmarshalErr),
types.ErrorCodeBadResponseBody,
http.StatusInternalServerError,
)
}
// Check base_resp status code
if minimaxResp.BaseResp.StatusCode != 0 {
return nil, types.NewErrorWithStatusCode(
fmt.Errorf("minimax TTS error: %d - %s", minimaxResp.BaseResp.StatusCode, minimaxResp.BaseResp.StatusMsg),
types.ErrorCodeBadResponse,
http.StatusBadRequest,
)
}
// Check if we have audio data
if minimaxResp.Data.Audio == "" {
return nil, types.NewErrorWithStatusCode(
fmt.Errorf("no audio data in minimax TTS response"),
types.ErrorCodeBadResponse,
http.StatusBadRequest,
)
}
if strings.HasPrefix(minimaxResp.Data.Audio, "http") {
c.Redirect(http.StatusFound, minimaxResp.Data.Audio)
} else {
// Handle hex-encoded audio data
audioData, decodeErr := hex.DecodeString(minimaxResp.Data.Audio)
if decodeErr != nil {
return nil, types.NewErrorWithStatusCode(
fmt.Errorf("failed to decode hex audio data: %w", decodeErr),
types.ErrorCodeBadResponse,
http.StatusInternalServerError,
)
}
// Determine content type - default to mp3
contentType := "audio/mpeg"
c.Data(http.StatusOK, contentType, audioData)
}
usage = &dto.Usage{
PromptTokens: info.GetEstimatePromptTokens(),
CompletionTokens: 0,
TotalTokens: int(minimaxResp.ExtraInfo.UsageCharacters),
}
return usage, nil
}
func handleChatCompletionResponse(c *gin.Context, resp *http.Response, info *relaycommon.RelayInfo) (usage any, err *types.NewAPIError) {
body, readErr := io.ReadAll(resp.Body)
if readErr != nil {
return nil, types.NewErrorWithStatusCode(
errors.New("failed to read minimax response"),
types.ErrorCodeReadResponseBodyFailed,
http.StatusInternalServerError,
)
}
defer resp.Body.Close()
// Set response headers
for key, values := range resp.Header {
if !service.ShouldCopyUpstreamHeader(c, key, values) {
continue
}
for _, value := range values {
c.Header(key, value)
}
}
c.Data(resp.StatusCode, "application/json", body)
return nil, nil
}
+93
View File
@@ -0,0 +1,93 @@
package mistral
import (
"errors"
"io"
"net/http"
"github.com/QuantumNous/new-api/relay/channel"
"github.com/QuantumNous/new-api/relay/channel/openai"
relaycommon "github.com/QuantumNous/new-api/relay/common"
"github.com/QuantumNous/new-api/relaykit/dto"
"github.com/QuantumNous/new-api/relaykit/types"
"github.com/gin-gonic/gin"
)
type Adaptor struct {
}
func (a *Adaptor) ConvertGeminiRequest(*gin.Context, *relaycommon.RelayInfo, *dto.GeminiChatRequest) (any, error) {
//TODO implement me
return nil, errors.New("not implemented")
}
func (a *Adaptor) ConvertClaudeRequest(*gin.Context, *relaycommon.RelayInfo, *dto.ClaudeRequest) (any, error) {
//TODO implement me
panic("implement me")
}
func (a *Adaptor) ConvertAudioRequest(c *gin.Context, info *relaycommon.RelayInfo, request dto.AudioRequest) (io.Reader, error) {
//TODO implement me
return nil, errors.New("not implemented")
}
func (a *Adaptor) ConvertImageRequest(c *gin.Context, info *relaycommon.RelayInfo, request dto.ImageRequest) (any, error) {
//TODO implement me
return nil, errors.New("not implemented")
}
func (a *Adaptor) Init(info *relaycommon.RelayInfo) {
}
func (a *Adaptor) GetRequestURL(info *relaycommon.RelayInfo) (string, error) {
return relaycommon.GetFullRequestURL(info.ChannelBaseUrl, info.RequestURLPath, info.ChannelType), nil
}
func (a *Adaptor) SetupRequestHeader(c *gin.Context, req *http.Header, info *relaycommon.RelayInfo) error {
channel.SetupApiRequestHeader(info, c, req)
req.Set("Authorization", "Bearer "+info.ApiKey)
return nil
}
func (a *Adaptor) ConvertOpenAIRequest(c *gin.Context, info *relaycommon.RelayInfo, request *dto.GeneralOpenAIRequest) (any, error) {
if request == nil {
return nil, errors.New("request is nil")
}
return requestOpenAI2Mistral(request), nil
}
func (a *Adaptor) ConvertRerankRequest(c *gin.Context, relayMode int, request dto.RerankRequest) (any, error) {
return nil, nil
}
func (a *Adaptor) ConvertEmbeddingRequest(c *gin.Context, info *relaycommon.RelayInfo, request dto.EmbeddingRequest) (any, error) {
//TODO implement me
return nil, errors.New("not implemented")
}
func (a *Adaptor) ConvertOpenAIResponsesRequest(c *gin.Context, info *relaycommon.RelayInfo, request dto.OpenAIResponsesRequest) (any, error) {
// TODO implement me
return nil, errors.New("not implemented")
}
func (a *Adaptor) DoRequest(c *gin.Context, info *relaycommon.RelayInfo, requestBody io.Reader) (any, error) {
return channel.DoApiRequest(a, c, info, requestBody)
}
func (a *Adaptor) DoResponse(c *gin.Context, resp *http.Response, info *relaycommon.RelayInfo) (usage any, err *types.NewAPIError) {
if info.IsStream {
usage, err = openai.OaiStreamHandler(c, info, resp)
} else {
usage, err = openai.OpenaiHandler(c, info, resp)
}
return
}
func (a *Adaptor) GetModelList() []string {
return ModelList
}
func (a *Adaptor) GetChannelName() string {
return ChannelName
}
+12
View File
@@ -0,0 +1,12 @@
package mistral
var ModelList = []string{
"open-mistral-7b",
"open-mixtral-8x7b",
"mistral-small-latest",
"mistral-medium-latest",
"mistral-large-latest",
"mistral-embed",
}
var ChannelName = "mistral"
+83
View File
@@ -0,0 +1,83 @@
package mistral
import (
"regexp"
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/relaykit/dto"
)
var mistralToolCallIdRegexp = regexp.MustCompile("^[a-zA-Z0-9]{9}$")
func requestOpenAI2Mistral(request *dto.GeneralOpenAIRequest) *dto.GeneralOpenAIRequest {
messages := make([]dto.Message, 0, len(request.Messages))
idMap := make(map[string]string)
for _, message := range request.Messages {
// 1. tool_calls.id
toolCalls := message.ParseToolCalls()
if toolCalls != nil {
for i := range toolCalls {
if !mistralToolCallIdRegexp.MatchString(toolCalls[i].ID) {
if newId, ok := idMap[toolCalls[i].ID]; ok {
toolCalls[i].ID = newId
} else {
newId, err := common.GenerateRandomCharsKey(9)
if err == nil {
idMap[toolCalls[i].ID] = newId
toolCalls[i].ID = newId
}
}
}
}
message.SetToolCalls(toolCalls)
}
// 2. tool_call_id
if message.ToolCallId != "" {
if newId, ok := idMap[message.ToolCallId]; ok {
message.ToolCallId = newId
} else {
if !mistralToolCallIdRegexp.MatchString(message.ToolCallId) {
newId, err := common.GenerateRandomCharsKey(9)
if err == nil {
idMap[message.ToolCallId] = newId
message.ToolCallId = newId
}
}
}
}
mediaMessages := message.ParseContent()
if message.Role == "assistant" && message.ToolCalls != nil && message.Content == "" {
mediaMessages = []dto.MediaContent{}
}
for j, mediaMessage := range mediaMessages {
if mediaMessage.Type == dto.ContentTypeImageURL {
imageUrl := mediaMessage.GetImageMedia()
mediaMessage.ImageUrl = imageUrl.Url
mediaMessages[j] = mediaMessage
}
}
message.SetMediaContent(mediaMessages)
messages = append(messages, dto.Message{
Role: message.Role,
Content: message.Content,
ToolCalls: message.ToolCalls,
ToolCallId: message.ToolCallId,
})
}
out := &dto.GeneralOpenAIRequest{
Model: request.Model,
Stream: request.Stream,
Messages: messages,
Temperature: request.Temperature,
TopP: request.TopP,
Tools: request.Tools,
ToolChoice: request.ToolChoice,
}
if request.MaxTokens != nil || request.MaxCompletionTokens != nil {
maxTokens := request.GetMaxTokens()
out.MaxTokens = &maxTokens
}
return out
}
+111
View File
@@ -0,0 +1,111 @@
package mokaai
import (
"errors"
"fmt"
"io"
"net/http"
"strings"
"github.com/QuantumNous/new-api/relay/channel"
relaycommon "github.com/QuantumNous/new-api/relay/common"
"github.com/QuantumNous/new-api/relay/constant"
"github.com/QuantumNous/new-api/relaykit/dto"
"github.com/QuantumNous/new-api/relaykit/types"
"github.com/gin-gonic/gin"
)
type Adaptor struct {
}
func (a *Adaptor) ConvertGeminiRequest(*gin.Context, *relaycommon.RelayInfo, *dto.GeminiChatRequest) (any, error) {
//TODO implement me
return nil, errors.New("not implemented")
}
func (a *Adaptor) ConvertClaudeRequest(*gin.Context, *relaycommon.RelayInfo, *dto.ClaudeRequest) (any, error) {
//TODO implement me
panic("implement me")
}
func (a *Adaptor) ConvertAudioRequest(c *gin.Context, info *relaycommon.RelayInfo, request dto.AudioRequest) (io.Reader, error) {
//TODO implement me
return nil, errors.New("not implemented")
}
func (a *Adaptor) ConvertImageRequest(c *gin.Context, info *relaycommon.RelayInfo, request dto.ImageRequest) (any, error) {
//TODO implement me
return nil, errors.New("not implemented")
}
func (a *Adaptor) ConvertEmbeddingRequest(c *gin.Context, info *relaycommon.RelayInfo, request dto.EmbeddingRequest) (any, error) {
//TODO implement me
return request, nil
}
func (a *Adaptor) Init(info *relaycommon.RelayInfo) {
}
func (a *Adaptor) GetRequestURL(info *relaycommon.RelayInfo) (string, error) {
// https://cloud.baidu.com/doc/WENXINWORKSHOP/s/clntwmv7t
suffix := "chat/"
if strings.HasPrefix(info.UpstreamModelName, "m3e") {
suffix = "embeddings"
}
fullRequestURL := fmt.Sprintf("%s/%s", info.ChannelBaseUrl, suffix)
return fullRequestURL, nil
}
func (a *Adaptor) SetupRequestHeader(c *gin.Context, req *http.Header, info *relaycommon.RelayInfo) error {
channel.SetupApiRequestHeader(info, c, req)
req.Set("Authorization", fmt.Sprintf("Bearer %s", info.ApiKey))
return nil
}
func (a *Adaptor) ConvertOpenAIRequest(c *gin.Context, info *relaycommon.RelayInfo, request *dto.GeneralOpenAIRequest) (any, error) {
if request == nil {
return nil, errors.New("request is nil")
}
switch info.RelayMode {
case constant.RelayModeEmbeddings:
baiduEmbeddingRequest := embeddingRequestOpenAI2Moka(*request)
return baiduEmbeddingRequest, nil
default:
return nil, errors.New("not implemented")
}
}
func (a *Adaptor) ConvertRerankRequest(c *gin.Context, relayMode int, request dto.RerankRequest) (any, error) {
return nil, nil
}
func (a *Adaptor) ConvertOpenAIResponsesRequest(c *gin.Context, info *relaycommon.RelayInfo, request dto.OpenAIResponsesRequest) (any, error) {
// TODO implement me
return nil, errors.New("not implemented")
}
func (a *Adaptor) DoRequest(c *gin.Context, info *relaycommon.RelayInfo, requestBody io.Reader) (any, error) {
return channel.DoApiRequest(a, c, info, requestBody)
}
func (a *Adaptor) DoResponse(c *gin.Context, resp *http.Response, info *relaycommon.RelayInfo) (usage any, err *types.NewAPIError) {
switch info.RelayMode {
case constant.RelayModeEmbeddings:
return mokaEmbeddingHandler(c, info, resp)
default:
// err, usage = mokaHandler(c, resp)
}
return
}
func (a *Adaptor) GetModelList() []string {
return ModelList
}
func (a *Adaptor) GetChannelName() string {
return ChannelName
}
+9
View File
@@ -0,0 +1,9 @@
package mokaai
var ModelList = []string{
"m3e-large",
"m3e-base",
"m3e-small",
}
var ChannelName = "mokaai"
+84
View File
@@ -0,0 +1,84 @@
package mokaai
import (
"encoding/json"
"io"
"net/http"
"github.com/QuantumNous/new-api/common"
relaycommon "github.com/QuantumNous/new-api/relay/common"
"github.com/QuantumNous/new-api/relaykit/dto"
"github.com/QuantumNous/new-api/relaykit/types"
"github.com/QuantumNous/new-api/service"
"github.com/gin-gonic/gin"
)
func embeddingRequestOpenAI2Moka(request dto.GeneralOpenAIRequest) *dto.EmbeddingRequest {
var input []string // Change input to []string
switch v := request.Input.(type) {
case string:
input = []string{v} // Convert string to []string
case []string:
input = v // Already a []string, no conversion needed
case []interface{}:
for _, part := range v {
if str, ok := part.(string); ok {
input = append(input, str) // Append each string to the slice
}
}
}
return &dto.EmbeddingRequest{
Input: input,
Model: request.Model,
}
}
func embeddingResponseMoka2OpenAI(response *dto.EmbeddingResponse) *dto.OpenAIEmbeddingResponse {
openAIEmbeddingResponse := dto.OpenAIEmbeddingResponse{
Object: "list",
Data: make([]dto.OpenAIEmbeddingResponseItem, 0, len(response.Data)),
Model: "baidu-embedding",
Usage: response.Usage,
}
for _, item := range response.Data {
openAIEmbeddingResponse.Data = append(openAIEmbeddingResponse.Data, dto.OpenAIEmbeddingResponseItem{
Object: item.Object,
Index: item.Index,
Embedding: item.Embedding,
})
}
return &openAIEmbeddingResponse
}
func mokaEmbeddingHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http.Response) (*dto.Usage, *types.NewAPIError) {
var baiduResponse dto.EmbeddingResponse
responseBody, err := io.ReadAll(resp.Body)
if err != nil {
return nil, types.NewError(err, types.ErrorCodeBadResponseBody)
}
service.CloseResponseBodyGracefully(resp)
err = json.Unmarshal(responseBody, &baiduResponse)
if err != nil {
return nil, types.NewError(err, types.ErrorCodeBadResponseBody)
}
// if baiduResponse.ErrorMsg != "" {
// return &dto.OpenAIErrorWithStatusCode{
// Error: dto.OpenAIError{
// Type: "baidu_error",
// Param: "",
// },
// StatusCode: resp.StatusCode,
// }, nil
// }
fullTextResponse := embeddingResponseMoka2OpenAI(&baiduResponse)
jsonResponse, err := common.Marshal(fullTextResponse)
if err != nil {
return nil, types.NewError(err, types.ErrorCodeBadResponseBody)
}
c.Writer.Header().Set("Content-Type", "application/json")
c.Writer.WriteHeader(resp.StatusCode)
service.IOCopyBytesGracefully(c, resp, jsonResponse)
return &fullTextResponse.Usage, nil
}
+135
View File
@@ -0,0 +1,135 @@
package moonshot
import (
"errors"
"fmt"
"io"
"net/http"
"strings"
"github.com/QuantumNous/new-api/common"
channelconstant "github.com/QuantumNous/new-api/constant"
"github.com/QuantumNous/new-api/relay/channel"
"github.com/QuantumNous/new-api/relay/channel/claude"
"github.com/QuantumNous/new-api/relay/channel/openai"
relaycommon "github.com/QuantumNous/new-api/relay/common"
"github.com/QuantumNous/new-api/relay/constant"
"github.com/QuantumNous/new-api/relaykit/dto"
"github.com/QuantumNous/new-api/relaykit/types"
"github.com/gin-gonic/gin"
)
type Adaptor struct {
}
func (a *Adaptor) ConvertGeminiRequest(*gin.Context, *relaycommon.RelayInfo, *dto.GeminiChatRequest) (any, error) {
//TODO implement me
return nil, errors.New("not implemented")
}
func (a *Adaptor) ConvertClaudeRequest(c *gin.Context, info *relaycommon.RelayInfo, req *dto.ClaudeRequest) (any, error) {
adaptor := claude.Adaptor{}
return adaptor.ConvertClaudeRequest(c, info, req)
}
func (a *Adaptor) ConvertAudioRequest(c *gin.Context, info *relaycommon.RelayInfo, request dto.AudioRequest) (io.Reader, error) {
//TODO implement me
return nil, errors.New("not supported")
}
func (a *Adaptor) ConvertImageRequest(c *gin.Context, info *relaycommon.RelayInfo, request dto.ImageRequest) (any, error) {
adaptor := openai.Adaptor{}
return adaptor.ConvertImageRequest(c, info, request)
}
func (a *Adaptor) Init(info *relaycommon.RelayInfo) {
}
func (a *Adaptor) GetRequestURL(info *relaycommon.RelayInfo) (string, error) {
baseURL := info.ChannelBaseUrl
if specialPlan, ok := channelconstant.ChannelSpecialBases[baseURL]; ok {
if info.RelayFormat == types.RelayFormatClaude {
return fmt.Sprintf("%s/v1/messages", specialPlan.ClaudeBaseURL), nil
}
if info.RelayFormat == types.RelayFormatOpenAI {
return fmt.Sprintf("%s/chat/completions", specialPlan.OpenAIBaseURL), nil
}
}
switch info.RelayFormat {
case types.RelayFormatClaude:
return fmt.Sprintf("%s/anthropic/v1/messages", info.ChannelBaseUrl), nil
default:
if info.RelayMode == constant.RelayModeRerank {
return fmt.Sprintf("%s/v1/rerank", info.ChannelBaseUrl), nil
} else if info.RelayMode == constant.RelayModeEmbeddings {
return fmt.Sprintf("%s/v1/embeddings", info.ChannelBaseUrl), nil
} else if info.RelayMode == constant.RelayModeChatCompletions {
return fmt.Sprintf("%s/v1/chat/completions", info.ChannelBaseUrl), nil
} else if info.RelayMode == constant.RelayModeCompletions {
return fmt.Sprintf("%s/v1/completions", info.ChannelBaseUrl), nil
}
return fmt.Sprintf("%s/v1/chat/completions", info.ChannelBaseUrl), nil
}
}
func (a *Adaptor) SetupRequestHeader(c *gin.Context, req *http.Header, info *relaycommon.RelayInfo) error {
channel.SetupApiRequestHeader(info, c, req)
req.Set("Authorization", fmt.Sprintf("Bearer %s", info.ApiKey))
return nil
}
func (a *Adaptor) ConvertOpenAIRequest(c *gin.Context, info *relaycommon.RelayInfo, request *dto.GeneralOpenAIRequest) (any, error) {
if request.Temperature != nil && isTemperatureOneOnlyModel(getUpstreamModelName(info, request.Model)) && *request.Temperature != 1.0 {
request.Temperature = common.GetPointer[float64](1.0)
}
return request, nil
}
func getUpstreamModelName(info *relaycommon.RelayInfo, fallback string) string {
if info != nil && info.ChannelMeta != nil && info.UpstreamModelName != "" {
return info.UpstreamModelName
}
return fallback
}
func isTemperatureOneOnlyModel(model string) bool {
return strings.EqualFold(model, "kimi-k2.6")
}
func (a *Adaptor) ConvertOpenAIResponsesRequest(c *gin.Context, info *relaycommon.RelayInfo, request dto.OpenAIResponsesRequest) (any, error) {
// TODO implement me
return nil, errors.New("not implemented")
}
func (a *Adaptor) DoRequest(c *gin.Context, info *relaycommon.RelayInfo, requestBody io.Reader) (any, error) {
return channel.DoApiRequest(a, c, info, requestBody)
}
func (a *Adaptor) ConvertRerankRequest(c *gin.Context, relayMode int, request dto.RerankRequest) (any, error) {
return request, nil
}
func (a *Adaptor) ConvertEmbeddingRequest(c *gin.Context, info *relaycommon.RelayInfo, request dto.EmbeddingRequest) (any, error) {
return request, nil
}
func (a *Adaptor) DoResponse(c *gin.Context, resp *http.Response, info *relaycommon.RelayInfo) (usage any, err *types.NewAPIError) {
switch info.RelayFormat {
case types.RelayFormatClaude:
adaptor := claude.Adaptor{}
return adaptor.DoResponse(c, resp, info)
default:
adaptor := openai.Adaptor{}
return adaptor.DoResponse(c, resp, info)
}
}
func (a *Adaptor) GetModelList() []string {
return ModelList
}
func (a *Adaptor) GetChannelName() string {
return ChannelName
}
+68
View File
@@ -0,0 +1,68 @@
package moonshot
import (
"testing"
"github.com/QuantumNous/new-api/common"
relaycommon "github.com/QuantumNous/new-api/relay/common"
"github.com/QuantumNous/new-api/relaykit/dto"
"github.com/stretchr/testify/require"
)
func TestConvertOpenAIRequestKimiK26UsesOnlyAllowedTemperature(t *testing.T) {
request := &dto.GeneralOpenAIRequest{
Model: "kimi-k2.6",
Temperature: common.GetPointer[float64](0.7),
}
info := &relaycommon.RelayInfo{
ChannelMeta: &relaycommon.ChannelMeta{
UpstreamModelName: "kimi-k2.6",
},
}
converted, err := (&Adaptor{}).ConvertOpenAIRequest(nil, info, request)
require.NoError(t, err)
convertedRequest, ok := converted.(*dto.GeneralOpenAIRequest)
require.True(t, ok)
require.NotNil(t, convertedRequest.Temperature)
require.Equal(t, 1.0, *convertedRequest.Temperature)
}
func TestConvertOpenAIRequestKimiK26KeepsOmittedTemperatureOmitted(t *testing.T) {
request := &dto.GeneralOpenAIRequest{
Model: "kimi-k2.6",
}
info := &relaycommon.RelayInfo{
ChannelMeta: &relaycommon.ChannelMeta{
UpstreamModelName: "kimi-k2.6",
},
}
converted, err := (&Adaptor{}).ConvertOpenAIRequest(nil, info, request)
require.NoError(t, err)
convertedRequest, ok := converted.(*dto.GeneralOpenAIRequest)
require.True(t, ok)
require.Nil(t, convertedRequest.Temperature)
}
func TestConvertOpenAIRequestOtherMoonshotModelKeepsTemperature(t *testing.T) {
request := &dto.GeneralOpenAIRequest{
Model: "kimi-k2.5",
Temperature: common.GetPointer[float64](0.7),
}
info := &relaycommon.RelayInfo{
ChannelMeta: &relaycommon.ChannelMeta{
UpstreamModelName: "kimi-k2.5",
},
}
converted, err := (&Adaptor{}).ConvertOpenAIRequest(nil, info, request)
require.NoError(t, err)
convertedRequest, ok := converted.(*dto.GeneralOpenAIRequest)
require.True(t, ok)
require.NotNil(t, convertedRequest.Temperature)
require.Equal(t, 0.7, *convertedRequest.Temperature)
}
+12
View File
@@ -0,0 +1,12 @@
package moonshot
var ModelList = []string{
"kimi-k3",
"kimi-k2.5",
"kimi-k2-0905-preview",
"kimi-k2-turbo-preview",
"kimi-k2-thinking",
"kimi-k2-thinking-turbo",
}
var ChannelName = "moonshot"
+121
View File
@@ -0,0 +1,121 @@
package newapi
import (
"errors"
"io"
"net/http"
"github.com/QuantumNous/new-api/relay/channel"
"github.com/QuantumNous/new-api/relay/channel/claude"
"github.com/QuantumNous/new-api/relay/channel/gemini"
"github.com/QuantumNous/new-api/relay/channel/openai"
relaycommon "github.com/QuantumNous/new-api/relay/common"
relayconstant "github.com/QuantumNous/new-api/relay/constant"
"github.com/QuantumNous/new-api/relaykit/dto"
"github.com/QuantumNous/new-api/relaykit/types"
"github.com/gin-gonic/gin"
)
type Adaptor struct {
openaiAdaptor openai.Adaptor
claudeAdaptor claude.Adaptor
geminiAdaptor gemini.Adaptor
}
func (a *Adaptor) Init(info *relaycommon.RelayInfo) {
a.openaiAdaptor.Init(info)
a.claudeAdaptor.Init(info)
a.geminiAdaptor.Init(info)
}
func (a *Adaptor) GetRequestURL(info *relaycommon.RelayInfo) (string, error) {
if info.RelayMode == relayconstant.RelayModeAlphaSearch {
return relaycommon.GetFullRequestURL(info.ChannelBaseUrl, "/v1/alpha/search", info.ChannelType), nil
}
return relaycommon.GetFullRequestURL(info.ChannelBaseUrl, info.RequestURLPath, info.ChannelType), nil
}
func (a *Adaptor) SetupRequestHeader(c *gin.Context, req *http.Header, info *relaycommon.RelayInfo) error {
channel.SetupApiRequestHeader(info, c, req)
req.Set("Authorization", "Bearer "+info.ApiKey)
switch info.RelayFormat {
case types.RelayFormatClaude:
req.Set("x-api-key", info.ApiKey)
if req.Get("anthropic-version") == "" {
anthropicVersion := c.Request.Header.Get("anthropic-version")
if anthropicVersion == "" {
anthropicVersion = "2023-06-01"
}
req.Set("anthropic-version", anthropicVersion)
}
case types.RelayFormatGemini:
req.Set("x-goog-api-key", info.ApiKey)
}
return nil
}
func (a *Adaptor) ConvertOpenAIRequest(c *gin.Context, info *relaycommon.RelayInfo, request *dto.GeneralOpenAIRequest) (any, error) {
if request == nil {
return nil, errors.New("request is nil")
}
return request, nil
}
func (a *Adaptor) ConvertOpenAIResponsesRequest(c *gin.Context, info *relaycommon.RelayInfo, request dto.OpenAIResponsesRequest) (any, error) {
return request, nil
}
func (a *Adaptor) ConvertEmbeddingRequest(c *gin.Context, info *relaycommon.RelayInfo, request dto.EmbeddingRequest) (any, error) {
return request, nil
}
func (a *Adaptor) ConvertClaudeRequest(c *gin.Context, info *relaycommon.RelayInfo, request *dto.ClaudeRequest) (any, error) {
if request == nil {
return nil, errors.New("request is nil")
}
return a.claudeAdaptor.ConvertClaudeRequest(c, info, request)
}
func (a *Adaptor) ConvertGeminiRequest(c *gin.Context, info *relaycommon.RelayInfo, request *dto.GeminiChatRequest) (any, error) {
if request == nil {
return nil, errors.New("request is nil")
}
return a.geminiAdaptor.ConvertGeminiRequest(c, info, request)
}
func (a *Adaptor) ConvertImageRequest(c *gin.Context, info *relaycommon.RelayInfo, request dto.ImageRequest) (any, error) {
return a.openaiAdaptor.ConvertImageRequest(c, info, request)
}
func (a *Adaptor) ConvertRerankRequest(c *gin.Context, relayMode int, request dto.RerankRequest) (any, error) {
return nil, errors.New("endpoint not supported")
}
func (a *Adaptor) ConvertAudioRequest(c *gin.Context, info *relaycommon.RelayInfo, request dto.AudioRequest) (io.Reader, error) {
return nil, errors.New("endpoint not supported")
}
func (a *Adaptor) DoRequest(c *gin.Context, info *relaycommon.RelayInfo, requestBody io.Reader) (any, error) {
return channel.DoApiRequest(a, c, info, requestBody)
}
func (a *Adaptor) DoResponse(c *gin.Context, resp *http.Response, info *relaycommon.RelayInfo) (usage any, err *types.NewAPIError) {
switch info.RelayFormat {
case types.RelayFormatClaude:
return a.claudeAdaptor.DoResponse(c, resp, info)
case types.RelayFormatGemini:
return a.geminiAdaptor.DoResponse(c, resp, info)
default:
return a.openaiAdaptor.DoResponse(c, resp, info)
}
}
func (a *Adaptor) GetModelList() []string {
return ModelList
}
func (a *Adaptor) GetChannelName() string {
return ChannelName
}
+6
View File
@@ -0,0 +1,6 @@
package newapi
const ChannelName = "newapi"
// ModelList is empty because models are fetched dynamically from upstream /v1/models.
var ModelList = []string{}
+134
View File
@@ -0,0 +1,134 @@
package ollama
import (
"errors"
"fmt"
"io"
"net/http"
"github.com/QuantumNous/new-api/relay/channel"
"github.com/QuantumNous/new-api/relay/channel/claude"
"github.com/QuantumNous/new-api/relay/channel/openai"
relaycommon "github.com/QuantumNous/new-api/relay/common"
relayconstant "github.com/QuantumNous/new-api/relay/constant"
"github.com/QuantumNous/new-api/relaykit/dto"
"github.com/QuantumNous/new-api/relaykit/types"
"github.com/gin-gonic/gin"
)
type Adaptor struct {
}
func (a *Adaptor) ConvertGeminiRequest(*gin.Context, *relaycommon.RelayInfo, *dto.GeminiChatRequest) (any, error) {
return nil, errors.New("not implemented")
}
func (a *Adaptor) ConvertClaudeRequest(c *gin.Context, info *relaycommon.RelayInfo, request *dto.ClaudeRequest) (any, error) {
adaptor := claude.Adaptor{}
return adaptor.ConvertClaudeRequest(c, info, request)
}
func (a *Adaptor) ConvertAudioRequest(c *gin.Context, info *relaycommon.RelayInfo, request dto.AudioRequest) (io.Reader, error) {
return nil, errors.New("not implemented")
}
func (a *Adaptor) ConvertImageRequest(c *gin.Context, info *relaycommon.RelayInfo, request dto.ImageRequest) (any, error) {
return nil, errors.New("not implemented")
}
func (a *Adaptor) Init(info *relaycommon.RelayInfo) {
}
func (a *Adaptor) GetRequestURL(info *relaycommon.RelayInfo) (string, error) {
switch info.RelayFormat {
case types.RelayFormatClaude:
return (&claude.Adaptor{}).GetRequestURL(info)
default:
switch info.RelayMode {
case relayconstant.RelayModeEmbeddings:
return fmt.Sprintf("%s/api/embed", info.ChannelBaseUrl), nil
case relayconstant.RelayModeResponses:
return fmt.Sprintf("%s/v1/responses", info.ChannelBaseUrl), nil
case relayconstant.RelayModeResponsesCompact:
return fmt.Sprintf("%s/v1/responses/compact", info.ChannelBaseUrl), nil
case relayconstant.RelayModeCompletions:
return fmt.Sprintf("%s/api/generate", info.ChannelBaseUrl), nil
default:
return fmt.Sprintf("%s/api/chat", info.ChannelBaseUrl), nil
}
}
}
func (a *Adaptor) SetupRequestHeader(c *gin.Context, req *http.Header, info *relaycommon.RelayInfo) error {
channel.SetupApiRequestHeader(info, c, req)
req.Set("Authorization", "Bearer "+info.ApiKey)
switch info.RelayFormat {
case types.RelayFormatClaude:
claude.CommonClaudeHeadersOperation(c, req, info)
anthropicVersion := c.Request.Header.Get("anthropic-version")
if anthropicVersion == "" {
anthropicVersion = "2023-06-01"
}
req.Set("anthropic-version", anthropicVersion)
}
return nil
}
func (a *Adaptor) ConvertOpenAIRequest(c *gin.Context, info *relaycommon.RelayInfo, request *dto.GeneralOpenAIRequest) (any, error) {
if request == nil {
return nil, errors.New("request is nil")
}
switch info.RelayMode {
case relayconstant.RelayModeCompletions:
return openAIToGenerate(c, request)
default:
return openAIChatToOllamaChat(c, request)
}
}
func (a *Adaptor) ConvertRerankRequest(c *gin.Context, relayMode int, request dto.RerankRequest) (any, error) {
return nil, nil
}
func (a *Adaptor) ConvertEmbeddingRequest(c *gin.Context, info *relaycommon.RelayInfo, request dto.EmbeddingRequest) (any, error) {
return requestOpenAI2Embeddings(request), nil
}
func (a *Adaptor) ConvertOpenAIResponsesRequest(c *gin.Context, info *relaycommon.RelayInfo, request dto.OpenAIResponsesRequest) (any, error) {
adaptor := openai.Adaptor{}
return adaptor.ConvertOpenAIResponsesRequest(c, info, request)
}
func (a *Adaptor) DoRequest(c *gin.Context, info *relaycommon.RelayInfo, requestBody io.Reader) (any, error) {
return channel.DoApiRequest(a, c, info, requestBody)
}
func (a *Adaptor) DoResponse(c *gin.Context, resp *http.Response, info *relaycommon.RelayInfo) (usage any, err *types.NewAPIError) {
switch info.RelayFormat {
case types.RelayFormatClaude:
adaptor := claude.Adaptor{}
return adaptor.DoResponse(c, resp, info)
default:
switch info.RelayMode {
case relayconstant.RelayModeEmbeddings:
return ollamaEmbeddingHandler(c, info, resp)
case relayconstant.RelayModeResponses, relayconstant.RelayModeResponsesCompact:
adaptor := openai.Adaptor{}
return adaptor.DoResponse(c, resp, info)
default:
if info.IsStream {
return ollamaStreamHandler(c, info, resp)
}
return ollamaChatHandler(c, info, resp)
}
}
}
func (a *Adaptor) GetModelList() []string {
return ModelList
}
func (a *Adaptor) GetChannelName() string {
return ChannelName
}
+7
View File
@@ -0,0 +1,7 @@
package ollama
var ModelList = []string{
"llama3-7b",
}
var ChannelName = "ollama"
+108
View File
@@ -0,0 +1,108 @@
package ollama
import (
"encoding/json"
)
type OllamaChatMessage struct {
Role string `json:"role"`
Content string `json:"content,omitempty"`
Images []string `json:"images,omitempty"`
ToolCalls []OllamaToolCall `json:"tool_calls,omitempty"`
ToolName string `json:"tool_name,omitempty"`
ToolCallID string `json:"tool_call_id,omitempty"`
Thinking json.RawMessage `json:"thinking,omitempty"`
}
type OllamaToolFunction struct {
Name string `json:"name"`
Description string `json:"description,omitempty"`
Parameters any `json:"parameters,omitempty"`
}
type OllamaTool struct {
Type string `json:"type"`
Function OllamaToolFunction `json:"function"`
}
type OllamaToolCall struct {
ID string `json:"id,omitempty"`
Function struct {
Name string `json:"name"`
Arguments any `json:"arguments"`
} `json:"function"`
}
type OllamaChatRequest struct {
Model string `json:"model"`
Messages []OllamaChatMessage `json:"messages"`
Tools any `json:"tools,omitempty"`
Format any `json:"format,omitempty"`
Stream bool `json:"stream"`
Options map[string]any `json:"options,omitempty"`
KeepAlive any `json:"keep_alive,omitempty"`
Think json.RawMessage `json:"think,omitempty"`
}
type OllamaGenerateRequest struct {
Model string `json:"model"`
Prompt string `json:"prompt,omitempty"`
Suffix string `json:"suffix,omitempty"`
Images []string `json:"images,omitempty"`
Format any `json:"format,omitempty"`
Stream bool `json:"stream"`
Options map[string]any `json:"options,omitempty"`
KeepAlive any `json:"keep_alive,omitempty"`
Think json.RawMessage `json:"think,omitempty"`
}
type OllamaEmbeddingRequest struct {
Model string `json:"model"`
Input any `json:"input"`
Options map[string]any `json:"options,omitempty"`
Dimensions int `json:"dimensions,omitempty"`
}
type OllamaEmbeddingResponse struct {
Error string `json:"error,omitempty"`
Model string `json:"model"`
Embeddings [][]float64 `json:"embeddings"`
PromptEvalCount int `json:"prompt_eval_count,omitempty"`
}
type OllamaTagsResponse struct {
Models []OllamaModel `json:"models"`
}
type OllamaModel struct {
Name string `json:"name"`
Size int64 `json:"size"`
Digest string `json:"digest,omitempty"`
ModifiedAt string `json:"modified_at"`
Details OllamaModelDetail `json:"details"`
}
type OllamaModelDetail struct {
ParentModel string `json:"parent_model,omitempty"`
Format string `json:"format,omitempty"`
Family string `json:"family,omitempty"`
Families []string `json:"families,omitempty"`
ParameterSize string `json:"parameter_size,omitempty"`
QuantizationLevel string `json:"quantization_level,omitempty"`
}
type OllamaPullRequest struct {
Name string `json:"name"`
Stream bool `json:"stream,omitempty"`
}
type OllamaPullResponse struct {
Status string `json:"status"`
Digest string `json:"digest,omitempty"`
Total int64 `json:"total,omitempty"`
Completed int64 `json:"completed,omitempty"`
}
type OllamaDeleteRequest struct {
Name string `json:"name"`
}
+575
View File
@@ -0,0 +1,575 @@
package ollama
import (
"fmt"
"io"
"net/http"
"strings"
"time"
"github.com/QuantumNous/new-api/common"
relaycommon "github.com/QuantumNous/new-api/relay/common"
"github.com/QuantumNous/new-api/relay/helper"
"github.com/QuantumNous/new-api/relaykit/dto"
"github.com/QuantumNous/new-api/relaykit/types"
"github.com/QuantumNous/new-api/service"
"github.com/gin-gonic/gin"
"github.com/samber/lo"
)
func toOllamaResponseFormat(responseFormat *dto.ResponseFormat) (any, error) {
if responseFormat == nil {
return nil, nil
}
switch responseFormat.Type {
case "json", "json_object":
return "json", nil
case "json_schema":
if len(responseFormat.JsonSchema) == 0 {
return nil, nil
}
var jsonSchema dto.FormatJsonSchema
if err := common.Unmarshal(responseFormat.JsonSchema, &jsonSchema); err != nil {
return nil, fmt.Errorf("invalid ollama response format: %w", err)
}
return jsonSchema.Schema, nil
default:
return nil, nil
}
}
func openAIChatToOllamaChat(c *gin.Context, r *dto.GeneralOpenAIRequest) (*OllamaChatRequest, error) {
think := r.Think
if len(think) == 0 {
effort := r.ReasoningEffort
if len(r.Reasoning) > 0 {
var reasoning dto.Reasoning
if err := common.Unmarshal(r.Reasoning, &reasoning); err != nil {
return nil, fmt.Errorf("invalid ollama reasoning: %w", err)
}
effort = lo.CoalesceOrEmpty(reasoning.Effort, effort)
}
if effort != "" {
var thinkValue any
switch effort {
case "none":
thinkValue = false
case "low", "medium", "high", "max":
thinkValue = effort
default:
return nil, fmt.Errorf("unsupported ollama reasoning effort %q", effort)
}
var err error
think, err = common.Marshal(thinkValue)
if err != nil {
return nil, fmt.Errorf("marshal ollama think: %w", err)
}
}
}
chatReq := &OllamaChatRequest{
Model: r.Model,
Stream: lo.FromPtrOr(r.Stream, false),
Options: map[string]any{},
Think: think,
}
format, err := toOllamaResponseFormat(r.ResponseFormat)
if err != nil {
return nil, err
}
chatReq.Format = format
// options mapping
if r.Temperature != nil {
chatReq.Options["temperature"] = r.Temperature
}
if r.TopP != nil {
chatReq.Options["top_p"] = lo.FromPtr(r.TopP)
}
if r.TopK != nil {
chatReq.Options["top_k"] = lo.FromPtr(r.TopK)
}
if r.FrequencyPenalty != nil {
chatReq.Options["frequency_penalty"] = lo.FromPtr(r.FrequencyPenalty)
}
if r.PresencePenalty != nil {
chatReq.Options["presence_penalty"] = lo.FromPtr(r.PresencePenalty)
}
if r.Seed != nil {
chatReq.Options["seed"] = int(lo.FromPtr(r.Seed))
}
if mt := r.GetMaxTokens(); mt != 0 {
chatReq.Options["num_predict"] = int(mt)
}
if r.Stop != nil {
switch v := r.Stop.(type) {
case string:
chatReq.Options["stop"] = []string{v}
case []string:
chatReq.Options["stop"] = v
case []any:
arr := lo.FilterMap(v, func(item any, _ int) (string, bool) {
value, ok := item.(string)
return value, ok
})
if len(arr) > 0 {
chatReq.Options["stop"] = arr
}
}
}
if len(r.Tools) > 0 {
chatReq.Tools = lo.Map(r.Tools, func(tool dto.ToolCallRequest, _ int) OllamaTool {
return OllamaTool{
Type: "function",
Function: OllamaToolFunction{
Name: tool.Function.Name,
Description: tool.Function.Description,
Parameters: tool.Function.Parameters,
},
}
})
}
chatReq.Messages = make([]OllamaChatMessage, 0, len(r.Messages))
toolNamesByCallID := make(map[string]string)
for _, m := range r.Messages {
var textBuilder strings.Builder
var images []string
if m.IsStringContent() {
textBuilder.WriteString(m.StringContent())
} else {
parts := m.ParseContent()
for _, part := range parts {
if part.Type == dto.ContentTypeImageURL {
source := part.ToFileSource()
if source != nil {
base64Data, _, err := service.GetBase64Data(c, source, "fetch image for ollama chat")
if err != nil {
return nil, err
}
if base64Data != "" {
images = append(images, base64Data)
}
}
} else if part.Type == dto.ContentTypeText {
textBuilder.WriteString(part.Text)
}
}
}
cm := OllamaChatMessage{Role: m.Role, Content: textBuilder.String()}
if len(images) > 0 {
cm.Images = images
}
if m.Role == "assistant" {
if reasoning, ok := lo.Coalesce(m.ReasoningContent, m.Reasoning); ok {
thinking, err := common.Marshal(*reasoning)
if err != nil {
return nil, fmt.Errorf("marshal ollama thinking: %w", err)
}
cm.Thinking = thinking
}
}
if m.Role == "tool" {
cm.ToolCallID = m.ToolCallId
cm.ToolName = lo.CoalesceOrEmpty(lo.FromPtr(m.Name), toolNamesByCallID[m.ToolCallId])
}
if m.ToolCalls != nil && len(m.ToolCalls) > 0 {
parsed := m.ParseToolCalls()
if len(parsed) > 0 {
calls := make([]OllamaToolCall, 0, len(parsed))
for _, tc := range parsed {
var args any
if tc.Function.Arguments != "" {
_ = common.Unmarshal([]byte(tc.Function.Arguments), &args)
}
if args == nil {
args = map[string]any{}
}
oc := OllamaToolCall{ID: tc.ID}
oc.Function.Name = tc.Function.Name
oc.Function.Arguments = args
calls = append(calls, oc)
if tc.ID != "" {
toolNamesByCallID[tc.ID] = tc.Function.Name
}
}
cm.ToolCalls = calls
}
}
chatReq.Messages = append(chatReq.Messages, cm)
}
return chatReq, nil
}
// openAIToGenerate converts OpenAI completions request to Ollama generate
func openAIToGenerate(c *gin.Context, r *dto.GeneralOpenAIRequest) (*OllamaGenerateRequest, error) {
gen := &OllamaGenerateRequest{
Model: r.Model,
Stream: lo.FromPtrOr(r.Stream, false),
Options: map[string]any{},
Think: r.Think,
}
// Prompt may be in r.Prompt (string or []any)
if r.Prompt != nil {
switch v := r.Prompt.(type) {
case string:
gen.Prompt = v
case []any:
var sb strings.Builder
for _, it := range v {
if s, ok := it.(string); ok {
sb.WriteString(s)
}
}
gen.Prompt = sb.String()
default:
gen.Prompt = fmt.Sprintf("%v", r.Prompt)
}
}
if r.Suffix != nil {
if s, ok := r.Suffix.(string); ok {
gen.Suffix = s
}
}
format, err := toOllamaResponseFormat(r.ResponseFormat)
if err != nil {
return nil, err
}
gen.Format = format
if r.Temperature != nil {
gen.Options["temperature"] = r.Temperature
}
if r.TopP != nil {
gen.Options["top_p"] = lo.FromPtr(r.TopP)
}
if r.TopK != nil {
gen.Options["top_k"] = lo.FromPtr(r.TopK)
}
if r.FrequencyPenalty != nil {
gen.Options["frequency_penalty"] = lo.FromPtr(r.FrequencyPenalty)
}
if r.PresencePenalty != nil {
gen.Options["presence_penalty"] = lo.FromPtr(r.PresencePenalty)
}
if r.Seed != nil {
gen.Options["seed"] = int(lo.FromPtr(r.Seed))
}
if mt := r.GetMaxTokens(); mt != 0 {
gen.Options["num_predict"] = int(mt)
}
if r.Stop != nil {
switch v := r.Stop.(type) {
case string:
gen.Options["stop"] = []string{v}
case []string:
gen.Options["stop"] = v
case []any:
arr := lo.FilterMap(v, func(item any, _ int) (string, bool) {
value, ok := item.(string)
return value, ok
})
if len(arr) > 0 {
gen.Options["stop"] = arr
}
}
}
return gen, nil
}
func requestOpenAI2Embeddings(r dto.EmbeddingRequest) *OllamaEmbeddingRequest {
opts := map[string]any{}
if r.Temperature != nil {
opts["temperature"] = r.Temperature
}
if r.TopP != nil {
opts["top_p"] = lo.FromPtr(r.TopP)
}
if r.FrequencyPenalty != nil {
opts["frequency_penalty"] = lo.FromPtr(r.FrequencyPenalty)
}
if r.PresencePenalty != nil {
opts["presence_penalty"] = lo.FromPtr(r.PresencePenalty)
}
if r.Seed != nil {
opts["seed"] = int(lo.FromPtr(r.Seed))
}
dimensions := lo.FromPtrOr(r.Dimensions, 0)
if r.Dimensions != nil {
opts["dimensions"] = dimensions
}
input := r.ParseInput()
if len(input) == 1 {
return &OllamaEmbeddingRequest{Model: r.Model, Input: input[0], Options: opts, Dimensions: dimensions}
}
return &OllamaEmbeddingRequest{Model: r.Model, Input: input, Options: opts, Dimensions: dimensions}
}
func ollamaEmbeddingHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http.Response) (*dto.Usage, *types.NewAPIError) {
var oResp OllamaEmbeddingResponse
body, err := io.ReadAll(resp.Body)
if err != nil {
return nil, types.NewOpenAIError(err, types.ErrorCodeBadResponseBody, http.StatusInternalServerError)
}
service.CloseResponseBodyGracefully(resp)
if err = common.Unmarshal(body, &oResp); err != nil {
return nil, types.NewOpenAIError(err, types.ErrorCodeBadResponseBody, http.StatusInternalServerError)
}
if oResp.Error != "" {
return nil, types.NewOpenAIError(fmt.Errorf("ollama error: %s", oResp.Error), types.ErrorCodeBadResponseBody, http.StatusInternalServerError)
}
data := make([]dto.OpenAIEmbeddingResponseItem, 0, len(oResp.Embeddings))
for i, emb := range oResp.Embeddings {
data = append(data, dto.OpenAIEmbeddingResponseItem{Index: i, Object: "embedding", Embedding: emb})
}
usage := &dto.Usage{PromptTokens: oResp.PromptEvalCount, CompletionTokens: 0, TotalTokens: oResp.PromptEvalCount}
embResp := &dto.OpenAIEmbeddingResponse{Object: "list", Data: data, Model: info.UpstreamModelName, Usage: *usage}
out, _ := common.Marshal(embResp)
service.IOCopyBytesGracefully(c, resp, out)
return usage, nil
}
func FetchOllamaModels(baseURL, apiKey string) ([]OllamaModel, error) {
url := fmt.Sprintf("%s/api/tags", baseURL)
client := &http.Client{}
request, err := http.NewRequest("GET", url, nil)
if err != nil {
return nil, fmt.Errorf("创建请求失败: %v", err)
}
// Ollama 通常不需要 Bearer token,但为了兼容性保留
if apiKey != "" {
request.Header.Set("Authorization", "Bearer "+apiKey)
}
response, err := client.Do(request)
if err != nil {
return nil, fmt.Errorf("请求失败: %v", err)
}
defer response.Body.Close()
if response.StatusCode != http.StatusOK {
body, _ := io.ReadAll(response.Body)
return nil, fmt.Errorf("服务器返回错误 %d: %s", response.StatusCode, string(body))
}
var tagsResponse OllamaTagsResponse
body, err := io.ReadAll(response.Body)
if err != nil {
return nil, fmt.Errorf("读取响应失败: %v", err)
}
err = common.Unmarshal(body, &tagsResponse)
if err != nil {
return nil, fmt.Errorf("解析响应失败: %v", err)
}
return tagsResponse.Models, nil
}
// 拉取 Ollama 模型 (非流式)
func PullOllamaModel(baseURL, apiKey, modelName string) error {
url := fmt.Sprintf("%s/api/pull", baseURL)
pullRequest := OllamaPullRequest{
Name: modelName,
Stream: false, // 非流式,简化处理
}
requestBody, err := common.Marshal(pullRequest)
if err != nil {
return fmt.Errorf("序列化请求失败: %v", err)
}
client := &http.Client{
Timeout: 30 * 60 * 1000 * time.Millisecond, // 30分钟超时,支持大模型
}
request, err := http.NewRequest("POST", url, strings.NewReader(string(requestBody)))
if err != nil {
return fmt.Errorf("创建请求失败: %v", err)
}
request.Header.Set("Content-Type", "application/json")
if apiKey != "" {
request.Header.Set("Authorization", "Bearer "+apiKey)
}
response, err := client.Do(request)
if err != nil {
return fmt.Errorf("请求失败: %v", err)
}
defer response.Body.Close()
if response.StatusCode != http.StatusOK {
body, _ := io.ReadAll(response.Body)
return fmt.Errorf("拉取模型失败 %d: %s", response.StatusCode, string(body))
}
return nil
}
// 流式拉取 Ollama 模型 (支持进度回调)
func PullOllamaModelStream(baseURL, apiKey, modelName string, progressCallback func(OllamaPullResponse)) error {
url := fmt.Sprintf("%s/api/pull", baseURL)
pullRequest := OllamaPullRequest{
Name: modelName,
Stream: true, // 启用流式
}
requestBody, err := common.Marshal(pullRequest)
if err != nil {
return fmt.Errorf("序列化请求失败: %v", err)
}
client := &http.Client{
Timeout: 60 * 60 * 1000 * time.Millisecond, // 1小时超时,支持超大模型
}
request, err := http.NewRequest("POST", url, strings.NewReader(string(requestBody)))
if err != nil {
return fmt.Errorf("创建请求失败: %v", err)
}
request.Header.Set("Content-Type", "application/json")
if apiKey != "" {
request.Header.Set("Authorization", "Bearer "+apiKey)
}
response, err := client.Do(request)
if err != nil {
return fmt.Errorf("请求失败: %v", err)
}
defer response.Body.Close()
if response.StatusCode != http.StatusOK {
body, _ := io.ReadAll(response.Body)
return fmt.Errorf("拉取模型失败 %d: %s", response.StatusCode, string(body))
}
// 读取流式响应
scanner := helper.NewStreamScanner(response.Body)
successful := false
for scanner.Scan() {
line := scanner.Text()
if strings.TrimSpace(line) == "" {
continue
}
var pullResponse OllamaPullResponse
if err := common.Unmarshal([]byte(line), &pullResponse); err != nil {
continue // 忽略解析失败的行
}
if progressCallback != nil {
progressCallback(pullResponse)
}
// 检查是否出现错误或完成
if strings.EqualFold(pullResponse.Status, "error") {
return fmt.Errorf("拉取模型失败: %s", strings.TrimSpace(line))
}
if strings.EqualFold(pullResponse.Status, "success") {
successful = true
break
}
}
if err := scanner.Err(); err != nil {
return fmt.Errorf("读取流式响应失败: %v", err)
}
if !successful {
return fmt.Errorf("拉取模型未完成: 未收到成功状态")
}
return nil
}
// 删除 Ollama 模型
func DeleteOllamaModel(baseURL, apiKey, modelName string) error {
url := fmt.Sprintf("%s/api/delete", baseURL)
deleteRequest := OllamaDeleteRequest{
Name: modelName,
}
requestBody, err := common.Marshal(deleteRequest)
if err != nil {
return fmt.Errorf("序列化请求失败: %v", err)
}
client := &http.Client{}
request, err := http.NewRequest("DELETE", url, strings.NewReader(string(requestBody)))
if err != nil {
return fmt.Errorf("创建请求失败: %v", err)
}
request.Header.Set("Content-Type", "application/json")
if apiKey != "" {
request.Header.Set("Authorization", "Bearer "+apiKey)
}
response, err := client.Do(request)
if err != nil {
return fmt.Errorf("请求失败: %v", err)
}
defer response.Body.Close()
if response.StatusCode != http.StatusOK {
body, _ := io.ReadAll(response.Body)
return fmt.Errorf("删除模型失败 %d: %s", response.StatusCode, string(body))
}
return nil
}
func FetchOllamaVersion(baseURL, apiKey string) (string, error) {
trimmedBase := strings.TrimRight(baseURL, "/")
if trimmedBase == "" {
return "", fmt.Errorf("baseURL 为空")
}
url := fmt.Sprintf("%s/api/version", trimmedBase)
client := &http.Client{Timeout: 10 * time.Second}
request, err := http.NewRequest("GET", url, nil)
if err != nil {
return "", fmt.Errorf("创建请求失败: %v", err)
}
if apiKey != "" {
request.Header.Set("Authorization", "Bearer "+apiKey)
}
response, err := client.Do(request)
if err != nil {
return "", fmt.Errorf("请求失败: %v", err)
}
defer response.Body.Close()
body, err := io.ReadAll(response.Body)
if err != nil {
return "", fmt.Errorf("读取响应失败: %v", err)
}
if response.StatusCode != http.StatusOK {
return "", fmt.Errorf("查询版本失败 %d: %s", response.StatusCode, string(body))
}
var versionResp struct {
Version string `json:"version"`
}
if err := common.Unmarshal(body, &versionResp); err != nil {
return "", fmt.Errorf("解析响应失败: %v", err)
}
if versionResp.Version == "" {
return "", fmt.Errorf("未返回版本信息")
}
return versionResp.Version, nil
}
+346
View File
@@ -0,0 +1,346 @@
package ollama
import (
"encoding/json"
"fmt"
"io"
"net/http"
"strings"
"time"
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/constant"
"github.com/QuantumNous/new-api/logger"
relaycommon "github.com/QuantumNous/new-api/relay/common"
"github.com/QuantumNous/new-api/relay/helper"
"github.com/QuantumNous/new-api/relaykit/dto"
"github.com/QuantumNous/new-api/relaykit/types"
"github.com/QuantumNous/new-api/service"
"github.com/gin-gonic/gin"
)
type ollamaChatStreamChunk struct {
Model string `json:"model"`
CreatedAt string `json:"created_at"`
// chat
Message *struct {
Role string `json:"role"`
Content string `json:"content"`
Thinking json.RawMessage `json:"thinking"`
ToolCalls []OllamaToolCall `json:"tool_calls"`
} `json:"message"`
// generate
Response string `json:"response"`
Done bool `json:"done"`
DoneReason string `json:"done_reason"`
TotalDuration int64 `json:"total_duration"`
LoadDuration int64 `json:"load_duration"`
PromptEvalCount int `json:"prompt_eval_count"`
EvalCount int `json:"eval_count"`
PromptEvalDuration int64 `json:"prompt_eval_duration"`
EvalDuration int64 `json:"eval_duration"`
}
func ollamaToolCallsToOpenAI(toolCalls []OllamaToolCall, startIndex int, includeIndex bool) ([]dto.ToolCallResponse, int) {
if len(toolCalls) == 0 {
return nil, startIndex
}
result := make([]dto.ToolCallResponse, 0, len(toolCalls))
for _, tc := range toolCalls {
var argBytes []byte
var err error
if tc.Function.Arguments == nil {
argBytes = []byte("{}")
} else {
argBytes, err = common.Marshal(tc.Function.Arguments)
if err != nil || len(argBytes) == 0 {
argBytes = []byte("{}")
}
}
toolCallID := tc.ID
if toolCallID == "" {
toolCallID = fmt.Sprintf("call_%d", startIndex)
}
tr := dto.ToolCallResponse{
ID: toolCallID,
Type: "function",
Function: dto.FunctionResponse{
Name: tc.Function.Name,
Arguments: string(argBytes),
},
}
if includeIndex {
tr.SetIndex(startIndex)
}
startIndex++
result = append(result, tr)
}
return result, startIndex
}
func toUnix(ts string) int64 {
if ts == "" {
return time.Now().Unix()
}
// try time.RFC3339 or with nanoseconds
t, err := time.Parse(time.RFC3339Nano, ts)
if err != nil {
t2, err2 := time.Parse(time.RFC3339, ts)
if err2 == nil {
return t2.Unix()
}
return time.Now().Unix()
}
return t.Unix()
}
func ollamaStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http.Response) (*dto.Usage, *types.NewAPIError) {
if resp == nil || resp.Body == nil {
return nil, types.NewOpenAIError(fmt.Errorf("empty response"), types.ErrorCodeBadResponse, http.StatusBadRequest)
}
defer service.CloseResponseBodyGracefully(resp)
helper.SetEventStreamHeaders(c)
scanner := helper.NewStreamScanner(resp.Body)
usage := &dto.Usage{}
var model = info.UpstreamModelName
var responseId = common.GetUUID()
var created = time.Now().Unix()
var toolCallIndex int
start := helper.GenerateStartEmptyResponse(responseId, created, model, nil)
if data, err := common.Marshal(start); err == nil {
_ = helper.StringData(c, string(data))
}
for scanner.Scan() {
line := scanner.Text()
line = strings.TrimSpace(line)
if line == "" {
continue
}
var chunk ollamaChatStreamChunk
if err := common.Unmarshal([]byte(line), &chunk); err != nil {
logger.LogError(c, "ollama stream json decode error: "+err.Error()+" line="+line)
return usage, types.NewOpenAIError(err, types.ErrorCodeBadResponseBody, http.StatusInternalServerError)
}
if chunk.Model != "" {
model = chunk.Model
}
created = toUnix(chunk.CreatedAt)
if !chunk.Done {
// delta content
var content string
if chunk.Message != nil {
content = chunk.Message.Content
} else {
content = chunk.Response
}
delta := dto.ChatCompletionsStreamResponse{
Id: responseId,
Object: "chat.completion.chunk",
Created: created,
Model: model,
Choices: []dto.ChatCompletionsStreamResponseChoice{{
Index: 0,
Delta: dto.ChatCompletionsStreamResponseChoiceDelta{Role: "assistant"},
}},
}
if content != "" {
delta.Choices[0].Delta.SetContentString(content)
}
if chunk.Message != nil && len(chunk.Message.Thinking) > 0 {
raw := strings.TrimSpace(string(chunk.Message.Thinking))
if raw != "" && raw != "null" {
// Unmarshal the JSON string to get the actual content without quotes
var thinkingContent string
if err := common.Unmarshal(chunk.Message.Thinking, &thinkingContent); err == nil {
delta.Choices[0].Delta.SetReasoningContent(thinkingContent)
} else {
// Fallback to raw string if it's not a JSON string
delta.Choices[0].Delta.SetReasoningContent(raw)
}
}
}
// tool calls
if chunk.Message != nil && len(chunk.Message.ToolCalls) > 0 {
delta.Choices[0].Delta.ToolCalls, toolCallIndex = ollamaToolCallsToOpenAI(chunk.Message.ToolCalls, toolCallIndex, true)
}
if data, err := common.Marshal(delta); err == nil {
_ = helper.StringData(c, string(data))
}
continue
}
// done frame
// finalize once and break loop
usage.PromptTokens = chunk.PromptEvalCount
usage.CompletionTokens = chunk.EvalCount
usage.TotalTokens = usage.PromptTokens + usage.CompletionTokens
finishReason := chunk.DoneReason
if finishReason == "" {
finishReason = "stop"
}
if toolCallIndex > 0 {
finishReason = constant.FinishReasonToolCalls
}
// emit stop delta
if stop := helper.GenerateStopResponse(responseId, created, model, finishReason); stop != nil {
if data, err := common.Marshal(stop); err == nil {
_ = helper.StringData(c, string(data))
}
}
// emit usage frame
if final := helper.GenerateFinalUsageResponse(responseId, created, model, *usage); final != nil {
if data, err := common.Marshal(final); err == nil {
_ = helper.StringData(c, string(data))
}
}
// send [DONE]
helper.Done(c)
break
}
if err := scanner.Err(); err != nil && err != io.EOF {
logger.LogError(c, "ollama stream scan error: "+err.Error())
}
return usage, nil
}
// non-stream handler for chat/generate
func ollamaChatHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http.Response) (*dto.Usage, *types.NewAPIError) {
body, err := io.ReadAll(resp.Body)
if err != nil {
return nil, types.NewOpenAIError(err, types.ErrorCodeReadResponseBodyFailed, http.StatusInternalServerError)
}
service.CloseResponseBodyGracefully(resp)
raw := string(body)
if common.DebugEnabled {
println("ollama non-stream raw resp:", raw)
}
lines := strings.Split(raw, "\n")
var (
aggContent strings.Builder
reasoningBuilder strings.Builder
lastChunk ollamaChatStreamChunk
parsedAny bool
toolCallIndex int
toolCalls []dto.ToolCallResponse
)
for _, ln := range lines {
ln = strings.TrimSpace(ln)
if ln == "" {
continue
}
var ck ollamaChatStreamChunk
if err := common.Unmarshal([]byte(ln), &ck); err != nil {
if len(lines) == 1 {
return nil, types.NewOpenAIError(err, types.ErrorCodeBadResponseBody, http.StatusInternalServerError)
}
continue
}
parsedAny = true
lastChunk = ck
if ck.Message != nil && len(ck.Message.Thinking) > 0 {
raw := strings.TrimSpace(string(ck.Message.Thinking))
if raw != "" && raw != "null" {
// Unmarshal the JSON string to get the actual content without quotes
var thinkingContent string
if err := common.Unmarshal(ck.Message.Thinking, &thinkingContent); err == nil {
reasoningBuilder.WriteString(thinkingContent)
} else {
// Fallback to raw string if it's not a JSON string
reasoningBuilder.WriteString(raw)
}
}
}
if ck.Message != nil && ck.Message.Content != "" {
aggContent.WriteString(ck.Message.Content)
} else if ck.Response != "" {
aggContent.WriteString(ck.Response)
}
if ck.Message != nil && len(ck.Message.ToolCalls) > 0 {
var converted []dto.ToolCallResponse
converted, toolCallIndex = ollamaToolCallsToOpenAI(ck.Message.ToolCalls, toolCallIndex, false)
toolCalls = append(toolCalls, converted...)
}
}
if !parsedAny {
var single ollamaChatStreamChunk
if err := common.Unmarshal(body, &single); err != nil {
return nil, types.NewOpenAIError(err, types.ErrorCodeBadResponseBody, http.StatusInternalServerError)
}
lastChunk = single
if single.Message != nil {
if len(single.Message.Thinking) > 0 {
raw := strings.TrimSpace(string(single.Message.Thinking))
if raw != "" && raw != "null" {
// Unmarshal the JSON string to get the actual content without quotes
var thinkingContent string
if err := common.Unmarshal(single.Message.Thinking, &thinkingContent); err == nil {
reasoningBuilder.WriteString(thinkingContent)
} else {
// Fallback to raw string if it's not a JSON string
reasoningBuilder.WriteString(raw)
}
}
}
aggContent.WriteString(single.Message.Content)
if len(single.Message.ToolCalls) > 0 {
var converted []dto.ToolCallResponse
converted, toolCallIndex = ollamaToolCallsToOpenAI(single.Message.ToolCalls, toolCallIndex, false)
toolCalls = append(toolCalls, converted...)
}
} else {
aggContent.WriteString(single.Response)
}
}
model := lastChunk.Model
if model == "" {
model = info.UpstreamModelName
}
created := toUnix(lastChunk.CreatedAt)
usage := &dto.Usage{PromptTokens: lastChunk.PromptEvalCount, CompletionTokens: lastChunk.EvalCount, TotalTokens: lastChunk.PromptEvalCount + lastChunk.EvalCount}
content := aggContent.String()
finishReason := lastChunk.DoneReason
if finishReason == "" {
finishReason = "stop"
}
if len(toolCalls) > 0 {
finishReason = constant.FinishReasonToolCalls
}
msg := dto.Message{Role: "assistant", Content: contentPtr(content)}
if len(toolCalls) > 0 {
if rawToolCalls, err := common.Marshal(toolCalls); err == nil {
msg.ToolCalls = rawToolCalls
}
}
if rc := reasoningBuilder.String(); rc != "" {
msg.ReasoningContent = &rc
}
full := dto.OpenAITextResponse{
Id: common.GetUUID(),
Model: model,
Object: "chat.completion",
Created: created,
Choices: []dto.OpenAITextResponseChoice{{
Index: 0,
Message: msg,
FinishReason: finishReason,
}},
Usage: *usage,
}
out, _ := common.Marshal(full)
service.IOCopyBytesGracefully(c, resp, out)
return usage, nil
}
func contentPtr(s string) *string {
if s == "" {
return nil
}
return &s
}
Loaded 100 of 2381 files, more files were not shown because too many files have changed in this diff. Show more