From 3cc5be35a8f01359f0ac64c3daa976059692c09f Mon Sep 17 00:00:00 2001 From: Andrey Avtomonov Date: Mon, 11 May 2026 12:01:48 +0200 Subject: [PATCH] feat: use managed runtime for agent semantic queries --- packages/cli/src/agent-runtime.test.ts | 44 +++++++++ packages/cli/src/agent-runtime.ts | 36 ++++++- packages/cli/src/agent.test.ts | 35 +++++++ packages/cli/src/agent.ts | 37 ++++---- packages/cli/src/commands/agent-commands.ts | 14 ++- packages/cli/src/index.test.ts | 100 ++++++++++++++++++++ 6 files changed, 245 insertions(+), 21 deletions(-) diff --git a/packages/cli/src/agent-runtime.test.ts b/packages/cli/src/agent-runtime.test.ts index a7634103..808ddac3 100644 --- a/packages/cli/src/agent-runtime.test.ts +++ b/packages/cli/src/agent-runtime.test.ts @@ -105,4 +105,48 @@ describe('agent runtime helpers', () => { queryExecutor, }); }); + + it('creates managed semantic compute when no test override is injected', async () => { + const project = { + projectDir: tempDir, + configPath: join(tempDir, 'ktx.yaml'), + config: { project: 'revenue', connections: {} }, + coreConfig: {}, + git: {}, + fileStore: {}, + } as never; + const ports = { semanticLayer: {} } as never; + const semanticLayerCompute = { query: vi.fn(), validateSources: vi.fn(), generateSources: vi.fn() }; + const loadProject = vi.fn(async () => project); + const createContextTools = vi.fn(() => ports); + const createManagedSemanticLayerCompute = vi.fn(async () => semanticLayerCompute); + const { io } = makeIo(); + + await expect( + createKtxAgentRuntime( + { + projectDir: tempDir, + enableSemanticCompute: true, + enableQueryExecution: false, + cliVersion: '0.2.0', + runtimeInstallPolicy: 'auto', + io, + }, + { + loadProject, + createContextTools, + createManagedSemanticLayerCompute, + }, + ), + ).resolves.toMatchObject({ project, ports, semanticLayerCompute }); + + expect(createManagedSemanticLayerCompute).toHaveBeenCalledWith({ + cliVersion: '0.2.0', + installPolicy: 'auto', + io, + }); + expect(createContextTools).toHaveBeenCalledWith(project, { + semanticLayerCompute, + }); + }); }); diff --git a/packages/cli/src/agent-runtime.ts b/packages/cli/src/agent-runtime.ts index 98ebcb3a..feccae7c 100644 --- a/packages/cli/src/agent-runtime.ts +++ b/packages/cli/src/agent-runtime.ts @@ -1,9 +1,13 @@ import { readFile } from 'node:fs/promises'; import { createDefaultLocalQueryExecutor, type KtxSqlQueryExecutorPort } from '@ktx/context/connections'; -import { createPythonSemanticLayerComputePort, type KtxSemanticLayerComputePort } from '@ktx/context/daemon'; +import type { KtxSemanticLayerComputePort } from '@ktx/context/daemon'; import { createLocalProjectMcpContextPorts, type KtxMcpContextPorts } from '@ktx/context/mcp'; import { type KtxLocalProject, loadKtxProject } from '@ktx/context/project'; import type { KtxCliIo } from './cli-runtime.js'; +import { + createManagedPythonSemanticLayerComputePort, + type KtxManagedPythonInstallPolicy, +} from './managed-python-command.js'; export const KTX_AGENT_MAX_ROWS_CAP = 1000; @@ -11,6 +15,9 @@ export interface KtxAgentRuntimeOptions { projectDir: string; enableSemanticCompute: boolean; enableQueryExecution: boolean; + cliVersion?: string; + runtimeInstallPolicy?: KtxManagedPythonInstallPolicy; + io?: KtxCliIo; } export interface KtxAgentRuntime { @@ -24,6 +31,7 @@ export interface KtxAgentRuntimeDeps { loadProject?: typeof loadKtxProject; createContextTools?: typeof createLocalProjectMcpContextPorts; createSemanticLayerCompute?: () => KtxSemanticLayerComputePort; + createManagedSemanticLayerCompute?: typeof createManagedPythonSemanticLayerComputePort; createQueryExecutor?: () => KtxSqlQueryExecutorPort; } @@ -57,14 +65,34 @@ export function parseAgentMaxRows(value: number | undefined): number { return value; } +async function createAgentSemanticLayerCompute( + options: KtxAgentRuntimeOptions, + deps: KtxAgentRuntimeDeps, +): Promise { + if (!options.enableSemanticCompute) { + return undefined; + } + if (deps.createSemanticLayerCompute) { + return deps.createSemanticLayerCompute(); + } + if (!options.cliVersion || !options.runtimeInstallPolicy || !options.io) { + throw new Error('Managed Python semantic compute requires cliVersion, runtimeInstallPolicy, and io.'); + } + const createManagedSemanticLayerCompute = + deps.createManagedSemanticLayerCompute ?? createManagedPythonSemanticLayerComputePort; + return createManagedSemanticLayerCompute({ + cliVersion: options.cliVersion, + installPolicy: options.runtimeInstallPolicy, + io: options.io, + }); +} + export async function createKtxAgentRuntime( options: KtxAgentRuntimeOptions, deps: KtxAgentRuntimeDeps = {}, ): Promise { const project = await (deps.loadProject ?? loadKtxProject)({ projectDir: options.projectDir }); - const semanticLayerCompute = options.enableSemanticCompute - ? (deps.createSemanticLayerCompute ?? createPythonSemanticLayerComputePort)() - : undefined; + const semanticLayerCompute = await createAgentSemanticLayerCompute(options, deps); const queryExecutor = options.enableQueryExecution ? (deps.createQueryExecutor ?? createDefaultLocalQueryExecutor)() : undefined; diff --git a/packages/cli/src/agent.test.ts b/packages/cli/src/agent.test.ts index a57e7d04..043bdddb 100644 --- a/packages/cli/src/agent.test.ts +++ b/packages/cli/src/agent.test.ts @@ -231,6 +231,8 @@ describe('runKtxAgent', () => { queryFile, execute: true, maxRows: 100, + cliVersion: '0.2.0', + runtimeInstallPolicy: 'never', }, io.io, { createRuntime: async () => runtime() }, @@ -240,6 +242,39 @@ describe('runKtxAgent', () => { expect(JSON.parse(io.stdout())).toMatchObject({ sql: 'select 1', rows: [[1]] }); }); + it('passes managed runtime options into default SL query runtime creation', async () => { + const queryFile = join(tempDir, 'sl-query.json'); + const io = makeIo(); + const createRuntime = vi.fn(async () => runtime()); + await writeFile(queryFile, '{"measures":["total_revenue"],"dimensions":[]}', 'utf-8'); + + await expect( + runKtxAgent( + { + command: 'sl-query', + projectDir: tempDir, + json: true, + connectionId: 'warehouse', + queryFile, + execute: false, + cliVersion: '0.2.0', + runtimeInstallPolicy: 'auto', + }, + io.io, + { createRuntime }, + ), + ).resolves.toBe(0); + + expect(createRuntime).toHaveBeenCalledWith({ + projectDir: tempDir, + enableSemanticCompute: true, + enableQueryExecution: false, + cliVersion: '0.2.0', + runtimeInstallPolicy: 'auto', + io: io.io, + }); + }); + it('executes read-only SQL from a SQL file with an explicit row limit', async () => { const sqlFile = join(tempDir, 'query.sql'); const fakeRuntime = runtime(); diff --git a/packages/cli/src/agent.ts b/packages/cli/src/agent.ts index ea2a224e..61d85b8c 100644 --- a/packages/cli/src/agent.ts +++ b/packages/cli/src/agent.ts @@ -17,6 +17,7 @@ import { noIndexedSourcesSlSearchReadiness, type KtxAgentSlSearchReadinessDetail, } from './agent-search-readiness.js'; +import type { KtxManagedPythonInstallPolicy } from './managed-python-command.js'; import { readKtxSetupStatus, type KtxSetupStatus } from './setup.js'; export type KtxAgentArgs = @@ -32,6 +33,8 @@ export type KtxAgentArgs = queryFile: string; execute: boolean; maxRows?: number; + cliVersion: string; + runtimeInstallPolicy: KtxManagedPythonInstallPolicy; } | { command: 'wiki-search'; projectDir: string; json: true; query: string; limit: number } | { command: 'wiki-read'; projectDir: string; json: true; pageId: string } @@ -42,6 +45,9 @@ export interface KtxAgentDeps extends KtxAgentRuntimeDeps { projectDir: string; enableSemanticCompute: boolean; enableQueryExecution: boolean; + cliVersion?: string; + runtimeInstallPolicy?: KtxManagedPythonInstallPolicy; + io?: KtxCliIo; }) => Promise; readSetupStatus?: ( projectDir: string, @@ -68,23 +74,22 @@ function writeAgentSlSearchReadinessError(io: KtxCliIo, detail: KtxAgentSlSearch writeAgentJsonError(io, detail.message, { code: detail.code, nextSteps: detail.nextSteps }); } -async function runtimeFor(args: KtxAgentArgs, deps: KtxAgentDeps): Promise { +async function runtimeFor(args: KtxAgentArgs, deps: KtxAgentDeps, io: KtxCliIo): Promise { const needsSemanticCompute = args.command === 'sl-query'; const needsQueryExecution = args.command === 'sql-execute' || (args.command === 'sl-query' && args.execute); - return deps.createRuntime - ? deps.createRuntime({ - projectDir: args.projectDir, - enableSemanticCompute: needsSemanticCompute, - enableQueryExecution: needsQueryExecution, - }) - : createKtxAgentRuntime( - { - projectDir: args.projectDir, - enableSemanticCompute: needsSemanticCompute, - enableQueryExecution: needsQueryExecution, - }, - deps, - ); + const runtimeOptions = { + projectDir: args.projectDir, + enableSemanticCompute: needsSemanticCompute, + enableQueryExecution: needsQueryExecution, + ...(args.command === 'sl-query' + ? { + cliVersion: args.cliVersion, + runtimeInstallPolicy: args.runtimeInstallPolicy, + io, + } + : {}), + }; + return deps.createRuntime ? deps.createRuntime(runtimeOptions) : createKtxAgentRuntime(runtimeOptions, deps); } function connectionIdForSource(runtime: KtxAgentRuntime, requested: string | undefined): string { @@ -101,7 +106,7 @@ export async function runKtxAgent(args: KtxAgentArgs, io: KtxCliIo, deps: KtxAge return 0; } - const runtime = await runtimeFor(args, deps); + const runtime = await runtimeFor(args, deps, io); if (args.command === 'context') { const [status, connections, semanticLayer] = await Promise.all([ diff --git a/packages/cli/src/commands/agent-commands.ts b/packages/cli/src/commands/agent-commands.ts index 57cc94c3..2593991a 100644 --- a/packages/cli/src/commands/agent-commands.ts +++ b/packages/cli/src/commands/agent-commands.ts @@ -2,6 +2,7 @@ import { Option, type Command } from '@commander-js/extra-typings'; import type { KtxAgentArgs } from '../agent.js'; import type { KtxCliCommandContext } from '../cli-program.js'; import { parsePositiveIntegerOption, resolveCommandProjectDir } from '../cli-program.js'; +import { runtimeInstallPolicyFromFlags } from '../managed-python-command.js'; async function runAgent(context: KtxCliCommandContext, args: KtxAgentArgs): Promise { const runner = context.deps.agent ?? (await import('../agent.js')).runKtxAgent; @@ -73,10 +74,19 @@ export function registerAgentCommands(program: Command, context: KtxCliCommandCo .requiredOption('--connection-id ', 'Connection id for execution') .requiredOption('--query-file ', 'JSON semantic-layer query file') .option('--execute', 'Execute the compiled query against the connection', false) + .option('--yes', 'Install the managed Python runtime without prompting when required', false) + .option('--no-input', 'Disable interactive managed runtime installation') .option('--max-rows ', 'Maximum rows to return when executing', parsePositiveIntegerOption) .action( async ( - options: { connectionId: string; queryFile: string; execute: boolean; maxRows?: number }, + options: { + connectionId: string; + queryFile: string; + execute: boolean; + maxRows?: number; + yes?: boolean; + input?: boolean; + }, command, ) => { await runAgent(context, { @@ -86,6 +96,8 @@ export function registerAgentCommands(program: Command, context: KtxCliCommandCo connectionId: options.connectionId, queryFile: options.queryFile, execute: options.execute, + cliVersion: context.packageInfo.version, + runtimeInstallPolicy: runtimeInstallPolicyFromFlags(options), ...(options.maxRows !== undefined ? { maxRows: options.maxRows } : {}), }); }, diff --git a/packages/cli/src/index.test.ts b/packages/cli/src/index.test.ts index 26fa00ae..ff61d91b 100644 --- a/packages/cli/src/index.test.ts +++ b/packages/cli/src/index.test.ts @@ -1401,6 +1401,8 @@ describe('runKtxCli', () => { queryFile: '/tmp/query.json', execute: true, maxRows: 100, + cliVersion: '0.0.0-private', + runtimeInstallPolicy: 'prompt', }, }, { @@ -1449,6 +1451,104 @@ describe('runKtxCli', () => { expect(helpIo.stdout()).not.toContain('agent '); }); + it('routes hidden agent SL query managed runtime policies', async () => { + const autoIo = makeIo(); + const neverIo = makeIo(); + const conflictIo = makeIo(); + const agent = vi.fn(async () => 0); + + await expect( + runKtxCli( + [ + '--project-dir', + tempDir, + 'agent', + 'sl', + 'query', + '--json', + '--connection-id', + 'warehouse', + '--query-file', + '/tmp/query.json', + '--yes', + ], + autoIo.io, + { agent }, + ), + ).resolves.toBe(0); + + await expect( + runKtxCli( + [ + '--project-dir', + tempDir, + 'agent', + 'sl', + 'query', + '--json', + '--connection-id', + 'warehouse', + '--query-file', + '/tmp/query.json', + '--no-input', + ], + neverIo.io, + { agent }, + ), + ).resolves.toBe(0); + + await expect( + runKtxCli( + [ + '--project-dir', + tempDir, + 'agent', + 'sl', + 'query', + '--json', + '--connection-id', + 'warehouse', + '--query-file', + '/tmp/query.json', + '--yes', + '--no-input', + ], + conflictIo.io, + { agent }, + ), + ).resolves.toBe(1); + + expect(agent).toHaveBeenNthCalledWith( + 1, + { + command: 'sl-query', + projectDir: tempDir, + json: true, + connectionId: 'warehouse', + queryFile: '/tmp/query.json', + execute: false, + cliVersion: '0.0.0-private', + runtimeInstallPolicy: 'auto', + }, + autoIo.io, + ); + expect(agent).toHaveBeenNthCalledWith( + 2, + { + command: 'sl-query', + projectDir: tempDir, + json: true, + connectionId: 'warehouse', + queryFile: '/tmp/query.json', + execute: false, + cliVersion: '0.0.0-private', + runtimeInstallPolicy: 'never', + }, + neverIo.io, + ); + expect(conflictIo.stderr()).toContain('Choose only one runtime install mode: --yes or --no-input'); + }); + it('prints semantic-layer hybrid search metadata from the hidden agent sl list command', async () => { const agent = vi.fn(async (args, io) => { expect(args).toEqual({