diff --git a/.changeset/direct-fim-providers.md b/.changeset/direct-fim-providers.md new file mode 100644 index 00000000000..d389af48ad1 --- /dev/null +++ b/.changeset/direct-fim-providers.md @@ -0,0 +1,5 @@ +--- +"kilo-code": patch +--- + +Use connected Mistral and Inception provider API keys for autocomplete FIM requests before falling back to Kilo Gateway. diff --git a/packages/kilo-docs/source-links.md b/packages/kilo-docs/source-links.md index bf6b1ebceef..11ceece8a47 100644 --- a/packages/kilo-docs/source-links.md +++ b/packages/kilo-docs/source-links.md @@ -1,13 +1,17 @@ # Source Code Links - + - +- + - +- + - - @@ -32,6 +36,8 @@ - +- + - - diff --git a/packages/kilo-vscode/package.json b/packages/kilo-vscode/package.json index ca633b3ea66..082adf9bd71 100644 --- a/packages/kilo-vscode/package.json +++ b/packages/kilo-vscode/package.json @@ -753,11 +753,15 @@ "type": "string", "enum": [ "mistralai/codestral-2508", - "inception/mercury-edit-2" + "inception/mercury-edit-2", + "mistral/codestral-2508", + "inception-direct/mercury-edit-2" ], "enumDescriptions": [ - "Codestral by Mistral AI (default)", - "Mercury Edit 2 by Inception" + "Codestral via Kilo Gateway (default)", + "Mercury Edit 2 via Kilo Gateway", + "Codestral via your connected Mistral provider API key", + "Mercury Edit 2 via your connected Inception provider API key" ], "description": "Model to use for inline autocomplete suggestions" }, diff --git a/packages/kilo-vscode/src/services/autocomplete/AutocompleteServiceManager.ts b/packages/kilo-vscode/src/services/autocomplete/AutocompleteServiceManager.ts index a272f677372..3484edf7ae6 100644 --- a/packages/kilo-vscode/src/services/autocomplete/AutocompleteServiceManager.ts +++ b/packages/kilo-vscode/src/services/autocomplete/AutocompleteServiceManager.ts @@ -319,7 +319,7 @@ export class AutocompleteServiceManager { } private getCurrentModelName(): string { - return this.inlineCompletionProvider.getModelId() + return getAutocompleteModel(this.inlineCompletionProvider.getModelId()).label } private getCurrentProviderName(): string { diff --git a/packages/kilo-vscode/src/services/autocomplete/classic-auto-complete/FillInTheMiddle.ts b/packages/kilo-vscode/src/services/autocomplete/classic-auto-complete/FillInTheMiddle.ts index 6a36ae3ab32..564423b2f18 100644 --- a/packages/kilo-vscode/src/services/autocomplete/classic-auto-complete/FillInTheMiddle.ts +++ b/packages/kilo-vscode/src/services/autocomplete/classic-auto-complete/FillInTheMiddle.ts @@ -65,36 +65,14 @@ export class FimPromptBuilder { signal?: AbortSignal, ): Promise { const { formattedPrefix, prunedSuffix, autocompleteInput } = prompt - let perflog = "" - const logtime = (() => { - let timestamp = performance.now() - return (msg: string) => { - const baseline = timestamp - timestamp = performance.now() - perflog += `${msg}: ${timestamp - baseline}\n` - } - })() - - logtime("snippets") - - console.log("[FIM] formattedPrefix:", formattedPrefix) - let response = "" const onChunk = (text: string) => { response += text } - logtime("prep fim") const usageInfo = await generateFim(connection, modelId, formattedPrefix, prunedSuffix, onChunk, signal) - logtime("fim network") - console.log("[FIM] response:", response) const fillInAtCursorSuggestion = processSuggestion(response) - if (fillInAtCursorSuggestion.text) { - console.info("Final FIM suggestion:", fillInAtCursorSuggestion) - } - logtime("processSuggestion") - console.log(perflog + `lengths: ${formattedPrefix.length + prunedSuffix.length}\n`) return { suggestion: fillInAtCursorSuggestion, cost: usageInfo.cost, diff --git a/packages/kilo-vscode/src/services/autocomplete/fim.ts b/packages/kilo-vscode/src/services/autocomplete/fim.ts index 8e5fcb429c6..2266f4aa055 100644 --- a/packages/kilo-vscode/src/services/autocomplete/fim.ts +++ b/packages/kilo-vscode/src/services/autocomplete/fim.ts @@ -2,6 +2,173 @@ import { ResponseMetaData } from "./types" import type { KiloConnectionService } from "../cli-backend" import { getAutocompleteModel } from "../../shared/autocomplete-models" +const FIM_MAX_TOKENS = 256 +const MISTRAL_FIM_URL = "https://api.mistral.ai/v1/fim/completions" +const CODESTRAL_FIM_URL = "https://codestral.mistral.ai/v1/fim/completions" +const INCEPTION_FIM_URL = "https://api.inceptionlabs.ai/v1/fim/completions" + +type FimProvider = "mistral" | "inception" + +interface DirectFimTarget { + provider: FimProvider + model: string + urls: string[] +} + +interface ProviderItem { + id: string + key?: string + options?: Record +} + +interface ProviderListResponse { + all: ProviderItem[] +} + +interface DirectFimChunk { + choices?: Array<{ + delta?: { content?: string } + text?: string + }> + usage?: { + prompt_tokens?: number + completion_tokens?: number + } +} + +interface DirectFimOptions { + apiKey: string + target: DirectFimTarget + prefix: string + suffix: string + temperature: number + onChunk: (text: string) => void + signal?: AbortSignal + fetchImpl?: typeof fetch +} + +export function getDirectFimTarget(model: string): DirectFimTarget | null { + const info = getAutocompleteModel(model) + if (info.directProvider === "mistral") { + return { provider: "mistral", model: info.requestModel, urls: [MISTRAL_FIM_URL, CODESTRAL_FIM_URL] } + } + if (info.directProvider === "inception") { + return { provider: "inception", model: info.requestModel, urls: [INCEPTION_FIM_URL] } + } + return null +} + +async function resolveProviderKey(connectionService: KiloConnectionService, provider: FimProvider): Promise { + const client = await connectionService.getClientAsync() + const result = await client.provider.list({}, { throwOnError: true }) + const data = result.data as ProviderListResponse | undefined + const item = data?.all.find((p) => p.id === provider) + const key = item?.options?.apiKey + return item?.key ?? (typeof key === "string" ? key : null) +} + +function extractDirectFimContent(chunk: DirectFimChunk): string { + const choice = chunk.choices?.[0] + return choice?.delta?.content ?? choice?.text ?? "" +} + +function parseDirectFimEvent(data: string): DirectFimChunk | null { + if (data === "[DONE]") return null + return JSON.parse(data) as DirectFimChunk +} + +function parseDirectFimLine(line: string): string | null { + const trimmed = line.trim() + if (!trimmed.startsWith("data:")) return null + return trimmed.slice("data:".length).trim() +} + +function handleDirectFimLine( + line: string, + onChunk: (text: string) => void, + usage: { inputTokens: number; outputTokens: number }, +) { + const data = parseDirectFimLine(line) + if (!data) return + const event = parseDirectFimEvent(data) + if (!event) return + const content = extractDirectFimContent(event) + if (content) onChunk(content) + usage.inputTokens = event.usage?.prompt_tokens ?? usage.inputTokens + usage.outputTokens = event.usage?.completion_tokens ?? usage.outputTokens +} + +export async function generateDirectFim(options: DirectFimOptions): Promise { + const urls = [...options.target.urls] + const [url] = urls + if (!url) throw new Error("FIM request failed: 500 missing provider endpoint") + return generateDirectFimWithUrl(options, url, urls.slice(1)) +} + +async function generateDirectFimWithUrl( + options: DirectFimOptions, + url: string, + fallbacks: string[], +): Promise { + const fetchImpl = options.fetchImpl ?? fetch + console.info(`[FIM] request provider=${options.target.provider} model=${options.target.model} url=${url}`) + const res = await fetchImpl(url, { + method: "POST", + headers: { + "Content-Type": "application/json", + Authorization: `Bearer ${options.apiKey}`, + }, + body: JSON.stringify({ + model: options.target.model, + prompt: options.prefix, + suffix: options.suffix, + max_tokens: FIM_MAX_TOKENS, + temperature: options.temperature, + stream: true, + }), + signal: options.signal, + }) + + if (!res.ok) { + const body = await res.text().catch(() => "") + const [next] = fallbacks + if (res.status === 401 && next) return generateDirectFimWithUrl(options, next, fallbacks.slice(1)) + throw new Error(`FIM request failed: ${res.status} ${res.statusText}: ${body}`) + } + + if (!res.body) throw new Error("FIM request failed: 500 empty response body") + + const reader = res.body.getReader() + const decoder = new TextDecoder() + const usage = { inputTokens: 0, outputTokens: 0 } + let pending = "" + + while (true) { + const chunk = await reader.read() + if (chunk.done) break + pending += decoder.decode(chunk.value, { stream: true }) + const lines = pending.split("\n") + pending = lines.pop() ?? "" + + for (const line of lines) { + handleDirectFimLine(line, options.onChunk, usage) + } + } + + pending += decoder.decode() + handleDirectFimLine(pending, options.onChunk, usage) + + reader.releaseLock() + + return { + cost: 0, + inputTokens: usage.inputTokens, + outputTokens: usage.outputTokens, + cacheWriteTokens: 0, + cacheReadTokens: 0, + } +} + /** * Generate a FIM (Fill-in-the-Middle) completion via the CLI backend. * Uses the SDK's kilo.fim() SSE endpoint which handles auth and streaming. @@ -17,6 +184,25 @@ export async function generateFim( signal?: AbortSignal, ): Promise { const client = await connectionService.getClientAsync() + const info = getAutocompleteModel(modelId) + const target = getDirectFimTarget(modelId) + const key = info.directProvider ? await resolveProviderKey(connectionService, info.directProvider).catch(() => null) : null + + if (target && !key) { + throw new Error(`FIM request failed: 401 Missing ${target.provider} provider API key`) + } + + if (target && key) { + return generateDirectFim({ + apiKey: key, + target, + prefix, + suffix, + temperature: info.temperature, + onChunk, + signal, + }) + } let cost = 0 let inputTokens = 0 @@ -27,14 +213,15 @@ export async function generateFim( // ends the stream. Without this, errors never reach ErrorBackoff. let sseError: Error | undefined - const temp = getAutocompleteModel(modelId).temperature + const temp = info.temperature + console.info(`[FIM] request provider=kilo model=${info.requestModel} url=/kilo/fim`) const { stream } = await client.kilo.fim( { prefix, suffix, - model: modelId, - maxTokens: 256, + model: info.requestModel, + maxTokens: FIM_MAX_TOKENS, temperature: temp, }, { diff --git a/packages/kilo-vscode/src/shared/autocomplete-models.ts b/packages/kilo-vscode/src/shared/autocomplete-models.ts index 6f46b56512b..23a805b6b0a 100644 --- a/packages/kilo-vscode/src/shared/autocomplete-models.ts +++ b/packages/kilo-vscode/src/shared/autocomplete-models.ts @@ -7,12 +7,20 @@ */ export interface AutocompleteModelDef { - /** Full model ID sent to the gateway, e.g. "mistralai/codestral-2508" */ + /** Stable setting value. */ readonly id: string + /** Model ID displayed under the selector provider group. */ + readonly modelID: string /** Human-readable label shown in the settings dropdown */ readonly label: string + /** Provider ID used by the selector group. */ + readonly providerID: string /** Provider display name for status bar / telemetry */ readonly provider: string + /** Full model ID sent to the FIM API. */ + readonly requestModel: string + /** Provider key to use for direct BYOK FIM. Empty means Kilo Gateway. */ + readonly directProvider?: "mistral" | "inception" /** FIM request temperature */ readonly temperature: number } @@ -20,14 +28,40 @@ export interface AutocompleteModelDef { const models: AutocompleteModelDef[] = [ { id: "mistralai/codestral-2508", - label: "Codestral (Mistral AI)", - provider: "Mistral AI", + modelID: "mistralai/codestral-2508", + label: "Codestral", + providerID: "kilo", + provider: "Kilo Gateway", + requestModel: "mistralai/codestral-2508", temperature: 0.2, }, { id: "inception/mercury-edit-2", - label: "Mercury Edit 2 (Inception)", + modelID: "inception/mercury-edit-2", + label: "Mercury Edit 2", + providerID: "kilo", + provider: "Kilo Gateway", + requestModel: "inception/mercury-edit-2", + temperature: 0, + }, + { + id: "mistral/codestral-2508", + modelID: "codestral-2508", + label: "Codestral", + providerID: "mistral", + provider: "Mistral", + requestModel: "codestral-2508", + directProvider: "mistral", + temperature: 0.2, + }, + { + id: "inception-direct/mercury-edit-2", + modelID: "mercury-edit-2", + label: "Mercury Edit 2", + providerID: "inception", provider: "Inception", + requestModel: "mercury-edit-2", + directProvider: "inception", temperature: 0, }, ] diff --git a/packages/kilo-vscode/tests/unit/autocomplete-fim-direct.test.ts b/packages/kilo-vscode/tests/unit/autocomplete-fim-direct.test.ts new file mode 100644 index 00000000000..8a0621ca9bf --- /dev/null +++ b/packages/kilo-vscode/tests/unit/autocomplete-fim-direct.test.ts @@ -0,0 +1,94 @@ +import { describe, expect, it } from "bun:test" +import { generateDirectFim, getDirectFimTarget } from "../../src/services/autocomplete/fim" + +function stream(text: string) { + return new ReadableStream({ + start(controller) { + controller.enqueue(new TextEncoder().encode(text)) + controller.close() + }, + }) +} + +describe("direct autocomplete FIM", () => { + it("maps autocomplete models to direct provider endpoints", () => { + expect(getDirectFimTarget("mistralai/codestral-2508")).toBeNull() + expect(getDirectFimTarget("inception/mercury-edit-2")).toBeNull() + expect(getDirectFimTarget("mistral/codestral-2508")).toEqual({ + provider: "mistral", + model: "codestral-2508", + urls: ["https://api.mistral.ai/v1/fim/completions", "https://codestral.mistral.ai/v1/fim/completions"], + }) + expect(getDirectFimTarget("inception-direct/mercury-edit-2")).toEqual({ + provider: "inception", + model: "mercury-edit-2", + urls: ["https://api.inceptionlabs.ai/v1/fim/completions"], + }) + expect(getDirectFimTarget("openai/gpt-5")).toBeNull() + }) + + it("streams provider FIM chunks and returns usage", async () => { + const chunks = [ + 'data: {"choices":[{"delta":{"content":"hel"}}]}\n\n', + 'data: {"choices":[{"delta":{"content":"lo"}}],"usage":{"prompt_tokens":3,"completion_tokens":2}}\n\n', + "data: [DONE]\n\n", + ].join("") + const calls: RequestInit[] = [] + const fetchImpl: typeof fetch = async (_url, init) => { + calls.push(init ?? {}) + return new Response(stream(chunks), { status: 200 }) + } + const text: string[] = [] + const usage = await generateDirectFim({ + apiKey: "test-key", + target: getDirectFimTarget("inception-direct/mercury-edit-2")!, + prefix: "const value = ", + suffix: "\n", + temperature: 0, + onChunk: (chunk) => text.push(chunk), + fetchImpl, + }) + + expect(text.join("")).toBe("hello") + expect(usage).toEqual({ + cost: 0, + inputTokens: 3, + outputTokens: 2, + cacheWriteTokens: 0, + cacheReadTokens: 0, + }) + expect(calls[0]?.headers).toEqual({ + "Content-Type": "application/json", + Authorization: "Bearer test-key", + }) + expect(JSON.parse(String(calls[0]?.body))).toEqual({ + model: "mercury-edit-2", + prompt: "const value = ", + suffix: "\n", + max_tokens: 256, + temperature: 0, + stream: true, + }) + }) + + it("retries Codestral-specific endpoint when Mistral rejects a Codestral key", async () => { + const urls: string[] = [] + const fetchImpl: typeof fetch = async (url) => { + urls.push(String(url)) + if (urls.length === 1) return new Response("unauthorized", { status: 401, statusText: "Unauthorized" }) + return new Response(stream("data: [DONE]\n\n"), { status: 200 }) + } + + await generateDirectFim({ + apiKey: "codestral-key", + target: getDirectFimTarget("mistral/codestral-2508")!, + prefix: "", + suffix: "", + temperature: 0.2, + onChunk: () => {}, + fetchImpl, + }) + + expect(urls).toEqual(["https://api.mistral.ai/v1/fim/completions", "https://codestral.mistral.ai/v1/fim/completions"]) + }) +}) diff --git a/packages/kilo-vscode/webview-ui/src/components/settings/ModelsTab.tsx b/packages/kilo-vscode/webview-ui/src/components/settings/ModelsTab.tsx index c961e2d943b..93105a2fa8c 100644 --- a/packages/kilo-vscode/webview-ui/src/components/settings/ModelsTab.tsx +++ b/packages/kilo-vscode/webview-ui/src/components/settings/ModelsTab.tsx @@ -7,7 +7,11 @@ import { parseModelString } from "../../../../src/shared/provider-model" import { DEFAULT_AUTOCOMPLETE_MODEL } from "../../../../src/shared/autocomplete-models" import { ModelSelectorBase } from "../shared/ModelSelector" import SettingsRow from "./SettingsRow" -import { AUTOCOMPLETE_PROVIDER_ID, AUTOCOMPLETE_SELECTOR_MODELS } from "./autocomplete-model-selector" +import { + AUTOCOMPLETE_SELECTOR_MODELS, + getAutocompleteSelection, + getAutocompleteSettingID, +} from "./autocomplete-model-selector" const ModelsTab: Component = () => { const { config, settings, updateConfig, updateSetting } = useConfig() @@ -39,8 +43,9 @@ const ModelsTab: Component = () => { } function handleAutocompleteModelSelect(providerID: string, modelID: string) { - if (providerID !== AUTOCOMPLETE_PROVIDER_ID || !modelID) return - updateSetting("autocomplete.model", modelID) + const id = getAutocompleteSettingID(providerID, modelID) + if (!id) return + updateSetting("autocomplete.model", id) } return ( @@ -77,7 +82,7 @@ const ModelsTab: Component = () => { last > m.providerID === providerID && m.modelID === modelID)?.id +} export const AUTOCOMPLETE_SELECTOR_MODELS: EnrichedModel[] = AUTOCOMPLETE_MODELS.map((m) => ({ - id: m.id, + id: m.modelID, name: m.label, - providerID: AUTOCOMPLETE_PROVIDER_ID, - providerName: AUTOCOMPLETE_PROVIDER_NAME, + providerID: m.providerID, + providerName: m.provider, }))