openai_gateway_count_tokens.go 25 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777
  1. package service
  2. import (
  3. "bytes"
  4. "context"
  5. "encoding/json"
  6. "fmt"
  7. "io"
  8. "net/http"
  9. "net/url"
  10. "strings"
  11. "github.com/Wei-Shaw/sub2api/internal/pkg/apicompat"
  12. "github.com/Wei-Shaw/sub2api/internal/pkg/logger"
  13. "github.com/gin-gonic/gin"
  14. "github.com/tidwall/gjson"
  15. "github.com/tiktoken-go/tokenizer"
  16. "go.uber.org/zap"
  17. )
  18. const (
  19. openAIResponsesInputItemTokenOverhead = 3
  20. openAIResponsesContentPartOverhead = 1
  21. openAIInputTokensFallbackMinimum = 1
  22. )
  23. type openAIInputTokensCountRequest struct {
  24. Model string `json:"model"`
  25. Instructions string `json:"instructions,omitempty"`
  26. Input json.RawMessage `json:"input,omitempty"`
  27. Tools []apicompat.ResponsesTool `json:"tools,omitempty"`
  28. ToolChoice json.RawMessage `json:"tool_choice,omitempty"`
  29. }
  30. type openAIInputTokensCountPrepared struct {
  31. Request openAIInputTokensCountRequest
  32. OriginalModel string
  33. NormalizedModel string
  34. BillingModel string
  35. UpstreamModel string
  36. }
  37. // ForwardResponsesInputTokens handles the native OpenAI
  38. // POST /v1/responses/input_tokens shape. Custom OpenAI-compatible relays often
  39. // implement /responses but not this preflight endpoint, so those accounts use
  40. // the local estimator instead of receiving a request that is known to fail.
  41. func (s *OpenAIGatewayService) ForwardResponsesInputTokens(
  42. ctx context.Context,
  43. c *gin.Context,
  44. account *Account,
  45. body []byte,
  46. ) error {
  47. if account == nil {
  48. writeOpenAIResponsesInputTokensError(c, http.StatusServiceUnavailable, "api_error", "No available OpenAI accounts")
  49. return fmt.Errorf("responses input_tokens: missing account")
  50. }
  51. prepared, err := prepareNativeOpenAIInputTokensCountRequest(body, account)
  52. if err != nil {
  53. writeOpenAIResponsesInputTokensError(c, http.StatusBadRequest, "invalid_request_error", "Failed to parse request body")
  54. return err
  55. }
  56. if shouldEstimateOpenAIInputTokensLocally(account) {
  57. writeOpenAIResponsesInputTokensFallback(c, account, prepared, 0, "custom_relay")
  58. return nil
  59. }
  60. token, _, err := s.GetAccessToken(ctx, account)
  61. if err != nil {
  62. writeOpenAIResponsesInputTokensError(c, http.StatusBadGateway, "upstream_error", "Failed to get access token")
  63. return fmt.Errorf("responses input_tokens: get access token: %w", err)
  64. }
  65. upstreamBody := ReplaceModelInBody(body, prepared.UpstreamModel)
  66. upstreamReq, err := s.buildInputTokensUpstreamRequest(ctx, c, account, upstreamBody, token)
  67. if err != nil {
  68. writeOpenAIResponsesInputTokensError(c, http.StatusInternalServerError, "api_error", "Failed to build request")
  69. return fmt.Errorf("responses input_tokens: build upstream request: %w", err)
  70. }
  71. proxyURL := ""
  72. if account.ProxyID != nil && account.Proxy != nil {
  73. proxyURL = account.Proxy.URL()
  74. }
  75. resp, err := s.doOpenAIUpstream(upstreamReq, proxyURL, account)
  76. if err != nil {
  77. safeErr := sanitizeUpstreamErrorMessage(err.Error())
  78. setOpsUpstreamError(c, 0, safeErr, "")
  79. writeOpenAIResponsesInputTokensError(c, http.StatusBadGateway, "upstream_error", "Upstream request failed")
  80. return fmt.Errorf("responses input_tokens: upstream request failed: %s", safeErr)
  81. }
  82. defer func() { _ = resp.Body.Close() }()
  83. respBody, err := io.ReadAll(resp.Body)
  84. if err != nil {
  85. writeOpenAIResponsesInputTokensError(c, http.StatusBadGateway, "upstream_error", "Failed to read response")
  86. return fmt.Errorf("responses input_tokens: read upstream response: %w", err)
  87. }
  88. if resp.StatusCode >= 400 {
  89. if isOpenAIResponsesInputTokensUnsupported(account, resp.StatusCode, respBody) {
  90. writeOpenAIResponsesInputTokensFallback(c, account, prepared, resp.StatusCode, "upstream_unsupported")
  91. return nil
  92. }
  93. if s.rateLimitService != nil {
  94. s.rateLimitService.HandleUpstreamError(ctx, account, resp.StatusCode, resp.Header, respBody)
  95. }
  96. upstreamMsg := sanitizeUpstreamErrorMessage(strings.TrimSpace(extractUpstreamErrorMessage(respBody)))
  97. setOpsUpstreamError(c, resp.StatusCode, upstreamMsg, "")
  98. writeOpenAIResponsesInputTokensError(c, resp.StatusCode, "upstream_error", "Upstream request failed")
  99. if upstreamMsg == "" {
  100. return fmt.Errorf("responses input_tokens: upstream error: %d", resp.StatusCode)
  101. }
  102. return fmt.Errorf("responses input_tokens: upstream error: %d message=%s", resp.StatusCode, upstreamMsg)
  103. }
  104. inputTokens := gjson.GetBytes(respBody, "input_tokens")
  105. if !inputTokens.Exists() {
  106. writeOpenAIResponsesInputTokensError(c, http.StatusBadGateway, "upstream_error", "Upstream response missing input_tokens")
  107. return fmt.Errorf("responses input_tokens: upstream response missing input_tokens")
  108. }
  109. contentType := strings.TrimSpace(resp.Header.Get("Content-Type"))
  110. if contentType == "" {
  111. contentType = "application/json"
  112. }
  113. c.Data(http.StatusOK, contentType, respBody)
  114. return nil
  115. }
  116. func prepareNativeOpenAIInputTokensCountRequest(body []byte, account *Account) (*openAIInputTokensCountPrepared, error) {
  117. var req openAIInputTokensCountRequest
  118. if err := json.Unmarshal(body, &req); err != nil {
  119. return nil, fmt.Errorf("parse responses input_tokens request: %w", err)
  120. }
  121. originalModel := strings.TrimSpace(req.Model)
  122. if originalModel == "" {
  123. return nil, fmt.Errorf("parse responses input_tokens request: model is required")
  124. }
  125. billingModel := resolveOpenAIForwardModel(account, originalModel, "")
  126. upstreamModel := normalizeOpenAIModelForUpstream(account, billingModel)
  127. req.Model = upstreamModel
  128. return &openAIInputTokensCountPrepared{
  129. Request: req,
  130. OriginalModel: originalModel,
  131. NormalizedModel: originalModel,
  132. BillingModel: billingModel,
  133. UpstreamModel: upstreamModel,
  134. }, nil
  135. }
  136. func shouldEstimateOpenAIInputTokensLocally(account *Account) bool {
  137. if account == nil || account.IsGrok() || account.IsCNProvider() || account.Type == AccountTypeUpstream {
  138. return true
  139. }
  140. if account.Type != AccountTypeAPIKey {
  141. return false
  142. }
  143. rawBaseURL := strings.TrimSpace(account.GetCredential("base_url"))
  144. if rawBaseURL == "" {
  145. return false
  146. }
  147. parsed, err := url.Parse(rawBaseURL)
  148. if err != nil {
  149. return true
  150. }
  151. return !strings.EqualFold(parsed.Hostname(), "api.openai.com")
  152. }
  153. func isOpenAIResponsesInputTokensUnsupported(account *Account, statusCode int, body []byte) bool {
  154. if statusCode == http.StatusNotFound {
  155. return true
  156. }
  157. return account != nil && account.Type == AccountTypeOAuth && isOpenAIOAuthInputTokensUnsupported(statusCode, body)
  158. }
  159. func writeOpenAIResponsesInputTokensFallback(c *gin.Context, account *Account, prepared *openAIInputTokensCountPrepared, statusCode int, reason string) {
  160. estimated := openAIInputTokensFallbackMinimum
  161. if prepared != nil {
  162. if got, err := estimateOpenAIInputTokens(prepared.Request); err == nil && got > 0 {
  163. estimated = got
  164. }
  165. }
  166. accountID := int64(0)
  167. upstreamModel := ""
  168. if account != nil {
  169. accountID = account.ID
  170. }
  171. if prepared != nil {
  172. upstreamModel = prepared.UpstreamModel
  173. }
  174. logger.L().Info("openai responses input_tokens: local estimate fallback",
  175. zap.Int64("account_id", accountID),
  176. zap.Int("upstream_status", statusCode),
  177. zap.Int("estimated_input_tokens", estimated),
  178. zap.String("upstream_model", upstreamModel),
  179. zap.String("reason", reason),
  180. )
  181. c.JSON(http.StatusOK, gin.H{
  182. "object": "response.input_tokens",
  183. "input_tokens": estimated,
  184. })
  185. }
  186. func writeOpenAIResponsesInputTokensError(c *gin.Context, status int, errType, message string) {
  187. c.JSON(status, gin.H{
  188. "error": gin.H{
  189. "type": errType,
  190. "message": message,
  191. },
  192. })
  193. }
  194. // EstimateGrokCountTokens estimates an Anthropic-compatible count_tokens request
  195. // locally. Grok does not expose a compatible token-counting endpoint, so this
  196. // path deliberately avoids account selection, credentials, and upstream calls.
  197. func EstimateGrokCountTokens(body []byte) (int, error) {
  198. return estimateAnthropicCountTokensLocally(body)
  199. }
  200. // estimateAnthropicCountTokensLocally 走 Anthropic→Responses→tiktoken 链本地估算
  201. // count_tokens,不发任何上游请求(上游无兼容端点的平台使用)。
  202. func estimateAnthropicCountTokensLocally(body []byte) (int, error) {
  203. var anthropicReq apicompat.AnthropicRequest
  204. if err := json.Unmarshal(body, &anthropicReq); err != nil {
  205. return 0, fmt.Errorf("parse anthropic count_tokens request: %w", err)
  206. }
  207. if strings.TrimSpace(anthropicReq.Model) == "" {
  208. return 0, fmt.Errorf("parse anthropic count_tokens request: model is required")
  209. }
  210. responsesReq, err := apicompat.AnthropicToResponses(&anthropicReq)
  211. if err != nil {
  212. return 0, fmt.Errorf("convert anthropic request to responses: %w", err)
  213. }
  214. estimated, err := estimateOpenAIInputTokens(openAIInputTokensCountRequest{
  215. Model: anthropicReq.Model,
  216. Instructions: responsesReq.Instructions,
  217. Input: responsesReq.Input,
  218. Tools: responsesReq.Tools,
  219. ToolChoice: responsesReq.ToolChoice,
  220. })
  221. if err != nil {
  222. return 0, fmt.Errorf("estimate input tokens: %w", err)
  223. }
  224. if estimated < openAIInputTokensFallbackMinimum {
  225. estimated = openAIInputTokensFallbackMinimum
  226. }
  227. return estimated, nil
  228. }
  229. // ForwardCountTokensAsAnthropic bridges Anthropic /v1/messages/count_tokens to
  230. // OpenAI POST /v1/responses/input_tokens and returns Anthropic-compatible output.
  231. func (s *OpenAIGatewayService) ForwardCountTokensAsAnthropic(
  232. ctx context.Context,
  233. c *gin.Context,
  234. account *Account,
  235. body []byte,
  236. defaultMappedModel string,
  237. ) error {
  238. if account == nil {
  239. writeAnthropicCountTokensError(c, http.StatusServiceUnavailable, "api_error", "No available OpenAI accounts")
  240. return fmt.Errorf("count_tokens: missing account")
  241. }
  242. // 国产供应商(全部协议,含 anthropic):一律本地估算,不发上游请求。
  243. // 依据(2026-08 核实):三家的 Anthropic 兼容层均未提供
  244. // /v1/messages/count_tokens——DeepSeek 官方 anthropic_api 文档无此端点
  245. // (且注明 anthropic-version 头被忽略),聚合网关 OpenModel 明确标注
  246. // count_tokens 为 "Anthropic only",Kimi/智谱亦无任何文档承诺。转发上游
  247. // 只会常态 404,且错误还会流入账号处置逻辑误伤整账号调度;Claude Code
  248. // 高频调用此端点,本地 tiktoken 估算是与 Grok 一致的既有方案。
  249. if account.IsCNProvider() {
  250. estimated, err := estimateAnthropicCountTokensLocally(body)
  251. if err != nil {
  252. writeAnthropicCountTokensError(c, http.StatusBadRequest, "invalid_request_error", "Failed to parse request body")
  253. return fmt.Errorf("count_tokens: estimate cn provider input tokens: %w", err)
  254. }
  255. logger.L().Debug("openai count_tokens: cn provider local estimate",
  256. zap.Int64("account_id", account.ID),
  257. zap.Int("estimated_input_tokens", estimated),
  258. )
  259. c.JSON(http.StatusOK, gin.H{
  260. "input_tokens": estimated,
  261. })
  262. return nil
  263. }
  264. prepared, err := prepareOpenAIInputTokensCountRequest(body, account, defaultMappedModel)
  265. if err != nil {
  266. writeAnthropicCountTokensError(c, http.StatusBadRequest, "invalid_request_error", "Failed to parse request body")
  267. return err
  268. }
  269. upstreamBody, err := marshalOpenAIUpstreamJSON(prepared.Request)
  270. if err != nil {
  271. writeAnthropicCountTokensError(c, http.StatusInternalServerError, "api_error", "Failed to build request")
  272. return fmt.Errorf("marshal openai input_tokens body: %w", err)
  273. }
  274. logger.L().Debug("openai count_tokens: model mapping applied",
  275. zap.Int64("account_id", account.ID),
  276. zap.String("original_model", prepared.OriginalModel),
  277. zap.String("normalized_model", prepared.NormalizedModel),
  278. zap.String("billing_model", prepared.BillingModel),
  279. zap.String("upstream_model", prepared.UpstreamModel),
  280. )
  281. token, _, err := s.GetAccessToken(ctx, account)
  282. if err != nil {
  283. writeAnthropicCountTokensError(c, http.StatusBadGateway, "upstream_error", "Failed to get access token")
  284. return fmt.Errorf("get access token: %w", err)
  285. }
  286. upstreamReq, err := s.buildInputTokensUpstreamRequest(ctx, c, account, upstreamBody, token)
  287. if err != nil {
  288. writeAnthropicCountTokensError(c, http.StatusInternalServerError, "api_error", "Failed to build request")
  289. return fmt.Errorf("build input_tokens request: %w", err)
  290. }
  291. proxyURL := ""
  292. if account.ProxyID != nil && account.Proxy != nil {
  293. proxyURL = account.Proxy.URL()
  294. }
  295. resp, err := s.doOpenAIUpstream(upstreamReq, proxyURL, account)
  296. if err != nil {
  297. safeErr := sanitizeUpstreamErrorMessage(err.Error())
  298. setOpsUpstreamError(c, 0, safeErr, "")
  299. writeAnthropicCountTokensError(c, http.StatusBadGateway, "upstream_error", "Upstream request failed")
  300. return fmt.Errorf("openai input_tokens upstream request failed: %s", safeErr)
  301. }
  302. defer func() { _ = resp.Body.Close() }()
  303. respBody, err := io.ReadAll(resp.Body)
  304. if err != nil {
  305. writeAnthropicCountTokensError(c, http.StatusBadGateway, "upstream_error", "Failed to read response")
  306. return fmt.Errorf("read input_tokens response: %w", err)
  307. }
  308. if resp.StatusCode >= 400 {
  309. upstreamMsg := sanitizeUpstreamErrorMessage(strings.TrimSpace(extractUpstreamErrorMessage(respBody)))
  310. if account.Type == AccountTypeOAuth && isOpenAIOAuthInputTokensUnsupported(resp.StatusCode, respBody) {
  311. writeOpenAIOAuthInputTokensFallback(c, account, prepared, resp.StatusCode)
  312. return nil
  313. }
  314. if s.rateLimitService != nil {
  315. s.rateLimitService.HandleUpstreamError(ctx, account, resp.StatusCode, resp.Header, respBody)
  316. }
  317. if isOpenAIInputTokensUnsupported(resp.StatusCode, respBody) {
  318. writeAnthropicCountTokensError(c, http.StatusNotFound, "not_found_error", "Token counting is not supported by upstream")
  319. return nil
  320. }
  321. upstreamDetail := ""
  322. if s.cfg != nil && s.cfg.Gateway.LogUpstreamErrorBody {
  323. maxBytes := s.cfg.Gateway.LogUpstreamErrorBodyMaxBytes
  324. if maxBytes <= 0 {
  325. maxBytes = 2048
  326. }
  327. upstreamDetail = truncateString(string(respBody), maxBytes)
  328. }
  329. setOpsUpstreamError(c, resp.StatusCode, upstreamMsg, upstreamDetail)
  330. errMsg := "Upstream request failed"
  331. switch resp.StatusCode {
  332. case 429:
  333. errMsg = "Rate limit exceeded"
  334. case 500, 502, 503, 504, 529:
  335. errMsg = "Upstream service temporarily unavailable"
  336. }
  337. writeAnthropicCountTokensError(c, resp.StatusCode, "upstream_error", errMsg)
  338. if upstreamMsg == "" {
  339. return fmt.Errorf("input_tokens upstream error: %d", resp.StatusCode)
  340. }
  341. return fmt.Errorf("input_tokens upstream error: %d message=%s", resp.StatusCode, upstreamMsg)
  342. }
  343. inputTokens := gjson.GetBytes(respBody, "input_tokens")
  344. if !inputTokens.Exists() {
  345. writeAnthropicCountTokensError(c, http.StatusBadGateway, "upstream_error", "Upstream response missing input_tokens")
  346. return fmt.Errorf("input_tokens response missing input_tokens field")
  347. }
  348. c.JSON(http.StatusOK, gin.H{
  349. "input_tokens": int(inputTokens.Int()),
  350. })
  351. return nil
  352. }
  353. func prepareOpenAIInputTokensCountRequest(
  354. body []byte,
  355. account *Account,
  356. defaultMappedModel string,
  357. ) (*openAIInputTokensCountPrepared, error) {
  358. var anthropicReq apicompat.AnthropicRequest
  359. if err := json.Unmarshal(body, &anthropicReq); err != nil {
  360. return nil, fmt.Errorf("parse anthropic count_tokens request: %w", err)
  361. }
  362. originalModel := anthropicReq.Model
  363. applyOpenAICompatModelNormalization(&anthropicReq)
  364. normalizedModel := anthropicReq.Model
  365. billingModel := resolveOpenAIForwardModel(account, normalizedModel, strings.TrimSpace(defaultMappedModel))
  366. upstreamModel := normalizeOpenAIModelForUpstream(account, billingModel)
  367. responsesReq, err := apicompat.AnthropicToResponses(&anthropicReq)
  368. if err != nil {
  369. return nil, fmt.Errorf("convert anthropic request to responses: %w", err)
  370. }
  371. return &openAIInputTokensCountPrepared{
  372. Request: openAIInputTokensCountRequest{
  373. Model: upstreamModel,
  374. Instructions: responsesReq.Instructions,
  375. Input: responsesReq.Input,
  376. Tools: responsesReq.Tools,
  377. ToolChoice: responsesReq.ToolChoice,
  378. },
  379. OriginalModel: originalModel,
  380. NormalizedModel: normalizedModel,
  381. BillingModel: billingModel,
  382. UpstreamModel: upstreamModel,
  383. }, nil
  384. }
  385. func (s *OpenAIGatewayService) buildInputTokensUpstreamRequest(
  386. ctx context.Context,
  387. c *gin.Context,
  388. account *Account,
  389. body []byte,
  390. token string,
  391. ) (*http.Request, error) {
  392. targetURL := openaiPlatformAPIInputTokensURL
  393. if account.Type == AccountTypeAPIKey {
  394. if baseURL := account.GetOpenAIBaseURL(); strings.TrimSpace(baseURL) != "" {
  395. validatedURL, err := s.validateUpstreamBaseURL(baseURL)
  396. if err != nil {
  397. return nil, err
  398. }
  399. targetURL = buildOpenAIResponsesInputTokensURL(validatedURL)
  400. }
  401. }
  402. req, err := http.NewRequestWithContext(ctx, http.MethodPost, targetURL, bytes.NewReader(body))
  403. if err != nil {
  404. return nil, err
  405. }
  406. req = req.WithContext(WithHTTPUpstreamProfile(req.Context(), HTTPUpstreamProfileOpenAI))
  407. authHeaders, err := s.buildOpenAIAuthenticationHeaders(ctx, account, token)
  408. if err != nil {
  409. return nil, err
  410. }
  411. for key, values := range authHeaders {
  412. for _, value := range values {
  413. req.Header.Add(key, value)
  414. }
  415. }
  416. req.Header.Set("content-type", "application/json")
  417. req.Header.Set("accept", "application/json")
  418. if c != nil && c.Request != nil {
  419. for key, values := range c.Request.Header {
  420. lower := strings.ToLower(strings.TrimSpace(key))
  421. if lower != "user-agent" && lower != "accept-language" {
  422. continue
  423. }
  424. for _, v := range values {
  425. req.Header.Add(key, v)
  426. }
  427. }
  428. }
  429. // 账号级请求头覆写(仅 openai api_key 账号启用时生效;OAuth 路径 no-op)
  430. account.ApplyHeaderOverrides(req.Header)
  431. return req, nil
  432. }
  433. func writeAnthropicCountTokensError(c *gin.Context, status int, errType, message string) {
  434. c.JSON(status, gin.H{
  435. "type": "error",
  436. "error": gin.H{
  437. "type": errType,
  438. "message": message,
  439. },
  440. })
  441. }
  442. func isOpenAIInputTokensUnsupported(statusCode int, body []byte) bool {
  443. if statusCode != http.StatusNotFound {
  444. return false
  445. }
  446. msg := strings.ToLower(strings.TrimSpace(extractUpstreamErrorMessage(body)))
  447. return strings.Contains(msg, "input_tokens") && strings.Contains(msg, "not found")
  448. }
  449. func writeOpenAIOAuthInputTokensFallback(c *gin.Context, account *Account, prepared *openAIInputTokensCountPrepared, statusCode int) {
  450. estimated := openAIInputTokensFallbackMinimum
  451. if got, err := estimateOpenAIInputTokens(prepared.Request); err == nil {
  452. if got > 0 {
  453. estimated = got
  454. }
  455. logger.L().Info("openai count_tokens: oauth fallback to local tiktoken estimate",
  456. zap.Int64("account_id", account.ID),
  457. zap.Int("upstream_status", statusCode),
  458. zap.Int("estimated_input_tokens", estimated),
  459. zap.String("upstream_model", prepared.UpstreamModel),
  460. )
  461. } else {
  462. logger.L().Warn("openai count_tokens: oauth local tiktoken fallback failed, using minimum estimate",
  463. zap.Int64("account_id", account.ID),
  464. zap.Int("upstream_status", statusCode),
  465. zap.Int("estimated_input_tokens", estimated),
  466. zap.String("upstream_model", prepared.UpstreamModel),
  467. zap.Error(err),
  468. )
  469. }
  470. c.JSON(http.StatusOK, gin.H{
  471. "input_tokens": estimated,
  472. })
  473. }
  474. func isOpenAIOAuthInputTokensUnsupported(statusCode int, body []byte) bool {
  475. switch statusCode {
  476. case http.StatusUnauthorized, http.StatusForbidden, http.StatusNotFound:
  477. default:
  478. return false
  479. }
  480. bodyLower := strings.ToLower(string(body))
  481. msg := strings.ToLower(strings.TrimSpace(extractUpstreamErrorMessage(body)))
  482. code := strings.ToLower(strings.TrimSpace(extractUpstreamErrorCode(body)))
  483. if code == "missing_scope" ||
  484. strings.Contains(bodyLower, "api.responses.write") ||
  485. strings.Contains(bodyLower, "missing scopes") ||
  486. strings.Contains(bodyLower, "insufficient_scope") {
  487. return true
  488. }
  489. if statusCode == http.StatusNotFound && isOpenAIInputTokensUnsupported(statusCode, body) {
  490. return true
  491. }
  492. // OAuth's platform endpoint can be blocked by an upstream proxy before it
  493. // reaches the API and return an HTML 403 page without a structured error.
  494. // Treat that endpoint-level response like the other unsupported cases so
  495. // count_tokens remains a local, non-health-affecting convenience request.
  496. if statusCode == http.StatusForbidden && isHTMLResponse(body) {
  497. return true
  498. }
  499. return strings.Contains(msg, "input_tokens") &&
  500. (strings.Contains(msg, "not found") ||
  501. strings.Contains(msg, "not supported") ||
  502. strings.Contains(msg, "unsupported"))
  503. }
  504. func isHTMLResponse(body []byte) bool {
  505. trimmed := strings.TrimSpace(strings.ToLower(string(body)))
  506. return strings.HasPrefix(trimmed, "<!doctype html") ||
  507. strings.HasPrefix(trimmed, "<html")
  508. }
  509. func estimateOpenAIInputTokens(req openAIInputTokensCountRequest) (int, error) {
  510. codec, err := openAIInputTokensCodecForModel(req.Model)
  511. if err != nil {
  512. return 0, err
  513. }
  514. total := 0
  515. addCount := func(text string) error {
  516. text = strings.TrimSpace(text)
  517. if text == "" {
  518. return nil
  519. }
  520. n, err := codec.Count(text)
  521. if err != nil {
  522. return err
  523. }
  524. total += n
  525. return nil
  526. }
  527. if err := addCount(req.Instructions); err != nil {
  528. return 0, err
  529. }
  530. inputTokens, err := estimateOpenAIInputTokensForInput(codec, req.Input)
  531. if err != nil {
  532. return 0, err
  533. }
  534. total += inputTokens
  535. for _, tool := range req.Tools {
  536. raw, err := marshalOpenAIUpstreamJSON(tool)
  537. if err != nil {
  538. return 0, err
  539. }
  540. if err := addCount(string(raw)); err != nil {
  541. return 0, err
  542. }
  543. }
  544. if len(req.ToolChoice) > 0 {
  545. compacted, err := compactOpenAIInputTokensJSON(req.ToolChoice)
  546. if err != nil {
  547. return 0, err
  548. }
  549. if err := addCount(compacted); err != nil {
  550. return 0, err
  551. }
  552. }
  553. if total < 0 {
  554. return 0, nil
  555. }
  556. return total, nil
  557. }
  558. func estimateOpenAIInputTokensForInput(codec tokenizer.Codec, raw json.RawMessage) (int, error) {
  559. if len(bytes.TrimSpace(raw)) == 0 {
  560. return 0, nil
  561. }
  562. var plainText string
  563. if err := json.Unmarshal(raw, &plainText); err == nil {
  564. return codec.Count(plainText)
  565. }
  566. var items []apicompat.ResponsesInputItem
  567. if err := json.Unmarshal(raw, &items); err == nil {
  568. return estimateOpenAIInputTokensForInputItems(codec, items)
  569. }
  570. compacted, err := compactOpenAIInputTokensJSON(raw)
  571. if err != nil {
  572. return 0, err
  573. }
  574. return codec.Count(compacted)
  575. }
  576. func estimateOpenAIInputTokensForInputItems(codec tokenizer.Codec, items []apicompat.ResponsesInputItem) (int, error) {
  577. total := 0
  578. countText := func(text string) error {
  579. text = strings.TrimSpace(text)
  580. if text == "" {
  581. return nil
  582. }
  583. n, err := codec.Count(text)
  584. if err != nil {
  585. return err
  586. }
  587. total += n
  588. return nil
  589. }
  590. for _, item := range items {
  591. total += openAIResponsesInputItemTokenOverhead
  592. if err := countText(item.Role); err != nil {
  593. return 0, err
  594. }
  595. if item.Type != "" && item.Type != "message" {
  596. if err := countText(item.Type); err != nil {
  597. return 0, err
  598. }
  599. }
  600. if err := countText(item.Name); err != nil {
  601. return 0, err
  602. }
  603. if err := countText(item.Arguments); err != nil {
  604. return 0, err
  605. }
  606. if err := countText(item.Output); err != nil {
  607. return 0, err
  608. }
  609. if err := countText(item.CallID); err != nil {
  610. return 0, err
  611. }
  612. if err := countText(item.ID); err != nil {
  613. return 0, err
  614. }
  615. if len(bytes.TrimSpace(item.Content)) == 0 {
  616. continue
  617. }
  618. var contentText string
  619. if err := json.Unmarshal(item.Content, &contentText); err == nil {
  620. if err := countText(contentText); err != nil {
  621. return 0, err
  622. }
  623. continue
  624. }
  625. var parts []apicompat.ResponsesContentPart
  626. if err := json.Unmarshal(item.Content, &parts); err == nil {
  627. for _, part := range parts {
  628. total += openAIResponsesContentPartOverhead
  629. switch part.Type {
  630. case "input_text", "output_text", "text":
  631. if err := countText(part.Text); err != nil {
  632. return 0, err
  633. }
  634. case "input_image":
  635. if err := countText(estimateOpenAIInputImageText(part.ImageURL)); err != nil {
  636. return 0, err
  637. }
  638. default:
  639. if err := countText(part.Type); err != nil {
  640. return 0, err
  641. }
  642. }
  643. }
  644. continue
  645. }
  646. compacted, err := compactOpenAIInputTokensJSON(item.Content)
  647. if err != nil {
  648. return 0, err
  649. }
  650. if err := countText(compacted); err != nil {
  651. return 0, err
  652. }
  653. }
  654. return total, nil
  655. }
  656. func estimateOpenAIInputImageText(imageURL string) string {
  657. trimmed := strings.TrimSpace(imageURL)
  658. if trimmed == "" {
  659. return ""
  660. }
  661. if strings.HasPrefix(strings.ToLower(trimmed), "data:") {
  662. if comma := strings.Index(trimmed, ","); comma > 0 {
  663. return trimmed[:comma]
  664. }
  665. }
  666. return trimmed
  667. }
  668. func compactOpenAIInputTokensJSON(raw json.RawMessage) (string, error) {
  669. if len(bytes.TrimSpace(raw)) == 0 {
  670. return "", nil
  671. }
  672. var buf bytes.Buffer
  673. if err := json.Compact(&buf, raw); err != nil {
  674. return "", err
  675. }
  676. return buf.String(), nil
  677. }
  678. func openAIInputTokensCodecForModel(model string) (tokenizer.Codec, error) {
  679. switch openAIInputTokensEncodingForModel(model) {
  680. case tokenizer.Cl100kBase:
  681. return tokenizer.Get(tokenizer.Cl100kBase)
  682. default:
  683. return tokenizer.Get(tokenizer.O200kBase)
  684. }
  685. }
  686. func openAIInputTokensEncodingForModel(model string) tokenizer.Encoding {
  687. normalized := strings.ToLower(strings.TrimSpace(model))
  688. switch {
  689. case strings.HasPrefix(normalized, "gpt-3.5"),
  690. (strings.HasPrefix(normalized, "gpt-4") &&
  691. !strings.HasPrefix(normalized, "gpt-4o") &&
  692. !strings.HasPrefix(normalized, "gpt-4.1")),
  693. strings.HasPrefix(normalized, "text-embedding-"):
  694. return tokenizer.Cl100kBase
  695. default:
  696. return tokenizer.O200kBase
  697. }
  698. }