Skip to content

Commit 726c3f2

Browse files
ctecte
committed
More progress
1 parent eb9c740 commit 726c3f2

File tree

10 files changed

+269
-290
lines changed

10 files changed

+269
-290
lines changed

‎src/api/transform/model-params.ts‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
1-
import type { ApiHandlerOptions, ModelInfo } from "../../shared/api"
1+
import type { ModelInfo, ProviderSettings } from "../../shared/api"
2+
23
import { ANTHROPIC_DEFAULT_MAX_TOKENS } from "../providers/constants"
34

45
export type ModelParams = {
@@ -19,7 +20,7 @@ export function getModelParams({
1920
defaultMaxTokens,
2021
defaultTemperature = 0,
2122
}: {
22-
options: ApiHandlerOptions
23+
options: ProviderSettings
2324
model: ModelInfo
2425
defaultMaxTokens?: number
2526
defaultTemperature?: number

‎src/api/transform/reasoning.ts‎

Lines changed: 5 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@ import { BetaThinkingConfigParam } from "@anthropic-ai/sdk/resources/beta"
22
import OpenAI from "openai"
33

44
import { ModelInfo, ProviderSettings } from "../../schemas"
5+
import { shouldUseReasoningBudget, shouldUseReasoningEffort } from "../../shared/api"
56

67
import type { ModelParams } from "./model-params"
78

@@ -21,21 +22,14 @@ export type GetModelResoningOptions = {
2122
settings: ProviderSettings
2223
}
2324

24-
const shouldUseReasoningBudget = (model: ModelInfo, settings: ProviderSettings, params: ModelParams) =>
25-
(model.requiredReasoningBudget || (model.supportsReasoningBudget && settings.enableReasoningEffort)) &&
26-
params.reasoningBudget
27-
28-
const shouldUseReasoningEffort = (model: ModelInfo, params: ModelParams) =>
29-
model.supportsReasoningEffort && params.reasoningEffort
30-
3125
export const getOpenRouterReasoning = ({
3226
model,
3327
params,
3428
settings,
3529
}: GetModelResoningOptions): OpenRouterReasoningParams | undefined =>
36-
shouldUseReasoningBudget(model, settings, params)
30+
shouldUseReasoningBudget({ model, settings })
3731
? { max_tokens: params.reasoningBudget }
38-
: shouldUseReasoningEffort(model, params)
32+
: shouldUseReasoningEffort({ model, settings })
3933
? { effort: params.reasoningEffort }
4034
: undefined
4135

@@ -44,7 +38,7 @@ export const getAnthropicReasoning = ({
4438
params,
4539
settings,
4640
}: GetModelResoningOptions): AnthropicReasoningParams | undefined =>
47-
shouldUseReasoningBudget(model, settings, params)
41+
shouldUseReasoningBudget({ model, settings })
4842
? { type: "enabled", budget_tokens: params.reasoningBudget! }
4943
: undefined
5044

@@ -53,4 +47,4 @@ export const getOpenAiReasoning = ({
5347
params,
5448
settings,
5549
}: GetModelResoningOptions): OpenAiReasoningParams | undefined =>
56-
shouldUseReasoningEffort(model, params) ? { reasoning_effort: params.reasoningEffort } : undefined
50+
shouldUseReasoningEffort({ model, settings }) ? { reasoning_effort: params.reasoningEffort } : undefined

‎src/shared/__tests__/api.test.ts‎

Lines changed: 154 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,154 @@
1+
// npx jest src/shared/__tests__/api.test.ts
2+
3+
import { type ModelInfo, ProviderSettings, getModelMaxOutputTokens } from "../api"
4+
5+
describe("getMaxTokensForModel", () => {
6+
/**
7+
* Testing the specific fix in commit cc79178f:
8+
* For thinking models, use apiConfig.modelMaxTokens if available,
9+
* otherwise fall back to 8192 (not modelInfo.maxTokens)
10+
*/
11+
12+
it("should return apiConfig.modelMaxTokens for thinking models when provided", () => {
13+
const model: ModelInfo = {
14+
contextWindow: 200_000,
15+
supportsPromptCache: true,
16+
requiredReasoningBudget: true,
17+
maxTokens: 8000,
18+
}
19+
20+
const settings: ProviderSettings = {
21+
modelMaxTokens: 4000,
22+
}
23+
24+
expect(getModelMaxOutputTokens({ model, settings })).toBe(4000)
25+
})
26+
27+
it("should return 16_384 for thinking models when modelMaxTokens not provided", () => {
28+
const model: ModelInfo = {
29+
contextWindow: 200_000,
30+
supportsPromptCache: true,
31+
requiredReasoningBudget: true,
32+
maxTokens: 8000,
33+
}
34+
35+
const settings = {}
36+
37+
expect(getModelMaxOutputTokens({ model, settings })).toBe(16_384)
38+
})
39+
40+
it("should return 16_384 for thinking models when apiConfig is undefined", () => {
41+
const model: ModelInfo = {
42+
contextWindow: 200_000,
43+
supportsPromptCache: true,
44+
requiredReasoningBudget: true,
45+
maxTokens: 8000,
46+
}
47+
48+
expect(getModelMaxOutputTokens({ model, settings: undefined })).toBe(16_384)
49+
})
50+
51+
it("should return modelInfo.maxTokens for non-thinking models", () => {
52+
const model: ModelInfo = {
53+
contextWindow: 200_000,
54+
supportsPromptCache: true,
55+
maxTokens: 8000,
56+
}
57+
58+
const settings: ProviderSettings = {
59+
modelMaxTokens: 4000,
60+
}
61+
62+
expect(getModelMaxOutputTokens({ model, settings })).toBe(8000)
63+
})
64+
65+
it("should return undefined for non-thinking models with undefined maxTokens", () => {
66+
const model: ModelInfo = {
67+
contextWindow: 200_000,
68+
supportsPromptCache: true,
69+
}
70+
71+
const settings: ProviderSettings = {
72+
modelMaxTokens: 4000,
73+
}
74+
75+
expect(getModelMaxOutputTokens({ model, settings })).toBeUndefined()
76+
})
77+
78+
test("should return maxTokens from modelInfo when thinking is false", () => {
79+
const model: ModelInfo = {
80+
contextWindow: 200_000,
81+
supportsPromptCache: true,
82+
maxTokens: 2048,
83+
}
84+
85+
const settings: ProviderSettings = {
86+
modelMaxTokens: 4096,
87+
}
88+
89+
const result = getModelMaxOutputTokens({ model, settings })
90+
expect(result).toBe(2048)
91+
})
92+
93+
test("should return modelMaxTokens from apiConfig when thinking is true", () => {
94+
const model: ModelInfo = {
95+
contextWindow: 200_000,
96+
supportsPromptCache: true,
97+
maxTokens: 2048,
98+
requiredReasoningBudget: true,
99+
}
100+
101+
const settings: ProviderSettings = {
102+
modelMaxTokens: 4096,
103+
}
104+
105+
const result = getModelMaxOutputTokens({ model, settings })
106+
expect(result).toBe(4096)
107+
})
108+
109+
test("should fallback to DEFAULT_THINKING_MODEL_MAX_TOKENS when thinking is true but apiConfig.modelMaxTokens is not defined", () => {
110+
const model: ModelInfo = {
111+
contextWindow: 200_000,
112+
supportsPromptCache: true,
113+
maxTokens: 2048,
114+
requiredReasoningBudget: true,
115+
}
116+
117+
const settings: ProviderSettings = {}
118+
119+
const result = getModelMaxOutputTokens({ model, settings: undefined })
120+
expect(result).toBe(16_384)
121+
})
122+
123+
test("should handle undefined inputs gracefully", () => {
124+
const modelInfoOnly: ModelInfo = {
125+
contextWindow: 200_000,
126+
supportsPromptCache: true,
127+
maxTokens: 2048,
128+
}
129+
130+
expect(getModelMaxOutputTokens({ model: modelInfoOnly, settings: undefined })).toBe(2048)
131+
})
132+
133+
test("should handle missing properties gracefully", () => {
134+
const modelInfoWithoutMaxTokens: ModelInfo = {
135+
contextWindow: 200_000,
136+
supportsPromptCache: true,
137+
requiredReasoningBudget: true,
138+
}
139+
140+
const settings: ProviderSettings = {
141+
modelMaxTokens: 4096,
142+
}
143+
144+
expect(getModelMaxOutputTokens({ model: modelInfoWithoutMaxTokens, settings })).toBe(4096)
145+
146+
const modelInfoWithoutThinking: ModelInfo = {
147+
contextWindow: 200_000,
148+
supportsPromptCache: true,
149+
maxTokens: 2048,
150+
}
151+
152+
expect(getModelMaxOutputTokens({ model: modelInfoWithoutThinking, settings: undefined })).toBe(2048)
153+
})
154+
})

‎src/shared/api.ts‎

Lines changed: 48 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,7 @@ import { ModelInfo, ProviderName, ProviderSettings } from "../schemas"
22

33
export type { ModelInfo, ProviderName, ProviderSettings }
44

5-
export type ApiHandlerOptions = Omit<ProviderSettings, "apiProvider" | "id">
5+
export type ApiHandlerOptions = Omit<ProviderSettings, "apiProvider">
66

77
// Anthropic
88
// https://docs.anthropic.com/en/docs/about-claude/models
@@ -1921,3 +1921,50 @@ export function toRouterName(value?: string): RouterName {
19211921
export type ModelRecord = Record<string, ModelInfo>
19221922

19231923
export type RouterModels = Record<RouterName, ModelRecord>
1924+
1925+
export const shouldUseReasoningBudget = ({
1926+
model,
1927+
settings,
1928+
}: {
1929+
model: ModelInfo
1930+
settings?: ProviderSettings
1931+
}): boolean => !!model.requiredReasoningBudget || (!!model.supportsReasoningBudget && !!settings?.enableReasoningEffort)
1932+
1933+
export const shouldUseReasoningEffort = ({
1934+
model,
1935+
settings,
1936+
}: {
1937+
model: ModelInfo
1938+
settings?: ProviderSettings
1939+
}): boolean => !!model.supportsReasoningEffort && (!!model.reasoningEffort || !!settings?.reasoningEffort)
1940+
1941+
export const DEFAULT_HYBRID_REASONING_MODEL_MAX_TOKENS = 16_384
1942+
export const DEFAULT_HYBRID_REASONING_MODEL_THINKING_TOKENS = 8_192
1943+
1944+
export const getModelMaxOutputTokens = ({
1945+
model,
1946+
settings,
1947+
}: {
1948+
model: ModelInfo
1949+
settings?: ProviderSettings
1950+
}): number | undefined => {
1951+
if (shouldUseReasoningBudget({ model, settings })) {
1952+
return settings?.modelMaxTokens || DEFAULT_HYBRID_REASONING_MODEL_MAX_TOKENS
1953+
}
1954+
1955+
return model.maxTokens ?? undefined
1956+
}
1957+
1958+
export const getModelMaxThinkingTokens = ({
1959+
model,
1960+
settings,
1961+
}: {
1962+
model: ModelInfo
1963+
settings?: ProviderSettings
1964+
}): number | undefined => {
1965+
if (shouldUseReasoningBudget({ model, settings })) {
1966+
return settings?.modelMaxThinkingTokens || DEFAULT_HYBRID_REASONING_MODEL_THINKING_TOKENS
1967+
}
1968+
1969+
return undefined
1970+
}

‎webview-ui/src/__mocks__/components/chat/TaskHeader.tsx‎

Lines changed: 1 addition & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -1,15 +1,3 @@
1-
import React from "react"
2-
// Import the actual utility instead of reimplementing it
3-
import { getMaxTokensForModel } from "@/utils/model-utils"
4-
5-
// Re-export the utility function to maintain the same interface
6-
export { getMaxTokensForModel }
7-
8-
/**
9-
* Mock version of the TaskHeader component
10-
*/
11-
const TaskHeader: React.FC<any> = () => {
12-
return <div data-testid="mocked-task-header">Mocked TaskHeader</div>
13-
}
1+
const TaskHeader = () => <div data-testid="mocked-task-header">Mocked TaskHeader</div>
142

153
export default TaskHeader

‎webview-ui/src/__tests__/getMaxTokensForModel.test.tsx‎

Lines changed: 0 additions & 81 deletions
This file was deleted.

0 commit comments

Comments
 (0)