Explorar el Código

Add Gemini 3.1 Pro (Preview) support and set as default Gemini model

Add Gemini 3.1 model entries for Gemini and Vertex providers and set
gemini-3.1-pro-preview as the default Gemini model.

Set supportsNativeTools/defaultToolProtocol for Gemini 3.1 entries to
preserve native tool-calling behavior.

Backfill the customtools alias only when it is missing and base model
info exists, so API-returned alias metadata is not overwritten.

Add regression tests for Gemini, Vertex, and Gemini fetcher alias
handling.

References:
- https://ai.google.dev/gemini-api/docs/models/gemini-3.1-pro-preview
- https://ai.google.dev/gemini-api/docs/pricing
- https://ai.google.dev/gemini-api/docs/thinking
- https://cloud.google.com/vertex-ai/generative-ai/docs/models/gemini/3-1-pro
- https://cloud.google.com/vertex-ai/generative-ai/pricing
- https://cloud.google.com/blog/products/ai-machine-learning/gemini-3-1-pro-on-gemini-cli-gemini-enterprise-and-vertex-ai
- https://deepmind.google/models/model-cards/gemini-3-1-pro/
- https://blog.google/innovation-and-ai/models-and-research/gemini-models/gemini-3-1-pro/
Peter Dave Hello hace 18 horas
padre
commit
34f7bc05d7

+ 5 - 0
.changeset/gemini-3-1-pro-default-model.md

@@ -0,0 +1,5 @@
+---
+"kilo-code": patch
+---
+
+Update Gemini default model metadata for Gemini 3.1 Pro and keep tool calling behavior consistent.

+ 63 - 1
packages/types/src/providers/gemini.ts

@@ -3,9 +3,71 @@ import type { ModelInfo } from "../model.js"
 // https://ai.google.dev/gemini-api/docs/models/gemini
 export type GeminiModelId = keyof typeof geminiModels
 
