Add reasoning effort

This commit is contained in:
dwrz
2026-07-07 15:14:07 +00:00
parent c730ecdc22
commit 27fa702071
9 changed files with 192 additions and 122 deletions

View File

@@ -25,6 +25,15 @@ type CompletionsResult struct {
Timings *Timings
}
// CompletionsParams holds parameters for Completions and CompletionsStream.
type CompletionsParams struct {
Model string
SystemMessage string
MaxIterations int
Messages []Message
ThinkingBudgetTokens *int
}
// Completions sends messages to the LLM using the specified model.
// If systemMessage is non-empty, it overrides the configured system message.
// maxIterations caps the number of agent loop iterations.
@@ -32,15 +41,14 @@ type CompletionsResult struct {
// Returns all messages generated during the completions request, including
// tool calls and tool results, along with usage and timings from the final response.
func (c *Client) Completions(
ctx context.Context,
model string,
systemMessage string,
maxIterations int,
messages []Message,
ctx context.Context, p CompletionsParams,
) (*CompletionsResult, error) {
ctx, cancel := context.WithTimeout(ctx, c.timeout)
defer cancel()
model, systemMessage, maxIterations, messages, thinkingBudgetTokens :=
p.Model, p.SystemMessage, p.MaxIterations, p.Messages, p.ThinkingBudgetTokens
// Use request system message if provided, config default otherwise.
smsg := c.systemMessage
if systemMessage != "" {
@@ -68,7 +76,7 @@ func (c *Client) Completions(
var finalUsage *Usage
var finalTimings *Timings
for i := 0; maxIterations == 0 || i < maxIterations; i++ {
cr, err := c.completions(ctx, model, messages)
cr, err := c.completions(ctx, model, thinkingBudgetTokens, messages)
if err != nil {
return nil, fmt.Errorf("completions: %w", err)
}
@@ -106,12 +114,18 @@ func (c *Client) Completions(
// completions sends a single non-streaming request and returns the
// full API response.
func (c *Client) completions(
ctx context.Context, model string, messages []Message,
ctx context.Context, model string, thinkingBudgetTokens *int, messages []Message,
) (*ChatResponse, error) {
kwargs := map[string]any{}
if thinkingBudgetTokens != nil {
kwargs["enable_thinking"] = *thinkingBudgetTokens != 0
}
body, err := json.Marshal(ChatRequest{
Model: model,
Messages: messages,
Tools: c.tools,
Model: model,
Messages: messages,
Tools: c.tools,
ThinkingBudgetTokens: thinkingBudgetTokens,
ChatTemplateKWARGS: kwargs,
})
if err != nil {
return nil, fmt.Errorf(
@@ -232,17 +246,15 @@ const (
// means unlimited (no cap).
// Returns ErrMaxIterations if the iteration limit is exceeded.
func (c *Client) CompletionsStream(
ctx context.Context,
model string,
systemMessage string,
maxIterations int,
messages []Message,
events chan<- StreamEvent,
ctx context.Context, p CompletionsParams, events chan<- StreamEvent,
) error {
defer close(events)
ctx, cancel := context.WithTimeout(ctx, c.timeout)
defer cancel()
model, systemMessage, maxIterations, messages, thinkingBudgetTokens :=
p.Model, p.SystemMessage, p.MaxIterations, p.Messages, p.ThinkingBudgetTokens
// Use request system message if provided, config default otherwise.
smsg := c.systemMessage
if systemMessage != "" {
@@ -265,7 +277,7 @@ func (c *Client) CompletionsStream(
// Loop for tool calls. maxIterations == 0 means unlimited.
for i := 0; maxIterations == 0 || i < maxIterations; i++ {
choices, err := c.completionsStream(ctx, messages, model, events)
choices, err := c.completionsStream(ctx, messages, model, thinkingBudgetTokens, events)
if err != nil {
return fmt.Errorf("completions stream: %w", err)
}
@@ -304,13 +316,19 @@ type streamAccum struct {
// response, and returns all assistant messages from the response choices.
// Sends each message to the events channel. Usage and timings are included
// on the done event for the last message.
func (c *Client) completionsStream(ctx context.Context, messages []Message, model string, events chan<- StreamEvent) ([]Message, error) {
func (c *Client) completionsStream(ctx context.Context, messages []Message, model string, thinkingBudgetTokens *int, events chan<- StreamEvent) ([]Message, error) {
kwargs := map[string]any{}
if thinkingBudgetTokens != nil {
kwargs["enable_thinking"] = *thinkingBudgetTokens != 0
}
body, err := json.Marshal(ChatRequest{
Model: model,
Messages: messages,
Stream: true,
Tools: c.tools,
StreamOptions: &StreamOptions{IncludeUsage: true},
Model: model,
Messages: messages,
Stream: true,
Tools: c.tools,
ThinkingBudgetTokens: thinkingBudgetTokens,
ChatTemplateKWARGS: kwargs,
StreamOptions: &StreamOptions{IncludeUsage: true},
})
if err != nil {
return nil, fmt.Errorf(

View File

@@ -35,18 +35,18 @@ func TestCompletionsStreamStreamsReasoningDeltas(t *testing.T) {
errCh := make(chan error, 1)
go func() {
errCh <- client.CompletionsStream(
context.Background(),
"test-model",
"",
0,
[]Message{{
Role: RoleUser,
ContentParts: []ContentPart{{
Type: ContentTypeText,
Text: "hello",
context.Background(), CompletionsParams{
Model: "test-model",
SystemMessage: "",
MaxIterations: 0,
Messages: []Message{{
Role: RoleUser,
ContentParts: []ContentPart{{
Type: ContentTypeText,
Text: "hello",
}},
}},
}},
events,
}, events,
)
}()
@@ -131,17 +131,18 @@ func TestCompletionsReturnsUsageAndTimings(t *testing.T) {
}
res, err := client.Completions(
context.Background(),
"test-model",
"",
0,
[]Message{{
Role: RoleUser,
ContentParts: []ContentPart{{
Type: ContentTypeText,
Text: "hi",
context.Background(), CompletionsParams{
Model: "test-model",
SystemMessage: "",
MaxIterations: 0,
Messages: []Message{{
Role: RoleUser,
ContentParts: []ContentPart{{
Type: ContentTypeText,
Text: "hi",
}},
}},
}},
},
)
if err != nil {
t.Fatalf("Completions: %v", err)
@@ -205,18 +206,18 @@ func TestCompletionsStreamEmitsUsageEvent(t *testing.T) {
errCh := make(chan error, 1)
go func() {
errCh <- client.CompletionsStream(
context.Background(),
"test-model",
"",
0,
[]Message{{
Role: RoleUser,
ContentParts: []ContentPart{{
Type: ContentTypeText,
Text: "hi",
context.Background(), CompletionsParams{
Model: "test-model",
SystemMessage: "",
MaxIterations: 0,
Messages: []Message{{
Role: RoleUser,
ContentParts: []ContentPart{{
Type: ContentTypeText,
Text: "hi",
}},
}},
}},
events,
}, events,
)
}()

View File

@@ -36,27 +36,28 @@ type APITool struct {
// ChatRequest is the request body for chat completions.
type ChatRequest struct {
Model string `json:"model"`
Messages []Message `json:"messages"`
Temperature *float32 `json:"temperature,omitempty"`
TopP *float32 `json:"top_p,omitempty"`
N int `json:"n,omitempty"`
Stream bool `json:"stream,omitempty"`
Stop []string `json:"stop,omitempty"`
MaxTokens int `json:"max_tokens,omitempty"`
MaxCompletionTokens int `json:"max_completion_tokens,omitempty"`
ResponseFormat *ResponseFormat `json:"response_format,omitempty"`
Seed *int `json:"seed,omitempty"`
PresencePenalty *float32 `json:"presence_penalty,omitempty"`
FrequencyPenalty *float32 `json:"frequency_penalty,omitempty"`
LogProbs bool `json:"logprobs,omitempty"`
TopLogProbs int `json:"top_logprobs,omitempty"`
Tools []APITool `json:"tools,omitempty"`
ToolChoice any `json:"tool_choice,omitempty"`
StreamOptions *StreamOptions `json:"stream_options,omitempty"`
ParallelToolCalls *bool `json:"parallel_tool_calls,omitempty"`
ReasoningEffort string `json:"reasoning_effort,omitempty"`
ServiceTier string `json:"service_tier,omitempty"`
Model string `json:"model"`
Messages []Message `json:"messages"`
Temperature *float32 `json:"temperature,omitempty"`
TopP *float32 `json:"top_p,omitempty"`
N int `json:"n,omitempty"`
Stream bool `json:"stream,omitempty"`
Stop []string `json:"stop,omitempty"`
MaxTokens int `json:"max_tokens,omitempty"`
MaxCompletionTokens int `json:"max_completion_tokens,omitempty"`
ResponseFormat *ResponseFormat `json:"response_format,omitempty"`
Seed *int `json:"seed,omitempty"`
PresencePenalty *float32 `json:"presence_penalty,omitempty"`
FrequencyPenalty *float32 `json:"frequency_penalty,omitempty"`
LogProbs bool `json:"logprobs,omitempty"`
TopLogProbs int `json:"top_logprobs,omitempty"`
Tools []APITool `json:"tools,omitempty"`
ToolChoice any `json:"tool_choice,omitempty"`
StreamOptions *StreamOptions `json:"stream_options,omitempty"`
ParallelToolCalls *bool `json:"parallel_tool_calls,omitempty"`
ChatTemplateKWARGS map[string]any `json:"chat_template_kwargs,omitempty"`
ThinkingBudgetTokens *int `json:"thinking_budget_tokens,omitempty"`
ServiceTier string `json:"service_tier,omitempty"`
}
// Usage represents token usage from the API response.

View File

@@ -77,17 +77,18 @@ func TestIntegration_Completions(t *testing.T) {
}
res, err := client.Completions(
context.Background(),
testModel(t),
"",
0,
[]Message{{
Role: RoleUser,
ContentParts: []ContentPart{{
Type: "text",
Text: "Say hello in three words.",
context.Background(), CompletionsParams{
Model: testModel(t),
SystemMessage: "",
MaxIterations: 0,
Messages: []Message{{
Role: RoleUser,
ContentParts: []ContentPart{{
Type: "text",
Text: "Say hello in three words.",
}},
}},
}},
},
)
if err != nil {
t.Fatalf("Completions: %v", err)
@@ -117,18 +118,18 @@ func TestIntegration_CompletionsStream(t *testing.T) {
errCh := make(chan error, 1)
go func() {
errCh <- client.CompletionsStream(
context.Background(),
testModel(t),
"",
0,
[]Message{{
Role: RoleUser,
ContentParts: []ContentPart{{
Type: "text",
Text: "Say hello in three words.",
context.Background(), CompletionsParams{
Model: testModel(t),
SystemMessage: "",
MaxIterations: 0,
Messages: []Message{{
Role: RoleUser,
ContentParts: []ContentPart{{
Type: "text",
Text: "Say hello in three words.",
}},
}},
}},
events,
}, events,
)
}()
@@ -182,17 +183,18 @@ func TestIntegration_ToolCalling(t *testing.T) {
}
res, err := client.Completions(
context.Background(),
testModel(t),
"",
10,
[]Message{{
Role: RoleUser,
ContentParts: []ContentPart{{
Type: "text",
Text: "Use the echo tool to say 'hello'",
context.Background(), CompletionsParams{
Model: testModel(t),
SystemMessage: "",
MaxIterations: 10,
Messages: []Message{{
Role: RoleUser,
ContentParts: []ContentPart{{
Type: "text",
Text: "Use the echo tool to say 'hello'",
}},
}},
}},
},
)
if err != nil {
t.Fatalf("Completions: %v", err)