Skip to content

Commit a2d441c

Browse files
ctecte
authored
Merge pull request RooCodeInc#1307 from RooVetGit/cte/dry-get-model
DRY up getModel
2 parents 2c0b162 + 9b646eb commit a2d441c

File tree

8 files changed

+341
-121
lines changed

8 files changed

+341
-121
lines changed

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

Lines changed: 257 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,257 @@
1+
// npx jest src/api/__tests__/index.test.ts
2+
3+
import { BetaThinkingConfigParam } from "@anthropic-ai/sdk/resources/beta/messages/index.mjs"
4+
5+
import { getModelParams } from "../index"
6+
import { ANTHROPIC_DEFAULT_MAX_TOKENS } from "../providers/constants"
7+
8+
describe("getModelParams", () => {
9+
it("should return default values when no custom values are provided", () => {
10+
const options = {}
11+
const model = {
12+
id: "test-model",
13+
contextWindow: 16000,
14+
supportsPromptCache: true,
15+
}
16+
17+
const result = getModelParams({
18+
options,
19+
model,
20+
defaultMaxTokens: 1000,
21+
defaultTemperature: 0.5,
22+
})
23+
24+
expect(result).toEqual({
25+
maxTokens: 1000,
26+
thinking: undefined,
27+
temperature: 0.5,
28+
})
29+
})
30+
31+
it("should use custom temperature from options when provided", () => {
32+
const options = { modelTemperature: 0.7 }
33+
const model = {
34+
id: "test-model",
35+
contextWindow: 16000,
36+
supportsPromptCache: true,
37+
}
38+
39+
const result = getModelParams({
40+
options,
41+
model,
42+
defaultMaxTokens: 1000,
43+
defaultTemperature: 0.5,
44+
})
45+
46+
expect(result).toEqual({
47+
maxTokens: 1000,
48+
thinking: undefined,
49+
temperature: 0.7,
50+
})
51+
})
52+
53+
it("should use model maxTokens when available", () => {
54+
const options = {}
55+
const model = {
56+
id: "test-model",
57+
maxTokens: 2000,
58+
contextWindow: 16000,
59+
supportsPromptCache: true,
60+
}
61+
62+
const result = getModelParams({
63+
options,
64+
model,
65+
defaultMaxTokens: 1000,
66+
})
67+
68+
expect(result).toEqual({
69+
maxTokens: 2000,
70+
thinking: undefined,
71+
temperature: 0,
72+
})
73+
})
74+
75+
it("should handle thinking models correctly", () => {
76+
const options = {}
77+
const model = {
78+
id: "test-model",
79+
thinking: true,
80+
maxTokens: 2000,
81+
contextWindow: 16000,
82+
supportsPromptCache: true,
83+
}
84+
85+
const result = getModelParams({
86+
options,
87+
model,
88+
})
89+
90+
const expectedThinking: BetaThinkingConfigParam = {
91+
type: "enabled",
92+
budget_tokens: 1600, // 80% of 2000
93+
}
94+
95+
expect(result).toEqual({
96+
maxTokens: 2000,
97+
thinking: expectedThinking,
98+
temperature: 1.0, // Thinking models require temperature 1.0.
99+
})
100+
})
101+
102+
it("should honor customMaxTokens for thinking models", () => {
103+
const options = { modelMaxTokens: 3000 }
104+
const model = {
105+
id: "test-model",
106+
thinking: true,
107+
contextWindow: 16000,
108+
supportsPromptCache: true,
109+
}
110+
111+
const result = getModelParams({
112+
options,
113+
model,
114+
defaultMaxTokens: 2000,
115+
})
116+
117+
const expectedThinking: BetaThinkingConfigParam = {
118+
type: "enabled",
119+
budget_tokens: 2400, // 80% of 3000
120+
}
121+
122+
expect(result).toEqual({
123+
maxTokens: 3000,
124+
thinking: expectedThinking,
125+
temperature: 1.0,
126+
})
127+
})
128+
129+
it("should honor customMaxThinkingTokens for thinking models", () => {
130+
const options = { modelMaxThinkingTokens: 1500 }
131+
const model = {
132+
id: "test-model",
133+
thinking: true,
134+
maxTokens: 4000,
135+
contextWindow: 16000,
136+
supportsPromptCache: true,
137+
}
138+
139+
const result = getModelParams({
140+
options,
141+
model,
142+
})
143+
144+
const expectedThinking: BetaThinkingConfigParam = {
145+
type: "enabled",
146+
budget_tokens: 1500, // Using the custom value
147+
}
148+
149+
expect(result).toEqual({
150+
maxTokens: 4000,
151+
thinking: expectedThinking,
152+
temperature: 1.0,
153+
})
154+
})
155+
156+
it("should not honor customMaxThinkingTokens for non-thinking models", () => {
157+
const options = { modelMaxThinkingTokens: 1500 }
158+
const model = {
159+
id: "test-model",
160+
maxTokens: 4000,
161+
contextWindow: 16000,
162+
supportsPromptCache: true,
163+
// Note: model.thinking is not set (so it's falsey).
164+
}
165+
166+
const result = getModelParams({
167+
options,
168+
model,
169+
})
170+
171+
expect(result).toEqual({
172+
maxTokens: 4000,
173+
thinking: undefined, // Should remain undefined despite customMaxThinkingTokens being set.
174+
temperature: 0, // Using default temperature.
175+
})
176+
})
177+
178+
it("should clamp thinking budget to at least 1024 tokens", () => {
179+
const options = { modelMaxThinkingTokens: 500 }
180+
const model = {
181+
id: "test-model",
182+
thinking: true,
183+
maxTokens: 2000,
184+
contextWindow: 16000,
185+
supportsPromptCache: true,
186+
}
187+
188+
const result = getModelParams({
189+
options,
190+
model,
191+
})
192+
193+
const expectedThinking: BetaThinkingConfigParam = {
194+
type: "enabled",
195+
budget_tokens: 1024, // Minimum is 1024
196+
}
197+
198+
expect(result).toEqual({
199+
maxTokens: 2000,
200+
thinking: expectedThinking,
201+
temperature: 1.0,
202+
})
203+
})
204+
205+
it("should clamp thinking budget to at most 80% of max tokens", () => {
206+
const options = { modelMaxThinkingTokens: 5000 }
207+
const model = {
208+
id: "test-model",
209+
thinking: true,
210+
maxTokens: 4000,
211+
contextWindow: 16000,
212+
supportsPromptCache: true,
213+
}
214+
215+
const result = getModelParams({
216+
options,
217+
model,
218+
})
219+
220+
const expectedThinking: BetaThinkingConfigParam = {
221+
type: "enabled",
222+
budget_tokens: 3200, // 80% of 4000
223+
}
224+
225+
expect(result).toEqual({
226+
maxTokens: 4000,
227+
thinking: expectedThinking,
228+
temperature: 1.0,
229+
})
230+
})
231+
232+
it("should use ANTHROPIC_DEFAULT_MAX_TOKENS when no maxTokens is provided for thinking models", () => {
233+
const options = {}
234+
const model = {
235+
id: "test-model",
236+
thinking: true,
237+
contextWindow: 16000,
238+
supportsPromptCache: true,
239+
}
240+
241+
const result = getModelParams({
242+
options,
243+
model,
244+
})
245+
246+
const expectedThinking: BetaThinkingConfigParam = {
247+
type: "enabled",
248+
budget_tokens: Math.floor(ANTHROPIC_DEFAULT_MAX_TOKENS * 0.8),
249+
}
250+
251+
expect(result).toEqual({
252+
maxTokens: undefined,
253+
thinking: expectedThinking,
254+
temperature: 1.0,
255+
})
256+
})
257+
})

