diff --git a/packages/kilo-docs/source-links.md b/packages/kilo-docs/source-links.md index 11ceece8a47..bf6b1ebceef 100644 --- a/packages/kilo-docs/source-links.md +++ b/packages/kilo-docs/source-links.md @@ -1,17 +1,13 @@ # Source Code Links - + - -- - - -- - - - @@ -36,8 +32,6 @@ - -- - - - diff --git a/packages/kilo-gateway/src/server/routes.ts b/packages/kilo-gateway/src/server/routes.ts index f92c00f3f31..381b2b184ef 100644 --- a/packages/kilo-gateway/src/server/routes.ts +++ b/packages/kilo-gateway/src/server/routes.ts @@ -29,6 +29,14 @@ type Auth = any type ModelCache = { clear: (providerID: string) => void } type Z = any +type FimProvider = "kilo" | "mistral" | "inception" + +interface FimTarget { + provider: FimProvider + model: string + urls: string[] +} + interface KiloRoutesDeps extends ImportDeps { Hono: new () => Hono describeRoute: DescribeRoute @@ -42,6 +50,20 @@ interface KiloRoutesDeps extends ImportDeps { } const FIM_TIMEOUT_MS = 30_000 +const KILO_FIM_URL = KILO_API_BASE + "/api/fim/completions" +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" + +export function resolveFimTarget(model?: string): FimTarget { + if (model === "mistral/codestral-2508") { + return { provider: "mistral", model: "codestral-2508", urls: [MISTRAL_FIM_URL, CODESTRAL_FIM_URL] } + } + if (model === "inception-direct/mercury-edit-2") { + return { provider: "inception", model: "mercury-edit-2", urls: [INCEPTION_FIM_URL] } + } + return { provider: "kilo", model: model ?? "mistralai/codestral-2501", urls: [KILO_FIM_URL] } +} /** * Create Kilo Gateway routes with OpenCode dependencies injected @@ -147,6 +169,52 @@ export function createKiloRoutes(deps: KiloRoutesDeps) { } } + const getProviderKey = async (provider: FimProvider) => { + const auth = await Auth.get(provider) + return auth?.type === "api" ? auth.key : undefined + } + + const fetchFim = async ( + target: FimTarget, + url: string, + fallbacks: string[], + key: string, + input: { + prefix: string + suffix: string + maxTokens: number + temperature: number + signal: AbortSignal + organizationId?: string + }, + ): Promise => { + console.info(`[FIM] request provider=${target.provider} model=${target.model} url=${url}`) + const response = await fetch(url, { + method: "POST", + headers: { + "Content-Type": "application/json", + Authorization: `Bearer ${key}`, + ...(target.provider === "kilo" + ? buildKiloHeaders(undefined, { kilocodeOrganizationId: input.organizationId }) + : {}), + ...(target.provider === "kilo" ? { [HEADER_FEATURE]: "autocomplete" } : {}), + }, + signal: input.signal, + body: JSON.stringify({ + model: target.model, + prompt: input.prefix, + suffix: input.suffix, + max_tokens: input.maxTokens, + temperature: input.temperature, + stream: true, + }), + }) + + const [next] = fallbacks + if (response.status === 401 && next) return fetchFim(target, next, fallbacks.slice(1), key, input) + return response + } + return new Hono() .get( "/profile", @@ -339,47 +407,38 @@ export function createKiloRoutes(deps: KiloRoutesDeps) { }), ), async (c: any) => { - const proxy = await getProxyAuth() + const { prefix, suffix, model, maxTokens, temperature } = c.req.valid("json") + const target = resolveFimTarget(model) + const fimMaxTokens = maxTokens ?? 256 + const fimTemperature = temperature ?? 0.2 + const proxy = target.provider === "kilo" ? await getProxyAuth() : undefined + const token = target.provider === "kilo" ? proxy?.token : await getProviderKey(target.provider) - if (!proxy.auth) { + if (target.provider === "kilo" && !proxy?.auth) { return c.json({ error: "Not authenticated with Kilo Gateway" }, 401) } - if (!proxy.token) { + if (target.provider === "kilo" && !token) { return c.json({ error: "No valid token found" }, 401) } - const { prefix, suffix, model, maxTokens, temperature } = c.req.valid("json") - const fimModel = model ?? "mistralai/codestral-2501" - const fimMaxTokens = maxTokens ?? 256 - const fimTemperature = temperature ?? 0.2 - - const baseApiUrl = KILO_API_BASE + "/api/" - const endpoint = new URL("fim/completions", baseApiUrl) - - const headers = { - "Content-Type": "application/json", - Authorization: `Bearer ${proxy.token}`, - ...buildKiloHeaders(undefined, { kilocodeOrganizationId: proxy.organizationId }), - [HEADER_FEATURE]: "autocomplete", + if (!token) { + return c.json({ error: `Missing ${target.provider} provider API key` }, 401) } const signal = AbortSignal.any([c.req.raw.signal, AbortSignal.timeout(FIM_TIMEOUT_MS)]) let response: Response try { - response = await fetch(endpoint, { - method: "POST", - headers, + const [url] = target.urls + if (!url) return c.json({ error: "No FIM endpoint configured" }, 500 as any) + response = await fetchFim(target, url, target.urls.slice(1), token, { + prefix, + suffix, + maxTokens: fimMaxTokens, + temperature: fimTemperature, signal, - body: JSON.stringify({ - model: fimModel, - prompt: prefix, - suffix, - max_tokens: fimMaxTokens, - temperature: fimTemperature, - stream: true, - }), + organizationId: proxy?.organizationId, }) } catch (err) { if (err instanceof DOMException && err.name === "TimeoutError") diff --git a/packages/kilo-gateway/test/fim.test.ts b/packages/kilo-gateway/test/fim.test.ts new file mode 100644 index 00000000000..9118506ad6e --- /dev/null +++ b/packages/kilo-gateway/test/fim.test.ts @@ -0,0 +1,30 @@ +import { describe, expect, test } from "bun:test" +import { resolveFimTarget } from "../src/server/routes" + +describe("FIM target resolution", () => { + test("keeps gateway autocomplete models on Kilo Gateway", () => { + expect(resolveFimTarget("mistralai/codestral-2508")).toEqual({ + provider: "kilo", + model: "mistralai/codestral-2508", + urls: ["https://api.kilo.ai/api/fim/completions"], + }) + expect(resolveFimTarget("inception/mercury-edit-2")).toEqual({ + provider: "kilo", + model: "inception/mercury-edit-2", + urls: ["https://api.kilo.ai/api/fim/completions"], + }) + }) + + test("routes explicit provider autocomplete models directly", () => { + expect(resolveFimTarget("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(resolveFimTarget("inception-direct/mercury-edit-2")).toEqual({ + provider: "inception", + model: "mercury-edit-2", + urls: ["https://api.inceptionlabs.ai/v1/fim/completions"], + }) + }) +}) diff --git a/packages/kilo-vscode/src/services/autocomplete/fim.ts b/packages/kilo-vscode/src/services/autocomplete/fim.ts index 2266f4aa055..d9f9a30a44c 100644 --- a/packages/kilo-vscode/src/services/autocomplete/fim.ts +++ b/packages/kilo-vscode/src/services/autocomplete/fim.ts @@ -3,171 +3,6 @@ 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. @@ -185,25 +20,6 @@ export async function generateFim( ): 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 let outputTokens = 0 @@ -213,16 +29,15 @@ export async function generateFim( // ends the stream. Without this, errors never reach ErrorBackoff. let sseError: Error | undefined - const temp = info.temperature - console.info(`[FIM] request provider=kilo model=${info.requestModel} url=/kilo/fim`) + console.info(`[FIM] request provider=${info.providerID} model=${info.requestModel} url=/kilo/fim`) const { stream } = await client.kilo.fim( { prefix, suffix, - model: info.requestModel, + model: info.id, maxTokens: FIM_MAX_TOKENS, - temperature: temp, + temperature: info.temperature, }, { signal, diff --git a/packages/kilo-vscode/tests/unit/autocomplete-fim-direct.test.ts b/packages/kilo-vscode/tests/unit/autocomplete-fim-direct.test.ts deleted file mode 100644 index 8a0621ca9bf..00000000000 --- a/packages/kilo-vscode/tests/unit/autocomplete-fim-direct.test.ts +++ /dev/null @@ -1,94 +0,0 @@ -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"]) - }) -})