mirror of
https://github.com/RooVetGit/Roo-Code.git
synced 2026-09-09 22:31:08 +00:00
feat: Adding more settings and control over Gemini
- with topP, topK, maxOutputTokens - allow users to enable URL context and Grounding Research
This commit is contained in:
Ton Hoang Nguyen (Bill)
parent
f18cf3d7ea
commit
5ff59936f9
4 changed files with 210 additions and 9 deletions
|
|
@ -151,6 +151,12 @@ const lmStudioSchema = baseProviderSettingsSchema.extend({
|
|||
const geminiSchema = apiModelIdProviderModelSchema.extend({
|
||||
geminiApiKey: z.string().optional(),
|
||||
googleGeminiBaseUrl: z.string().optional(),
|
||||
topP: z.number().optional(),
|
||||
topK: z.number().optional(),
|
||||
maxOutputTokens: z.number().optional(),
|
||||
enableUrlContext: z.boolean().optional(),
|
||||
enableGrounding: z.boolean().optional(),
|
||||
contextLimit: z.number().optional(),
|
||||
})
|
||||
|
||||
const openAiNativeSchema = apiModelIdProviderModelSchema.extend({
|
||||
|
|
|
|||
72
src/api/providers/__tests__/gemini-handler.spec.ts
Normal file
72
src/api/providers/__tests__/gemini-handler.spec.ts
Normal file
|
|
@ -0,0 +1,72 @@
|
|||
import { describe, it, expect, vi } from "vitest"
|
||||
import { GeminiHandler } from "../gemini"
|
||||
import type { ApiHandlerOptions } from "../../../shared/api"
|
||||
import type { Anthropic } from "@anthropic-ai/sdk"
|
||||
|
||||
describe("GeminiHandler backend support", () => {
|
||||
it("slices messages when contextLimit is set", async () => {
|
||||
const options = { apiProvider: "gemini", contextLimit: 1 } as ApiHandlerOptions
|
||||
const handler = new GeminiHandler(options)
|
||||
const stub = vi.fn().mockReturnValue((async function* () {})())
|
||||
// @ts-ignore access private client
|
||||
handler["client"].models.generateContentStream = stub
|
||||
const messages = [
|
||||
{ role: "user", content: [{ type: "text", text: "first" }] },
|
||||
{ role: "assistant", content: [{ type: "text", text: "second" }] },
|
||||
] as Anthropic.Messages.MessageParam[]
|
||||
for await (const _ of handler.createMessage("instr", messages)) {
|
||||
}
|
||||
expect(stub).toHaveBeenCalledOnce()
|
||||
const params = stub.mock.calls[0][0]
|
||||
expect(params.contents).toHaveLength(1)
|
||||
})
|
||||
|
||||
it("passes maxOutputTokens, topP, topK, and tools for URL context and grounding in config", async () => {
|
||||
const options = {
|
||||
apiProvider: "gemini",
|
||||
maxOutputTokens: 5,
|
||||
topP: 0.5,
|
||||
topK: 10,
|
||||
enableUrlContext: true,
|
||||
enableGrounding: true,
|
||||
} as ApiHandlerOptions
|
||||
const handler = new GeminiHandler(options)
|
||||
const stub = vi.fn().mockReturnValue((async function* () {})())
|
||||
// @ts-ignore access private client
|
||||
handler["client"].models.generateContentStream = stub
|
||||
await handler.createMessage("instr", [] as any).next()
|
||||
const config = stub.mock.calls[0][0].config
|
||||
expect(config.maxOutputTokens).toBe(5)
|
||||
expect(config.topP).toBe(0.5)
|
||||
expect(config.topK).toBe(10)
|
||||
expect(config.tools).toEqual([{ urlContext: {} }, { googleSearch: {} }])
|
||||
})
|
||||
|
||||
it("completePrompt passes config overrides without tools when URL context and grounding disabled", async () => {
|
||||
const options = {
|
||||
apiProvider: "gemini",
|
||||
maxOutputTokens: 7,
|
||||
topP: 0.7,
|
||||
topK: 3,
|
||||
enableUrlContext: false,
|
||||
enableGrounding: false,
|
||||
} as ApiHandlerOptions
|
||||
const handler = new GeminiHandler(options)
|
||||
const stub = vi.fn().mockResolvedValue({ text: "ok" })
|
||||
// @ts-ignore access private client
|
||||
handler["client"].models.generateContent = stub
|
||||
const res = await handler.completePrompt("hi")
|
||||
expect(res).toBe("ok")
|
||||
expect(stub).toHaveBeenCalledWith(
|
||||
expect.objectContaining({
|
||||
config: expect.objectContaining({
|
||||
maxOutputTokens: 7,
|
||||
topP: 0.7,
|
||||
topK: 3,
|
||||
}),
|
||||
}),
|
||||
)
|
||||
const promptConfig = stub.mock.calls[0][0].config
|
||||
expect(promptConfig.tools).toBeUndefined()
|
||||
})
|
||||
})
|
||||
|
|
@ -65,15 +65,27 @@ export class GeminiHandler extends BaseProvider implements SingleCompletionHandl
|
|||
): ApiStream {
|
||||
const { id: model, info, reasoning: thinkingConfig, maxTokens } = this.getModel()
|
||||
|
||||
const contents = messages.map(convertAnthropicMessageToGemini)
|
||||
const limitedMessages = this.options.contextLimit ? messages.slice(-this.options.contextLimit) : messages
|
||||
const contents = limitedMessages.map(convertAnthropicMessageToGemini)
|
||||
|
||||
const config: GenerateContentConfig = {
|
||||
const tools: Array<Record<string, object>> = []
|
||||
if (this.options.enableUrlContext) {
|
||||
tools.push({ urlContext: {} })
|
||||
}
|
||||
if (this.options.enableGrounding) {
|
||||
tools.push({ googleSearch: {} })
|
||||
}
|
||||
const rawConfig = {
|
||||
systemInstruction,
|
||||
httpOptions: this.options.googleGeminiBaseUrl ? { baseUrl: this.options.googleGeminiBaseUrl } : undefined,
|
||||
thinkingConfig,
|
||||
maxOutputTokens: this.options.modelMaxTokens ?? maxTokens ?? undefined,
|
||||
maxOutputTokens: this.options.maxOutputTokens ?? this.options.modelMaxTokens ?? maxTokens ?? undefined,
|
||||
temperature: this.options.modelTemperature ?? 0,
|
||||
topP: this.options.topP,
|
||||
topK: this.options.topK,
|
||||
...(tools.length > 0 ? { tools } : {}),
|
||||
}
|
||||
const config = rawConfig as unknown as GenerateContentConfig
|
||||
|
||||
const params: GenerateContentParameters = { model, contents, config }
|
||||
|
||||
|
|
@ -146,15 +158,29 @@ export class GeminiHandler extends BaseProvider implements SingleCompletionHandl
|
|||
try {
|
||||
const { id: model } = this.getModel()
|
||||
|
||||
const tools: Array<Record<string, object>> = []
|
||||
if (this.options.enableUrlContext) {
|
||||
tools.push({ urlContext: {} })
|
||||
}
|
||||
if (this.options.enableGrounding) {
|
||||
tools.push({ googleSearch: {} })
|
||||
}
|
||||
const rawPromptConfig = {
|
||||
httpOptions: this.options.googleGeminiBaseUrl
|
||||
? { baseUrl: this.options.googleGeminiBaseUrl }
|
||||
: undefined,
|
||||
temperature: this.options.modelTemperature ?? 0,
|
||||
maxOutputTokens: this.options.maxOutputTokens ?? this.options.modelMaxTokens,
|
||||
topP: this.options.topP,
|
||||
topK: this.options.topK,
|
||||
...(tools.length > 0 ? { tools } : {}),
|
||||
}
|
||||
const promptConfig = rawPromptConfig as unknown as GenerateContentConfig
|
||||
|
||||
const result = await this.client.models.generateContent({
|
||||
model,
|
||||
contents: [{ role: "user", parts: [{ text: prompt }] }],
|
||||
config: {
|
||||
httpOptions: this.options.googleGeminiBaseUrl
|
||||
? { baseUrl: this.options.googleGeminiBaseUrl }
|
||||
: undefined,
|
||||
temperature: this.options.modelTemperature ?? 0,
|
||||
},
|
||||
config: promptConfig,
|
||||
})
|
||||
|
||||
return result.text ?? ""
|
||||
|
|
|
|||
|
|
@ -1,6 +1,7 @@
|
|||
import { useCallback, useState } from "react"
|
||||
import { Checkbox } from "vscrui"
|
||||
import { VSCodeTextField } from "@vscode/webview-ui-toolkit/react"
|
||||
import { Slider } from "@src/components/ui"
|
||||
|
||||
import type { ProviderSettings } from "@roo-code/types"
|
||||
|
||||
|
|
@ -72,6 +73,102 @@ export const Gemini = ({ apiConfiguration, setApiConfigurationField }: GeminiPro
|
|||
/>
|
||||
)}
|
||||
</div>
|
||||
<div className="mt-4">
|
||||
<label className="block font-medium mb-1">{t("settings:providers.topP")}</label>
|
||||
<div className="flex items-center space-x-2">
|
||||
<Slider
|
||||
min={0}
|
||||
max={1}
|
||||
step={0.01}
|
||||
value={[apiConfiguration.topP ?? 0]}
|
||||
onValueChange={(values: number[]) => setApiConfigurationField("topP", values[0])}
|
||||
className="flex-grow"
|
||||
/>
|
||||
<span className="w-10 text-right">{(apiConfiguration.topP ?? 0).toFixed(2)}</span>
|
||||
</div>
|
||||
<div className="text-sm text-vscode-descriptionForeground">
|
||||
{t("settings:providers.topPDescription")}
|
||||
</div>
|
||||
</div>
|
||||
<div className="mt-4">
|
||||
<label className="block font-medium mb-1">{t("settings:providers.topK")}</label>
|
||||
<div className="flex items-center space-x-2">
|
||||
<Slider
|
||||
min={0}
|
||||
max={100}
|
||||
step={1}
|
||||
value={[apiConfiguration.topK ?? 0]}
|
||||
onValueChange={(values: number[]) => setApiConfigurationField("topK", values[0])}
|
||||
className="flex-grow"
|
||||
/>
|
||||
<span className="w-10 text-right">{apiConfiguration.topK ?? 0}</span>
|
||||
</div>
|
||||
<div className="text-sm text-vscode-descriptionForeground">
|
||||
{t("settings:providers.topKDescription")}
|
||||
</div>
|
||||
</div>
|
||||
<div className="mt-4">
|
||||
<label className="block font-medium mb-1">{t("settings:providers.maxOutputTokens")}</label>
|
||||
<div className="flex items-center space-x-2">
|
||||
<Slider
|
||||
min={0}
|
||||
max={2048}
|
||||
step={1}
|
||||
value={[apiConfiguration.maxOutputTokens ?? 0]}
|
||||
onValueChange={(values: number[]) => setApiConfigurationField("maxOutputTokens", values[0])}
|
||||
className="flex-grow"
|
||||
/>
|
||||
<VSCodeTextField
|
||||
value={(apiConfiguration.maxOutputTokens ?? 0).toString()}
|
||||
type="text"
|
||||
inputMode="numeric"
|
||||
onInput={handleInputChange("maxOutputTokens", (e) => parseInt((e as any).target.value, 10))}
|
||||
className="w-16"
|
||||
/>
|
||||
</div>
|
||||
<div className="text-sm text-vscode-descriptionForeground">
|
||||
{t("settings:providers.maxOutputTokensDescription")}
|
||||
</div>
|
||||
</div>
|
||||
<Checkbox
|
||||
checked={!!apiConfiguration.enableUrlContext}
|
||||
onChange={(checked: boolean) => setApiConfigurationField("enableUrlContext", checked)}>
|
||||
{t("settings:providers.enableUrlContext")}
|
||||
</Checkbox>
|
||||
<div className="text-sm text-vscode-descriptionForeground mb-2">
|
||||
{t("settings:providers.enableUrlContextDescription")}
|
||||
</div>
|
||||
<Checkbox
|
||||
checked={!!apiConfiguration.enableGrounding}
|
||||
onChange={(checked: boolean) => setApiConfigurationField("enableGrounding", checked)}>
|
||||
{t("settings:providers.enableGrounding")}
|
||||
</Checkbox>
|
||||
<div className="text-sm text-vscode-descriptionForeground mb-2">
|
||||
{t("settings:providers.enableGroundingDescription")}
|
||||
</div>
|
||||
<div className="mt-4">
|
||||
<label className="block font-medium mb-1">{t("settings:providers.contextLimit")}</label>
|
||||
<div className="flex items-center space-x-2">
|
||||
<Slider
|
||||
min={0}
|
||||
max={2048}
|
||||
step={1}
|
||||
value={[apiConfiguration.contextLimit ?? 0]}
|
||||
onValueChange={(values: number[]) => setApiConfigurationField("contextLimit", values[0])}
|
||||
className="flex-grow"
|
||||
/>
|
||||
<VSCodeTextField
|
||||
value={(apiConfiguration.contextLimit ?? 0).toString()}
|
||||
type="text"
|
||||
inputMode="numeric"
|
||||
onInput={handleInputChange("contextLimit", (e) => parseInt((e as any).target.value, 10))}
|
||||
className="w-16"
|
||||
/>
|
||||
</div>
|
||||
<div className="text-sm text-vscode-descriptionForeground">
|
||||
{t("settings:providers.contextLimitDescription")}
|
||||
</div>
|
||||
</div>
|
||||
</>
|
||||
)
|
||||
}
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue