Like-api/relay/controller/text.go

package controller

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io"
	"net/http"

	"github.com/gin-gonic/gin"

	"github.com/songquanpeng/one-api/common/logger"
	"github.com/songquanpeng/one-api/relay"
	"github.com/songquanpeng/one-api/relay/adaptor"
	"github.com/songquanpeng/one-api/relay/adaptor/openai"
	"github.com/songquanpeng/one-api/relay/billing"
	billingratio "github.com/songquanpeng/one-api/relay/billing/ratio"
	"github.com/songquanpeng/one-api/relay/meta"
	"github.com/songquanpeng/one-api/relay/model"
)

func RelayTextHelper(c *gin.Context) *model.ErrorWithStatusCode {
	ctx := c.Request.Context()
	meta := meta.GetByContext(c)
	// get & validate textRequest
	textRequest, err := getAndValidateTextRequest(c, meta.Mode)
	if err != nil {
		logger.Errorf(ctx, "getAndValidateTextRequest failed: %s", err.Error())
		return openai.ErrorWrapper(err, "invalid_text_request", http.StatusBadRequest)
	}
	meta.IsStream = textRequest.Stream

	// map model name
	meta.OriginModelName = textRequest.Model
	textRequest.Model, _ = getMappedModelName(textRequest.Model, meta.ModelMapping)
	meta.ActualModelName = textRequest.Model
	// set system prompt if not empty
	systemPromptReset := setSystemPrompt(ctx, textRequest, meta.ForcedSystemPrompt)
	// get model ratio & group ratio
	modelRatio := billingratio.GetModelRatio(textRequest.Model, meta.ChannelType)
	groupRatio := billingratio.GetGroupRatio(meta.Group)
	ratio := modelRatio * groupRatio
	// pre-consume quota
	promptTokens := getPromptTokens(textRequest, meta.Mode)
	meta.PromptTokens = promptTokens
	preConsumedQuota, bizErr := preConsumeQuota(ctx, textRequest, promptTokens, ratio, meta)
	if bizErr != nil {
		logger.Warnf(ctx, "preConsumeQuota failed: %+v", *bizErr)
		return bizErr
	}

	adaptor := relay.GetAdaptor(meta.APIType)
	if adaptor == nil {
		return openai.ErrorWrapper(fmt.Errorf("invalid api type: %d", meta.APIType), "invalid_api_type", http.StatusBadRequest)
	}
	adaptor.Init(meta)

	// get request body
	requestBody, err := getRequestBody(c, meta, textRequest, adaptor)
	if err != nil {
		return openai.ErrorWrapper(err, "convert_request_failed", http.StatusInternalServerError)
	}

	// do request
	resp, err := adaptor.DoRequest(c, meta, requestBody)
	if err != nil {
		logger.Errorf(ctx, "DoRequest failed: %s", err.Error())
		return openai.ErrorWrapper(err, "do_request_failed", http.StatusInternalServerError)
	}
	if isErrorHappened(meta, resp) {
		billing.ReturnPreConsumedQuota(ctx, preConsumedQuota, meta.TokenId)
		return RelayErrorHandler(resp)
	}

	// do response
	usage, respErr := adaptor.DoResponse(c, resp, meta)
	if respErr != nil {
		logger.Errorf(ctx, "respErr is not nil: %+v", respErr)
		billing.ReturnPreConsumedQuota(ctx, preConsumedQuota, meta.TokenId)
		return respErr
	}
	// If the channel has a UsageAPIURL configured and the adaptor captured a
	// generation ID, fetch accurate multimodal token counts now (works for any
	// channel type that sets meta.GenerationId during DoResponse).
	if meta.Config.UsageAPIURL != "" && meta.GenerationId != "" {
		if fetchedUsage, fetchErr := openai.FetchUsageFromMetadataURL(meta.Config.UsageAPIURL, meta.GenerationId); fetchErr == nil {
			usage = fetchedUsage
		} else {
			logger.Warnf(ctx, "failed to fetch usage from metadata URL: %s", fetchErr.Error())
		}
	}
	// post-consume quota
	go postConsumeQuota(ctx, usage, meta, textRequest, ratio, preConsumedQuota, modelRatio, groupRatio, systemPromptReset)
	return nil
}

func getRequestBody(c *gin.Context, meta *meta.Meta, textRequest *model.GeneralOpenAIRequest, adaptor adaptor.Adaptor) (io.Reader, error) {
	// Always convert request through adaptor before sending upstream.
	// This avoids passthrough edge cases where strict-compat sanitization may be bypassed.
	// The slight performance cost is acceptable compared to malformed request risk.

	// get request body
	var requestBody io.Reader
	convertedRequest, err := adaptor.ConvertRequest(c, meta.Mode, textRequest)
	if err != nil {
		logger.Debugf(c.Request.Context(), "converted request failed: %s\n", err.Error())
		return nil, err
	}
	jsonData, err := json.Marshal(convertedRequest)
	if err != nil {
		logger.Debugf(c.Request.Context(), "converted request json_marshal_failed: %s\n", err.Error())
		return nil, err
	}
	jsonData = sanitizeStrictToolChoiceJSON(jsonData)
	logger.Debugf(c.Request.Context(), "converted request: \n%s", string(jsonData))
	requestBody = bytes.NewBuffer(jsonData)
	return requestBody, nil
}

// Final outbound guard for strict OpenAI-compatible gateways:
// - remove any disable_parallel_tool_use key at any depth
// - coerce object/array tool_choice to OpenAI-compatible string form
func sanitizeStrictToolChoiceJSON(data []byte) []byte {
	var payload any
	if err := json.Unmarshal(data, &payload); err != nil {
		return data
	}

	root, ok := payload.(map[string]any)
	if !ok {
		return data
	}

	cleaned := stripDisableParallelToolUseAny(root)
	root, ok = cleaned.(map[string]any)
	if !ok {
		return data
	}

	hasTools := false
	if tools, ok := root["tools"].([]any); ok && len(tools) > 0 {
		hasTools = true
	}
	if functions, ok := root["functions"]; ok && functions != nil {
		hasTools = true
	}

	if tc, ok := root["tool_choice"]; ok {
		switch tc.(type) {
		case map[string]any, []any:
			if hasTools {
				root["tool_choice"] = "auto"
			} else {
				delete(root, "tool_choice")
			}
		}
	}

	out, err := json.Marshal(root)
	if err != nil {
		return data
	}
	return out
}

func stripDisableParallelToolUseAny(v any) any {
	switch val := v.(type) {
	case map[string]any:
		delete(val, "disable_parallel_tool_use")
		for k, child := range val {
			val[k] = stripDisableParallelToolUseAny(child)
		}
		return val
	case []any:
		for i, child := range val {
			val[i] = stripDisableParallelToolUseAny(child)
		}
		return val
	default:
		return v
	}
}