Skip to content

feat: add custom context window support for AI providers #7210

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
wants to merge 1 commit into from
Closed

feat: add custom context window support for AI providers #7210

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
2 changes: 2 additions & 0 deletions packages/types/src/provider-settings.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -147,6 +147,7 @@ const vertexSchema = apiModelIdProviderModelSchema.extend({
vertexJsonCredentials: z.string().optional(),
vertexProjectId: z.string().optional(),
vertexRegion: z.string().optional(),
vertexCustomModelInfo: modelInfoSchema.nullish(),
})

const openAiSchema = baseProviderSettingsSchema.extend({
Expand Down Expand Up @@ -248,6 +249,7 @@ const xaiSchema = apiModelIdProviderModelSchema.extend({

const groqSchema = apiModelIdProviderModelSchema.extend({
groqApiKey: z.string().optional(),
groqCustomModelInfo: modelInfoSchema.nullish(),
})

const huggingFaceSchema = baseProviderSettingsSchema.extend({
Expand Down
24 changes: 23 additions & 1 deletion src/api/providers/base-openai-compatible-provider.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -130,6 +130,28 @@ export abstract class BaseOpenAiCompatibleProvider<ModelName extends string>
? (this.options.apiModelId as ModelName)
: this.defaultProviderModelId

return { id, info: this.providerModels[id] }
const defaultInfo = this.providerModels[id]

// Check if there's custom model info for this provider
// This allows Groq and other providers to override context window
const customModelInfo = this.getCustomModelInfo()

const info: ModelInfo = customModelInfo
? {
...defaultInfo,
...customModelInfo,
// Ensure required fields are present
maxTokens: customModelInfo.maxTokens ?? defaultInfo.maxTokens,
contextWindow: customModelInfo.contextWindow ?? defaultInfo.contextWindow,
supportsPromptCache: customModelInfo.supportsPromptCache ?? defaultInfo.supportsPromptCache,
}
: defaultInfo

return { id, info }
}

protected getCustomModelInfo(): ModelInfo | undefined {
Copy link
Contributor Author

@roomote roomote bot Aug 19, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Is this the best approach? I notice that OpenAI provider directly accesses this.options.openAiCustomModelInfo while here we're introducing a getCustomModelInfo() method pattern. Should we standardize on one approach across all providers for consistency?

Also, this method lacks JSDoc documentation explaining its purpose and when subclasses should override it.

// Override in subclasses to provide custom model info
return undefined
}
}
6 changes: 5 additions & 1 deletion src/api/providers/groq.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
import { type GroqModelId, groqDefaultModelId, groqModels } from "@roo-code/types"
import { type GroqModelId, type ModelInfo, groqDefaultModelId, groqModels } from "@roo-code/types"

import type { ApiHandlerOptions } from "../../shared/api"

Expand All @@ -16,4 +16,8 @@ export class GroqHandler extends BaseOpenAiCompatibleProvider<GroqModelId> {
defaultTemperature: 0.5,
})
}

protected override getCustomModelInfo(): ModelInfo | undefined {
return this.options.groqCustomModelInfo ?? undefined
Copy link
Contributor Author

@roomote roomote bot Aug 19, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Missing test coverage for this new functionality. Could we add tests to verify that custom context window values properly override the defaults?

}
}
16 changes: 15 additions & 1 deletion src/api/providers/vertex.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,7 +15,21 @@ export class VertexHandler extends GeminiHandler implements SingleCompletionHand
override getModel() {
const modelId = this.options.apiModelId
let id = modelId && modelId in vertexModels ? (modelId as VertexModelId) : vertexDefaultModelId
const info: ModelInfo = vertexModels[id]
const defaultInfo: ModelInfo = vertexModels[id]

// Apply custom model info if provided
const customModelInfo = this.options.vertexCustomModelInfo
Copy link
Contributor Author

@roomote roomote bot Aug 19, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I notice this implementation directly accesses this.options.vertexCustomModelInfo instead of using the getCustomModelInfo() pattern like GroqHandler. Should we be consistent?

Also missing test coverage here.

const info: ModelInfo = customModelInfo
? {
...defaultInfo,
...customModelInfo,
// Ensure required fields are present
maxTokens: customModelInfo.maxTokens ?? defaultInfo.maxTokens,
contextWindow: customModelInfo.contextWindow ?? defaultInfo.contextWindow,
supportsPromptCache: customModelInfo.supportsPromptCache ?? defaultInfo.supportsPromptCache,
}
: defaultInfo

const params = getModelParams({ format: "gemini", modelId: id, model: info, settings: this.options })

// The `:thinking` suffix indicates that the model is a "Hybrid"
Expand Down
61 changes: 61 additions & 0 deletions webview-ui/src/components/common/ContextWindow.tsx
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,61 @@
import { useCallback } from "react"
import { VSCodeTextField } from "@vscode/webview-ui-toolkit/react"

import type { ModelInfo } from "@roo-code/types"

type ContextWindowProps = {
customModelInfo?: ModelInfo | null
defaultContextWindow?: number
onContextWindowChange: (contextWindow: number | undefined) => void
label?: string
placeholder?: string
helperText?: string
}

const inputEventTransform = (event: any) => (event as { target: HTMLInputElement })?.target?.value

export const ContextWindow = ({
Copy link
Contributor Author

@roomote roomote bot Aug 19, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Missing test coverage for this new component. We should add tests for:

  • Rendering with different props
  • Input validation behavior
  • Callback invocation with correct values
  • Edge cases (empty input, invalid numbers, etc.)

customModelInfo,
defaultContextWindow,
onContextWindowChange,
label,
placeholder,
helperText,
}: ContextWindowProps) => {
const handleContextWindowChange = useCallback(
(event: any) => {
const value = inputEventTransform(event)?.trim()

if (value === "") {
// Clear custom context window
onContextWindowChange(undefined)
} else {
const numValue = parseInt(value, 10)
if (!isNaN(numValue) && numValue > 0) {
Copy link
Contributor Author

@roomote roomote bot Aug 19, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The validation only checks if the number is positive, but doesn't validate reasonable upper bounds. Could we add a warning for very large values that might impact performance? The issue description mentions performance degradation with large context windows.

Also consider adding a max value constraint or at least a warning when values exceed certain thresholds (e.g., > 1M tokens).

onContextWindowChange(numValue)
}
}
},
[onContextWindowChange],
)

const currentValue = customModelInfo?.contextWindow?.toString() || ""
const placeholderText = placeholder || defaultContextWindow?.toString() || "128000"
const labelText = label || "Context Window Size"
Copy link
Contributor

@ellipsis-dev ellipsis-dev bot Aug 19, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Use localization instead of hardcoded English strings. Consider passing a translated label and helper text (or use t()) rather than "Context Window Size" and the helper text below.

This comment was generated because it violated a code review rule: irule_C0ez7Rji6ANcGkkX.

Copy link
Contributor Author

@roomote roomote bot Aug 19, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

These hardcoded English strings should use i18n keys for consistency with the rest of the application. The label and helper text should come from translation files.

const helperTextContent = helperText || "Custom context window size in tokens (leave empty to use default)"

return (
<>
<VSCodeTextField
value={currentValue}
onInput={handleContextWindowChange}
placeholder={placeholderText}
className="w-full">
<label className="block font-medium mb-1">{labelText}</label>
</VSCodeTextField>
{helperTextContent && (
<div className="text-sm text-vscode-descriptionForeground -mt-2">{helperTextContent}</div>
)}
</>
)
}
44 changes: 43 additions & 1 deletion webview-ui/src/components/settings/providers/Groq.tsx
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,10 +1,11 @@
import { useCallback } from "react"
import { VSCodeTextField } from "@vscode/webview-ui-toolkit/react"

import type { ProviderSettings } from "@roo-code/types"
import type { ProviderSettings, ModelInfo } from "@roo-code/types"

import { useAppTranslation } from "@src/i18n/TranslationContext"
import { VSCodeButtonLink } from "@src/components/common/VSCodeButtonLink"
import { ContextWindow } from "@src/components/common/ContextWindow"

import { inputEventTransform } from "../transforms"

Expand All @@ -27,6 +28,42 @@ export const Groq = ({ apiConfiguration, setApiConfigurationField }: GroqProps)
[setApiConfigurationField],
)

const handleContextWindowChange = useCallback(
Copy link
Contributor Author

@roomote roomote bot Aug 19, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This entire block (lines 31-65) is duplicated in Vertex.tsx. Could we extract this logic into a shared utility function like createCustomModelInfo(contextWindow, currentModelInfo) to avoid duplication and make maintenance easier?

(contextWindow: number | undefined) => {
const currentModelInfo = apiConfiguration?.groqCustomModelInfo
const updatedModelInfo: ModelInfo | undefined = contextWindow
? {
maxTokens: currentModelInfo?.maxTokens ?? null,
contextWindow,
supportsPromptCache: currentModelInfo?.supportsPromptCache ?? false,
// Preserve other fields if they exist
...(currentModelInfo && {
maxThinkingTokens: currentModelInfo.maxThinkingTokens,
supportsImages: currentModelInfo.supportsImages,
supportsComputerUse: currentModelInfo.supportsComputerUse,
supportsVerbosity: currentModelInfo.supportsVerbosity,
supportsReasoningBudget: currentModelInfo.supportsReasoningBudget,
requiredReasoningBudget: currentModelInfo.requiredReasoningBudget,
supportsReasoningEffort: currentModelInfo.supportsReasoningEffort,
supportedParameters: currentModelInfo.supportedParameters,
inputPrice: currentModelInfo.inputPrice,
outputPrice: currentModelInfo.outputPrice,
cacheWritesPrice: currentModelInfo.cacheWritesPrice,
cacheReadsPrice: currentModelInfo.cacheReadsPrice,
description: currentModelInfo.description,
reasoningEffort: currentModelInfo.reasoningEffort,
minTokensPerCachePoint: currentModelInfo.minTokensPerCachePoint,
maxCachePoints: currentModelInfo.maxCachePoints,
cachableFields: currentModelInfo.cachableFields,
tiers: currentModelInfo.tiers,
}),
}
: undefined
setApiConfigurationField("groqCustomModelInfo", updatedModelInfo)
},
[apiConfiguration?.groqCustomModelInfo, setApiConfigurationField],
)

return (
<>
<VSCodeTextField
Expand All @@ -45,6 +82,11 @@ export const Groq = ({ apiConfiguration, setApiConfigurationField }: GroqProps)
{t("settings:providers.getGroqApiKey")}
</VSCodeButtonLink>
)}
<ContextWindow
customModelInfo={apiConfiguration?.groqCustomModelInfo}
defaultContextWindow={128000}
onContextWindowChange={handleContextWindowChange}
/>
</>
)
}
44 changes: 43 additions & 1 deletion webview-ui/src/components/settings/providers/Vertex.tsx
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,10 +1,11 @@
import { useCallback } from "react"
import { VSCodeLink, VSCodeTextField } from "@vscode/webview-ui-toolkit/react"

