| 123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777 |
- package service
- import (
- "bytes"
- "context"
- "encoding/json"
- "fmt"
- "io"
- "net/http"
- "net/url"
- "strings"
- "github.com/Wei-Shaw/sub2api/internal/pkg/apicompat"
- "github.com/Wei-Shaw/sub2api/internal/pkg/logger"
- "github.com/gin-gonic/gin"
- "github.com/tidwall/gjson"
- "github.com/tiktoken-go/tokenizer"
- "go.uber.org/zap"
- )
- const (
- openAIResponsesInputItemTokenOverhead = 3
- openAIResponsesContentPartOverhead = 1
- openAIInputTokensFallbackMinimum = 1
- )
- type openAIInputTokensCountRequest struct {
- Model string `json:"model"`
- Instructions string `json:"instructions,omitempty"`
- Input json.RawMessage `json:"input,omitempty"`
- Tools []apicompat.ResponsesTool `json:"tools,omitempty"`
- ToolChoice json.RawMessage `json:"tool_choice,omitempty"`
- }
- type openAIInputTokensCountPrepared struct {
- Request openAIInputTokensCountRequest
- OriginalModel string
- NormalizedModel string
- BillingModel string
- UpstreamModel string
- }
- // ForwardResponsesInputTokens handles the native OpenAI
- // POST /v1/responses/input_tokens shape. Custom OpenAI-compatible relays often
- // implement /responses but not this preflight endpoint, so those accounts use
- // the local estimator instead of receiving a request that is known to fail.
- func (s *OpenAIGatewayService) ForwardResponsesInputTokens(
- ctx context.Context,
- c *gin.Context,
- account *Account,
- body []byte,
- ) error {
- if account == nil {
- writeOpenAIResponsesInputTokensError(c, http.StatusServiceUnavailable, "api_error", "No available OpenAI accounts")
- return fmt.Errorf("responses input_tokens: missing account")
- }
- prepared, err := prepareNativeOpenAIInputTokensCountRequest(body, account)
- if err != nil {
- writeOpenAIResponsesInputTokensError(c, http.StatusBadRequest, "invalid_request_error", "Failed to parse request body")
- return err
- }
- if shouldEstimateOpenAIInputTokensLocally(account) {
- writeOpenAIResponsesInputTokensFallback(c, account, prepared, 0, "custom_relay")
- return nil
- }
- token, _, err := s.GetAccessToken(ctx, account)
- if err != nil {
- writeOpenAIResponsesInputTokensError(c, http.StatusBadGateway, "upstream_error", "Failed to get access token")
- return fmt.Errorf("responses input_tokens: get access token: %w", err)
- }
- upstreamBody := ReplaceModelInBody(body, prepared.UpstreamModel)
- upstreamReq, err := s.buildInputTokensUpstreamRequest(ctx, c, account, upstreamBody, token)
- if err != nil {
- writeOpenAIResponsesInputTokensError(c, http.StatusInternalServerError, "api_error", "Failed to build request")
- return fmt.Errorf("responses input_tokens: build upstream request: %w", err)
- }
- proxyURL := ""
- if account.ProxyID != nil && account.Proxy != nil {
- proxyURL = account.Proxy.URL()
- }
- resp, err := s.doOpenAIUpstream(upstreamReq, proxyURL, account)
- if err != nil {
- safeErr := sanitizeUpstreamErrorMessage(err.Error())
- setOpsUpstreamError(c, 0, safeErr, "")
- writeOpenAIResponsesInputTokensError(c, http.StatusBadGateway, "upstream_error", "Upstream request failed")
- return fmt.Errorf("responses input_tokens: upstream request failed: %s", safeErr)
- }
- defer func() { _ = resp.Body.Close() }()
- respBody, err := io.ReadAll(resp.Body)
- if err != nil {
- writeOpenAIResponsesInputTokensError(c, http.StatusBadGateway, "upstream_error", "Failed to read response")
- return fmt.Errorf("responses input_tokens: read upstream response: %w", err)
- }
- if resp.StatusCode >= 400 {
- if isOpenAIResponsesInputTokensUnsupported(account, resp.StatusCode, respBody) {
- writeOpenAIResponsesInputTokensFallback(c, account, prepared, resp.StatusCode, "upstream_unsupported")
- return nil
- }
- if s.rateLimitService != nil {
- s.rateLimitService.HandleUpstreamError(ctx, account, resp.StatusCode, resp.Header, respBody)
- }
- upstreamMsg := sanitizeUpstreamErrorMessage(strings.TrimSpace(extractUpstreamErrorMessage(respBody)))
- setOpsUpstreamError(c, resp.StatusCode, upstreamMsg, "")
- writeOpenAIResponsesInputTokensError(c, resp.StatusCode, "upstream_error", "Upstream request failed")
- if upstreamMsg == "" {
- return fmt.Errorf("responses input_tokens: upstream error: %d", resp.StatusCode)
- }
- return fmt.Errorf("responses input_tokens: upstream error: %d message=%s", resp.StatusCode, upstreamMsg)
- }
- inputTokens := gjson.GetBytes(respBody, "input_tokens")
- if !inputTokens.Exists() {
- writeOpenAIResponsesInputTokensError(c, http.StatusBadGateway, "upstream_error", "Upstream response missing input_tokens")
- return fmt.Errorf("responses input_tokens: upstream response missing input_tokens")
- }
- contentType := strings.TrimSpace(resp.Header.Get("Content-Type"))
- if contentType == "" {
- contentType = "application/json"
- }
- c.Data(http.StatusOK, contentType, respBody)
- return nil
- }
- func prepareNativeOpenAIInputTokensCountRequest(body []byte, account *Account) (*openAIInputTokensCountPrepared, error) {
- var req openAIInputTokensCountRequest
- if err := json.Unmarshal(body, &req); err != nil {
- return nil, fmt.Errorf("parse responses input_tokens request: %w", err)
- }
- originalModel := strings.TrimSpace(req.Model)
- if originalModel == "" {
- return nil, fmt.Errorf("parse responses input_tokens request: model is required")
- }
- billingModel := resolveOpenAIForwardModel(account, originalModel, "")
- upstreamModel := normalizeOpenAIModelForUpstream(account, billingModel)
- req.Model = upstreamModel
- return &openAIInputTokensCountPrepared{
- Request: req,
- OriginalModel: originalModel,
- NormalizedModel: originalModel,
- BillingModel: billingModel,
- UpstreamModel: upstreamModel,
- }, nil
- }
- func shouldEstimateOpenAIInputTokensLocally(account *Account) bool {
- if account == nil || account.IsGrok() || account.IsCNProvider() || account.Type == AccountTypeUpstream {
- return true
- }
- if account.Type != AccountTypeAPIKey {
- return false
- }
- rawBaseURL := strings.TrimSpace(account.GetCredential("base_url"))
- if rawBaseURL == "" {
- return false
- }
- parsed, err := url.Parse(rawBaseURL)
- if err != nil {
- return true
- }
- return !strings.EqualFold(parsed.Hostname(), "api.openai.com")
- }
- func isOpenAIResponsesInputTokensUnsupported(account *Account, statusCode int, body []byte) bool {
- if statusCode == http.StatusNotFound {
- return true
- }
- return account != nil && account.Type == AccountTypeOAuth && isOpenAIOAuthInputTokensUnsupported(statusCode, body)
- }
- func writeOpenAIResponsesInputTokensFallback(c *gin.Context, account *Account, prepared *openAIInputTokensCountPrepared, statusCode int, reason string) {
- estimated := openAIInputTokensFallbackMinimum
- if prepared != nil {
- if got, err := estimateOpenAIInputTokens(prepared.Request); err == nil && got > 0 {
- estimated = got
- }
- }
- accountID := int64(0)
- upstreamModel := ""
- if account != nil {
- accountID = account.ID
- }
- if prepared != nil {
- upstreamModel = prepared.UpstreamModel
- }
- logger.L().Info("openai responses input_tokens: local estimate fallback",
- zap.Int64("account_id", accountID),
- zap.Int("upstream_status", statusCode),
- zap.Int("estimated_input_tokens", estimated),
- zap.String("upstream_model", upstreamModel),
- zap.String("reason", reason),
- )
- c.JSON(http.StatusOK, gin.H{
- "object": "response.input_tokens",
- "input_tokens": estimated,
- })
- }
- func writeOpenAIResponsesInputTokensError(c *gin.Context, status int, errType, message string) {
- c.JSON(status, gin.H{
- "error": gin.H{
- "type": errType,
- "message": message,
- },
- })
- }
- // EstimateGrokCountTokens estimates an Anthropic-compatible count_tokens request
- // locally. Grok does not expose a compatible token-counting endpoint, so this
- // path deliberately avoids account selection, credentials, and upstream calls.
- func EstimateGrokCountTokens(body []byte) (int, error) {
- return estimateAnthropicCountTokensLocally(body)
- }
- // estimateAnthropicCountTokensLocally 走 Anthropic→Responses→tiktoken 链本地估算
- // count_tokens,不发任何上游请求(上游无兼容端点的平台使用)。
- func estimateAnthropicCountTokensLocally(body []byte) (int, error) {
- var anthropicReq apicompat.AnthropicRequest
- if err := json.Unmarshal(body, &anthropicReq); err != nil {
- return 0, fmt.Errorf("parse anthropic count_tokens request: %w", err)
- }
- if strings.TrimSpace(anthropicReq.Model) == "" {
- return 0, fmt.Errorf("parse anthropic count_tokens request: model is required")
- }
- responsesReq, err := apicompat.AnthropicToResponses(&anthropicReq)
- if err != nil {
- return 0, fmt.Errorf("convert anthropic request to responses: %w", err)
- }
- estimated, err := estimateOpenAIInputTokens(openAIInputTokensCountRequest{
- Model: anthropicReq.Model,
- Instructions: responsesReq.Instructions,
- Input: responsesReq.Input,
- Tools: responsesReq.Tools,
- ToolChoice: responsesReq.ToolChoice,
- })
- if err != nil {
- return 0, fmt.Errorf("estimate input tokens: %w", err)
- }
- if estimated < openAIInputTokensFallbackMinimum {
- estimated = openAIInputTokensFallbackMinimum
- }
- return estimated, nil
- }
- // ForwardCountTokensAsAnthropic bridges Anthropic /v1/messages/count_tokens to
- // OpenAI POST /v1/responses/input_tokens and returns Anthropic-compatible output.
- func (s *OpenAIGatewayService) ForwardCountTokensAsAnthropic(
- ctx context.Context,
- c *gin.Context,
- account *Account,
- body []byte,
- defaultMappedModel string,
- ) error {
- if account == nil {
- writeAnthropicCountTokensError(c, http.StatusServiceUnavailable, "api_error", "No available OpenAI accounts")
- return fmt.Errorf("count_tokens: missing account")
- }
- // 国产供应商(全部协议,含 anthropic):一律本地估算,不发上游请求。
- // 依据(2026-08 核实):三家的 Anthropic 兼容层均未提供
- // /v1/messages/count_tokens——DeepSeek 官方 anthropic_api 文档无此端点
- // (且注明 anthropic-version 头被忽略),聚合网关 OpenModel 明确标注
- // count_tokens 为 "Anthropic only",Kimi/智谱亦无任何文档承诺。转发上游
- // 只会常态 404,且错误还会流入账号处置逻辑误伤整账号调度;Claude Code
- // 高频调用此端点,本地 tiktoken 估算是与 Grok 一致的既有方案。
- if account.IsCNProvider() {
- estimated, err := estimateAnthropicCountTokensLocally(body)
- if err != nil {
- writeAnthropicCountTokensError(c, http.StatusBadRequest, "invalid_request_error", "Failed to parse request body")
- return fmt.Errorf("count_tokens: estimate cn provider input tokens: %w", err)
- }
- logger.L().Debug("openai count_tokens: cn provider local estimate",
- zap.Int64("account_id", account.ID),
- zap.Int("estimated_input_tokens", estimated),
- )
- c.JSON(http.StatusOK, gin.H{
- "input_tokens": estimated,
- })
- return nil
- }
- prepared, err := prepareOpenAIInputTokensCountRequest(body, account, defaultMappedModel)
- if err != nil {
- writeAnthropicCountTokensError(c, http.StatusBadRequest, "invalid_request_error", "Failed to parse request body")
- return err
- }
- upstreamBody, err := marshalOpenAIUpstreamJSON(prepared.Request)
- if err != nil {
- writeAnthropicCountTokensError(c, http.StatusInternalServerError, "api_error", "Failed to build request")
- return fmt.Errorf("marshal openai input_tokens body: %w", err)
- }
- logger.L().Debug("openai count_tokens: model mapping applied",
- zap.Int64("account_id", account.ID),
- zap.String("original_model", prepared.OriginalModel),
- zap.String("normalized_model", prepared.NormalizedModel),
- zap.String("billing_model", prepared.BillingModel),
- zap.String("upstream_model", prepared.UpstreamModel),
- )
- token, _, err := s.GetAccessToken(ctx, account)
- if err != nil {
- writeAnthropicCountTokensError(c, http.StatusBadGateway, "upstream_error", "Failed to get access token")
- return fmt.Errorf("get access token: %w", err)
- }
- upstreamReq, err := s.buildInputTokensUpstreamRequest(ctx, c, account, upstreamBody, token)
- if err != nil {
- writeAnthropicCountTokensError(c, http.StatusInternalServerError, "api_error", "Failed to build request")
- return fmt.Errorf("build input_tokens request: %w", err)
- }
- proxyURL := ""
- if account.ProxyID != nil && account.Proxy != nil {
- proxyURL = account.Proxy.URL()
- }
- resp, err := s.doOpenAIUpstream(upstreamReq, proxyURL, account)
- if err != nil {
- safeErr := sanitizeUpstreamErrorMessage(err.Error())
- setOpsUpstreamError(c, 0, safeErr, "")
- writeAnthropicCountTokensError(c, http.StatusBadGateway, "upstream_error", "Upstream request failed")
- return fmt.Errorf("openai input_tokens upstream request failed: %s", safeErr)
- }
- defer func() { _ = resp.Body.Close() }()
- respBody, err := io.ReadAll(resp.Body)
- if err != nil {
- writeAnthropicCountTokensError(c, http.StatusBadGateway, "upstream_error", "Failed to read response")
- return fmt.Errorf("read input_tokens response: %w", err)
- }
- if resp.StatusCode >= 400 {
- upstreamMsg := sanitizeUpstreamErrorMessage(strings.TrimSpace(extractUpstreamErrorMessage(respBody)))
- if account.Type == AccountTypeOAuth && isOpenAIOAuthInputTokensUnsupported(resp.StatusCode, respBody) {
- writeOpenAIOAuthInputTokensFallback(c, account, prepared, resp.StatusCode)
- return nil
- }
- if s.rateLimitService != nil {
- s.rateLimitService.HandleUpstreamError(ctx, account, resp.StatusCode, resp.Header, respBody)
- }
- if isOpenAIInputTokensUnsupported(resp.StatusCode, respBody) {
- writeAnthropicCountTokensError(c, http.StatusNotFound, "not_found_error", "Token counting is not supported by upstream")
- return nil
- }
- upstreamDetail := ""
- if s.cfg != nil && s.cfg.Gateway.LogUpstreamErrorBody {
- maxBytes := s.cfg.Gateway.LogUpstreamErrorBodyMaxBytes
- if maxBytes <= 0 {
- maxBytes = 2048
- }
- upstreamDetail = truncateString(string(respBody), maxBytes)
- }
- setOpsUpstreamError(c, resp.StatusCode, upstreamMsg, upstreamDetail)
- errMsg := "Upstream request failed"
- switch resp.StatusCode {
- case 429:
- errMsg = "Rate limit exceeded"
- case 500, 502, 503, 504, 529:
- errMsg = "Upstream service temporarily unavailable"
- }
- writeAnthropicCountTokensError(c, resp.StatusCode, "upstream_error", errMsg)
- if upstreamMsg == "" {
- return fmt.Errorf("input_tokens upstream error: %d", resp.StatusCode)
- }
- return fmt.Errorf("input_tokens upstream error: %d message=%s", resp.StatusCode, upstreamMsg)
- }
- inputTokens := gjson.GetBytes(respBody, "input_tokens")
- if !inputTokens.Exists() {
- writeAnthropicCountTokensError(c, http.StatusBadGateway, "upstream_error", "Upstream response missing input_tokens")
- return fmt.Errorf("input_tokens response missing input_tokens field")
- }
- c.JSON(http.StatusOK, gin.H{
- "input_tokens": int(inputTokens.Int()),
- })
- return nil
- }
- func prepareOpenAIInputTokensCountRequest(
- body []byte,
- account *Account,
- defaultMappedModel string,
- ) (*openAIInputTokensCountPrepared, error) {
- var anthropicReq apicompat.AnthropicRequest
- if err := json.Unmarshal(body, &anthropicReq); err != nil {
- return nil, fmt.Errorf("parse anthropic count_tokens request: %w", err)
- }
- originalModel := anthropicReq.Model
- applyOpenAICompatModelNormalization(&anthropicReq)
- normalizedModel := anthropicReq.Model
- billingModel := resolveOpenAIForwardModel(account, normalizedModel, strings.TrimSpace(defaultMappedModel))
- upstreamModel := normalizeOpenAIModelForUpstream(account, billingModel)
- responsesReq, err := apicompat.AnthropicToResponses(&anthropicReq)
- if err != nil {
- return nil, fmt.Errorf("convert anthropic request to responses: %w", err)
- }
- return &openAIInputTokensCountPrepared{
- Request: openAIInputTokensCountRequest{
- Model: upstreamModel,
- Instructions: responsesReq.Instructions,
- Input: responsesReq.Input,
- Tools: responsesReq.Tools,
- ToolChoice: responsesReq.ToolChoice,
- },
- OriginalModel: originalModel,
- NormalizedModel: normalizedModel,
- BillingModel: billingModel,
- UpstreamModel: upstreamModel,
- }, nil
- }
- func (s *OpenAIGatewayService) buildInputTokensUpstreamRequest(
- ctx context.Context,
- c *gin.Context,
- account *Account,
- body []byte,
- token string,
- ) (*http.Request, error) {
- targetURL := openaiPlatformAPIInputTokensURL
- if account.Type == AccountTypeAPIKey {
- if baseURL := account.GetOpenAIBaseURL(); strings.TrimSpace(baseURL) != "" {
- validatedURL, err := s.validateUpstreamBaseURL(baseURL)
- if err != nil {
- return nil, err
- }
- targetURL = buildOpenAIResponsesInputTokensURL(validatedURL)
- }
- }
- req, err := http.NewRequestWithContext(ctx, http.MethodPost, targetURL, bytes.NewReader(body))
- if err != nil {
- return nil, err
- }
- req = req.WithContext(WithHTTPUpstreamProfile(req.Context(), HTTPUpstreamProfileOpenAI))
- authHeaders, err := s.buildOpenAIAuthenticationHeaders(ctx, account, token)
- if err != nil {
- return nil, err
- }
- for key, values := range authHeaders {
- for _, value := range values {
- req.Header.Add(key, value)
- }
- }
- req.Header.Set("content-type", "application/json")
- req.Header.Set("accept", "application/json")
- if c != nil && c.Request != nil {
- for key, values := range c.Request.Header {
- lower := strings.ToLower(strings.TrimSpace(key))
- if lower != "user-agent" && lower != "accept-language" {
- continue
- }
- for _, v := range values {
- req.Header.Add(key, v)
- }
- }
- }
- // 账号级请求头覆写(仅 openai api_key 账号启用时生效;OAuth 路径 no-op)
- account.ApplyHeaderOverrides(req.Header)
- return req, nil
- }
- func writeAnthropicCountTokensError(c *gin.Context, status int, errType, message string) {
- c.JSON(status, gin.H{
- "type": "error",
- "error": gin.H{
- "type": errType,
- "message": message,
- },
- })
- }
- func isOpenAIInputTokensUnsupported(statusCode int, body []byte) bool {
- if statusCode != http.StatusNotFound {
- return false
- }
- msg := strings.ToLower(strings.TrimSpace(extractUpstreamErrorMessage(body)))
- return strings.Contains(msg, "input_tokens") && strings.Contains(msg, "not found")
- }
- func writeOpenAIOAuthInputTokensFallback(c *gin.Context, account *Account, prepared *openAIInputTokensCountPrepared, statusCode int) {
- estimated := openAIInputTokensFallbackMinimum
- if got, err := estimateOpenAIInputTokens(prepared.Request); err == nil {
- if got > 0 {
- estimated = got
- }
- logger.L().Info("openai count_tokens: oauth fallback to local tiktoken estimate",
- zap.Int64("account_id", account.ID),
- zap.Int("upstream_status", statusCode),
- zap.Int("estimated_input_tokens", estimated),
- zap.String("upstream_model", prepared.UpstreamModel),
- )
- } else {
- logger.L().Warn("openai count_tokens: oauth local tiktoken fallback failed, using minimum estimate",
- zap.Int64("account_id", account.ID),
- zap.Int("upstream_status", statusCode),
- zap.Int("estimated_input_tokens", estimated),
- zap.String("upstream_model", prepared.UpstreamModel),
- zap.Error(err),
- )
- }
- c.JSON(http.StatusOK, gin.H{
- "input_tokens": estimated,
- })
- }
- func isOpenAIOAuthInputTokensUnsupported(statusCode int, body []byte) bool {
- switch statusCode {
- case http.StatusUnauthorized, http.StatusForbidden, http.StatusNotFound:
- default:
- return false
- }
- bodyLower := strings.ToLower(string(body))
- msg := strings.ToLower(strings.TrimSpace(extractUpstreamErrorMessage(body)))
- code := strings.ToLower(strings.TrimSpace(extractUpstreamErrorCode(body)))
- if code == "missing_scope" ||
- strings.Contains(bodyLower, "api.responses.write") ||
- strings.Contains(bodyLower, "missing scopes") ||
- strings.Contains(bodyLower, "insufficient_scope") {
- return true
- }
- if statusCode == http.StatusNotFound && isOpenAIInputTokensUnsupported(statusCode, body) {
- return true
- }
- // OAuth's platform endpoint can be blocked by an upstream proxy before it
- // reaches the API and return an HTML 403 page without a structured error.
- // Treat that endpoint-level response like the other unsupported cases so
- // count_tokens remains a local, non-health-affecting convenience request.
- if statusCode == http.StatusForbidden && isHTMLResponse(body) {
- return true
- }
- return strings.Contains(msg, "input_tokens") &&
- (strings.Contains(msg, "not found") ||
- strings.Contains(msg, "not supported") ||
- strings.Contains(msg, "unsupported"))
- }
- func isHTMLResponse(body []byte) bool {
- trimmed := strings.TrimSpace(strings.ToLower(string(body)))
- return strings.HasPrefix(trimmed, "<!doctype html") ||
- strings.HasPrefix(trimmed, "<html")
- }
- func estimateOpenAIInputTokens(req openAIInputTokensCountRequest) (int, error) {
- codec, err := openAIInputTokensCodecForModel(req.Model)
- if err != nil {
- return 0, err
- }
- total := 0
- addCount := func(text string) error {
- text = strings.TrimSpace(text)
- if text == "" {
- return nil
- }
- n, err := codec.Count(text)
- if err != nil {
- return err
- }
- total += n
- return nil
- }
- if err := addCount(req.Instructions); err != nil {
- return 0, err
- }
- inputTokens, err := estimateOpenAIInputTokensForInput(codec, req.Input)
- if err != nil {
- return 0, err
- }
- total += inputTokens
- for _, tool := range req.Tools {
- raw, err := marshalOpenAIUpstreamJSON(tool)
- if err != nil {
- return 0, err
- }
- if err := addCount(string(raw)); err != nil {
- return 0, err
- }
- }
- if len(req.ToolChoice) > 0 {
- compacted, err := compactOpenAIInputTokensJSON(req.ToolChoice)
- if err != nil {
- return 0, err
- }
- if err := addCount(compacted); err != nil {
- return 0, err
- }
- }
- if total < 0 {
- return 0, nil
- }
- return total, nil
- }
- func estimateOpenAIInputTokensForInput(codec tokenizer.Codec, raw json.RawMessage) (int, error) {
- if len(bytes.TrimSpace(raw)) == 0 {
- return 0, nil
- }
- var plainText string
- if err := json.Unmarshal(raw, &plainText); err == nil {
- return codec.Count(plainText)
- }
- var items []apicompat.ResponsesInputItem
- if err := json.Unmarshal(raw, &items); err == nil {
- return estimateOpenAIInputTokensForInputItems(codec, items)
- }
- compacted, err := compactOpenAIInputTokensJSON(raw)
- if err != nil {
- return 0, err
- }
- return codec.Count(compacted)
- }
- func estimateOpenAIInputTokensForInputItems(codec tokenizer.Codec, items []apicompat.ResponsesInputItem) (int, error) {
- total := 0
- countText := func(text string) error {
- text = strings.TrimSpace(text)
- if text == "" {
- return nil
- }
- n, err := codec.Count(text)
- if err != nil {
- return err
- }
- total += n
- return nil
- }
- for _, item := range items {
- total += openAIResponsesInputItemTokenOverhead
- if err := countText(item.Role); err != nil {
- return 0, err
- }
- if item.Type != "" && item.Type != "message" {
- if err := countText(item.Type); err != nil {
- return 0, err
- }
- }
- if err := countText(item.Name); err != nil {
- return 0, err
- }
- if err := countText(item.Arguments); err != nil {
- return 0, err
- }
- if err := countText(item.Output); err != nil {
- return 0, err
- }
- if err := countText(item.CallID); err != nil {
- return 0, err
- }
- if err := countText(item.ID); err != nil {
- return 0, err
- }
- if len(bytes.TrimSpace(item.Content)) == 0 {
- continue
- }
- var contentText string
- if err := json.Unmarshal(item.Content, &contentText); err == nil {
- if err := countText(contentText); err != nil {
- return 0, err
- }
- continue
- }
- var parts []apicompat.ResponsesContentPart
- if err := json.Unmarshal(item.Content, &parts); err == nil {
- for _, part := range parts {
- total += openAIResponsesContentPartOverhead
- switch part.Type {
- case "input_text", "output_text", "text":
- if err := countText(part.Text); err != nil {
- return 0, err
- }
- case "input_image":
- if err := countText(estimateOpenAIInputImageText(part.ImageURL)); err != nil {
- return 0, err
- }
- default:
- if err := countText(part.Type); err != nil {
- return 0, err
- }
- }
- }
- continue
- }
- compacted, err := compactOpenAIInputTokensJSON(item.Content)
- if err != nil {
- return 0, err
- }
- if err := countText(compacted); err != nil {
- return 0, err
- }
- }
- return total, nil
- }
- func estimateOpenAIInputImageText(imageURL string) string {
- trimmed := strings.TrimSpace(imageURL)
- if trimmed == "" {
- return ""
- }
- if strings.HasPrefix(strings.ToLower(trimmed), "data:") {
- if comma := strings.Index(trimmed, ","); comma > 0 {
- return trimmed[:comma]
- }
- }
- return trimmed
- }
- func compactOpenAIInputTokensJSON(raw json.RawMessage) (string, error) {
- if len(bytes.TrimSpace(raw)) == 0 {
- return "", nil
- }
- var buf bytes.Buffer
- if err := json.Compact(&buf, raw); err != nil {
- return "", err
- }
- return buf.String(), nil
- }
- func openAIInputTokensCodecForModel(model string) (tokenizer.Codec, error) {
- switch openAIInputTokensEncodingForModel(model) {
- case tokenizer.Cl100kBase:
- return tokenizer.Get(tokenizer.Cl100kBase)
- default:
- return tokenizer.Get(tokenizer.O200kBase)
- }
- }
- func openAIInputTokensEncodingForModel(model string) tokenizer.Encoding {
- normalized := strings.ToLower(strings.TrimSpace(model))
- switch {
- case strings.HasPrefix(normalized, "gpt-3.5"),
- (strings.HasPrefix(normalized, "gpt-4") &&
- !strings.HasPrefix(normalized, "gpt-4o") &&
- !strings.HasPrefix(normalized, "gpt-4.1")),
- strings.HasPrefix(normalized, "text-embedding-"):
- return tokenizer.Cl100kBase
- default:
- return tokenizer.O200kBase
- }
- }
|