-export const geminiDefaultModelId: GeminiModelId = "gemini-3-pro-preview"
+export const geminiDefaultModelId: GeminiModelId = "gemini-3.1-pro-preview"
 
 export const geminiModels = {
+	"gemini-3.1-pro-preview": {
+		maxTokens: 65_536,
+		contextWindow: 1_048_576,
+		supportsImages: true,
+		supportsNativeTools: true, // kilocode_change
+		defaultToolProtocol: "native", // kilocode_change
+		supportsPromptCache: true,
+		supportsReasoningEffort: ["low", "medium", "high"],
+		reasoningEffort: "low",
+
+		supportsTemperature: true,
+		defaultTemperature: 1,
+		inputPrice: 4.0,
+		outputPrice: 18.0,
+		cacheReadsPrice: 0.4,
+		cacheWritesPrice: 4.5,
+		tiers: [
+			{
+				contextWindow: 200_000,
+				inputPrice: 2.0,
+				outputPrice: 12.0,
+				cacheReadsPrice: 0.2,
+			},
+			{
+				contextWindow: Infinity,
+				inputPrice: 4.0,
+				outputPrice: 18.0,
+				cacheReadsPrice: 0.4,
+			},
+		],
+	},
+	"gemini-3.1-pro-preview-customtools": {
+		maxTokens: 65_536,
+		contextWindow: 1_048_576,
+		supportsImages: true,
+		supportsNativeTools: true, // kilocode_change
+		defaultToolProtocol: "native", // kilocode_change
+		supportsPromptCache: true,
+		supportsReasoningEffort: ["low", "medium", "high"],
+		reasoningEffort: "low",
+
+		supportsTemperature: true,
+		defaultTemperature: 1,
+		inputPrice: 4.0,
+		outputPrice: 18.0,
+		cacheReadsPrice: 0.4,
+		cacheWritesPrice: 4.5,
+		tiers: [
+			{
+				contextWindow: 200_000,
+				inputPrice: 2.0,
+				outputPrice: 12.0,
+				cacheReadsPrice: 0.2,
+			},
+			{
+				contextWindow: Infinity,
+				inputPrice: 4.0,
+				outputPrice: 18.0,
+				cacheReadsPrice: 0.4,
+			},
+		],
+	},
 	"gemini-3-pro-preview": {
 		maxTokens: 65_536,
 		contextWindow: 1_048_576,

+ 31 - 0
packages/types/src/providers/vertex.ts

@@ -22,6 +22,37 @@ export const vertexModels = {
 		supportsVerbosity: ["low", "medium", "high", "max"],
 	},
 	// kilocode_change end
+	"gemini-3.1-pro-preview": {
+		maxTokens: 65_536,
+		contextWindow: 1_048_576,
+		supportsImages: true,
+		supportsNativeTools: true, // kilocode_change
+		defaultToolProtocol: "native", // kilocode_change
+		supportsPromptCache: true,
+		supportsReasoningEffort: ["low", "medium", "high"],
+		reasoningEffort: "low",
+
+		supportsTemperature: true,
+		defaultTemperature: 1,
+		inputPrice: 4.0,
+		outputPrice: 18.0,
+		cacheReadsPrice: 0.4,
+		cacheWritesPrice: 4.5,
+		tiers: [
+			{
+				contextWindow: 200_000,
+				inputPrice: 2.0,
+				outputPrice: 12.0,
+				cacheReadsPrice: 0.2,
+			},
+			{
+				contextWindow: Infinity,
+				inputPrice: 4.0,
+				outputPrice: 18.0,
+				cacheReadsPrice: 0.4,
+			},
+		],
+	},
 	"gemini-3-pro-preview": {
 		maxTokens: 65_536,
 		contextWindow: 1_048_576,

+ 55 - 0
src/api/providers/__tests__/gemini.spec.ts

@@ -181,6 +181,8 @@ describe("GeminiHandler", () => {
 			const modelInfo = handler.getModel()
 			expect(modelInfo.id).toBe(GEMINI_MODEL_NAME)
 			expect(modelInfo.info).toBeDefined()
+			expect(modelInfo.info.supportsNativeTools).toBe(true) // kilocode_change
+			expect(modelInfo.info.defaultToolProtocol).toBe("native") // kilocode_change
 		})
 
 		it("should return default model if invalid model specified", () => {
@@ -191,6 +193,59 @@ describe("GeminiHandler", () => {
 			const modelInfo = invalidHandler.getModel()
 			expect(modelInfo.id).toBe(geminiDefaultModelId) // Default model
 		})
+
+		// kilocode_change start
+		it("should preserve customtools model alias after dynamic model loading", async () => {
+			const customToolsModelId = "gemini-3.1-pro-preview-customtools"
+			getGeminiModelsMock.mockResolvedValueOnce({
+				"gemini-3.1-pro-preview": {
+					maxTokens: 65_536,
+					contextWindow: 1_048_576,
+					supportsNativeTools: true,
+					defaultToolProtocol: "native",
+					supportsPromptCache: true,
+				},
+				[customToolsModelId]: {
+					maxTokens: 65_536,
+					contextWindow: 1_048_576,
+					supportsNativeTools: true,
+					defaultToolProtocol: "native",
+					supportsPromptCache: true,
+				},
+			})
+
+			const customToolsHandler = new GeminiHandler({
+				apiModelId: customToolsModelId,
+				geminiApiKey: "test-key",
+			})
+
+			const mockGenerateContentStream = vitest.fn().mockResolvedValue({
+				[Symbol.asyncIterator]: async function* () {
+					yield { text: "ok" }
+					yield { usageMetadata: { promptTokenCount: 1, candidatesTokenCount: 1 } }
+				},
+			})
+
+			customToolsHandler["client"] = {
+				models: {
+					generateContentStream: mockGenerateContentStream,
+					generateContent: vitest.fn(),
+					getGenerativeModel: vitest.fn(),
+				},
+			} as any
+
+			const stream = customToolsHandler.createMessage("sys", [{ role: "user", content: "hello" }])
+			for await (const _chunk of stream) {
+				// Drain stream
+			}
+
+			expect(mockGenerateContentStream).toHaveBeenCalledWith(
+				expect.objectContaining({
+					model: customToolsModelId,
+				}),
+			)
+		})
+		// kilocode_change end
 	})
 
 	describe("calculateCost", () => {

+ 15 - 0
src/api/providers/__tests__/vertex.spec.ts

@@ -149,5 +149,20 @@ describe("VertexHandler", () => {
 			expect(modelInfo.info.maxTokens).toBe(8192)
 			expect(modelInfo.info.contextWindow).toBe(1048576)
 		})
+
+		// kilocode_change start
+		it("should expose native tools metadata for gemini-3.1-pro-preview", () => {
+			const testHandler = new VertexHandler({
+				apiModelId: "gemini-3.1-pro-preview",
+				vertexProjectId: "test-project",
+				vertexRegion: "us-central1",
+			})
+
+			const modelInfo = testHandler.getModel()
+			expect(modelInfo.id).toBe("gemini-3.1-pro-preview")
+			expect(modelInfo.info.supportsNativeTools).toBe(true)
+			expect(modelInfo.info.defaultToolProtocol).toBe("native")
+		})
+		// kilocode_change end
 	})
 })

+ 101 - 0
src/api/providers/fetchers/__tests__/gemini.spec.ts

@@ -0,0 +1,101 @@
+// kilocode_change - new file
+import { beforeEach, describe, expect, it, vi } from "vitest"
+
+const listModelsMock = vi.hoisted(() => vi.fn())
+const googleGenAICtorMock = vi.hoisted(() =>
+	vi.fn(() => ({
+		models: {
+			list: listModelsMock,
+		},
+	})),
+)
+
+vi.mock("@google/genai", () => ({
+	GoogleGenAI: googleGenAICtorMock,
+}))
+
+import { getGeminiModels } from "../gemini"
+
+const createPager = (models: Array<Record<string, any>>) => ({
+	async *[Symbol.asyncIterator]() {
+		for (const model of models) {
+			yield model
+		}
+	},
+})
+
+describe("getGeminiModels", () => {
+	beforeEach(() => {
+		vi.clearAllMocks()
+	})
+
+	// kilocode_change start
+	it("adds customtools alias when gemini-3.1-pro-preview is available", async () => {
+		listModelsMock.mockResolvedValueOnce(
+			createPager([
+				{
+					name: "models/gemini-3.1-pro-preview",
+					inputTokenLimit: 222_222,
+					outputTokenLimit: 11_111,
+					displayName: "Gemini 3.1 Pro",
+				},
+			]),
+		)
+
+		const models = await getGeminiModels({ apiKey: "test-key" })
+
+		expect(googleGenAICtorMock).toHaveBeenCalledWith({ apiKey: "test-key" })
+		expect(models["gemini-3.1-pro-preview"]).toBeDefined()
+		expect(models["gemini-3.1-pro-preview-customtools"]).toMatchObject({
+			supportsNativeTools: true,
+			defaultToolProtocol: "native",
+			contextWindow: 222_222,
+			maxTokens: 11_111,
+		})
+	})
+
+	it("does not add customtools alias when gemini-3.1-pro-preview is unavailable", async () => {
+		listModelsMock.mockResolvedValueOnce(
+			createPager([
+				{
+					name: "models/gemini-2.5-pro",
+					inputTokenLimit: 1_048_576,
+					outputTokenLimit: 65_536,
+				},
+			]),
+		)
+
+		const models = await getGeminiModels({ apiKey: "test-key" })
+
+		expect(models["gemini-3.1-pro-preview-customtools"]).toBeUndefined()
+	})
+
+	it("does not overwrite API metadata when customtools alias is returned by API", async () => {
+		listModelsMock.mockResolvedValueOnce(
+			createPager([
+				{
+					name: "models/gemini-3.1-pro-preview",
+					inputTokenLimit: 1_048_576,
+					outputTokenLimit: 65_536,
+				},
+				{
+					name: "models/gemini-3.1-pro-preview-customtools",
+					inputTokenLimit: 333_333,
+					outputTokenLimit: 12_345,
+					displayName: "Gemini 3.1 Pro Custom Tools from API",
+					description: "API alias description",
+				},
+			]),
+		)
+
+		const models = await getGeminiModels({ apiKey: "test-key" })
+		const alias = models["gemini-3.1-pro-preview-customtools"]
+
+		expect(alias).toBeDefined()
+		expect(alias.displayName).toBe("Gemini 3.1 Pro Custom Tools from API")
+		expect(alias.description).toBe("API alias description")
+		expect(alias.maxTokens).toBe(12_345)
+		expect(alias.contextWindow).toBe(333_333)
+	})
+	// kilocode_change end
+})

+ 26 - 0
src/api/providers/fetchers/gemini.ts

@@ -10,6 +10,12 @@ const STATIC_MODELS: Record<string, ModelInfo> = geminiModels as Record<string,
 
 const SUPPORTED_MODEL_PREFIXES = ["gemini-", "learnlm-"]
 
+// kilocode_change start
+const MODEL_ALIASES: ReadonlyArray<{ alias: string; base: string }> = [
+	{ alias: "gemini-3.1-pro-preview-customtools", base: "gemini-3.1-pro-preview" },
+]
+// kilocode_change end
+
 interface GeminiFetcherOptions {
 	apiKey?: string
 	baseUrl?: string
@@ -123,6 +129,26 @@ export const getGeminiModels = async ({ apiKey, baseUrl }: GeminiFetcherOptions
 			}
 		}
 
+		// kilocode_change start
+		// Include static aliases (e.g. *-customtools) when their base model is available.
+		for (const { alias, base } of MODEL_ALIASES) {
+			const staticInfo = STATIC_MODELS[alias]
+			if (!staticInfo) {
+				continue
+			}
+
+			const baseInfo = models[base]
+			if (!models[alias] && baseInfo) {
+				models[alias] = createModelInfo({
+					id: alias,
+					staticInfo,
+					inputTokenLimit: typeof baseInfo?.contextWindow === "number" ? baseInfo.contextWindow : undefined,
+					outputTokenLimit: typeof baseInfo?.maxTokens === "number" ? baseInfo.maxTokens : undefined,
+				})
+			}
+		}
+		// kilocode_change end
+
 		if (!Object.keys(models).length) {
 			console.debug("[getGeminiModels] No models returned from API, falling back to static list")
 			return { ...STATIC_MODELS }