import { type ProviderSettings, VERTEX_REGIONS } from "@roo-code/types"
import { type ProviderSettings, type ModelInfo, VERTEX_REGIONS } from "@roo-code/types"

import { useAppTranslation } from "@src/i18n/TranslationContext"
import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from "@src/components/ui"
import { ContextWindow } from "@src/components/common/ContextWindow"

import { inputEventTransform } from "../transforms"

Expand All @@ -27,6 +28,42 @@ export const Vertex = ({ apiConfiguration, setApiConfigurationField }: VertexPro
[setApiConfigurationField],
)

const handleContextWindowChange = useCallback(
Copy link
Contributor Author

@roomote roomote bot Aug 19, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is duplicate code from Groq.tsx. Should be extracted into a shared utility to maintain DRY principles.

(contextWindow: number | undefined) => {
const currentModelInfo = apiConfiguration?.vertexCustomModelInfo
const updatedModelInfo: ModelInfo | undefined = contextWindow
? {
maxTokens: currentModelInfo?.maxTokens ?? null,
contextWindow,
supportsPromptCache: currentModelInfo?.supportsPromptCache ?? false,
// Preserve other fields if they exist
...(currentModelInfo && {
maxThinkingTokens: currentModelInfo.maxThinkingTokens,
supportsImages: currentModelInfo.supportsImages,
supportsComputerUse: currentModelInfo.supportsComputerUse,
supportsVerbosity: currentModelInfo.supportsVerbosity,
supportsReasoningBudget: currentModelInfo.supportsReasoningBudget,
requiredReasoningBudget: currentModelInfo.requiredReasoningBudget,
supportsReasoningEffort: currentModelInfo.supportsReasoningEffort,
supportedParameters: currentModelInfo.supportedParameters,
inputPrice: currentModelInfo.inputPrice,
outputPrice: currentModelInfo.outputPrice,
cacheWritesPrice: currentModelInfo.cacheWritesPrice,
cacheReadsPrice: currentModelInfo.cacheReadsPrice,
description: currentModelInfo.description,
reasoningEffort: currentModelInfo.reasoningEffort,
minTokensPerCachePoint: currentModelInfo.minTokensPerCachePoint,
maxCachePoints: currentModelInfo.maxCachePoints,
cachableFields: currentModelInfo.cachableFields,
tiers: currentModelInfo.tiers,
}),
}
: undefined
setApiConfigurationField("vertexCustomModelInfo", updatedModelInfo)
},
[apiConfiguration?.vertexCustomModelInfo, setApiConfigurationField],
)

return (
<>
<div className="text-sm text-vscode-descriptionForeground">
Expand Down Expand Up @@ -91,6 +128,11 @@ export const Vertex = ({ apiConfiguration, setApiConfigurationField }: VertexPro
</SelectContent>
</Select>
</div>
<ContextWindow
customModelInfo={apiConfiguration?.vertexCustomModelInfo}
defaultContextWindow={128000}
onContextWindowChange={handleContextWindowChange}
/>
</>
)
}
Loading