feat(provider): add retry logic and vendor prefix handling in OpenAI provider

This commit is contained in:
tong3jie 2026-03-11 14:09:07 +08:00
parent 4e3769e989
commit 297f64103b
4 changed files with 88 additions and 26 deletions

2
go.mod
View file

@ -47,6 +47,7 @@ require (
github.com/elliotchance/orderedmap/v3 v3.1.0 // indirect github.com/elliotchance/orderedmap/v3 v3.1.0 // indirect
github.com/gdamore/encoding v1.0.1 // indirect github.com/gdamore/encoding v1.0.1 // indirect
github.com/godbus/dbus/v5 v5.1.0 // indirect github.com/godbus/dbus/v5 v5.1.0 // indirect
github.com/hashicorp/go-cleanhttp v0.5.2 // indirect
github.com/inconshreveable/mousetrap v1.1.0 // indirect github.com/inconshreveable/mousetrap v1.1.0 // indirect
github.com/lucasb-eyer/go-colorful v1.3.0 // indirect github.com/lucasb-eyer/go-colorful v1.3.0 // indirect
github.com/mattn/go-colorable v0.1.14 // indirect github.com/mattn/go-colorable v0.1.14 // indirect
@ -81,6 +82,7 @@ require (
github.com/gogo/protobuf v1.3.2 // indirect github.com/gogo/protobuf v1.3.2 // indirect
github.com/google/jsonschema-go v0.4.2 // indirect github.com/google/jsonschema-go v0.4.2 // indirect
github.com/grbit/go-json v0.11.0 // 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/compress v1.18.4 // indirect
github.com/klauspost/cpuid/v2 v2.3.0 // indirect github.com/klauspost/cpuid/v2 v2.3.0 // indirect
github.com/tidwall/gjson v1.18.0 // indirect github.com/tidwall/gjson v1.18.0 // indirect

4
go.sum
View file

@ -108,6 +108,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/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 h1:FKkx9QbD7HR/zjK1Ia5XiBsq9zdLi5Kf3zGyFTAFkGg=
github.com/h2non/filetype v1.1.3/go.mod h1:319b3zT68BvV+WRj7cwy856M2ehB3HqNOt6sy1HndBY= 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 h1:a+bsQ5rvGLjzHuww6tVxozPZFVghXaHOwFs4luLUK2k=
github.com/hashicorp/golang-lru/v2 v2.0.7/go.mod h1:QeFd9opnmA6QUJc5vARoKUSoFhyfM2/ZepoAG6RGpeM= 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= github.com/hpcloud/tail v1.0.0/go.mod h1:ab1qPbhIpdTxEkNHXyeSf5vhxWSCs/tWer42PpOxQnU=

View file

@ -36,7 +36,12 @@ type (
ReasoningDetail = protocoltypes.ReasoningDetail ReasoningDetail = protocoltypes.ReasoningDetail
) )
const DefaultRequestTimeout = 120 * time.Second const (
DefaultRequestTimeout = 120 * time.Second
// set response size limit to prevent OOM if server returns a huge response (e.g., an HTML error page instead of JSON)
maxResponseSize = 10 * 1024 * 1024
)
// NewHTTPClient creates an *http.Client with an optional proxy and the default timeout. // NewHTTPClient creates an *http.Client with an optional proxy and the default timeout.
func NewHTTPClient(proxy string) *http.Client { func NewHTTPClient(proxy string) *http.Client {
@ -316,7 +321,9 @@ func HandleErrorResponse(resp *http.Response, apiBase string) error {
// then parses the JSON response into an LLMResponse. // then parses the JSON response into an LLMResponse.
func ReadAndParseResponse(resp *http.Response, apiBase string) (*LLMResponse, error) { func ReadAndParseResponse(resp *http.Response, apiBase string) (*LLMResponse, error) {
contentType := resp.Header.Get("Content-Type") contentType := resp.Header.Get("Content-Type")
reader := bufio.NewReader(resp.Body)
safeReader := io.LimitReader(resp.Body, maxResponseSize)
reader := bufio.NewReader(safeReader)
prefix, err := reader.Peek(256) prefix, err := reader.Peek(256)
if err != nil && err != io.EOF && err != bufio.ErrBufferFull { if err != nil && err != io.EOF && err != bufio.ErrBufferFull {
return nil, fmt.Errorf("failed to inspect response: %w", err) return nil, fmt.Errorf("failed to inspect response: %w", err)
@ -326,6 +333,10 @@ func ReadAndParseResponse(resp *http.Response, apiBase string) (*LLMResponse, er
} }
out, err := ParseResponse(reader) out, err := ParseResponse(reader)
if err != nil { 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, apiBase)
}
return nil, fmt.Errorf("failed to parse JSON response: %w", err) return nil, fmt.Errorf("failed to parse JSON response: %w", err)
} }
return out, nil return out, nil

View file

@ -8,11 +8,14 @@ import (
"fmt" "fmt"
"io" "io"
"log" "log"
"maps"
"net/http" "net/http"
"net/url" "net/url"
"strings" "strings"
"time" "time"
"github.com/hashicorp/go-retryablehttp"
"github.com/sipeed/picoclaw/pkg/providers/common" "github.com/sipeed/picoclaw/pkg/providers/common"
"github.com/sipeed/picoclaw/pkg/providers/protocoltypes" "github.com/sipeed/picoclaw/pkg/providers/protocoltypes"
) )
@ -30,6 +33,21 @@ type (
ReasoningDetail = protocoltypes.ReasoningDetail ReasoningDetail = protocoltypes.ReasoningDetail
) )
var vendorPrefixes = map[string]struct{}{
"litellm": {},
"moonshot": {},
"nvidia": {},
"groq": {},
"ollama": {},
"deepseek": {},
"google": {},
"openrouter": {},
"zhipu": {},
"mistral": {},
"vivgrid": {},
"minimax": {},
}
type Provider struct { type Provider struct {
apiKey string apiKey string
apiBase string apiBase string
@ -62,11 +80,46 @@ func WithExtraBody(extraBody map[string]any) Option {
} }
} }
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 { 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{ p := &Provider{
apiKey: apiKey, apiKey: apiKey,
apiBase: strings.TrimRight(apiBase, "/"), apiBase: strings.TrimRight(apiBase, "/"),
httpClient: common.NewHTTPClient(proxy), httpClient: retryClient.StandardClient(),
} }
for _, opt := range opts { for _, opt := range opts {
@ -75,6 +128,10 @@ func NewProvider(apiKey, apiBase, proxy string, opts ...Option) *Provider {
} }
} }
if p.maxTokensField == "" {
p.maxTokensField = "max_tokens"
}
return p return p
} }
@ -115,17 +172,7 @@ func (p *Provider) buildRequestBody(
} }
if maxTokens, ok := common.AsInt(options["max_tokens"]); ok { if maxTokens, ok := common.AsInt(options["max_tokens"]); ok {
fieldName := p.maxTokensField requestBody[p.maxTokensField] = maxTokens
if fieldName == "" {
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"
}
}
requestBody[fieldName] = maxTokens
} }
if temperature, ok := common.AsFloat(options["temperature"]); ok { if temperature, ok := common.AsFloat(options["temperature"]); ok {
@ -149,8 +196,8 @@ func (p *Provider) buildRequestBody(
// Merge extra body fields configured per-provider/model. // Merge extra body fields configured per-provider/model.
// These are injected last so they take precedence over defaults. // These are injected last so they take precedence over defaults.
for k, v := range p.extraBody { if p.extraBody != nil {
requestBody[k] = v maps.Copy(requestBody, p.extraBody)
} }
return requestBody return requestBody
@ -190,6 +237,7 @@ func (p *Provider) Chat(
} }
defer resp.Body.Close() defer resp.Body.Close()
// Non-200: read a prefix to tell HTML error page apart from JSON error body.
if resp.StatusCode != http.StatusOK { if resp.StatusCode != http.StatusOK {
return nil, common.HandleErrorResponse(resp, p.apiBase) return nil, common.HandleErrorResponse(resp, p.apiBase)
} }
@ -387,23 +435,20 @@ func parseStreamResponse(
} }
func normalizeModel(model, apiBase string) string { func normalizeModel(model, apiBase string) string {
if strings.Contains(strings.ToLower(apiBase), "openrouter.ai") {
return model
}
before, after, ok := strings.Cut(model, "/") before, after, ok := strings.Cut(model, "/")
if !ok { if !ok {
return model return model
} }
if strings.Contains(strings.ToLower(apiBase), "openrouter.ai") { if _, exists := vendorPrefixes[strings.ToLower(before)]; exists {
return model return after
} }
prefix := strings.ToLower(before) return model
switch prefix {
case "litellm", "moonshot", "nvidia", "groq", "ollama", "deepseek", "google",
"openrouter", "zhipu", "mistral", "vivgrid", "minimax", "novita":
return after
default:
return model
}
} }
func buildToolsList(tools []ToolDefinition, nativeSearch bool) []any { func buildToolsList(tools []ToolDefinition, nativeSearch bool) []any {