‎src/api/index.ts‎

Lines changed: 42 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,9 @@
11
import { Anthropic } from "@anthropic-ai/sdk"
2+
import { BetaThinkingConfigParam } from "@anthropic-ai/sdk/resources/beta/messages/index.mjs"
3+
4+
import { ApiConfiguration, ModelInfo, ApiHandlerOptions } from "../shared/api"
5+
import { ANTHROPIC_DEFAULT_MAX_TOKENS } from "./providers/constants"
26
import { GlamaHandler } from "./providers/glama"
3-
import { ApiConfiguration, ModelInfo } from "../shared/api"
47
import { AnthropicHandler } from "./providers/anthropic"
58
import { AwsBedrockHandler } from "./providers/bedrock"
69
import { OpenRouterHandler } from "./providers/openrouter"
@@ -63,3 +66,41 @@ export function buildApiHandler(configuration: ApiConfiguration): ApiHandler {
6366
return new AnthropicHandler(options)
6467
}
6568
}
69+
70+
export function getModelParams({
71+
options,
72+
model,
73+
defaultMaxTokens,
74+
defaultTemperature = 0,
75+
}: {
76+
options: ApiHandlerOptions
77+
model: ModelInfo
78+
defaultMaxTokens?: number
79+
defaultTemperature?: number
80+
}) {
81+
const {
82+
modelMaxTokens: customMaxTokens,
83+
modelMaxThinkingTokens: customMaxThinkingTokens,
84+
modelTemperature: customTemperature,
85+
} = options
86+
87+
let maxTokens = model.maxTokens ?? defaultMaxTokens
88+
let thinking: BetaThinkingConfigParam | undefined = undefined
89+
let temperature = customTemperature ?? defaultTemperature
90+
91+
if (model.thinking) {
92+
// Only honor `customMaxTokens` for thinking models.
93+
maxTokens = customMaxTokens ?? maxTokens
94+
95+
// Clamp the thinking budget to be at most 80% of max tokens and at
96+
// least 1024 tokens.
97+
const maxBudgetTokens = Math.floor((maxTokens || ANTHROPIC_DEFAULT_MAX_TOKENS) * 0.8)
98+
const budgetTokens = Math.max(Math.min(customMaxThinkingTokens ?? maxBudgetTokens, maxBudgetTokens), 1024)
99+
thinking = { type: "enabled", budget_tokens: budgetTokens }
100+
101+
// Anthropic "Thinking" models require a temperature of 1.0.
102+
temperature = 1.0
103+
}
104+
105+
return { maxTokens, thinking, temperature }
106+
}

0 commit comments

Comments
 (0)