From aa5ee02507a9f2bb54daf97c4aebfd6438afa98e Mon Sep 17 00:00:00 2001 From: tong3jie <14191774+tong3jie@users.noreply.github.com> Date: Wed, 11 Mar 2026 14:09:07 +0800 Subject: [PATCH] feat(provider): add retry logic and vendor prefix handling in OpenAI provider --- go.mod | 2 + go.sum | 4 + pkg/providers/openai_compat/provider.go | 119 +++++++++++++++++------- 3 files changed, 91 insertions(+), 34 deletions(-) diff --git a/go.mod b/go.mod index f60be046f..3539f863c 100644 --- a/go.mod +++ b/go.mod @@ -40,6 +40,7 @@ require ( github.com/dustin/go-humanize v1.0.1 // indirect github.com/elliotchance/orderedmap/v3 v3.1.0 // indirect github.com/gdamore/encoding v1.0.1 // indirect + github.com/hashicorp/go-cleanhttp v0.5.2 // indirect github.com/inconshreveable/mousetrap v1.1.0 // indirect github.com/lucasb-eyer/go-colorful v1.3.0 // indirect github.com/mattn/go-colorable v0.1.14 // indirect @@ -77,6 +78,7 @@ require ( github.com/gogo/protobuf v1.3.2 // indirect github.com/google/jsonschema-go v0.4.2 // indirect github.com/grbit/go-json v0.11.0 // indirect + github.com/hashicorp/go-retryablehttp v0.7.8 github.com/klauspost/compress v1.18.4 // indirect github.com/klauspost/cpuid/v2 v2.3.0 // indirect github.com/tidwall/gjson v1.18.0 // indirect diff --git a/go.sum b/go.sum index 4060997f8..1e20d30da 100644 --- a/go.sum +++ b/go.sum @@ -102,6 +102,10 @@ github.com/grbit/go-json v0.11.0 h1:bAbyMdYrYl/OjYsSqLH99N2DyQ291mHy726Mx+sYrnc= github.com/grbit/go-json v0.11.0/go.mod h1:IYpHsdybQ386+6g3VE6AXQ3uTGa5mquBme5/ZWmtzek= github.com/h2non/filetype v1.1.3 h1:FKkx9QbD7HR/zjK1Ia5XiBsq9zdLi5Kf3zGyFTAFkGg= github.com/h2non/filetype v1.1.3/go.mod h1:319b3zT68BvV+WRj7cwy856M2ehB3HqNOt6sy1HndBY= +github.com/hashicorp/go-cleanhttp v0.5.2 h1:035FKYIWjmULyFRBKPs8TBQoi0x6d9G4xc9neXJWAZQ= +github.com/hashicorp/go-cleanhttp v0.5.2/go.mod h1:kO/YDlP8L1346E6Sodw+PrpBSV4/SoxCXGY6BqNFT48= +github.com/hashicorp/go-retryablehttp v0.7.8 h1:ylXZWnqa7Lhqpk0L1P1LzDtGcCR0rPVUrx/c8Unxc48= +github.com/hashicorp/go-retryablehttp v0.7.8/go.mod h1:rjiScheydd+CxvumBsIrFKlx3iS0jrZ7LvzFGFmuKbw= github.com/hashicorp/golang-lru/v2 v2.0.7 h1:a+bsQ5rvGLjzHuww6tVxozPZFVghXaHOwFs4luLUK2k= github.com/hashicorp/golang-lru/v2 v2.0.7/go.mod h1:QeFd9opnmA6QUJc5vARoKUSoFhyfM2/ZepoAG6RGpeM= github.com/hpcloud/tail v1.0.0/go.mod h1:ab1qPbhIpdTxEkNHXyeSf5vhxWSCs/tWer42PpOxQnU= diff --git a/pkg/providers/openai_compat/provider.go b/pkg/providers/openai_compat/provider.go index 0e8db7409..2b9f40e17 100644 --- a/pkg/providers/openai_compat/provider.go +++ b/pkg/providers/openai_compat/provider.go @@ -13,6 +13,7 @@ import ( "strings" "time" + "github.com/hashicorp/go-retryablehttp" "github.com/sipeed/picoclaw/pkg/providers/protocoltypes" ) @@ -29,6 +30,21 @@ type ( ReasoningDetail = protocoltypes.ReasoningDetail ) +var vendorPrefixes = map[string]struct{}{ + "litellm": {}, + "moonshot": {}, + "nvidia": {}, + "groq": {}, + "ollama": {}, + "deepseek": {}, + "google": {}, + "openrouter": {}, + "zhipu": {}, + "mistral": {}, + "vivgrid": {}, + "minimax": {}, +} + type Provider struct { apiKey string apiBase string @@ -54,26 +70,48 @@ func WithRequestTimeout(timeout time.Duration) Option { } } -func NewProvider(apiKey, apiBase, proxy string, opts ...Option) *Provider { - client := &http.Client{ - Timeout: defaultRequestTimeout, - } - - if proxy != "" { - parsed, err := url.Parse(proxy) - if err == nil { - client.Transport = &http.Transport{ - Proxy: http.ProxyURL(parsed), +func WithRetry(maxRetries int, minWait, maxWait time.Duration) Option { + return func(p *Provider) { + if rc, ok := p.httpClient.Transport.(*retryablehttp.RoundTripper); ok { + if maxRetries >= 0 { + rc.Client.RetryMax = maxRetries } + if minWait > 0 { + rc.Client.RetryWaitMin = minWait + } + if maxWait > 0 { + rc.Client.RetryWaitMax = maxWait + } + } + } +} + +func NewProvider(apiKey, apiBase, proxy string, opts ...Option) *Provider { + + retryClient := retryablehttp.NewClient() + retryClient.RetryMax = 3 + retryClient.RetryWaitMin = 1 * time.Second + retryClient.RetryWaitMax = 30 * time.Second + retryClient.Backoff = retryablehttp.LinearJitterBackoff + retryClient.Logger = nil + + transport := &http.Transport{} + if proxy != "" { + + if parsed, err := url.Parse(proxy); err == nil { + transport.Proxy = http.ProxyURL(parsed) } else { log.Printf("openai_compat: invalid proxy URL %q: %v", proxy, err) } } + retryClient.HTTPClient.Transport = transport + retryClient.HTTPClient.Timeout = defaultRequestTimeout + p := &Provider{ apiKey: apiKey, apiBase: strings.TrimRight(apiBase, "/"), - httpClient: client, + httpClient: retryClient.StandardClient(), } for _, opt := range opts { @@ -126,18 +164,7 @@ func (p *Provider) Chat( } if maxTokens, ok := asInt(options["max_tokens"]); ok { - // Use configured maxTokensField if specified, otherwise fallback to model-based detection - fieldName := p.maxTokensField - if fieldName == "" { - // Fallback: detect from model name for backward compatibility - lowerModel := strings.ToLower(model) - if strings.Contains(lowerModel, "glm") || strings.Contains(lowerModel, "o1") || - strings.Contains(lowerModel, "gpt-5") { - fieldName = "max_completion_tokens" - } else { - fieldName = "max_tokens" - } - } + fieldName := p.resolveMaxTokenField(model) requestBody[fieldName] = maxTokens } @@ -184,6 +211,13 @@ func (p *Provider) Chat( } defer resp.Body.Close() + select { + case <-ctx.Done(): + return nil, ctx.Err() + default: + // continue processing response + } + contentType := resp.Header.Get("Content-Type") // Non-200: read a prefix to tell HTML error page apart from JSON error body. @@ -202,8 +236,12 @@ func (p *Provider) Chat( ) } + //set response size limit to prevent OOM if server returns a huge response (e.g., an HTML error page instead of JSON) + const maxResponseSize = 10 * 1024 * 1024 + // Peek without consuming so the full stream reaches the JSON decoder. - reader := bufio.NewReader(resp.Body) + safeReader := io.LimitReader(resp.Body, maxResponseSize) + reader := bufio.NewReader(safeReader) prefix, err := reader.Peek(256) // io.EOF/ErrBufferFull are normal; only real errors abort if err != nil && err != io.EOF && err != bufio.ErrBufferFull { return nil, fmt.Errorf("failed to inspect response: %w", err) @@ -214,12 +252,28 @@ func (p *Provider) Chat( out, err := parseResponse(reader) if err != nil { + // some APIs return 200 with an HTML error page, so check for that before giving up on JSON parsing + if looksLikeHTML(prefix, contentType) { + return nil, wrapHTMLResponseError(resp.StatusCode, prefix, contentType, p.apiBase) + } return nil, fmt.Errorf("failed to parse JSON response: %w", err) } return out, nil } +func (p *Provider) resolveMaxTokenField(model string) string { + if p.maxTokensField != "" { + return p.maxTokensField + } + + lowerModel := strings.ToLower(model) + if strings.Contains(lowerModel, "o1") || strings.Contains(lowerModel, "glm-4") || strings.Contains(lowerModel, "gpt-5") { + return "max_completion_tokens" + } + return "max_tokens" +} + func wrapHTMLResponseError(statusCode int, body []byte, contentType, apiBase string) error { respPreview := responsePreview(body, 128) return fmt.Errorf( @@ -428,23 +482,20 @@ func serializeMessages(messages []Message) []any { } func normalizeModel(model, apiBase string) string { + if strings.Contains(strings.ToLower(apiBase), "openrouter.ai") { + return model + } + before, after, ok := strings.Cut(model, "/") if !ok { return model } - if strings.Contains(strings.ToLower(apiBase), "openrouter.ai") { - return model + if _, exists := vendorPrefixes[strings.ToLower(before)]; exists { + return after } - prefix := strings.ToLower(before) - switch prefix { - case "litellm", "moonshot", "nvidia", "groq", "ollama", "deepseek", "google", - "openrouter", "zhipu", "mistral", "vivgrid", "minimax": - return after - default: - return model - } + return model } func asInt(v any) (int, bool) {