From 4973ca562f560c3a5bb4bdafbe42012a06a16a3e Mon Sep 17 00:00:00 2001 From: Luca Martial <48870843+luca-martial@users.noreply.github.com> Date: Wed, 13 May 2026 08:42:38 -0400 Subject: [PATCH] Restore Vertex AI LLM setup (#56) * feat(context): resolve Vertex AI config references * feat(cli): restore Vertex AI LLM setup --------- Co-authored-by: Andrey Avtomonov --- packages/cli/src/commands/setup-commands.ts | 30 + packages/cli/src/index.test.ts | 41 ++ packages/cli/src/setup-models.test.ts | 359 ++++++++++- packages/cli/src/setup-models.ts | 600 +++++++++++++++++- packages/cli/src/setup.test.ts | 64 +- packages/cli/src/setup.ts | 13 +- packages/context/src/llm/local-config.test.ts | 46 ++ packages/context/src/llm/local-config.ts | 19 +- 8 files changed, 1113 insertions(+), 59 deletions(-) diff --git a/packages/cli/src/commands/setup-commands.ts b/packages/cli/src/commands/setup-commands.ts index 90251ae1..6a215651 100644 --- a/packages/cli/src/commands/setup-commands.ts +++ b/packages/cli/src/commands/setup-commands.ts @@ -2,6 +2,7 @@ import { type Command, InvalidArgumentError, Option } from '@commander-js/extra- import type { KtxCliCommandContext } from '../cli-program.js'; import { resolveCommandProjectDir } from '../cli-program.js'; import type { KtxSetupDatabaseDriver } from '../setup-databases.js'; +import type { KtxSetupLlmBackend } from '../setup-models.js'; import type { KtxSetupSourceType } from '../setup-sources.js'; async function runSetupArgs( @@ -27,6 +28,13 @@ function embeddingBackend(value: string): 'openai' | 'sentence-transformers' { throw new InvalidArgumentError(`invalid choice '${value}'`); } +function llmBackend(value: string): KtxSetupLlmBackend { + if (value === 'anthropic' || value === 'vertex') { + return value; + } + throw new InvalidArgumentError(`invalid choice '${value}'`); +} + function databaseDriver(value: string): KtxSetupDatabaseDriver { if ( value === 'sqlite' || @@ -93,9 +101,12 @@ function shouldShowSetupEntryMenu( skipAgents?: boolean; yes?: boolean; input?: boolean; + llmBackend?: KtxSetupLlmBackend; anthropicApiKeyEnv?: string; anthropicApiKeyFile?: string; anthropicModel?: string; + vertexProject?: string; + vertexLocation?: string; skipLlm?: boolean; embeddingBackend?: string; embeddingApiKeyEnv?: string; @@ -166,9 +177,12 @@ function shouldShowSetupEntryMenu( 'skipAgents', 'yes', 'input', + 'llmBackend', 'anthropicApiKeyEnv', 'anthropicApiKeyFile', 'anthropicModel', + 'vertexProject', + 'vertexLocation', 'skipLlm', 'embeddingBackend', 'embeddingApiKeyEnv', @@ -227,9 +241,12 @@ export function registerSetupCommands(program: Command, context: KtxCliCommandCo .option('--skip-agents', 'Leave agent integration incomplete for now', false) .option('--yes', 'Accept safe defaults in non-interactive setup', false) .option('--no-input', 'Disable interactive terminal input') + .addOption(new Option('--llm-backend ', 'LLM backend').argParser(llmBackend)) .option('--anthropic-api-key-env ', 'Environment variable containing the Anthropic API key') .option('--anthropic-api-key-file ', 'File containing the Anthropic API key') .option('--anthropic-model ', 'Anthropic model ID to validate and save') + .option('--vertex-project ', 'Google Vertex AI project ID, env:NAME, or file:/path') + .option('--vertex-location ', 'Google Vertex AI location, env:NAME, or file:/path') .addOption(new Option('--skip-llm', 'Leave LLM setup incomplete for now').hideHelp().default(false)) .addOption(new Option('--embedding-backend ', 'Embedding backend').argParser(embeddingBackend)) .option('--embedding-api-key-env ', 'Environment variable containing the embedding provider API key') @@ -325,6 +342,16 @@ export function registerSetupCommands(program: Command, context: KtxCliCommandCo context.setExitCode(1); return; } + if (options.llmBackend === 'vertex' && (options.anthropicApiKeyEnv || options.anthropicApiKeyFile)) { + context.io.stderr.write('Anthropic API key flags are only valid with --llm-backend anthropic.\n'); + context.setExitCode(1); + return; + } + if (options.llmBackend === 'anthropic' && (options.vertexProject || options.vertexLocation)) { + context.io.stderr.write('Vertex AI flags are only valid with --llm-backend vertex.\n'); + context.setExitCode(1); + return; + } if (options.embeddingApiKeyEnv && options.embeddingApiKeyFile) { context.io.stderr.write( 'Choose only one embedding credential source: --embedding-api-key-env or --embedding-api-key-file.\n', @@ -364,9 +391,12 @@ export function registerSetupCommands(program: Command, context: KtxCliCommandCo inputMode: options.input === false ? 'disabled' : 'auto', yes: options.yes === true, cliVersion: context.packageInfo.version, + ...(options.llmBackend ? { llmBackend: options.llmBackend } : {}), ...(options.anthropicApiKeyEnv ? { anthropicApiKeyEnv: options.anthropicApiKeyEnv } : {}), ...(options.anthropicApiKeyFile ? { anthropicApiKeyFile: options.anthropicApiKeyFile } : {}), ...(options.anthropicModel ? { anthropicModel: options.anthropicModel } : {}), + ...(options.vertexProject ? { vertexProject: options.vertexProject } : {}), + ...(options.vertexLocation ? { vertexLocation: options.vertexLocation } : {}), skipLlm: options.skipLlm === true, ...(options.embeddingBackend ? { embeddingBackend: options.embeddingBackend } : {}), ...(options.embeddingApiKeyEnv ? { embeddingApiKeyEnv: options.embeddingApiKeyEnv } : {}), diff --git a/packages/cli/src/index.test.ts b/packages/cli/src/index.test.ts index 1e69c590..58848dab 100644 --- a/packages/cli/src/index.test.ts +++ b/packages/cli/src/index.test.ts @@ -853,6 +853,47 @@ describe('runKtxCli', () => { ); }); + it('dispatches Vertex AI setup flags to the setup runner', async () => { + const setup = vi.fn(async () => 0); + const setupIo = makeIo(); + + await expect( + runKtxCli( + [ + '--project-dir', + tempDir, + 'setup', + '--no-input', + '--llm-backend', + 'vertex', + '--vertex-project', + 'local-gcp-project', + '--vertex-location', + 'us-east5', + '--anthropic-model', + 'claude-sonnet-4-6', + ], + setupIo.io, + { setup }, + ), + ).resolves.toBe(0); + + expect(setup).toHaveBeenCalledWith( + expect.objectContaining({ + command: 'run', + projectDir: tempDir, + inputMode: 'disabled', + cliVersion: '0.0.0-private', + llmBackend: 'vertex', + vertexProject: 'local-gcp-project', + vertexLocation: 'us-east5', + anthropicModel: 'claude-sonnet-4-6', + skipLlm: false, + }), + setupIo.io, + ); + }); + it('rejects conflicting Anthropic credential setup flags', async () => { const setup = vi.fn(async () => 0); const setupIo = makeIo(); diff --git a/packages/cli/src/setup-models.test.ts b/packages/cli/src/setup-models.test.ts index fb8acb47..2e83ade2 100644 --- a/packages/cli/src/setup-models.test.ts +++ b/packages/cli/src/setup-models.test.ts @@ -7,8 +7,10 @@ import { BUNDLED_ANTHROPIC_MODELS, fetchAnthropicModels, type KtxSetupModelPromptAdapter, + runKtxSetupGcloudApplicationDefaultAuth, runKtxSetupAnthropicModelStep, } from './setup-models.js'; +import type { KtxCliIo } from './cli-runtime.js'; function makeIo() { let stdout = ''; @@ -33,6 +35,7 @@ function makeIo() { } function makePromptAdapter(options: { + providerChoice?: string; selectValues?: string[]; credentialChoice?: string; modelChoice?: string; @@ -43,8 +46,20 @@ function makePromptAdapter(options: { const selectValues = [...(options.selectValues ?? [])]; const textValues = [...(options.textValues ?? [])]; const passwordValues = [...(options.passwordValues ?? [])]; + let providerPromptCount = 0; return { select: vi.fn(async ({ message }) => { + if (message.includes('LLM provider')) { + providerPromptCount += 1; + const nextProviderChoice = selectValues[0]; + if (nextProviderChoice === 'anthropic' || nextProviderChoice === 'vertex' || nextProviderChoice === 'back') { + return selectValues.shift() ?? nextProviderChoice; + } + if (options.credentialChoice === 'back' && providerPromptCount > 1) { + return 'back'; + } + return options.providerChoice ?? 'anthropic'; + } const nextValue = selectValues.shift(); if (nextValue) { return nextValue; @@ -55,7 +70,10 @@ function makePromptAdapter(options: { return options.modelChoice ?? 'claude-sonnet-4-6'; }), text: vi.fn(async () => textValues.shift() ?? ''), - password: vi.fn(async () => (passwordValues.length > 0 ? passwordValues.shift() : options.passwordValue ?? 'sk-ant-pasted')), + password: vi.fn( + async () => + passwordValues.length > 0 ? passwordValues.shift() : options.passwordValue ?? 'sk-ant-pasted', // pragma: allowlist secret + ), cancel: vi.fn(), }; } @@ -89,7 +107,7 @@ describe('setup Anthropic model step', () => { ), ); - await expect(fetchAnthropicModels('sk-ant-test', fetchModels)).resolves.toEqual([ + await expect(fetchAnthropicModels('sk-ant-test', fetchModels)).resolves.toEqual([ // pragma: allowlist secret { id: 'claude-sonnet-4-6', label: 'Claude Sonnet 4.6', recommended: true }, { id: 'claude-opus-4-6', label: 'Claude Opus 4.6', recommended: false }, { id: 'claude-haiku-4-5', label: 'Claude Haiku 4.5', recommended: false }, @@ -107,7 +125,7 @@ describe('setup Anthropic model step', () => { makeIo().io, { prompts, - env: { ANTHROPIC_API_KEY: 'sk-ant-test' }, + env: { ANTHROPIC_API_KEY: 'sk-ant-test' }, // pragma: allowlist secret listModels: vi.fn(async () => [ { id: 'claude-sonnet-4', label: 'Claude Sonnet 4', recommended: true }, { id: 'claude-opus-4', label: 'Claude Opus 4', recommended: false }, @@ -132,19 +150,58 @@ describe('setup Anthropic model step', () => { ); }); + it('offers Vertex AI as an Anthropic model provider option', async () => { + const prompts = makePromptAdapter({ providerChoice: 'back' }); + + const result = await runKtxSetupAnthropicModelStep( + { projectDir: tempDir, inputMode: 'auto', skipLlm: false }, + makeIo().io, + { prompts, env: {} }, + ); + + expect(result.status).toBe('back'); + expect(prompts.select).toHaveBeenCalledWith( + expect.objectContaining({ + message: expect.stringContaining('Which LLM provider should KTX use?'), + options: expect.arrayContaining([ + { value: 'vertex', label: 'Google Vertex AI for Anthropic Claude' }, + { value: 'back', label: 'Back' }, + ]), + }), + ); + }); + + it('returns from Anthropic credential Back to provider selection', async () => { + const prompts = makePromptAdapter({ selectValues: ['anthropic', 'back', 'back'] }); + + const result = await runKtxSetupAnthropicModelStep( + { projectDir: tempDir, inputMode: 'auto', skipLlm: false }, + makeIo().io, + { prompts, env: {} }, + ); + + expect(result.status).toBe('back'); + expect(prompts.select).toHaveBeenNthCalledWith( + 3, + expect.objectContaining({ + message: expect.stringContaining('Which LLM provider should KTX use?'), + }), + ); + }); + it('configures env credentials, selected model, prompt caching, and llm completion state', async () => { const io = makeIo(); const result = await runKtxSetupAnthropicModelStep( { projectDir: tempDir, inputMode: 'disabled', - anthropicApiKeyEnv: 'ANTHROPIC_API_KEY', + anthropicApiKeyEnv: 'ANTHROPIC_API_KEY', // pragma: allowlist secret anthropicModel: 'claude-sonnet-4-6', skipLlm: false, }, io.io, { - env: { ANTHROPIC_API_KEY: 'sk-ant-test' }, + env: { ANTHROPIC_API_KEY: 'sk-ant-test' }, // pragma: allowlist secret healthCheck: vi.fn(async () => ({ ok: true as const })), }, ); @@ -154,7 +211,7 @@ describe('setup Anthropic model step', () => { expect(config.llm).toMatchObject({ provider: { backend: 'anthropic', - anthropic: { api_key: 'env:ANTHROPIC_API_KEY' }, + anthropic: { api_key: 'env:ANTHROPIC_API_KEY' }, // pragma: allowlist secret }, models: { default: 'claude-sonnet-4-6' }, promptCaching: { enabled: true }, @@ -166,10 +223,258 @@ describe('setup Anthropic model step', () => { expect(io.stdout()).not.toContain('sk-ant-test'); }); + it('configures Vertex AI provider, selected model, prompt caching, and llm completion state', async () => { + const io = makeIo(); + const healthCheck = vi.fn(async () => ({ ok: true as const })); + + const result = await runKtxSetupAnthropicModelStep( + { + projectDir: tempDir, + inputMode: 'disabled', + llmBackend: 'vertex', + vertexProject: 'local-gcp-project', + vertexLocation: 'us-east5', + anthropicModel: 'claude-sonnet-4-6', + skipLlm: false, + }, + io.io, + { env: {}, healthCheck }, + ); + + expect(result.status).toBe('ready'); + expect(healthCheck).toHaveBeenCalledWith({ + backend: 'vertex', + vertex: { project: 'local-gcp-project', location: 'us-east5' }, + modelSlots: { default: 'claude-sonnet-4-6' }, + promptCaching: { enabled: true, vertexFallbackTo5m: true }, + }); + const config = parseKtxProjectConfig(await readFile(join(tempDir, 'ktx.yaml'), 'utf-8')); + expect(config.llm).toMatchObject({ + provider: { + backend: 'vertex', + vertex: { project: 'local-gcp-project', location: 'us-east5' }, + }, + models: { default: 'claude-sonnet-4-6' }, + promptCaching: { enabled: true, vertexFallbackTo5m: true }, + }); + expect(config.scan.enrichment.mode).toBe('llm'); + expect(await readFile(join(tempDir, 'ktx.yaml'), 'utf-8')).not.toContain('completed_steps:'); + expect((await readKtxSetupState(tempDir)).completed_steps).toContain('llm'); + expect(io.stdout()).toContain('LLM ready: yes (claude-sonnet-4-6)'); + }); + + it('can run gcloud auth for Vertex AI and infer project and default location', async () => { + const io = makeIo(); + const prompts = makePromptAdapter({ selectValues: ['vertex', 'gcloud', 'local-gcp-project', 'claude-sonnet-4-6'] }); + const runGcloudAuth = vi.fn(async () => ({ ok: true as const })); + const readGcloudProject = vi.fn(async () => 'local-gcp-project'); + const listGcloudProjects = vi.fn(async () => [ + { projectId: 'local-gcp-project', name: 'Local project' }, + { projectId: 'other-gcp-project', name: 'Other project' }, + ]); + const healthCheck = vi.fn(async () => ({ ok: true as const })); + + const result = await runKtxSetupAnthropicModelStep( + { projectDir: tempDir, inputMode: 'auto', skipLlm: false }, + io.io, + { + prompts, + env: {}, + runGcloudAuth, + readGcloudProject, + listGcloudProjects, + healthCheck, + }, + ); + + expect(result.status).toBe('ready'); + expect(runGcloudAuth).toHaveBeenCalledWith(io.io); + expect(readGcloudProject).toHaveBeenCalled(); + expect(listGcloudProjects).toHaveBeenCalled(); + expect(prompts.text).not.toHaveBeenCalled(); + expect(prompts.select).toHaveBeenCalledWith( + expect.objectContaining({ + message: expect.stringContaining('Which Google Cloud project should KTX use for Vertex AI?'), + options: [ + { value: 'local-gcp-project', label: 'local-gcp-project - Local project (current gcloud project)' }, + { value: 'other-gcp-project', label: 'other-gcp-project - Other project' }, + { value: 'manual', label: 'Enter a project ID manually' }, + { value: 'back', label: 'Back' }, + ], + }), + ); + expect(healthCheck).toHaveBeenCalledWith({ + backend: 'vertex', + vertex: { project: 'local-gcp-project', location: 'us-east5' }, + modelSlots: { default: 'claude-sonnet-4-6' }, + promptCaching: { enabled: true, vertexFallbackTo5m: true }, + }); + const config = parseKtxProjectConfig(await readFile(join(tempDir, 'ktx.yaml'), 'utf-8')); + expect(config.llm.provider).toMatchObject({ + backend: 'vertex', + vertex: { project: 'local-gcp-project', location: 'us-east5' }, + }); + }); + + it('lets users choose a different visible gcloud project for Vertex AI', async () => { + const io = makeIo(); + const prompts = makePromptAdapter({ selectValues: ['vertex', 'existing', 'other-gcp-project', 'claude-sonnet-4-6'] }); + const healthCheck = vi.fn(async () => ({ ok: true as const })); + + const result = await runKtxSetupAnthropicModelStep( + { projectDir: tempDir, inputMode: 'auto', skipLlm: false }, + io.io, + { + prompts, + env: {}, + readGcloudProject: vi.fn(async () => 'current-gcp-project'), + listGcloudProjects: vi.fn(async () => [ + { projectId: 'current-gcp-project', name: 'Current project' }, + { projectId: 'other-gcp-project', name: 'Other project' }, + ]), + healthCheck, + }, + ); + + expect(result.status).toBe('ready'); + expect(healthCheck).toHaveBeenCalledWith({ + backend: 'vertex', + vertex: { project: 'other-gcp-project', location: 'us-east5' }, + modelSlots: { default: 'claude-sonnet-4-6' }, + promptCaching: { enabled: true, vertexFallbackTo5m: true }, + }); + const config = parseKtxProjectConfig(await readFile(join(tempDir, 'ktx.yaml'), 'utf-8')); + expect(config.llm.provider).toMatchObject({ + backend: 'vertex', + vertex: { project: 'other-gcp-project', location: 'us-east5' }, + }); + }); + + it('allows manual Vertex AI project entry when gcloud project listing is empty', async () => { + const io = makeIo(); + const prompts = makePromptAdapter({ + selectValues: ['vertex', 'existing', 'manual', 'claude-sonnet-4-6'], + textValues: ['manual-gcp-project'], + }); + const healthCheck = vi.fn(async () => ({ ok: true as const })); + + const result = await runKtxSetupAnthropicModelStep( + { projectDir: tempDir, inputMode: 'auto', skipLlm: false }, + io.io, + { + prompts, + env: {}, + readGcloudProject: vi.fn(async () => undefined), + listGcloudProjects: vi.fn(async () => []), + healthCheck, + }, + ); + + expect(result.status).toBe('ready'); + expect(prompts.select).toHaveBeenCalledWith( + expect.objectContaining({ + message: expect.stringContaining('Which Google Cloud project should KTX use for Vertex AI?'), + options: [ + { value: 'manual', label: 'Enter a project ID manually' }, + { value: 'back', label: 'Back' }, + ], + }), + ); + expect(prompts.text).toHaveBeenCalledWith( + expect.objectContaining({ + message: 'Google Cloud project ID\n│ Press Escape to go back.\n│', + }), + ); + expect(healthCheck).toHaveBeenCalledWith( + expect.objectContaining({ + vertex: { project: 'manual-gcp-project', location: 'us-east5' }, + }), + ); + }); + + it('returns from Vertex AI project selection Back to provider selection', async () => { + const prompts = makePromptAdapter({ selectValues: ['vertex', 'existing', 'back', 'back'] }); + + const result = await runKtxSetupAnthropicModelStep( + { projectDir: tempDir, inputMode: 'auto', skipLlm: false }, + makeIo().io, + { + prompts, + env: {}, + readGcloudProject: vi.fn(async () => 'current-gcp-project'), + listGcloudProjects: vi.fn(async () => [{ projectId: 'current-gcp-project', name: 'Current project' }]), + }, + ); + + expect(result.status).toBe('back'); + expect(prompts.select).toHaveBeenNthCalledWith( + 4, + expect.objectContaining({ + message: expect.stringContaining('Which LLM provider should KTX use?'), + }), + ); + }); + + it('runs only gcloud application-default login for Vertex AI auth', async () => { + const io = makeIo(); + const runGcloud = vi.fn(async () => ({ ok: true as const })); + + await expect(runKtxSetupGcloudApplicationDefaultAuth(io.io, runGcloud)).resolves.toEqual({ ok: true }); + + expect(runGcloud).toHaveBeenCalledTimes(1); + expect(runGcloud).toHaveBeenCalledWith(['auth', 'application-default', 'login'], expect.anything()); + expect(runGcloud).not.toHaveBeenCalledWith(['auth', 'login'], expect.anything()); + expect(io.stdout()).toContain('gcloud auth application-default login'); + expect(io.stdout()).not.toContain('gcloud auth login'); + }); + + it('indents gcloud auth output inside the setup gutter', async () => { + const io = makeIo(); + const runGcloud = vi.fn(async (_args: string[], commandIo: KtxCliIo) => { + commandIo.stdout.write('Your browser has been opened to visit:\n\n https://accounts.example/auth\n'); + commandIo.stderr.write('Credentials saved to file: [/tmp/application_default_credentials.json]\n'); + return { ok: true as const }; + }); + + await expect(runKtxSetupGcloudApplicationDefaultAuth(io.io, runGcloud)).resolves.toEqual({ ok: true }); + + expect(io.stdout()).toContain('│ Your browser has been opened to visit:'); + expect(io.stdout()).toContain('│ https://accounts.example/auth'); + expect(io.stderr()).toContain('│ Credentials saved to file: [/tmp/application_default_credentials.json]'); + expect(io.stdout()).not.toContain('\nYour browser has been opened'); + }); + + it('explains common Vertex AI Forbidden health-check causes', async () => { + const io = makeIo(); + + const result = await runKtxSetupAnthropicModelStep( + { + projectDir: tempDir, + inputMode: 'disabled', + llmBackend: 'vertex', + vertexProject: 'kaelio-orbit-looker-20260430', + vertexLocation: 'us-east5', + anthropicModel: 'claude-sonnet-4-6', + skipLlm: false, + }, + io.io, + { + env: {}, + healthCheck: vi.fn(async () => ({ ok: false as const, message: 'Forbidden' })), + }, + ); + + expect(result.status).toBe('failed'); + expect(io.stderr()).toContain('project kaelio-orbit-looker-20260430'); + expect(io.stderr()).toContain('Vertex AI API is enabled'); + expect(io.stderr()).toContain('Anthropic Claude model access'); + expect(io.stderr()).toContain('roles/aiplatform.user'); + }); + it('resolves --anthropic-api-key-file for health checks and stores a file reference', async () => { const io = makeIo(); const secretPath = join(tempDir, 'anthropic-api-key'); - await writeFile(secretPath, 'sk-ant-file', 'utf-8'); + await writeFile(secretPath, 'sk-ant-file', 'utf-8'); // pragma: allowlist secret const healthCheck = vi.fn(async () => ({ ok: true as const })); const result = await runKtxSetupAnthropicModelStep( @@ -187,7 +492,7 @@ describe('setup Anthropic model step', () => { expect(result.status).toBe('ready'); expect(healthCheck).toHaveBeenCalledWith( expect.objectContaining({ - anthropic: { apiKey: 'sk-ant-file' }, + anthropic: { apiKey: 'sk-ant-file' }, // pragma: allowlist secret modelSlots: { default: 'claude-sonnet-4-6' }, }), ); @@ -195,7 +500,7 @@ describe('setup Anthropic model step', () => { expect(config.llm).toMatchObject({ provider: { backend: 'anthropic', - anthropic: { api_key: `file:${secretPath}` }, + anthropic: { api_key: `file:${secretPath}` }, // pragma: allowlist secret }, models: { default: 'claude-sonnet-4-6' }, }); @@ -249,11 +554,11 @@ describe('setup Anthropic model step', () => { { projectDir: tempDir, inputMode: 'disabled', - anthropicApiKeyEnv: 'ANTHROPIC_API_KEY', + anthropicApiKeyEnv: 'ANTHROPIC_API_KEY', // pragma: allowlist secret skipLlm: false, }, io.io, - { env: { ANTHROPIC_API_KEY: 'sk-ant-test' }, healthCheck }, + { env: { ANTHROPIC_API_KEY: 'sk-ant-test' }, healthCheck }, // pragma: allowlist secret ); expect(result.status).toBe('missing-input'); @@ -267,7 +572,7 @@ describe('setup Anthropic model step', () => { const prompts = makePromptAdapter({ credentialChoice: 'paste', modelChoice: 'claude-sonnet-4-6', - passwordValue: 'sk-ant-pasted', + passwordValue: 'sk-ant-pasted', // pragma: allowlist secret }); const result = await runKtxSetupAnthropicModelStep( @@ -282,7 +587,7 @@ describe('setup Anthropic model step', () => { ); expect(result.status).toBe('ready'); - await expect(readFile(join(tempDir, '.ktx/secrets/anthropic-api-key'), 'utf-8')).resolves.toBe('sk-ant-pasted\n'); + await expect(readFile(join(tempDir, '.ktx/secrets/anthropic-api-key'), 'utf-8')).resolves.toBe('sk-ant-pasted\n'); // pragma: allowlist secret if (process.platform !== 'win32') { expect((await stat(join(tempDir, '.ktx/secrets/anthropic-api-key'))).mode & 0o777).toBe(0o600); } @@ -295,7 +600,7 @@ describe('setup Anthropic model step', () => { it('opens pasted key entry directly and tells users Escape goes back', async () => { const prompts = makePromptAdapter({ selectValues: ['paste', 'claude-sonnet-4-6'], - passwordValue: 'sk-ant-pasted', + passwordValue: 'sk-ant-pasted', // pragma: allowlist secret }); const result = await runKtxSetupAnthropicModelStep( @@ -370,7 +675,7 @@ describe('setup Anthropic model step', () => { makeIo().io, { prompts, - env: { ANTHROPIC_API_KEY: 'sk-ant-test' }, + env: { ANTHROPIC_API_KEY: 'sk-ant-test' }, // pragma: allowlist secret listModels: vi.fn(async () => [{ id: 'claude-sonnet-4-6', label: 'Claude Sonnet 4.6', recommended: true }]), }, ); @@ -401,7 +706,7 @@ describe('setup Anthropic model step', () => { io.io, { prompts, - env: { ANTHROPIC_API_KEY: 'sk-ant-test' }, + env: { ANTHROPIC_API_KEY: 'sk-ant-test' }, // pragma: allowlist secret listModels: vi.fn(async () => [{ id: 'claude-sonnet-4-6', label: 'Claude Sonnet 4.6', recommended: true }]), healthCheck: vi.fn(async () => ({ ok: true as const })), }, @@ -424,7 +729,7 @@ describe('setup Anthropic model step', () => { await expect( runKtxSetupAnthropicModelStep({ projectDir: tempDir, inputMode: 'auto', skipLlm: false }, io.io, { prompts, - env: { ANTHROPIC_API_KEY: 'sk-ant-test' }, + env: { ANTHROPIC_API_KEY: 'sk-ant-test' }, // pragma: allowlist secret listModels: vi.fn(async () => { throw new Error('network unavailable'); }), @@ -444,7 +749,7 @@ describe('setup Anthropic model step', () => { io.io, { prompts, - env: { ANTHROPIC_API_KEY: 'sk-ant-test' }, + env: { ANTHROPIC_API_KEY: 'sk-ant-test' }, // pragma: allowlist secret listModels: vi.fn(async () => { throw new Error('network unavailable'); }), @@ -504,13 +809,13 @@ describe('setup Anthropic model step', () => { { projectDir: tempDir, inputMode: 'disabled', - anthropicApiKeyEnv: 'ANTHROPIC_API_KEY', + anthropicApiKeyEnv: 'ANTHROPIC_API_KEY', // pragma: allowlist secret anthropicModel: 'claude-sonnet-4-6', skipLlm: false, }, io.io, { - env: { ANTHROPIC_API_KEY: 'sk-ant-test' }, + env: { ANTHROPIC_API_KEY: 'sk-ant-test' }, // pragma: allowlist secret healthCheck: vi.fn(async () => ({ ok: false as const, message: '401 invalid x-api-key [redacted]' })), }, ); @@ -536,7 +841,7 @@ describe('setup Anthropic model step', () => { io.io, { prompts, - env: { ANTHROPIC_API_KEY: 'sk-ant-test' }, + env: { ANTHROPIC_API_KEY: 'sk-ant-test' }, // pragma: allowlist secret listModels: vi.fn(async () => [ { id: 'claude-haiku-3-5', label: 'Claude Haiku 3.5', recommended: false }, { id: 'claude-sonnet-4-6', label: 'Claude Sonnet 4.6', recommended: true }, @@ -547,7 +852,7 @@ describe('setup Anthropic model step', () => { expect(result.status).toBe('ready'); expect(healthCheck).toHaveBeenCalledTimes(2); - expect(prompts.select).toHaveBeenCalledTimes(4); + expect(prompts.select).toHaveBeenCalledTimes(5); expect(io.stderr()).toContain('Anthropic model health check failed: model not found'); expect(io.stderr()).toContain('Choose a different credential source or model, or Back.'); const config = parseKtxProjectConfig(await readFile(join(tempDir, 'ktx.yaml'), 'utf-8')); @@ -583,7 +888,7 @@ describe('setup Anthropic model step', () => { it('returns from model selection Back to credential selection instead of exiting setup', async () => { const prompts = makePromptAdapter({ selectValues: ['paste', 'back', 'back'], - passwordValue: 'sk-ant-pasted', + passwordValue: 'sk-ant-pasted', // pragma: allowlist secret }); const result = await runKtxSetupAnthropicModelStep( @@ -599,7 +904,7 @@ describe('setup Anthropic model step', () => { expect(result.status).toBe('back'); expect(prompts.select).toHaveBeenNthCalledWith( - 3, + 4, expect.objectContaining({ message: expect.stringContaining('How should KTX find your Anthropic API key?'), }), @@ -635,7 +940,7 @@ describe('setup Anthropic model step', () => { const config = parseKtxProjectConfig(await readFile(join(tempDir, 'ktx.yaml'), 'utf-8')); expect(config.llm.provider).toMatchObject({ backend: 'anthropic', - anthropic: { api_key: 'env:ANTHROPIC_API_KEY' }, + anthropic: { api_key: 'env:ANTHROPIC_API_KEY' }, // pragma: allowlist secret }); }); @@ -653,7 +958,7 @@ describe('setup Anthropic model step', () => { ' provider:', ' backend: anthropic', ' anthropic:', - ' api_key: env:ANTHROPIC_API_KEY', + ' api_key: env:ANTHROPIC_API_KEY', // pragma: allowlist secret ' models:', ' default: claude-sonnet-4-6', 'ingest:', @@ -669,7 +974,7 @@ describe('setup Anthropic model step', () => { const healthCheck = vi.fn(async () => ({ ok: true as const })); await expect( runKtxSetupAnthropicModelStep({ projectDir: tempDir, inputMode: 'disabled', skipLlm: false }, makeIo().io, { - env: { ANTHROPIC_API_KEY: 'sk-ant-test' }, + env: { ANTHROPIC_API_KEY: 'sk-ant-test' }, // pragma: allowlist secret healthCheck, }), ).resolves.toMatchObject({ status: 'ready' }); diff --git a/packages/cli/src/setup-models.ts b/packages/cli/src/setup-models.ts index 221dbd14..37ebdeec 100644 --- a/packages/cli/src/setup-models.ts +++ b/packages/cli/src/setup-models.ts @@ -1,4 +1,6 @@ +import { execFile, spawn } from 'node:child_process'; import { writeFile } from 'node:fs/promises'; +import { promisify } from 'node:util'; import { cancel, isCancel, password, select, text } from '@clack/prompts'; import { resolveLocalKtxLlmConfig } from '@ktx/context'; import { resolveKtxConfigReference } from '@ktx/context/core'; @@ -18,9 +20,12 @@ import { envCredentialReference, writeProjectLocalSecretReference } from './setu export interface KtxSetupModelArgs { projectDir: string; inputMode: 'auto' | 'disabled'; + llmBackend?: KtxSetupLlmBackend; anthropicApiKeyEnv?: string; anthropicApiKeyFile?: string; anthropicModel?: string; + vertexProject?: string; + vertexLocation?: string; forcePrompt?: boolean; showPromptInstructions?: boolean; skipLlm: boolean; @@ -39,6 +44,8 @@ export interface AnthropicModelChoice { recommended: boolean; } +export type KtxSetupLlmBackend = 'anthropic' | 'vertex'; + export interface KtxSetupModelPromptAdapter { select(options: { message: string; options: Array<{ value: string; label: string }> }): Promise; text(options: { message: string; placeholder?: string }): Promise; @@ -52,6 +59,9 @@ export interface KtxSetupModelDeps { prompts?: KtxSetupModelPromptAdapter; listModels?: (apiKey: string) => Promise; healthCheck?: (config: KtxLlmConfig) => Promise; + runGcloudAuth?: (io: KtxCliIo) => Promise; + readGcloudProject?: () => Promise; + listGcloudProjects?: () => Promise; } export const BUNDLED_ANTHROPIC_MODEL_REGISTRY_VERSION = '2026-05-07'; @@ -78,6 +88,16 @@ const ANTHROPIC_MODEL_PROMPT_CONTEXT = 'KTX uses this as the default model for ingest agents that turn schemas, SQL, BI metadata, and docs ' + 'into semantic-layer sources and wiki context.'; +const VERTEX_AUTH_PROMPT_CONTEXT = + 'KTX can use Google Cloud Application Default Credentials for local Vertex AI access. This opens the normal ' + + 'gcloud browser login flow and does not store Google credentials in ktx.yaml.'; +const VERTEX_PROJECT_PROMPT_CONTEXT = + 'KTX stores the selected Google Cloud project ID in ktx.yaml and uses Application Default Credentials for ' + + 'access. Project visibility depends on the signed-in Google account and organization permissions.'; +const DEFAULT_VERTEX_LOCATION = 'us-east5'; + +const execFileAsync = promisify(execFile); + type AnthropicModelDiscoveryErrorReason = 'authentication' | 'http' | 'empty-response'; export class AnthropicModelDiscoveryError extends Error { @@ -103,6 +123,27 @@ type ChooseModelResult = | { status: 'ready'; model: string } | { status: 'back' | 'missing-input' | 'invalid-credential' }; +type ChooseBackendResult = + | { status: 'ready'; backend: KtxSetupLlmBackend; prompted: boolean } + | { status: 'back' }; + +type VertexConfigChoice = + | { + status: 'ready'; + refs: { project?: string; location: string }; + values: { project?: string; location: string }; + } + | { status: 'back' | 'missing-input' }; + +type VertexAuthChoice = { status: 'ready' } | { status: 'back' | 'missing-input' }; + +export type GcloudAuthResult = { ok: true } | { ok: false; message: string }; +interface GcloudProjectChoice { + projectId: string; + name?: string; +} +type GcloudCommandRunner = (args: string[], io: KtxCliIo) => Promise; + function createPromptAdapter(): KtxSetupModelPromptAdapter { return { async select(options) { @@ -131,6 +172,122 @@ function createPromptAdapter(): KtxSetupModelPromptAdapter { }; } +function createIndentedCommandIo(io: KtxCliIo): KtxCliIo { + const indentedWriter = (write: (chunk: string) => void) => { + let atLineStart = true; + return (chunk: string) => { + for (const char of chunk) { + if (atLineStart) { + write('│ '); + atLineStart = false; + } + write(char); + if (char === '\n') { + atLineStart = true; + } + } + }; + }; + + return { + stdout: { + isTTY: io.stdout.isTTY, + columns: io.stdout.columns, + write: indentedWriter((chunk) => io.stdout.write(chunk)), + }, + stderr: { + write: indentedWriter((chunk) => io.stderr.write(chunk)), + }, + }; +} + +function runInteractiveGcloud(args: string[], io: KtxCliIo): Promise { + return new Promise((resolve) => { + let settled = false; + const child = spawn('gcloud', args, { stdio: ['inherit', 'pipe', 'pipe'] }); + child.stdout?.on('data', (chunk: Buffer) => { + io.stdout.write(chunk.toString('utf8')); + }); + child.stderr?.on('data', (chunk: Buffer) => { + io.stderr.write(chunk.toString('utf8')); + }); + child.on('error', (error: NodeJS.ErrnoException) => { + if (settled) { + return; + } + settled = true; + if (error.code === 'ENOENT') { + resolve({ ok: false, message: 'gcloud CLI was not found on PATH.' }); + return; + } + resolve({ ok: false, message: error.message }); + }); + child.on('close', (code, signal) => { + if (settled) { + return; + } + settled = true; + if (code === 0) { + resolve({ ok: true }); + return; + } + resolve({ + ok: false, + message: signal ? `gcloud exited after signal ${signal}.` : `gcloud exited with code ${code ?? 'unknown'}.`, + }); + }); + }); +} + +export async function runKtxSetupGcloudApplicationDefaultAuth( + io: KtxCliIo, + runGcloud: GcloudCommandRunner = runInteractiveGcloud, +): Promise { + io.stdout.write('│ Running gcloud auth application-default login...\n'); + return await runGcloud(['auth', 'application-default', 'login'], createIndentedCommandIo(io)); +} + +async function defaultReadGcloudProject(): Promise { + try { + const { stdout } = await execFileAsync('gcloud', ['config', 'get-value', 'project'], { encoding: 'utf8' }); + const value = stdout.trim(); + return value && value !== '(unset)' ? value : undefined; + } catch { + return undefined; + } +} + +async function defaultListGcloudProjects(): Promise { + try { + const { stdout } = await execFileAsync('gcloud', ['projects', 'list', '--format=json(projectId,name)'], { + encoding: 'utf8', + }); + const parsed = JSON.parse(stdout.trim() || '[]') as unknown; + if (!Array.isArray(parsed)) { + return []; + } + + return parsed + .map((item): GcloudProjectChoice | undefined => { + if (!item || typeof item !== 'object') { + return undefined; + } + const record = item as { projectId?: unknown; name?: unknown }; + if (typeof record.projectId !== 'string' || !record.projectId.trim()) { + return undefined; + } + const name = typeof record.name === 'string' && record.name.trim() ? record.name.trim() : undefined; + return { + projectId: record.projectId.trim(), + ...(name ? { name } : {}), + }; + }) + .filter((project): project is GcloudProjectChoice => Boolean(project)); + } catch { + return []; + } +} + export async function fetchAnthropicModels( apiKey: string, fetchFn: typeof fetch = fetch, @@ -195,20 +352,33 @@ function hasUsableConfiguredLlm(config: KtxProjectConfig): boolean { function buildProjectLlmConfig( existing: KtxProjectLlmConfig, - credentialRef: string, + provider: + | { backend: 'anthropic'; credentialRef: string } + | { backend: 'vertex'; vertex: { project?: string; location: string } }, model: string, ): KtxProjectLlmConfig { + if (provider.backend === 'vertex') { + return { + provider: { + backend: 'vertex', + vertex: provider.vertex, + }, + models: { ...existing.models, default: model }, + promptCaching: { ...(existing.promptCaching ?? {}), enabled: true, vertexFallbackTo5m: true }, + }; + } + return { provider: { backend: 'anthropic', - anthropic: { api_key: credentialRef }, + anthropic: { api_key: provider.credentialRef }, }, models: { ...existing.models, default: model }, promptCaching: { ...(existing.promptCaching ?? {}), enabled: true }, }; } -function buildHealthConfig(credentialValue: string, model: string): KtxLlmConfig { +function buildAnthropicHealthConfig(credentialValue: string, model: string): KtxLlmConfig { return { backend: 'anthropic', anthropic: { apiKey: credentialValue }, @@ -217,6 +387,28 @@ function buildHealthConfig(credentialValue: string, model: string): KtxLlmConfig }; } +function buildVertexHealthConfig(vertex: { project?: string; location: string }, model: string): KtxLlmConfig { + return { + backend: 'vertex', + vertex, + modelSlots: { default: model }, + promptCaching: { enabled: true, vertexFallbackTo5m: true }, + }; +} + +function formatVertexHealthFailure(message: string, vertex: { project?: string; location: string }): string { + const trimmed = message.trim() || 'unknown error'; + if (!/(forbidden|permission|permission_denied|403)/i.test(trimmed)) { + return trimmed; + } + + return ( + `${trimmed}. Check that Vertex AI API is enabled for project ${vertex.project ?? '(unknown)'}, ` + + `Anthropic Claude model access is enabled for location ${vertex.location}, and that your Application Default ` + + 'Credentials principal has Vertex AI User (roles/aiplatform.user) or equivalent permissions.' + ); +} + async function chooseCredentialRef( args: KtxSetupModelArgs, io: KtxCliIo, @@ -298,6 +490,266 @@ async function chooseCredentialRef( } } +function requestedBackend(args: KtxSetupModelArgs): KtxSetupLlmBackend | undefined { + if (args.llmBackend) { + return args.llmBackend; + } + if (args.vertexProject || args.vertexLocation) { + return 'vertex'; + } + if (args.anthropicApiKeyEnv || args.anthropicApiKeyFile || args.anthropicModel) { + return 'anthropic'; + } + return undefined; +} + +async function chooseBackend( + args: KtxSetupModelArgs, + io: KtxCliIo, + deps: KtxSetupModelDeps, +): Promise { + const explicit = requestedBackend(args); + if (explicit) { + return { status: 'ready', backend: explicit, prompted: false }; + } + if (args.inputMode === 'disabled') { + return { status: 'ready', backend: 'anthropic', prompted: false }; + } + + const prompts = deps.prompts ?? createPromptAdapter(); + if (args.showPromptInstructions !== false) { + io.stdout.write( + '│ Use Up/Down to move, Enter to confirm the current selection, choose Back to return to the previous step, Ctrl+C to exit.\n', + ); + } + const choice = await prompts.select({ + message: 'Which LLM provider should KTX use?', + options: [ + { value: 'anthropic', label: 'Anthropic API' }, + { value: 'vertex', label: 'Google Vertex AI for Anthropic Claude' }, + { value: 'back', label: 'Back' }, + ], + }); + if (choice === 'back') { + return { status: 'back' }; + } + return { status: 'ready', backend: choice === 'vertex' ? 'vertex' : 'anthropic', prompted: true }; +} + +async function chooseVertexAuth( + args: KtxSetupModelArgs, + io: KtxCliIo, + deps: KtxSetupModelDeps, +): Promise { + if (args.inputMode === 'disabled' || args.vertexProject || args.vertexLocation) { + return { status: 'ready' }; + } + + const prompts = deps.prompts ?? createPromptAdapter(); + const choice = await prompts.select({ + message: `How should KTX authenticate with Google Vertex AI?\n\n${VERTEX_AUTH_PROMPT_CONTEXT}`, + options: [ + { value: 'gcloud', label: 'Run gcloud Application Default Credentials login' }, + { value: 'existing', label: 'Use existing gcloud/Application Default Credentials' }, + { value: 'back', label: 'Back' }, + ], + }); + if (choice === 'back') { + return { status: 'back' }; + } + if (choice !== 'gcloud') { + return { status: 'ready' }; + } + + const result = await (deps.runGcloudAuth ?? runKtxSetupGcloudApplicationDefaultAuth)(io); + if (!result.ok) { + io.stderr.write(`gcloud authentication failed: ${result.message}\n`); + return { status: 'missing-input' }; + } + return { status: 'ready' }; +} + +function resolveProvidedVertexRef( + label: 'project' | 'location', + ref: string, + env: NodeJS.ProcessEnv, + io: KtxCliIo, +): { status: 'ready'; ref: string; value: string } | { status: 'missing-input' } { + let value: string | undefined; + try { + value = resolveKtxConfigReference(ref, env); + } catch { + value = undefined; + } + if (!value) { + io.stderr.write(`Missing Vertex AI ${label}: ${ref} could not be resolved.\n`); + return { status: 'missing-input' }; + } + return { status: 'ready', ref, value }; +} + +function normalizeGcloudProjectId(projectId: string | undefined): string | undefined { + const trimmed = projectId?.trim(); + return trimmed ? trimmed : undefined; +} + +function orderGcloudProjects(projects: GcloudProjectChoice[], currentProject: string | undefined): GcloudProjectChoice[] { + const ordered: GcloudProjectChoice[] = []; + const seen = new Set(); + const addProject = (project: GcloudProjectChoice) => { + const projectId = normalizeGcloudProjectId(project.projectId); + if (!projectId || seen.has(projectId)) { + return; + } + seen.add(projectId); + const name = normalizeGcloudProjectId(project.name); + ordered.push({ + projectId, + ...(name ? { name } : {}), + }); + }; + + if (currentProject) { + addProject(projects.find((project) => project.projectId.trim() === currentProject) ?? { projectId: currentProject }); + } + for (const project of projects) { + addProject(project); + } + return ordered; +} + +function formatGcloudProjectLabel(project: GcloudProjectChoice, currentProject: string | undefined): string { + const name = project.name && project.name !== project.projectId ? ` - ${project.name}` : ''; + const current = project.projectId === currentProject ? ' (current gcloud project)' : ''; + return `${project.projectId}${name}${current}`; +} + +async function chooseInteractiveVertexProject( + currentProject: string | undefined, + io: KtxCliIo, + deps: KtxSetupModelDeps, +): Promise<{ status: 'ready'; ref: string; value: string } | { status: 'back' | 'missing-input' }> { + const prompts = deps.prompts ?? createPromptAdapter(); + let projects: GcloudProjectChoice[] = []; + try { + projects = await (deps.listGcloudProjects ?? defaultListGcloudProjects)(); + } catch { + io.stderr.write('Could not list Google Cloud projects with gcloud. Enter a project ID manually or choose Back.\n'); + } + + const orderedProjects = orderGcloudProjects(projects, currentProject); + if (orderedProjects.length === 0) { + io.stdout.write('│ gcloud did not return any visible Google Cloud projects. Enter a project ID manually or choose Back.\n'); + } + + const choice = await prompts.select({ + message: `Which Google Cloud project should KTX use for Vertex AI?\n\n${VERTEX_PROJECT_PROMPT_CONTEXT}`, + options: [ + ...orderedProjects.map((project) => ({ + value: project.projectId, + label: formatGcloudProjectLabel(project, currentProject), + })), + { value: 'manual', label: 'Enter a project ID manually' }, + { value: 'back', label: 'Back' }, + ], + }); + if (choice === 'back') { + return { status: 'back' }; + } + if (choice === 'manual') { + const manual = await prompts.text({ + message: withTextInputNavigation('Google Cloud project ID'), + placeholder: currentProject ?? orderedProjects[0]?.projectId, + }); + if (manual === undefined) { + return { status: 'back' }; + } + const project = normalizeGcloudProjectId(manual); + return project ? { status: 'ready', ref: project, value: project } : { status: 'missing-input' }; + } + + return { status: 'ready', ref: choice, value: choice }; +} + +async function chooseVertexConfig( + args: KtxSetupModelArgs, + io: KtxCliIo, + deps: KtxSetupModelDeps, +): Promise { + const env = deps.env ?? process.env; + let projectRef: string | undefined; + let projectValue: string | undefined; + let gcloudProject: string | undefined; + + if (args.vertexProject) { + const project = resolveProvidedVertexRef('project', args.vertexProject, env, io); + if (project.status !== 'ready') { + return { status: project.status }; + } + projectRef = project.ref; + projectValue = project.value; + } else if (env.GOOGLE_VERTEX_PROJECT?.trim()) { + projectRef = envCredentialReference('GOOGLE_VERTEX_PROJECT'); + projectValue = env.GOOGLE_VERTEX_PROJECT.trim(); + } else { + gcloudProject = normalizeGcloudProjectId(await (deps.readGcloudProject ?? defaultReadGcloudProject)()); + if (args.inputMode === 'disabled') { + if (gcloudProject) { + projectRef = gcloudProject; + projectValue = gcloudProject; + } + } else { + const project = await chooseInteractiveVertexProject(gcloudProject, io, deps); + if (project.status !== 'ready') { + return { status: project.status }; + } + projectRef = project.ref; + projectValue = project.value; + } + } + + let locationRef: string | undefined; + let locationValue: string | undefined; + if (args.vertexLocation) { + const location = resolveProvidedVertexRef('location', args.vertexLocation, env, io); + if (location.status !== 'ready') { + return { status: location.status }; + } + locationRef = location.ref; + locationValue = location.value; + } else if (env.GOOGLE_VERTEX_LOCATION?.trim()) { + locationRef = envCredentialReference('GOOGLE_VERTEX_LOCATION'); + locationValue = env.GOOGLE_VERTEX_LOCATION.trim(); + } else { + locationRef = DEFAULT_VERTEX_LOCATION; + locationValue = DEFAULT_VERTEX_LOCATION; + } + + if (!projectRef || !projectValue) { + io.stderr.write( + 'Missing Vertex AI project: run `gcloud config set project PROJECT_ID`, pass --vertex-project, or set GOOGLE_VERTEX_PROJECT.\n', + ); + return { status: 'missing-input' }; + } + + if (!locationRef || !locationValue) { + io.stderr.write('Missing Vertex AI location: pass --vertex-location.\n'); + return { status: 'missing-input' }; + } + + return { + status: 'ready', + refs: { + ...(projectRef ? { project: projectRef } : {}), + location: locationRef, + }, + values: { + ...(projectValue ? { project: projectValue } : {}), + location: locationValue, + }, + }; +} + async function chooseModel( args: KtxSetupModelArgs, credentialValue: string, @@ -359,28 +811,73 @@ async function chooseModel( return { status: 'ready', model: choice }; } -async function persistLlmConfig(projectDir: string, credentialRef: string, model: string): Promise { +async function chooseVertexModel(args: KtxSetupModelArgs, io: KtxCliIo, deps: KtxSetupModelDeps): Promise { + if (args.anthropicModel) { + return { status: 'ready', model: args.anthropicModel }; + } + if (args.inputMode === 'disabled') { + io.stderr.write('Missing Anthropic model: pass --anthropic-model.\n'); + return { status: 'missing-input' }; + } + + const selectableModels = BUNDLED_ANTHROPIC_MODELS.filter(isSelectableAnthropicModel); + const prompts = deps.prompts ?? createPromptAdapter(); + const choice = await prompts.select({ + message: `Which Anthropic model should KTX use?\n\n${ANTHROPIC_MODEL_PROMPT_CONTEXT}`, + options: [ + ...selectableModels.map((model) => ({ + value: model.id, + label: `${model.label || model.id}${model.recommended ? ' (recommended)' : ''}`, + })), + { value: 'manual', label: 'Enter a model ID manually' }, + { value: 'back', label: 'Back' }, + ], + }); + if (choice === 'back') { + return { status: 'back' }; + } + if (choice === 'manual') { + const manual = await prompts.text({ + message: withTextInputNavigation('Anthropic model ID'), + placeholder: selectableModels.find((model) => model.recommended)?.id ?? selectableModels[0]?.id, + }); + if (manual === undefined) { + return { status: 'back' }; + } + return manual.trim() ? { status: 'ready', model: manual.trim() } : { status: 'missing-input' }; + } + return { status: 'ready', model: choice }; +} + +async function persistLlmConfig( + projectDir: string, + provider: + | { backend: 'anthropic'; credentialRef: string } + | { backend: 'vertex'; vertex: { project?: string; location: string } }, + model: string, +): Promise { const project = await loadKtxProject({ projectDir }); const config = { ...project.config, - llm: buildProjectLlmConfig(project.config.llm, credentialRef, model), + llm: buildProjectLlmConfig(project.config.llm, provider, model), scan: { ...project.config.scan, enrichment: { - ...project.config.scan.enrichment, - mode: 'llm' as const, - }, + ...project.config.scan.enrichment, + mode: 'llm' as const, }, - }; + }, + }; await writeFile(project.configPath, serializeKtxProjectConfig(config), 'utf-8'); await markKtxSetupStateStepComplete(projectDir, 'llm'); } -function buildInteractiveRetryArgs(args: KtxSetupModelArgs): KtxSetupModelArgs { +function buildInteractiveRetryArgs(args: KtxSetupModelArgs, backend?: KtxSetupLlmBackend): KtxSetupModelArgs { return { projectDir: args.projectDir, inputMode: args.inputMode, - ...(args.showPromptInstructions !== undefined ? { showPromptInstructions: args.showPromptInstructions } : {}), + ...(backend ?? args.llmBackend ? { llmBackend: backend ?? args.llmBackend } : {}), + showPromptInstructions: false, skipLlm: args.skipLlm, }; } @@ -399,9 +896,12 @@ export async function runKtxSetupAnthropicModelStep( if ( args.forcePrompt !== true && hasUsableConfiguredLlm(project.config) && + !args.llmBackend && !args.anthropicApiKeyEnv && !args.anthropicApiKeyFile && - !args.anthropicModel + !args.anthropicModel && + !args.vertexProject && + !args.vertexLocation ) { io.stdout.write(`│ LLM ready: yes (${project.config.llm.models.default})\n`); return { status: 'ready', projectDir: args.projectDir }; @@ -411,31 +911,91 @@ export async function runKtxSetupAnthropicModelStep( let attemptArgs = args; while (true) { - const credential = await chooseCredentialRef(attemptArgs, io, deps); + const backendChoice = await chooseBackend(attemptArgs, io, deps); + if (backendChoice.status !== 'ready') { + return { status: backendChoice.status, projectDir: args.projectDir }; + } + + const backendArgs = backendChoice.prompted + ? ({ ...attemptArgs, llmBackend: backendChoice.backend, showPromptInstructions: false } satisfies KtxSetupModelArgs) + : attemptArgs; + + if (backendChoice.backend === 'vertex') { + const auth = await chooseVertexAuth(backendArgs, io, deps); + if (auth.status === 'back' && backendChoice.prompted) { + attemptArgs = buildInteractiveRetryArgs(args); + continue; + } + if (auth.status !== 'ready') { + return { status: auth.status, projectDir: args.projectDir }; + } + + const vertex = await chooseVertexConfig(backendArgs, io, deps); + if (vertex.status === 'back' && backendChoice.prompted) { + attemptArgs = buildInteractiveRetryArgs(args); + continue; + } + if (vertex.status !== 'ready') { + return { status: vertex.status, projectDir: args.projectDir }; + } + + const model = await chooseVertexModel(backendArgs, io, deps); + if (model.status === 'back' && !backendArgs.vertexLocation) { + attemptArgs = buildInteractiveRetryArgs(args, backendChoice.backend); + continue; + } + if (model.status === 'invalid-credential') { + return { status: 'failed', projectDir: args.projectDir }; + } + if (model.status !== 'ready') { + return { status: model.status, projectDir: args.projectDir }; + } + + const health = await healthCheck(buildVertexHealthConfig(vertex.values, model.model)); + if (health.ok) { + await persistLlmConfig(args.projectDir, { backend: 'vertex', vertex: vertex.refs }, model.model); + io.stdout.write(`│ LLM ready: yes (${model.model})\n`); + return { status: 'ready', projectDir: args.projectDir }; + } + + io.stderr.write(`Vertex AI Anthropic model health check failed: ${formatVertexHealthFailure(health.message, vertex.values)}\n`); + if (args.inputMode === 'disabled') { + return { status: 'failed', projectDir: args.projectDir }; + } + io.stderr.write('Choose a different Vertex AI project, location, or model, or Back.\n'); + attemptArgs = buildInteractiveRetryArgs(args, backendChoice.backend); + continue; + } + + const credential = await chooseCredentialRef(backendArgs, io, deps); + if (credential.status === 'back' && backendChoice.prompted) { + attemptArgs = buildInteractiveRetryArgs(args); + continue; + } if (credential.status !== 'ready') { return { status: credential.status, projectDir: args.projectDir }; } - const model = await chooseModel(attemptArgs, credential.value, io, deps); + const model = await chooseModel(backendArgs, credential.value, io, deps); if (model.status === 'invalid-credential') { if (args.inputMode === 'disabled') { return { status: 'failed', projectDir: args.projectDir }; } io.stderr.write('Choose a different credential source or Back.\n'); - attemptArgs = buildInteractiveRetryArgs(args); + attemptArgs = buildInteractiveRetryArgs(args, backendChoice.backend); continue; } - if (model.status === 'back' && !attemptArgs.anthropicApiKeyEnv && !attemptArgs.anthropicApiKeyFile) { - attemptArgs = buildInteractiveRetryArgs(args); + if (model.status === 'back' && !backendArgs.anthropicApiKeyEnv && !backendArgs.anthropicApiKeyFile) { + attemptArgs = buildInteractiveRetryArgs(args, backendChoice.backend); continue; } if (model.status !== 'ready') { return { status: model.status, projectDir: args.projectDir }; } - const health = await healthCheck(buildHealthConfig(credential.value, model.model)); + const health = await healthCheck(buildAnthropicHealthConfig(credential.value, model.model)); if (health.ok) { - await persistLlmConfig(args.projectDir, credential.ref, model.model); + await persistLlmConfig(args.projectDir, { backend: 'anthropic', credentialRef: credential.ref }, model.model); io.stdout.write(`│ LLM ready: yes (${model.model})\n`); return { status: 'ready', projectDir: args.projectDir }; } @@ -445,6 +1005,6 @@ export async function runKtxSetupAnthropicModelStep( return { status: 'failed', projectDir: args.projectDir }; } io.stderr.write('Choose a different credential source or model, or Back.\n'); - attemptArgs = buildInteractiveRetryArgs(args); + attemptArgs = buildInteractiveRetryArgs(args, backendChoice.backend); } } diff --git a/packages/cli/src/setup.test.ts b/packages/cli/src/setup.test.ts index 0cad3ebc..dd134fce 100644 --- a/packages/cli/src/setup.test.ts +++ b/packages/cli/src/setup.test.ts @@ -73,7 +73,7 @@ describe('setup status', () => { ' provider:', ' backend: anthropic', ' anthropic:', - ' api_key: env:ANTHROPIC_API_KEY', + ' api_key: env:ANTHROPIC_API_KEY', // pragma: allowlist secret ' models:', ' default: claude-sonnet-4-6', 'ingest:', @@ -144,7 +144,7 @@ describe('setup status', () => { ' model: text-embedding-3-small', ' dimensions: 1536', ' openai:', - ' api_key: env:OPENAI_API_KEY', + ' api_key: env:OPENAI_API_KEY', // pragma: allowlist secret ].join('\n'), 'utf-8', ); @@ -908,7 +908,7 @@ describe('setup status', () => { inputMode: 'disabled', yes: false, cliVersion: '0.2.0', - anthropicApiKeyEnv: 'ANTHROPIC_API_KEY', + anthropicApiKeyEnv: 'ANTHROPIC_API_KEY', // pragma: allowlist secret anthropicModel: 'claude-sonnet-4-6', skipLlm: false, skipEmbeddings: true, @@ -925,7 +925,51 @@ describe('setup status', () => { expect.objectContaining({ projectDir: tempDir, inputMode: 'disabled', - anthropicApiKeyEnv: 'ANTHROPIC_API_KEY', + anthropicApiKeyEnv: 'ANTHROPIC_API_KEY', // pragma: allowlist secret + anthropicModel: 'claude-sonnet-4-6', + skipLlm: false, + }), + testIo.io, + ); + }); + + it('passes Vertex AI model setup args after project selection succeeds', async () => { + const testIo = makeIo(); + const model = vi.fn(async () => ({ status: 'ready' as const, projectDir: tempDir })); + + await expect( + runKtxSetup( + { + command: 'run', + projectDir: tempDir, + mode: 'new', + agents: false, + skipAgents: true, + inputMode: 'disabled', + yes: false, + cliVersion: '0.2.0', + llmBackend: 'vertex', + vertexProject: 'local-gcp-project', + vertexLocation: 'us-east5', + anthropicModel: 'claude-sonnet-4-6', + skipLlm: false, + skipEmbeddings: true, + databaseSchemas: [], + skipDatabases: true, + skipSources: true, + }, + testIo.io, + { model }, + ), + ).resolves.toBe(0); + + expect(model).toHaveBeenCalledWith( + expect.objectContaining({ + projectDir: tempDir, + inputMode: 'disabled', + llmBackend: 'vertex', + vertexProject: 'local-gcp-project', + vertexLocation: 'us-east5', anthropicModel: 'claude-sonnet-4-6', skipLlm: false, }), @@ -949,11 +993,11 @@ describe('setup status', () => { inputMode: 'disabled', yes: true, cliVersion: '0.2.0', - anthropicApiKeyEnv: 'ANTHROPIC_API_KEY', + anthropicApiKeyEnv: 'ANTHROPIC_API_KEY', // pragma: allowlist secret anthropicModel: 'claude-sonnet-4-6', skipLlm: false, embeddingBackend: 'openai', - embeddingApiKeyEnv: 'OPENAI_API_KEY', + embeddingApiKeyEnv: 'OPENAI_API_KEY', // pragma: allowlist secret skipEmbeddings: false, databaseSchemas: [], skipDatabases: true, @@ -971,7 +1015,7 @@ describe('setup status', () => { cliVersion: '0.2.0', runtimeInstallPolicy: 'auto', embeddingBackend: 'openai', - embeddingApiKeyEnv: 'OPENAI_API_KEY', + embeddingApiKeyEnv: 'OPENAI_API_KEY', // pragma: allowlist secret skipEmbeddings: false, }), testIo.io, @@ -1169,11 +1213,11 @@ describe('setup status', () => { inputMode: 'disabled', yes: false, cliVersion: '0.2.0', - anthropicApiKeyEnv: 'ANTHROPIC_API_KEY', + anthropicApiKeyEnv: 'ANTHROPIC_API_KEY', // pragma: allowlist secret anthropicModel: 'claude-sonnet-4-6', skipLlm: false, embeddingBackend: 'openai', - embeddingApiKeyEnv: 'OPENAI_API_KEY', + embeddingApiKeyEnv: 'OPENAI_API_KEY', // pragma: allowlist secret skipEmbeddings: false, databaseDrivers: ['postgres'], databaseConnectionId: 'warehouse', @@ -2020,7 +2064,7 @@ describe('setup status', () => { inputMode: 'disabled', yes: false, cliVersion: '0.2.0', - anthropicApiKeyEnv: 'ANTHROPIC_API_KEY', + anthropicApiKeyEnv: 'ANTHROPIC_API_KEY', // pragma: allowlist secret anthropicModel: 'claude-sonnet-4-6', skipLlm: false, skipEmbeddings: false, diff --git a/packages/cli/src/setup.ts b/packages/cli/src/setup.ts index 064da729..1ab48f0b 100644 --- a/packages/cli/src/setup.ts +++ b/packages/cli/src/setup.ts @@ -24,7 +24,12 @@ import { runKtxSetupDatabasesStep, } from './setup-databases.js'; import { type KtxSetupEmbeddingsDeps, runKtxSetupEmbeddingsStep } from './setup-embeddings.js'; -import { type KtxSetupModelDeps, isKtxSetupLlmConfigReady, runKtxSetupAnthropicModelStep } from './setup-models.js'; +import { + type KtxSetupLlmBackend, + type KtxSetupModelDeps, + isKtxSetupLlmConfigReady, + runKtxSetupAnthropicModelStep, +} from './setup-models.js'; import { type KtxSetupProjectDeps, runKtxSetupProjectStep } from './setup-project.js'; import { isKtxPreAgentSetupReady, @@ -65,9 +70,12 @@ export type KtxSetupArgs = inputMode: 'auto' | 'disabled'; yes: boolean; cliVersion: string; + llmBackend?: KtxSetupLlmBackend; anthropicApiKeyEnv?: string; anthropicApiKeyFile?: string; anthropicModel?: string; + vertexProject?: string; + vertexLocation?: string; skipLlm: boolean; embeddingBackend?: 'openai' | 'sentence-transformers'; embeddingApiKeyEnv?: string; @@ -578,9 +586,12 @@ async function runKtxSetupInner(args: KtxSetupArgs, io: KtxCliIo, deps: KtxSetup { projectDir: projectResult.projectDir, inputMode: args.inputMode, + ...(args.llmBackend ? { llmBackend: args.llmBackend } : {}), ...(args.anthropicApiKeyEnv ? { anthropicApiKeyEnv: args.anthropicApiKeyEnv } : {}), ...(args.anthropicApiKeyFile ? { anthropicApiKeyFile: args.anthropicApiKeyFile } : {}), ...(args.anthropicModel ? { anthropicModel: args.anthropicModel } : {}), + ...(args.vertexProject ? { vertexProject: args.vertexProject } : {}), + ...(args.vertexLocation ? { vertexLocation: args.vertexLocation } : {}), forcePrompt: forcePromptSteps.has('models') || runOnly === 'models', showPromptInstructions, skipLlm: args.skipLlm || !shouldRunModels, diff --git a/packages/context/src/llm/local-config.test.ts b/packages/context/src/llm/local-config.test.ts index ffb00b36..5d114e12 100644 --- a/packages/context/src/llm/local-config.test.ts +++ b/packages/context/src/llm/local-config.test.ts @@ -37,6 +37,52 @@ describe('local KTX LLM config', () => { }); }); + it('resolves Vertex AI env references into a KtxLlmConfig', () => { + const config: KtxProjectLlmConfig = { + provider: { + backend: 'vertex', + vertex: { project: 'env:GOOGLE_VERTEX_PROJECT', location: 'env:GOOGLE_VERTEX_LOCATION' }, + }, + models: { default: 'env:KTX_MODEL' }, + promptCaching: { enabled: true, vertexFallbackTo5m: true }, + }; + + expect( + resolveLocalKtxLlmConfig(config, { + GOOGLE_VERTEX_PROJECT: 'local-gcp-project', + GOOGLE_VERTEX_LOCATION: 'us-east5', + KTX_MODEL: 'claude-sonnet-4-6', + }), + ).toEqual({ + backend: 'vertex', + vertex: { project: 'local-gcp-project', location: 'us-east5' }, + modelSlots: { default: 'claude-sonnet-4-6' }, + promptCaching: { enabled: true, vertexFallbackTo5m: true }, + }); + }); + + it('ignores inactive Vertex AI references for non-Vertex backends', () => { + const config: KtxProjectLlmConfig = { + provider: { + backend: 'anthropic', + anthropic: { api_key: 'env:ANTHROPIC_API_KEY' }, // pragma: allowlist secret + vertex: { location: 'env:MISSING_VERTEX_LOCATION' }, + }, + models: { default: 'claude-sonnet-4-6' }, + }; + + expect( + resolveLocalKtxLlmConfig(config, { + ANTHROPIC_API_KEY: 'sk-ant-test', // pragma: allowlist secret + }), + ).toEqual({ + backend: 'anthropic', + anthropic: { apiKey: 'sk-ant-test' }, // pragma: allowlist secret + modelSlots: { default: 'claude-sonnet-4-6' }, + promptCaching: undefined, + }); + }); + it('returns null when the local LLM backend is disabled', () => { expect( createLocalKtxLlmProviderFromConfig({ diff --git a/packages/context/src/llm/local-config.ts b/packages/context/src/llm/local-config.ts index 76f1905f..2709c4b7 100644 --- a/packages/context/src/llm/local-config.ts +++ b/packages/context/src/llm/local-config.ts @@ -67,16 +67,33 @@ function resolvedProviderConfig( }; } +function resolvedVertexConfig( + config: { project?: string; location?: string } | undefined, + env: NodeJS.ProcessEnv, +): { project?: string; location: string } | undefined { + if (!config) { + return undefined; + } + + const project = resolveOptional(config.project, env); + const location = resolveRequired(config.location, env, 'llm.provider.vertex.location is required'); + return { + ...(project ? { project } : {}), + location, + }; +} + export function resolveLocalKtxLlmConfig(config: KtxProjectLlmConfig, env: NodeJS.ProcessEnv): KtxLlmConfig | null { if (config.provider.backend === 'none') { return null; } const modelSlots = resolveModelSlots(config.models, env); + const vertex = config.provider.backend === 'vertex' ? resolvedVertexConfig(config.provider.vertex, env) : undefined; const anthropic = resolvedProviderConfig(config.provider.anthropic, env); const gateway = resolvedProviderConfig(config.provider.gateway, env); return { backend: config.provider.backend, - ...(config.provider.vertex ? { vertex: config.provider.vertex } : {}), + ...(vertex ? { vertex } : {}), ...(anthropic ? { anthropic } : {}), ...(gateway ? { gateway } : {}), modelSlots,