feat(provider): add retry logic and vendor prefix handling in OpenAI provider
This commit is contained in:
tong3jie
parent
d5cbf198b2
commit
aa5ee02507
3 changed files with 91 additions and 34 deletions
2
go.mod
2
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
|
||||
|
|
|
|||
4
go.sum
4
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=
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
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 != "" {
|
||||
parsed, err := url.Parse(proxy)
|
||||
if err == nil {
|
||||
client.Transport = &http.Transport{
|
||||
Proxy: http.ProxyURL(parsed),
|
||||
}
|
||||
|
||||
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,24 +482,21 @@ 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
|
||||
}
|
||||
}
|
||||
|
||||
func asInt(v any) (int, bool) {
|
||||
switch val := v.(type) {
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue