Skip to content
Open
Show file tree
Hide file tree
Changes from 1 commit
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 4 additions & 4 deletions sdk/go/ai/client.go
Original file line number Diff line number Diff line change
Expand Up @@ -82,8 +82,8 @@ func (c *Client) CompleteWithMessages(ctx context.Context, messages []Message, o
}

func (c *Client) doRequest(ctx context.Context, req *Request) (*Response, error) {
// Marshal request
body, err := json.Marshal(req)
// Apply provider-specific parameter rewrites before serialization.
body, err := c.marshalRequest(req)
if err != nil {
return nil, fmt.Errorf("marshal request: %w", err)
}
Expand Down Expand Up @@ -176,8 +176,8 @@ func (c *Client) StreamComplete(ctx context.Context, prompt string, opts ...Opti
}
}

// Marshal request
body, err := json.Marshal(req)
// Apply provider-specific parameter rewrites before serialization.
body, err := c.marshalRequest(req)
if err != nil {
errCh <- fmt.Errorf("marshal request: %w", err)
return
Expand Down
95 changes: 95 additions & 0 deletions sdk/go/ai/model_params.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,95 @@
package ai

import (
"encoding/json"
"strings"
)

// needsMaxCompletionTokens returns true if the model requires
// max_completion_tokens instead of max_tokens. This applies to OpenAI's
// newer model families (o1, o3, gpt-4o, etc.) which dropped support for
// the legacy max_tokens parameter.
//
// Reference: https://platform.openai.com/docs/api-reference/chat/create
func needsMaxCompletionTokens(model string) bool {
m := strings.ToLower(strings.TrimSpace(model))

// Strip provider prefix if present (e.g. "openai/gpt-4o" → "gpt-4o")
if idx := strings.LastIndex(m, "/"); idx >= 0 {
m = m[idx+1:]
}

// o1, o3, o4 series always use max_completion_tokens
if strings.HasPrefix(m, "o1") || strings.HasPrefix(m, "o3") || strings.HasPrefix(m, "o4") {
return true
}

// gpt-4o and gpt-4o-mini series use max_completion_tokens
if strings.HasPrefix(m, "gpt-4o") {
return true
}

// gpt-4.1, gpt-4.5 etc. (newer dot-release models)
if strings.HasPrefix(m, "gpt-4.") {
return true
}

return false
}

// isOpenAICompatible returns true if the base URL points to OpenAI or an
// OpenAI-compatible endpoint (not Anthropic, Cohere, etc.).
func isOpenAICompatible(baseURL string) bool {
lower := strings.ToLower(baseURL)
// OpenAI's own endpoint
if strings.Contains(lower, "openai.com") {
return true
}
// OpenRouter proxies to OpenAI models
if strings.Contains(lower, "openrouter.ai") {
return true
}
// Azure OpenAI
if strings.Contains(lower, "openai.azure.com") {
return true
}
// For unknown/custom endpoints, assume OpenAI-compatible since that's
// the most common case for the chat/completions API shape.
return true
}
Comment thread
7vignesh marked this conversation as resolved.

// marshalRequest serializes the request, applying provider-specific parameter
// rewrites. For OpenAI-compatible endpoints with newer models, max_tokens is
// rewritten to max_completion_tokens.
func (c *Client) marshalRequest(req *Request) ([]byte, error) {
model := req.Model
if model == "" {
model = c.config.Model
}

// If the model needs max_completion_tokens and we have a max_tokens value,
// serialize with the rewritten field name.
if req.MaxTokens != nil && needsMaxCompletionTokens(model) && isOpenAICompatible(c.config.BaseURL) {
return marshalWithMaxCompletionTokens(req)
}

return json.Marshal(req)
}

// marshalWithMaxCompletionTokens serializes the request with max_completion_tokens
// instead of max_tokens. We use a shadow struct to avoid modifying the original.
func marshalWithMaxCompletionTokens(req *Request) ([]byte, error) {
type requestAlias Request

wire := struct {
*requestAlias
MaxTokens *int `json:"max_tokens,omitempty"`
MaxCompletionTokens *int `json:"max_completion_tokens,omitempty"`
}{
requestAlias: (*requestAlias)(req),
MaxTokens: nil, // suppress max_tokens
MaxCompletionTokens: req.MaxTokens,
}

return json.Marshal(wire)
}
230 changes: 230 additions & 0 deletions sdk/go/ai/model_params_test.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,230 @@
package ai

import (
"encoding/json"
"testing"
)

func TestNeedsMaxCompletionTokens(t *testing.T) {
tests := []struct {
model string
expected bool
}{
// o1 series
{"o1-mini", true},
{"o1-preview", true},
{"o1", true},
// o3 series
{"o3-mini", true},
{"o3", true},
// o4 series
{"o4-mini", true},
// gpt-4o series
{"gpt-4o", true},
{"gpt-4o-mini", true},
{"gpt-4o-2024-05-13", true},
// gpt-4.x dot releases
{"gpt-4.1", true},
{"gpt-4.5-preview", true},
// With provider prefix
{"openai/gpt-4o", true},
{"openai/o1-mini", true},
{"openrouter/openai/gpt-4o-mini", true},
// Case insensitive
{"GPT-4O", true},
{"O1-Mini", true},
// Legacy models — should NOT rewrite
{"gpt-4", false},
{"gpt-4-turbo", false},
{"gpt-3.5-turbo", false},
{"gpt-4-0613", false},
// Non-OpenAI models
{"claude-3-opus", false},
{"mistral-large", false},
{"llama-3-70b", false},
// Empty
{"", false},
}

for _, tt := range tests {
got := needsMaxCompletionTokens(tt.model)
if got != tt.expected {
t.Errorf("needsMaxCompletionTokens(%q) = %v, want %v", tt.model, got, tt.expected)
}
}
}

func TestMarshalRequest_RewritesMaxTokensForNewModels(t *testing.T) {
maxTokens := 1000
cfg := DefaultConfig()
cfg.APIKey = "test-key"
cfg.Model = "gpt-4o"

client, _ := NewClient(cfg)

req := &Request{
Model: "gpt-4o",
MaxTokens: &maxTokens,
Messages: []Message{
{Role: "user", Content: []ContentPart{{Type: "text", Text: "hello"}}},
},
}

body, err := client.marshalRequest(req)
if err != nil {
t.Fatalf("marshalRequest error: %v", err)
}

var raw map[string]interface{}
if err := json.Unmarshal(body, &raw); err != nil {
t.Fatalf("unmarshal error: %v", err)
}

if _, ok := raw["max_tokens"]; ok {
t.Error("expected max_tokens to be absent for gpt-4o model")
}
if val, ok := raw["max_completion_tokens"]; !ok {
t.Error("expected max_completion_tokens to be present for gpt-4o model")
} else if int(val.(float64)) != 1000 {
t.Errorf("expected max_completion_tokens=1000, got %v", val)
}
}

func TestMarshalRequest_KeepsMaxTokensForLegacyModels(t *testing.T) {
maxTokens := 2000
cfg := DefaultConfig()
cfg.APIKey = "test-key"
cfg.Model = "gpt-3.5-turbo"

client, _ := NewClient(cfg)

req := &Request{
Model: "gpt-3.5-turbo",
MaxTokens: &maxTokens,
Messages: []Message{
{Role: "user", Content: []ContentPart{{Type: "text", Text: "hello"}}},
},
}

body, err := client.marshalRequest(req)
if err != nil {
t.Fatalf("marshalRequest error: %v", err)
}

var raw map[string]interface{}
if err := json.Unmarshal(body, &raw); err != nil {
t.Fatalf("unmarshal error: %v", err)
}

if _, ok := raw["max_completion_tokens"]; ok {
t.Error("expected max_completion_tokens to be absent for gpt-3.5-turbo model")
}
if val, ok := raw["max_tokens"]; !ok {
t.Error("expected max_tokens to be present for gpt-3.5-turbo model")
} else if int(val.(float64)) != 2000 {
t.Errorf("expected max_tokens=2000, got %v", val)
}
}

func TestMarshalRequest_O1MiniUsesMaxCompletionTokens(t *testing.T) {
maxTokens := 500
cfg := DefaultConfig()
cfg.APIKey = "test-key"
cfg.Model = "o1-mini"

client, _ := NewClient(cfg)

req := &Request{
Model: "o1-mini",
MaxTokens: &maxTokens,
Messages: []Message{
{Role: "user", Content: []ContentPart{{Type: "text", Text: "hello"}}},
},
}

body, err := client.marshalRequest(req)
if err != nil {
t.Fatalf("marshalRequest error: %v", err)
}

var raw map[string]interface{}
if err := json.Unmarshal(body, &raw); err != nil {
t.Fatalf("unmarshal error: %v", err)
}

if _, ok := raw["max_tokens"]; ok {
t.Error("expected max_tokens to be absent for o1-mini")
}
if val, ok := raw["max_completion_tokens"]; !ok {
t.Error("expected max_completion_tokens to be present for o1-mini")
} else if int(val.(float64)) != 500 {
t.Errorf("expected max_completion_tokens=500, got %v", val)
}
}

func TestMarshalRequest_NilMaxTokensOmitted(t *testing.T) {
cfg := DefaultConfig()
cfg.APIKey = "test-key"
cfg.Model = "gpt-4o"

client, _ := NewClient(cfg)

req := &Request{
Model: "gpt-4o",
MaxTokens: nil,
Messages: []Message{
{Role: "user", Content: []ContentPart{{Type: "text", Text: "hello"}}},
},
}

body, err := client.marshalRequest(req)
if err != nil {
t.Fatalf("marshalRequest error: %v", err)
}

var raw map[string]interface{}
if err := json.Unmarshal(body, &raw); err != nil {
t.Fatalf("unmarshal error: %v", err)
}

if _, ok := raw["max_tokens"]; ok {
t.Error("expected max_tokens to be absent when nil")
}
if _, ok := raw["max_completion_tokens"]; ok {
t.Error("expected max_completion_tokens to be absent when MaxTokens is nil")
}
}

func TestMarshalRequest_UsesRequestModelOverClientModel(t *testing.T) {
maxTokens := 800
cfg := DefaultConfig()
cfg.APIKey = "test-key"
cfg.Model = "gpt-3.5-turbo" // client default is legacy

client, _ := NewClient(cfg)

req := &Request{
Model: "o1-preview", // request overrides to new model
MaxTokens: &maxTokens,
Messages: []Message{
{Role: "user", Content: []ContentPart{{Type: "text", Text: "hello"}}},
},
}

body, err := client.marshalRequest(req)
if err != nil {
t.Fatalf("marshalRequest error: %v", err)
}

var raw map[string]interface{}
if err := json.Unmarshal(body, &raw); err != nil {
t.Fatalf("unmarshal error: %v", err)
}

if _, ok := raw["max_tokens"]; ok {
t.Error("expected max_tokens to be absent for o1-preview")
}
if _, ok := raw["max_completion_tokens"]; !ok {
t.Error("expected max_completion_tokens to be present for o1-preview")
}
}
Loading