| 1 | package provider |
| 2 | |
| 3 | import ( |
| 4 | "encoding/json" |
| 5 | "errors" |
| 6 | "net/http" |
| 7 | "regexp" |
| 8 | "strconv" |
| 9 | "strings" |
| 10 | ) |
| 11 | |
| 12 | // ContextLimitError is a trusted shared-window overflow from a provider HTTP |
| 13 | // 400/413/422. Unwrap returns the original APIError so localization, trace IDs, |
| 14 | // and telemetry keep working. The body is never persisted or replayed. |
| 15 | type ContextLimitError struct { |
| 16 | APIError *APIError |
| 17 | WindowTokens int |
| 18 | RequestedTokens int |
| 19 | PromptTokens int |
| 20 | CompletionTokens int |
| 21 | } |
| 22 | |
| 23 | // OutputLimitError is a provider-reported completion-token ceiling. It is |
| 24 | // separate from ContextLimitError because the request may fit the model |
| 25 | // context window while exceeding the route's output-only limit. |
| 26 | type OutputLimitError struct { |
| 27 | APIError *APIError |
| 28 | RequestedTokens int |
| 29 | MaxOutputTokens int |
| 30 | } |
| 31 | |
| 32 | func (e *OutputLimitError) Error() string { |
| 33 | if e == nil { |
| 34 | return "output token limit exceeded" |
| 35 | } |
| 36 | if e.APIError != nil { |
| 37 | return e.APIError.Error() |
| 38 | } |
| 39 | return "output token limit exceeded" |
| 40 | } |
| 41 | |
| 42 | func (e *OutputLimitError) Unwrap() error { |
| 43 | if e == nil { |
| 44 | return nil |
| 45 | } |
| 46 | return e.APIError |
| 47 | } |
| 48 | |
| 49 | func (e *ContextLimitError) Error() string { |
| 50 | if e == nil { |
| 51 | return "context limit exceeded" |
| 52 | } |
| 53 | if e.APIError != nil { |
| 54 | return e.APIError.Error() |
| 55 | } |
| 56 | return "context limit exceeded" |
| 57 | } |
| 58 | |
| 59 | func (e *ContextLimitError) Unwrap() error { |
| 60 | if e == nil { |
| 61 | return nil |
| 62 | } |
| 63 | return e.APIError |
| 64 | } |
| 65 | |
| 66 | var ( |
| 67 | contextLimitEnglishRe = regexp.MustCompile(`(?i)maximum context length is (\d+) tokens?\.?\s*however,\s*you requested (\d+) tokens? \((\d+) in the (?:messages|prompt), (\d+) in the completion\)`) |
| 68 | contextLimitPromptRe = regexp.MustCompile(`(?i)prompt is too long:\s*(\d+) tokens? > (\d+) maximum`) |
| 69 | contextLimitInputRe = regexp.MustCompile(`(?i)\binput length (\d+) exceeds the maximum length (\d+)\b`) |
| 70 | contextLimitSumRe = regexp.MustCompile("(?i)input length and [`']?max_tokens[`']? exceed context limit:\\s*(\\d+)\\s*\\+\\s*(\\d+)\\s*>\\s*(\\d+)") |
| 71 | outputLimitRe = regexp.MustCompile(`(?i)max_tokens\s*(?:is\s+too\s+large|too\s+large)\s*[:=]?\s*(\d+).*?(?:supports?|maximum|at\s+most)[^\d]*(\d+)`) |
| 72 | ) |
| 73 | |
| 74 | func contextLimitStatusOK(status int) bool { |
| 75 | return status == http.StatusBadRequest || status == http.StatusRequestEntityTooLarge || status == http.StatusUnprocessableEntity |
| 76 | } |
| 77 | |
| 78 | func positiveToken(n int) bool { return n > 0 } |
| 79 | |
| 80 | func contextLimitInvariant(window, requested, prompt, completion int) bool { |
| 81 | if !positiveToken(window) { |
| 82 | return false |
| 83 | } |
| 84 | if positiveToken(prompt) && positiveToken(completion) { |
| 85 | if prompt+completion <= window { |
| 86 | return false |
| 87 | } |
| 88 | if requested > 0 && requested != prompt+completion { |
| 89 | return false |
| 90 | } |
| 91 | return true |
| 92 | } |
| 93 | if requested > window && (prompt > 0 || completion > 0 || requested > 0) { |
| 94 | return requested > window |
| 95 | } |
| 96 | return false |
| 97 | } |
| 98 | |
| 99 | func completeContextLimit(window, requested, prompt, completion int) (int, int, int, int, bool) { |
| 100 | if window <= 0 { |
| 101 | return 0, 0, 0, 0, false |
| 102 | } |
| 103 | if prompt > 0 && completion > 0 && requested <= 0 { |
| 104 | requested = prompt + completion |
| 105 | } |
| 106 | if requested > 0 && prompt > 0 && completion <= 0 && requested > prompt { |
| 107 | completion = requested - prompt |
| 108 | } |
| 109 | if requested > 0 && completion > 0 && prompt <= 0 && requested > completion { |
| 110 | prompt = requested - completion |
| 111 | } |
| 112 | if !contextLimitInvariant(window, requested, prompt, completion) { |
| 113 | return 0, 0, 0, 0, false |
| 114 | } |
| 115 | if requested <= 0 { |
| 116 | requested = prompt + completion |
| 117 | } |
| 118 | return window, requested, prompt, completion, true |
| 119 | } |
| 120 | |
| 121 | type contextLimitJSON struct { |
| 122 | Error *struct { |
| 123 | Message string `json:"message"` |
| 124 | ContextLength int `json:"context_length"` |
| 125 | MaxContextLength int `json:"max_context_length"` |
| 126 | MaxTokens int `json:"max_tokens"` |
| 127 | RequestedTokens int `json:"requested_tokens"` |
| 128 | Requested int `json:"requested"` |
| 129 | PromptTokens int `json:"prompt_tokens"` |
| 130 | InputTokens int `json:"input_tokens"` |
| 131 | CompletionTokens int `json:"completion_tokens"` |
| 132 | OutputTokens int `json:"output_tokens"` |
| 133 | } `json:"error"` |
| 134 | ContextLength int `json:"context_length"` |
| 135 | MaxContextLength int `json:"max_context_length"` |
| 136 | RequestedTokens int `json:"requested_tokens"` |
| 137 | PromptTokens int `json:"prompt_tokens"` |
| 138 | InputTokens int `json:"input_tokens"` |
| 139 | CompletionTokens int `json:"completion_tokens"` |
| 140 | OutputTokens int `json:"output_tokens"` |
| 141 | Usage *struct { |
| 142 | InputTokens int `json:"input_tokens"` |
| 143 | OutputTokens int `json:"output_tokens"` |
| 144 | PromptTokens int `json:"prompt_tokens"` |
| 145 | CompletionTokens int `json:"completion_tokens"` |
| 146 | } `json:"usage"` |
| 147 | } |
| 148 | |
| 149 | func firstPositive(values ...int) int { |
| 150 | for _, n := range values { |
| 151 | if n > 0 { |
| 152 | return n |
| 153 | } |
| 154 | } |
| 155 | return 0 |
| 156 | } |
| 157 | |
| 158 | func parseContextLimitJSON(body string) (window, requested, prompt, completion int, message string, ok bool) { |
| 159 | var parsed contextLimitJSON |
| 160 | if json.Unmarshal([]byte(body), &parsed) != nil { |
| 161 | return 0, 0, 0, 0, "", false |
| 162 | } |
| 163 | if parsed.Error != nil { |
| 164 | message = parsed.Error.Message |
| 165 | window = firstPositive(parsed.Error.ContextLength, parsed.Error.MaxContextLength) |
| 166 | requested = firstPositive(parsed.Error.RequestedTokens, parsed.Error.Requested) |
| 167 | prompt = firstPositive(parsed.Error.PromptTokens, parsed.Error.InputTokens) |
| 168 | completion = firstPositive(parsed.Error.CompletionTokens, parsed.Error.OutputTokens) |
| 169 | } |
| 170 | window = firstPositive(window, parsed.ContextLength, parsed.MaxContextLength) |
| 171 | requested = firstPositive(requested, parsed.RequestedTokens) |
| 172 | prompt = firstPositive(prompt, parsed.PromptTokens, parsed.InputTokens) |
| 173 | completion = firstPositive(completion, parsed.CompletionTokens, parsed.OutputTokens) |
| 174 | if parsed.Usage != nil { |
| 175 | prompt = firstPositive(prompt, parsed.Usage.PromptTokens, parsed.Usage.InputTokens) |
| 176 | completion = firstPositive(completion, parsed.Usage.CompletionTokens, parsed.Usage.OutputTokens) |
| 177 | } |
| 178 | if window, requested, prompt, completion, ok = completeContextLimit(window, requested, prompt, completion); ok { |
| 179 | return window, requested, prompt, completion, message, true |
| 180 | } |
| 181 | return 0, 0, 0, 0, message, false |
| 182 | } |
| 183 | |
| 184 | func parseContextLimitText(text string) (window, requested, prompt, completion int, ok bool) { |
| 185 | text = strings.TrimSpace(text) |
| 186 | if text == "" { |
| 187 | return 0, 0, 0, 0, false |
| 188 | } |
| 189 | if m := contextLimitEnglishRe.FindStringSubmatch(text); len(m) == 5 { |
| 190 | return completeContextLimit(atoiStrict(m[1]), atoiStrict(m[2]), atoiStrict(m[3]), atoiStrict(m[4])) |
| 191 | } |
| 192 | if m := contextLimitSumRe.FindStringSubmatch(text); len(m) == 4 { |
| 193 | return completeContextLimit(atoiStrict(m[3]), 0, atoiStrict(m[1]), atoiStrict(m[2])) |
| 194 | } |
| 195 | for _, re := range []*regexp.Regexp{contextLimitPromptRe, contextLimitInputRe} { |
| 196 | if m := re.FindStringSubmatch(text); len(m) == 3 { |
| 197 | prompt = atoiStrict(m[1]) |
| 198 | window = atoiStrict(m[2]) |
| 199 | if prompt > 0 && window > 0 && prompt > window { |
| 200 | return window, prompt, prompt, 0, true |
| 201 | } |
| 202 | } |
| 203 | } |
| 204 | return 0, 0, 0, 0, false |
| 205 | } |
| 206 | |
| 207 | func atoiStrict(s string) int { |
| 208 | n, err := strconv.Atoi(strings.TrimSpace(s)) |
| 209 | if err != nil || n <= 0 { |
| 210 | return 0 |
| 211 | } |
| 212 | return n |
| 213 | } |
| 214 | |
| 215 | // ParseContextLimitError extracts a trusted overflow from an APIError. |
| 216 | // Unparseable, non-context, or invariant-breaking bodies return nil. |
| 217 | func ParseContextLimitError(apiErr *APIError) *ContextLimitError { |
| 218 | if apiErr == nil || !contextLimitStatusOK(apiErr.Status) { |
| 219 | return nil |
| 220 | } |
| 221 | window, requested, prompt, completion, message, jsonOK := parseContextLimitJSON(apiErr.Body) |
| 222 | if !jsonOK { |
| 223 | if w, r, p, c, ok := parseContextLimitText(apiErr.Body); ok { |
| 224 | window, requested, prompt, completion = w, r, p, c |
| 225 | } else if w, r, p, c, ok := parseContextLimitText(message); ok { |
| 226 | window, requested, prompt, completion = w, r, p, c |
| 227 | } else { |
| 228 | // A bare overflow with no token numbers (Zhipu GLM 1261) is still |
| 229 | // provider-confirmed: trust it with an unknown window so consumers |
| 230 | // fall back to the configured window instead of resending as-is. |
| 231 | if isUnnumberedPromptTooLong(message, apiErr.Body) { |
| 232 | return &ContextLimitError{APIError: apiErr} |
| 233 | } |
| 234 | return nil |
| 235 | } |
| 236 | } |
| 237 | if !contextLimitInvariant(window, requested, prompt, completion) && |
| 238 | !(window > 0 && requested > window && prompt > 0) { |
| 239 | if isUnnumberedPromptTooLong(message, apiErr.Body) { |
| 240 | return &ContextLimitError{APIError: apiErr} |
| 241 | } |
| 242 | return nil |
| 243 | } |
| 244 | if requested <= 0 { |
| 245 | requested = prompt + completion |
| 246 | } |
| 247 | return &ContextLimitError{ |
| 248 | APIError: apiErr, |
| 249 | WindowTokens: window, |
| 250 | RequestedTokens: requested, |
| 251 | PromptTokens: prompt, |
| 252 | CompletionTokens: completion, |
| 253 | } |
| 254 | } |
| 255 | |
| 256 | // isUnnumberedPromptTooLong matches provider overflow errors that carry no |
| 257 | // token numbers at all. Canonical shape — Zhipu GLM 1261: |
| 258 | // |
| 259 | // {"error":{"code":"1261","message":"Prompt exceeds max length"}} |
| 260 | // |
| 261 | // The message (or the whole body, when the JSON shape differs) is matched |
| 262 | // case-insensitively; code 1261 is not matched directly so sibling GLM codes |
| 263 | // that reuse the message stay covered and numeric codes never false-positive. |
| 264 | func isUnnumberedPromptTooLong(message, body string) bool { |
| 265 | for _, s := range []string{message, body} { |
| 266 | if s == "" { |
| 267 | continue |
| 268 | } |
| 269 | if strings.Contains(strings.ToLower(s), "prompt exceeds max length") { |
| 270 | return true |
| 271 | } |
| 272 | } |
| 273 | return false |
| 274 | } |
| 275 | |
| 276 | // AsContextLimitError unwraps err to a trusted overflow, if any. |
| 277 | func AsContextLimitError(err error) *ContextLimitError { |
| 278 | var limit *ContextLimitError |
| 279 | if err != nil && errors.As(err, &limit) { |
| 280 | return limit |
| 281 | } |
| 282 | return nil |
| 283 | } |
| 284 | |
| 285 | // ParseOutputLimitError extracts a completion-only ceiling from a 400/413/422 |
| 286 | // API error. The parser is intentionally conservative: it only accepts text |
| 287 | // that names both the requested max_tokens and a smaller supported maximum. |
| 288 | func ParseOutputLimitError(apiErr *APIError) *OutputLimitError { |
| 289 | if apiErr == nil || !contextLimitStatusOK(apiErr.Status) { |
| 290 | return nil |
| 291 | } |
| 292 | text := strings.TrimSpace(apiErr.Body) |
| 293 | if text == "" { |
| 294 | return nil |
| 295 | } |
| 296 | m := outputLimitRe.FindStringSubmatch(text) |
| 297 | if len(m) != 3 { |
| 298 | return nil |
| 299 | } |
| 300 | requested, maxOutput := atoiStrict(m[1]), atoiStrict(m[2]) |
| 301 | if requested <= 0 || maxOutput <= 0 || requested <= maxOutput { |
| 302 | return nil |
| 303 | } |
| 304 | return &OutputLimitError{APIError: apiErr, RequestedTokens: requested, MaxOutputTokens: maxOutput} |
| 305 | } |
| 306 | |
| 307 | // AsOutputLimitError unwraps err to a trusted output ceiling, if any. |
| 308 | func AsOutputLimitError(err error) *OutputLimitError { |
| 309 | var limit *OutputLimitError |
| 310 | if err != nil && errors.As(err, &limit) { |
| 311 | return limit |
| 312 | } |
| 313 | return nil |
| 314 | } |
| 315 |