From 954a64dd76dd27de6cdc662f1fa0ddb4493f8e6f Mon Sep 17 00:00:00 2001 From: TimothyDing Date: Tue, 14 Jul 2026 14:35:26 +0800 Subject: [PATCH 01/14] feat(llm): add openai-compatible backend runtime Add an openai-compatible LLM backend that talks to any OpenAI-compatible endpoint (base URL + API key) via @ai-sdk/openai-compatible. Wire the backend through the config schema, provider factory, local-config resolution, rate-limit classification, and health-check redaction. Prompt caching stays disabled for non-Anthropic models via the existing isAnthropicProtocolModel gate. --- packages/cli/package.json | 1 + .../cli/src/context/llm/ai-sdk-runtime.ts | 11 ++- packages/cli/src/context/llm/local-config.ts | 2 + .../src/context/llm/rate-limit-governor.ts | 2 +- packages/cli/src/context/project/config.ts | 5 +- packages/cli/src/llm/model-health.ts | 2 +- packages/cli/src/llm/model-provider.ts | 16 ++++ packages/cli/src/llm/types.ts | 3 +- .../test/context/llm/ai-sdk-runtime.test.ts | 19 ++++- .../cli/test/context/llm/local-config.test.ts | 66 ++++++++++++++++ .../cli/test/context/project/config.test.ts | 24 +++++- packages/cli/test/llm/model-health.test.ts | 26 +++++++ packages/cli/test/llm/model-provider.test.ts | 75 +++++++++++++++++++ pnpm-lock.yaml | 3 + 14 files changed, 246 insertions(+), 9 deletions(-) diff --git a/packages/cli/package.json b/packages/cli/package.json index af4bf65cf..c5fad800c 100644 --- a/packages/cli/package.json +++ b/packages/cli/package.json @@ -50,6 +50,7 @@ "@ai-sdk/anthropic": "3.0.78", "@ai-sdk/devtools": "0.0.18", "@ai-sdk/google-vertex": "^4.0.134", + "@ai-sdk/openai-compatible": "^2.0.47", "@anthropic-ai/claude-agent-sdk": "0.3.146", "@aws-sdk/client-athena": "^3.1068.0", "@aws-sdk/client-glue": "^3.1068.0", diff --git a/packages/cli/src/context/llm/ai-sdk-runtime.ts b/packages/cli/src/context/llm/ai-sdk-runtime.ts index 7787279d5..a62222732 100644 --- a/packages/cli/src/context/llm/ai-sdk-runtime.ts +++ b/packages/cli/src/context/llm/ai-sdk-runtime.ts @@ -49,9 +49,16 @@ function hasTools(tools: Record): boolean { return Object.keys(tools).length > 0; } -function modelProviderName(model: unknown): RateLimitProvider { +/** @internal */ +export function modelProviderName(model: unknown): RateLimitProvider { const provider = (model as { provider?: string }).provider ?? ''; - return provider.includes('vertex') || provider.includes('google') ? 'vertex' : 'anthropic-api'; + if (provider.includes('vertex') || provider.includes('google')) { + return 'vertex'; + } + if (provider.includes('openai')) { + return 'openai-compatible'; + } + return 'anthropic-api'; } interface HeaderLimitPair { diff --git a/packages/cli/src/context/llm/local-config.ts b/packages/cli/src/context/llm/local-config.ts index 0d64a27fe..d864368b0 100644 --- a/packages/cli/src/context/llm/local-config.ts +++ b/packages/cli/src/context/llm/local-config.ts @@ -102,11 +102,13 @@ export function resolveLocalKtxLlmConfig(config: KtxProjectLlmConfig, env: NodeJ 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); + const openaiCompatible = resolvedProviderConfig(config.provider.openaiCompatible, env); return { backend: config.provider.backend, ...(vertex ? { vertex } : {}), ...(anthropic ? { anthropic } : {}), ...(gateway ? { gateway } : {}), + ...(openaiCompatible ? { openaiCompatible } : {}), modelSlots, promptCaching: config.promptCaching, }; diff --git a/packages/cli/src/context/llm/rate-limit-governor.ts b/packages/cli/src/context/llm/rate-limit-governor.ts index 909e4c44a..6e764cee5 100644 --- a/packages/cli/src/context/llm/rate-limit-governor.ts +++ b/packages/cli/src/context/llm/rate-limit-governor.ts @@ -1,6 +1,6 @@ import { createAbortError, throwIfAborted } from '../core/abort.js'; -export type RateLimitProvider = 'claude-subscription' | 'anthropic-api' | 'vertex' | 'codex'; +export type RateLimitProvider = 'claude-subscription' | 'anthropic-api' | 'vertex' | 'codex' | 'openai-compatible'; type RateLimitSignalStatus = 'allowed' | 'warning' | 'rejected'; export interface RateLimitSignal { diff --git a/packages/cli/src/context/project/config.ts b/packages/cli/src/context/project/config.ts index 92377b39d..6843ee922 100644 --- a/packages/cli/src/context/project/config.ts +++ b/packages/cli/src/context/project/config.ts @@ -3,7 +3,7 @@ import YAML from 'yaml'; import * as z from 'zod'; import { connectionConfigSchema } from './driver-schemas.js'; -const KTX_LLM_BACKENDS = ['none', 'anthropic', 'vertex', 'gateway', 'claude-code', 'codex'] as const; +const KTX_LLM_BACKENDS = ['none', 'anthropic', 'vertex', 'gateway', 'openai-compatible', 'claude-code', 'codex'] as const; const KTX_EMBEDDING_BACKENDS = ['none', 'openai', 'sentence-transformers'] as const; const KTX_PROMPT_CACHE_TTLS = ['5m', '1h'] as const; const KTX_ENRICHMENT_MODES = ['none', 'deterministic', 'llm'] as const; @@ -38,11 +38,12 @@ const llmProviderSchema = z .enum(KTX_LLM_BACKENDS) .default('none') .describe( - 'LLM provider backend. "none" disables LLM features; "anthropic" / "vertex" / "gateway" require the matching nested credentials block; "claude-code" uses the local Claude Code session; "codex" uses the local Codex session.', + 'LLM provider backend. "none" disables LLM features; "anthropic" / "vertex" / "gateway" / "openai-compatible" require the matching nested credentials block; "claude-code" uses the local Claude Code session; "codex" uses the local Codex session.', ), vertex: vertexProviderSchema.optional().describe('Vertex AI credentials, used when backend is "vertex".'), anthropic: apiCredentialsSchema.optional().describe('Anthropic API credentials, used when backend is "anthropic".'), gateway: apiCredentialsSchema.optional().describe('AI Gateway credentials, used when backend is "gateway".'), + openaiCompatible: apiCredentialsSchema.optional().describe('OpenAI-compatible endpoint credentials, used when backend is "openai-compatible".'), }) .describe('LLM provider selection and credentials.'); diff --git a/packages/cli/src/llm/model-health.ts b/packages/cli/src/llm/model-health.ts index 919204783..3c0d5bf2e 100644 --- a/packages/cli/src/llm/model-health.ts +++ b/packages/cli/src/llm/model-health.ts @@ -15,7 +15,7 @@ export interface KtxLlmHealthCheckOptions { } function redactHealthCheckMessage(message: string, config: KtxLlmConfig): string { - const secrets = [config.anthropic?.apiKey, config.gateway?.apiKey].filter( + const secrets = [config.anthropic?.apiKey, config.gateway?.apiKey, config.openaiCompatible?.apiKey].filter( (value): value is string => typeof value === 'string' && value.length > 0, ); return secrets.reduce((current, secret) => current.split(secret).join('[redacted]'), message); diff --git a/packages/cli/src/llm/model-provider.ts b/packages/cli/src/llm/model-provider.ts index 8322c35cd..363883601 100644 --- a/packages/cli/src/llm/model-provider.ts +++ b/packages/cli/src/llm/model-provider.ts @@ -1,6 +1,7 @@ import { createAnthropic } from '@ai-sdk/anthropic'; import { devToolsMiddleware } from '@ai-sdk/devtools'; import { createVertexAnthropic } from '@ai-sdk/google-vertex/anthropic'; +import { createOpenAICompatible } from '@ai-sdk/openai-compatible'; import { createGateway, generateText, wrapLanguageModel, type LanguageModel } from 'ai'; import { createKtxToolCallRepairHandler } from './repair.js'; import type { @@ -16,11 +17,13 @@ type AnthropicFactory = typeof createAnthropic; type AnthropicModelFactory = (modelId: string) => LanguageModel; type VertexAnthropicFactory = (options?: Parameters[0]) => AnthropicModelFactory; type GatewayFactory = (options?: Parameters[0]) => AnthropicModelFactory; +type OpenAICompatibleFactory = (options: Parameters[0]) => AnthropicModelFactory; export interface KtxLlmProviderFactoryDeps { createAnthropic?: (options?: Parameters[0]) => AnthropicModelFactory; createVertexAnthropic?: VertexAnthropicFactory; createGateway?: GatewayFactory; + createOpenAICompatible?: OpenAICompatibleFactory; generateText?: typeof generateText; devtoolsEnabled?: boolean; wrapLanguageModel?: typeof wrapLanguageModel; @@ -175,6 +178,19 @@ class DefaultKtxLlmProvider implements KtxLlmProvider { return (modelId) => vertex(modelId); } + if (config.backend === 'openai-compatible') { + const baseURL = config.openaiCompatible?.baseURL; + if (!baseURL) { + throw new Error('openaiCompatible.baseURL is required when ktx LLM backend is openai-compatible'); + } + const openaiCompatible = (deps.createOpenAICompatible ?? createOpenAICompatible)({ + name: 'openai-compatible', + baseURL, + ...(config.openaiCompatible?.apiKey ? { apiKey: config.openaiCompatible.apiKey } : {}), + }); + return (modelId) => openaiCompatible(modelId); + } + if (config.backend === 'gateway') { const gateway = (deps.createGateway ?? createGateway)({ ...(config.gateway?.apiKey ? { apiKey: config.gateway.apiKey } : {}), diff --git a/packages/cli/src/llm/types.ts b/packages/cli/src/llm/types.ts index a190b1c00..a1cdf7fcc 100644 --- a/packages/cli/src/llm/types.ts +++ b/packages/cli/src/llm/types.ts @@ -3,7 +3,7 @@ import type { LanguageModel, TelemetrySettings, ToolCallRepairFunction, ToolSet export const KTX_MODEL_ROLES = ['default', 'triage', 'candidateExtraction', 'curator', 'reconcile', 'repair'] as const; export type KtxModelRole = (typeof KTX_MODEL_ROLES)[number]; -type KtxLlmBackend = 'anthropic' | 'vertex' | 'gateway' | 'claude-code' | 'codex'; +type KtxLlmBackend = 'anthropic' | 'vertex' | 'gateway' | 'openai-compatible' | 'claude-code' | 'codex'; export type KtxPromptCacheTtl = '5m' | '1h'; type KtxJsonValue = @@ -40,6 +40,7 @@ export interface KtxLlmConfig { vertex?: { project?: string; location: string }; anthropic?: { apiKey?: string; baseURL?: string }; gateway?: { baseURL?: string; apiKey?: string }; + openaiCompatible?: { apiKey?: string; baseURL?: string }; modelSlots: { default: string } & Partial>; promptCaching?: Partial; telemetry?: { diff --git a/packages/cli/test/context/llm/ai-sdk-runtime.test.ts b/packages/cli/test/context/llm/ai-sdk-runtime.test.ts index 6c0bbe1d8..480a267db 100644 --- a/packages/cli/test/context/llm/ai-sdk-runtime.test.ts +++ b/packages/cli/test/context/llm/ai-sdk-runtime.test.ts @@ -7,7 +7,7 @@ vi.mock('ai', () => ({ })); import { generateText } from 'ai'; -import { AiSdkKtxLlmRuntime } from '../../../src/context/llm/ai-sdk-runtime.js'; +import { AiSdkKtxLlmRuntime, modelProviderName } from '../../../src/context/llm/ai-sdk-runtime.js'; describe('AiSdkKtxLlmRuntime.runAgentLoop', () => { let runtime: AiSdkKtxLlmRuntime; @@ -566,3 +566,20 @@ describe('AiSdkKtxLlmRuntime.runAgentLoop', () => { expect(serialized).not.toContain('SECRET TOOL DESCRIPTION'); }); }); + +describe('modelProviderName', () => { + it('classifies OpenAI-compatible providers for rate limiting', () => { + expect(modelProviderName({ provider: 'openai-compatible.chat' })).toBe('openai-compatible'); + expect(modelProviderName({ provider: 'openai' })).toBe('openai-compatible'); + }); + + it('classifies Vertex and Google providers as vertex', () => { + expect(modelProviderName({ provider: 'google-vertex' })).toBe('vertex'); + expect(modelProviderName({ provider: 'google.generative-ai' })).toBe('vertex'); + }); + + it('defaults unknown or Anthropic providers to anthropic-api', () => { + expect(modelProviderName({ provider: 'anthropic' })).toBe('anthropic-api'); + expect(modelProviderName({})).toBe('anthropic-api'); + }); +}); diff --git a/packages/cli/test/context/llm/local-config.test.ts b/packages/cli/test/context/llm/local-config.test.ts index d13ec4706..35b1354bb 100644 --- a/packages/cli/test/context/llm/local-config.test.ts +++ b/packages/cli/test/context/llm/local-config.test.ts @@ -82,6 +82,72 @@ describe('local ktx LLM config', () => { }); }); + it('resolves openai-compatible env references into a KtxLlmConfig', () => { + const config: KtxProjectLlmConfig = { + provider: { + backend: 'openai-compatible', + openaiCompatible: { api_key: 'env:MAAS_KEY', base_url: 'https://maas.example/compatible-mode/v1' }, // pragma: allowlist secret + }, + models: { default: 'qwen3.7-max' }, + promptCaching: { enabled: false }, + }; + + expect( + resolveLocalKtxLlmConfig(config, { MAAS_KEY: 'sk-maas' }), // pragma: allowlist secret + ).toEqual({ + backend: 'openai-compatible', + openaiCompatible: { apiKey: 'sk-maas', baseURL: 'https://maas.example/compatible-mode/v1' }, // pragma: allowlist secret + modelSlots: { default: 'qwen3.7-max' }, + promptCaching: { enabled: false }, + }); + }); + + it('creates a non-null provider for the openai-compatible backend', () => { + const createKtxLlmProvider = vi.fn(() => ({ getModel: vi.fn() }) as never); + const provider = createLocalKtxLlmProviderFromConfig( + { + provider: { backend: 'openai-compatible', openaiCompatible: { base_url: 'https://maas.example/v1' } }, + models: { default: 'qwen3.7-max' }, + }, + { createKtxLlmProvider }, + ); + + expect(provider).not.toBeNull(); + expect(createKtxLlmProvider).toHaveBeenCalledWith( + expect.objectContaining({ backend: 'openai-compatible', openaiCompatible: { baseURL: 'https://maas.example/v1' } }), + ); + }); + + it('routes the openai-compatible backend through the AI SDK runtime', () => { + const createAiSdkRuntime = vi.fn(() => ({ + generateText: vi.fn(), + generateObject: vi.fn(), + runAgentLoop: vi.fn(), + subprocessForkSpec: vi.fn(() => null), + })); + const createKtxLlmProvider = vi.fn(() => ({ getModel: vi.fn() }) as never); + + const runtime = createLocalKtxLlmRuntimeFromConfig( + { + provider: { + backend: 'openai-compatible', + openaiCompatible: { api_key: 'env:MAAS_KEY', base_url: 'https://maas.example/v1' }, // pragma: allowlist secret + }, + models: { default: 'qwen3.7-max' }, + }, + { env: { MAAS_KEY: 'sk-maas' }, createAiSdkRuntime, createKtxLlmProvider }, // pragma: allowlist secret + ); + + expect(runtime).not.toBeNull(); + expect(createAiSdkRuntime).toHaveBeenCalledTimes(1); + expect(createKtxLlmProvider).toHaveBeenCalledWith( + expect.objectContaining({ + backend: 'openai-compatible', + openaiCompatible: { apiKey: 'sk-maas', baseURL: 'https://maas.example/v1' }, // pragma: allowlist secret + }), + ); + }); + it('returns null when the local LLM backend is disabled', () => { expect( createLocalKtxLlmProviderFromConfig({ diff --git a/packages/cli/test/context/project/config.test.ts b/packages/cli/test/context/project/config.test.ts index cd4be37ed..b77f2ec4e 100644 --- a/packages/cli/test/context/project/config.test.ts +++ b/packages/cli/test/context/project/config.test.ts @@ -413,6 +413,28 @@ scan: expect(config.scan.enrichment.embeddings?.dimensions).toBe(1536); }); + it('parses and round-trips an openai-compatible LLM provider block', () => { + const yaml = ` +llm: + provider: + backend: openai-compatible + openaiCompatible: + api_key: env:MAAS_KEY + base_url: https://maas.example/compatible-mode/v1 + models: + default: qwen3.7-max +`; + const config = parseKtxProjectConfig(yaml); + expect(config.llm.provider).toMatchObject({ + backend: 'openai-compatible', + openaiCompatible: { api_key: 'env:MAAS_KEY', base_url: 'https://maas.example/compatible-mode/v1' }, + }); + expect(config.llm.models.default).toBe('qwen3.7-max'); + + const roundTripped = parseKtxProjectConfig(serializeKtxProjectConfig(config)); + expect(roundTripped.llm.provider).toEqual(config.llm.provider); + }); + it('parses scan relationship settings', () => { const config = parseKtxProjectConfig(` scan: @@ -689,7 +711,7 @@ describe('generateKtxProjectConfigJsonSchema', () => { const llm = (schema.properties as Record }>).llm; const provider = llm?.properties?.provider as { properties?: Record }; const backend = provider?.properties?.backend as { enum?: readonly string[] }; - expect(backend?.enum).toEqual(['none', 'anthropic', 'vertex', 'gateway', 'claude-code', 'codex']); + expect(backend?.enum).toEqual(['none', 'anthropic', 'vertex', 'gateway', 'openai-compatible', 'claude-code', 'codex']); const storage = (schema.properties as Record }>).storage; const state = storage?.properties?.state as { enum?: readonly string[] }; diff --git a/packages/cli/test/llm/model-health.test.ts b/packages/cli/test/llm/model-health.test.ts index 5b0b0cc48..70b22f7aa 100644 --- a/packages/cli/test/llm/model-health.test.ts +++ b/packages/cli/test/llm/model-health.test.ts @@ -62,6 +62,32 @@ describe('ktx LLM health check', () => { }); }); + it('redacts the openai-compatible api key from failed health-check messages', async () => { + const openaiModel = { modelId: 'qwen3.7-max', provider: 'openai-compatible.chat' } as never; + const generateText = vi.fn(async () => { + throw new Error('401 Unauthorized bearer sk-maas-secret'); + }); + + await expect( + runKtxLlmHealthCheck( + { + backend: 'openai-compatible', + openaiCompatible: { apiKey: 'sk-maas-secret', baseURL: 'https://maas.example/v1' }, // pragma: allowlist secret + modelSlots: { default: 'qwen3.7-max' }, + }, + { + deps: { + createOpenAICompatible: vi.fn(() => vi.fn(() => openaiModel)), + generateText, + }, + }, + ), + ).resolves.toEqual({ + ok: false, + message: '401 Unauthorized bearer [redacted]', + }); + }); + it('reports claude-code as unsupported by the AI SDK health check', async () => { const result = await runKtxLlmHealthCheck({ backend: 'claude-code', diff --git a/packages/cli/test/llm/model-provider.test.ts b/packages/cli/test/llm/model-provider.test.ts index 17d47c6a1..7f049badc 100644 --- a/packages/cli/test/llm/model-provider.test.ts +++ b/packages/cli/test/llm/model-provider.test.ts @@ -206,6 +206,81 @@ describe('createKtxLlmProvider', () => { expect(gateway).toHaveBeenCalledWith('anthropic/claude-sonnet-4-6'); }); + it('uses an OpenAI-compatible endpoint with name, baseURL, and apiKey and no Anthropic beta header', () => { + const openaiModel = languageModel('qwen3.7-max', 'openai-compatible.chat'); + const openaiCompatible = vi.fn(() => openaiModel); + const createOpenAICompatible = vi.fn(() => openaiCompatible); + + const provider = createKtxLlmProvider( + { + backend: 'openai-compatible', + openaiCompatible: { apiKey: 'sk-openai-compatible', baseURL: 'https://maas.test/compatible-mode/v1' }, // pragma: allowlist secret + modelSlots: { default: 'qwen3.7-max' }, + promptCaching: { enabled: false }, + }, + { createOpenAICompatible, devtoolsEnabled: false }, + ); + + expect(provider.getModel('default')).toBe(openaiModel); + expect(createOpenAICompatible).toHaveBeenCalledWith({ + name: 'openai-compatible', + baseURL: 'https://maas.test/compatible-mode/v1', + apiKey: 'sk-openai-compatible', // pragma: allowlist secret + }); + expect(openaiCompatible).toHaveBeenCalledWith('qwen3.7-max'); + }); + + it('omits the apiKey for an OpenAI-compatible endpoint that needs none', () => { + const openaiCompatible = vi.fn((modelId: string) => languageModel(modelId, 'openai-compatible.chat')); + const createOpenAICompatible = vi.fn(() => openaiCompatible); + + const provider = createKtxLlmProvider( + { + backend: 'openai-compatible', + openaiCompatible: { baseURL: 'https://maas.test/v1' }, + modelSlots: { default: 'qwen3.7-max' }, + promptCaching: { enabled: false }, + }, + { createOpenAICompatible, devtoolsEnabled: false }, + ); + + provider.getModel('default'); + expect(createOpenAICompatible).toHaveBeenCalledWith({ + name: 'openai-compatible', + baseURL: 'https://maas.test/v1', + }); + }); + + it('throws when the OpenAI-compatible backend is missing a base URL', () => { + expect(() => + createKtxLlmProvider( + { + backend: 'openai-compatible', + openaiCompatible: { apiKey: 'sk-openai-compatible' }, // pragma: allowlist secret + modelSlots: { default: 'qwen3.7-max' }, + promptCaching: { enabled: false }, + }, + { createOpenAICompatible: vi.fn(() => vi.fn()) }, + ), + ).toThrow('openaiCompatible.baseURL is required'); + }); + + it('disables cache markers for OpenAI-compatible (non-Anthropic) models', () => { + const provider = createKtxLlmProvider( + { + backend: 'openai-compatible', + openaiCompatible: { baseURL: 'https://maas.test/v1' }, + modelSlots: { default: 'qwen3.7-max' }, + promptCaching: { enabled: true }, + }, + { + createOpenAICompatible: vi.fn(() => vi.fn((modelId: string) => languageModel(modelId, 'openai-compatible.chat'))), + }, + ); + + expect(provider.cacheMarker('1h', 'qwen3.7-max')).toBeUndefined(); + }); + it('uses explicit role overrides before default', () => { const anthropic = vi.fn((modelId: string) => languageModel(modelId, 'anthropic')); diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index 6e5a17104..a38e6b274 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -134,6 +134,9 @@ importers: '@ai-sdk/google-vertex': specifier: ^4.0.134 version: 4.0.134(zod@4.4.3) + '@ai-sdk/openai-compatible': + specifier: ^2.0.47 + version: 2.0.47(zod@4.4.3) '@anthropic-ai/claude-agent-sdk': specifier: 0.3.146 version: 0.3.146(@anthropic-ai/sdk@0.97.1(zod@4.4.3))(@modelcontextprotocol/sdk@1.29.0(zod@4.4.3))(zod@4.4.3) From b1d212bfb759d11384aaf5d4706279472d21477a Mon Sep 17 00:00:00 2001 From: TimothyDing Date: Tue, 14 Jul 2026 14:52:19 +0800 Subject: [PATCH 02/14] feat(setup): add openai-compatible LLM setup flow and flags Add the openai-compatible backend to the interactive setup wizard and the scriptable CLI flags (--llm-base-url, --llm-model, --llm-api-key-env, --llm-api-key-file). The wizard prompts for base URL, an optional API key (paste / env var / none), and a single model id applied to every ktx role. ktx status reports the new backend, and the flags validate that they only apply to --llm-backend openai-compatible. --- packages/cli/src/commands/setup-commands.ts | 40 +++ packages/cli/src/setup-models.ts | 283 +++++++++++++++++++- packages/cli/src/setup.ts | 8 + packages/cli/src/status-project.ts | 14 + packages/cli/test/index.test.ts | 75 +++++- packages/cli/test/setup-models.test.ts | 181 +++++++++++++ 6 files changed, 594 insertions(+), 7 deletions(-) diff --git a/packages/cli/src/commands/setup-commands.ts b/packages/cli/src/commands/setup-commands.ts index d838e4723..8fe661e01 100644 --- a/packages/cli/src/commands/setup-commands.ts +++ b/packages/cli/src/commands/setup-commands.ts @@ -101,6 +101,10 @@ function shouldShowSetupEntryMenu( anthropicApiKeyFile?: string; vertexProject?: string; vertexLocation?: string; + llmBaseUrl?: string; + llmModel?: string; + llmApiKeyEnv?: string; + llmApiKeyFile?: string; skipLlm?: boolean; embeddingBackend?: string; embeddingApiKeyEnv?: string; @@ -175,6 +179,10 @@ function shouldShowSetupEntryMenu( 'anthropicApiKeyFile', 'vertexProject', 'vertexLocation', + 'llmBaseUrl', + 'llmModel', + 'llmApiKeyEnv', + 'llmApiKeyFile', 'skipLlm', 'embeddingBackend', 'embeddingApiKeyEnv', @@ -244,6 +252,16 @@ export function registerSetupCommands(program: Command, context: KtxCliCommandCo ) .addOption(new Option('--vertex-project ', 'Google Vertex AI project ID, env:NAME, or file:/path').hideHelp()) .addOption(new Option('--vertex-location ', 'Google Vertex AI location, env:NAME, or file:/path').hideHelp()) + .addOption( + new Option('--llm-base-url ', 'OpenAI-compatible LLM base URL, env:NAME, or file:/path').hideHelp(), + ) + .addOption(new Option('--llm-model ', 'OpenAI-compatible model id used for every ktx role').hideHelp()) + .addOption( + new Option('--llm-api-key-env ', 'Environment variable containing the OpenAI-compatible API key').hideHelp(), + ) + .addOption( + new Option('--llm-api-key-file ', 'File containing the OpenAI-compatible API key').hideHelp(), + ) .addOption(new Option('--skip-llm', 'Leave LLM setup incomplete for now').hideHelp().default(false)) .addOption(new Option('--embedding-backend ', 'Embedding backend').argParser(embeddingBackend).hideHelp()) .addOption( @@ -381,6 +399,24 @@ export function registerSetupCommands(program: Command, context: KtxCliCommandCo context.setExitCode(1); return; } + if ( + options.llmBackend && + options.llmBackend !== 'openai-compatible' && + (options.llmBaseUrl || options.llmModel || options.llmApiKeyEnv || options.llmApiKeyFile) + ) { + context.io.stderr.write( + 'OpenAI-compatible flags (--llm-base-url, --llm-model, --llm-api-key-env, --llm-api-key-file) are only valid with --llm-backend openai-compatible.\n', + ); + context.setExitCode(1); + return; + } + if (options.llmApiKeyEnv && options.llmApiKeyFile) { + context.io.stderr.write( + 'Choose only one OpenAI-compatible credential source: --llm-api-key-env or --llm-api-key-file.\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', @@ -454,6 +490,10 @@ export function registerSetupCommands(program: Command, context: KtxCliCommandCo ...(options.anthropicApiKeyFile ? { anthropicApiKeyFile: options.anthropicApiKeyFile } : {}), ...(options.vertexProject ? { vertexProject: options.vertexProject } : {}), ...(options.vertexLocation ? { vertexLocation: options.vertexLocation } : {}), + ...(options.llmBaseUrl ? { llmBaseUrl: options.llmBaseUrl } : {}), + ...(options.llmModel ? { llmModel: options.llmModel } : {}), + ...(options.llmApiKeyEnv ? { llmApiKeyEnv: options.llmApiKeyEnv } : {}), + ...(options.llmApiKeyFile ? { llmApiKeyFile: options.llmApiKeyFile } : {}), skipLlm: options.skipLlm === true, ...(options.embeddingBackend ? { embeddingBackend: options.embeddingBackend } : {}), ...(options.embeddingApiKeyEnv ? { embeddingApiKeyEnv: options.embeddingApiKeyEnv } : {}), diff --git a/packages/cli/src/setup-models.ts b/packages/cli/src/setup-models.ts index 8b7af13eb..31029068d 100644 --- a/packages/cli/src/setup-models.ts +++ b/packages/cli/src/setup-models.ts @@ -39,6 +39,10 @@ export interface KtxSetupModelArgs { anthropicApiKeyFile?: string; vertexProject?: string; vertexLocation?: string; + llmBaseUrl?: string; + llmModel?: string; + llmApiKeyEnv?: string; + llmApiKeyFile?: string; forcePrompt?: boolean; showPromptInstructions?: boolean; skipLlm: boolean; @@ -55,7 +59,7 @@ export type KtxSetupModelResult = // The CLI arg parser, the interactive prompt, and the missing-backend error all // derive from this list, so adding a backend is one edit. Order is the prompt's // preference order (subscription backends first). -const KTX_SETUP_LLM_BACKENDS = ['claude-code', 'codex', 'anthropic', 'vertex'] as const; +const KTX_SETUP_LLM_BACKENDS = ['claude-code', 'codex', 'openai-compatible', 'anthropic', 'vertex'] as const; export type KtxSetupLlmBackend = (typeof KTX_SETUP_LLM_BACKENDS)[number]; /** Validates a raw CLI or prompt value against the setup-selectable LLM backends. */ @@ -69,6 +73,7 @@ export function isKtxSetupLlmBackend(value: string): value is KtxSetupLlmBackend const KTX_SETUP_LLM_BACKEND_LABELS: Record = { 'claude-code': 'Claude subscription (Pro/Max)', codex: 'Codex subscription', + 'openai-compatible': 'OpenAI-compatible endpoint (base URL + API key)', anthropic: 'Anthropic API key', vertex: 'Google Vertex AI for Anthropic Claude', }; @@ -145,12 +150,22 @@ const MODEL_PRESETS = { vertex: ANTHROPIC_PRESET, 'claude-code': CLAUDE_CODE_PRESET, codex: CODEX_PRESET, -} satisfies Record; +} satisfies Record, KtxSetupModelPreset>; -function presetForBackend(backend: KtxSetupLlmBackend): KtxSetupModelPreset { +function presetForBackend(backend: Exclude): KtxSetupModelPreset { return MODEL_PRESETS[backend]; } +// OpenAI-compatible endpoints expose arbitrary model ids, so the user names one +// model and every role points at it. Per-role overrides remain hand-editable in +// ktx.yaml. +function singleModelPreset(model: string): KtxSetupModelPreset { + return KTX_MODEL_ROLES.reduce((preset, role) => { + preset[role] = model; + return preset; + }, {} as KtxSetupModelPreset); +} + const execFileAsync = promisify(execFile); type ChooseBackendResult = @@ -166,7 +181,8 @@ type ChooseBackendResult = const MISSING_LLM_BACKEND_MESSAGE = `Missing LLM backend: pass --llm-backend with one of ${KTX_SETUP_LLM_BACKENDS.join(', ')}. ` + 'claude-code and codex use local CLI authentication; anthropic also needs --anthropic-api-key-env or ' + - '--anthropic-api-key-file, and vertex also needs --vertex-project.'; + '--anthropic-api-key-file, vertex also needs --vertex-project, and openai-compatible also needs --llm-base-url ' + + 'and --llm-model (add --llm-api-key-env or --llm-api-key-file when the endpoint requires a key).'; type VertexConfigChoice = | { @@ -237,6 +253,10 @@ export function isKtxSetupLlmConfigReady(config: KtxProjectLlmConfig): boolean { return typeof resolved.vertex?.location === 'string' && resolved.vertex.location.trim().length > 0; } + if (resolved.backend === 'openai-compatible') { + return typeof resolved.openaiCompatible?.baseURL === 'string' && resolved.openaiCompatible.baseURL.trim().length > 0; + } + return ( resolved.backend === 'anthropic' || resolved.backend === 'gateway' || @@ -254,6 +274,7 @@ function buildProjectLlmConfig( provider: | { backend: 'anthropic'; credentialRef: string } | { backend: 'vertex'; vertex: { project?: string; location: string } } + | { backend: 'openai-compatible'; credentialRef?: string; baseUrl: string } | { backend: 'claude-code' } | { backend: 'codex' }, models: KtxSetupModelPreset, @@ -285,6 +306,20 @@ function buildProjectLlmConfig( }; } + if (provider.backend === 'openai-compatible') { + return { + provider: { + backend: 'openai-compatible', + openaiCompatible: { + ...(provider.credentialRef ? { api_key: provider.credentialRef } : {}), + base_url: provider.baseUrl, + }, + }, + models, + promptCaching: existing.promptCaching, + }; + } + return { provider: { backend: 'anthropic', @@ -313,7 +348,16 @@ function buildVertexHealthConfig(vertex: { project?: string; location: string }, }; } -type LlmCheckProvider = 'Anthropic API' | 'Vertex AI' | 'Claude subscription' | 'Codex'; +function buildOpenAiCompatibleHealthConfig(baseUrl: string, apiKeyValue: string | undefined, model: string): KtxLlmConfig { + return { + backend: 'openai-compatible', + openaiCompatible: { baseURL: baseUrl, ...(apiKeyValue ? { apiKey: apiKeyValue } : {}) }, + modelSlots: { default: model }, + promptCaching: { enabled: false }, + }; +} + +type LlmCheckProvider = 'Anthropic API' | 'Vertex AI' | 'Claude subscription' | 'Codex' | 'OpenAI-compatible'; function llmCheckStartText(provider: LlmCheckProvider, model: string): string { return `Checking ${provider} LLM (${model}).`; @@ -450,6 +494,171 @@ async function chooseCredentialRef( } } +const OPENAI_COMPATIBLE_PROMPT_CONTEXT = + 'ktx sends chat/completions requests to this OpenAI-compatible endpoint to verify model access now and to run ' + + 'ingest agents that turn schemas, SQL, BI metadata, and docs into semantic-layer sources and wiki context. ' + + 'ktx.yaml stores the base URL and an env: or file: reference for the key, not the raw key.'; + +async function chooseOpenAiCompatibleBaseUrl( + args: KtxSetupModelArgs, + io: KtxCliIo, + deps: KtxSetupModelDeps, +): Promise<{ status: 'ready'; ref: string; value: string } | { status: 'back' | 'missing-input' }> { + const env = deps.env ?? process.env; + if (args.llmBaseUrl) { + let value: string | undefined; + try { + value = resolveKtxConfigReference(args.llmBaseUrl, env); + } catch { + value = undefined; + } + if (!value) { + io.stderr.write(`Missing OpenAI-compatible base URL: ${args.llmBaseUrl} could not be resolved.\n`); + return { status: 'missing-input' }; + } + return { status: 'ready', ref: args.llmBaseUrl, value }; + } + if (args.inputMode === 'disabled') { + io.stderr.write('Missing OpenAI-compatible base URL: pass --llm-base-url.\n'); + return { status: 'missing-input' }; + } + + const prompts = deps.prompts ?? createPromptAdapter(); + const url = await prompts.text({ + message: withTextInputNavigation('OpenAI-compatible base URL (for example https://host/v1)'), + }); + if (url === undefined) { + return { status: 'back' }; + } + const trimmed = url.trim(); + if (!trimmed) { + return { status: 'missing-input' }; + } + return { status: 'ready', ref: trimmed, value: trimmed }; +} + +async function chooseOpenAiCompatibleCredentialRef( + args: KtxSetupModelArgs, + io: KtxCliIo, + deps: KtxSetupModelDeps, +): Promise<{ status: 'ready'; ref?: string; value?: string } | { status: 'back' | 'missing-input' }> { + const env = deps.env ?? process.env; + if (args.llmApiKeyEnv) { + const ref = envCredentialReference(args.llmApiKeyEnv); + const value = resolveKtxConfigReference(ref, env); + if (!value) { + io.stderr.write(`Missing OpenAI-compatible API key: ${args.llmApiKeyEnv} is not set.\n`); + return { status: 'missing-input' }; + } + return { status: 'ready', ref, value }; + } + if (args.llmApiKeyFile) { + const ref = `file:${args.llmApiKeyFile}`; + let value: string | undefined; + try { + value = resolveKtxConfigReference(ref, env); + } catch { + value = undefined; + } + if (!value) { + io.stderr.write(`Missing OpenAI-compatible API key file: ${args.llmApiKeyFile}\n`); + return { status: 'missing-input' }; + } + return { status: 'ready', ref, value }; + } + if (args.inputMode === 'disabled') { + // Endpoints that need no key are valid; a scripted run without a key flag proceeds keyless. + return { status: 'ready' }; + } + + const prompts = deps.prompts ?? createPromptAdapter(); + while (true) { + const choice = await prompts.select({ + message: `How should ktx authenticate to the OpenAI-compatible endpoint?\n\n${OPENAI_COMPATIBLE_PROMPT_CONTEXT}`, + options: [ + { value: 'paste', label: 'Paste an API key and save it as a local secret file' }, + { value: 'env', label: 'Use an API key from an environment variable' }, + { value: 'none', label: 'No API key required' }, + { value: 'back', label: 'Back' }, + ], + }); + if (choice === 'back') { + return { status: 'back' }; + } + if (choice === 'none') { + return { status: 'ready' }; + } + if (choice === 'paste') { + io.stdout.write( + '│ ktx will save the key in .ktx/secrets/openai-compatible-api-key with local file permissions, then write a file: reference in ktx.yaml.\n', + ); + const value = await prompts.password({ message: withTextInputNavigation('OpenAI-compatible API key') }); + if (value === undefined) { + continue; + } + if (!value.trim()) { + return { status: 'missing-input' }; + } + const ref = await writeProjectLocalSecretReference({ + projectDir: args.projectDir, + fileName: 'openai-compatible-api-key', + value, + }); + return { status: 'ready', ref, value: value.trim() }; + } + + const envName = await prompts.text({ + message: withTextInputNavigation('Environment variable holding the API key (for example DASHSCOPE_API_KEY)'), + }); + if (envName === undefined) { + continue; + } + const trimmedName = envName.trim(); + if (!trimmedName) { + return { status: 'missing-input' }; + } + const ref = envCredentialReference(trimmedName); + const value = resolveKtxConfigReference(ref, env); + if (!value) { + io.stderr.write(`Missing OpenAI-compatible API key: ${trimmedName} is not set.\n`); + return { status: 'missing-input' }; + } + return { status: 'ready', ref, value }; + } +} + +async function chooseOpenAiCompatibleModel( + args: KtxSetupModelArgs, + io: KtxCliIo, + deps: KtxSetupModelDeps, +): Promise<{ status: 'ready'; model: string } | { status: 'back' | 'missing-input' }> { + if (args.llmModel) { + const trimmed = args.llmModel.trim(); + if (!trimmed) { + io.stderr.write('Missing OpenAI-compatible model: --llm-model was empty.\n'); + return { status: 'missing-input' }; + } + return { status: 'ready', model: trimmed }; + } + if (args.inputMode === 'disabled') { + io.stderr.write('Missing OpenAI-compatible model: pass --llm-model.\n'); + return { status: 'missing-input' }; + } + + const prompts = deps.prompts ?? createPromptAdapter(); + const model = await prompts.text({ + message: withTextInputNavigation('Model id used for every ktx role (for example qwen3.7-max)'), + }); + if (model === undefined) { + return { status: 'back' }; + } + const trimmed = model.trim(); + if (!trimmed) { + return { status: 'missing-input' }; + } + return { status: 'ready', model: trimmed }; +} + function requestedBackend(args: KtxSetupModelArgs): KtxSetupLlmBackend | undefined { if (args.llmBackend) { return args.llmBackend; @@ -460,6 +669,9 @@ function requestedBackend(args: KtxSetupModelArgs): KtxSetupLlmBackend | undefin if (args.anthropicApiKeyEnv || args.anthropicApiKeyFile) { return 'anthropic'; } + if (args.llmBaseUrl || args.llmModel || args.llmApiKeyEnv || args.llmApiKeyFile) { + return 'openai-compatible'; + } return undefined; } @@ -718,6 +930,7 @@ async function persistLlmConfig( provider: | { backend: 'anthropic'; credentialRef: string } | { backend: 'vertex'; vertex: { project?: string; location: string } } + | { backend: 'openai-compatible'; credentialRef?: string; baseUrl: string } | { backend: 'claude-code' } | { backend: 'codex' }, models: KtxSetupModelPreset, @@ -922,6 +1135,66 @@ export async function runKtxSetupAnthropicModelStep( return { status: 'ready', projectDir: args.projectDir }; } + if (backendChoice.backend === 'openai-compatible') { + const baseUrl = await chooseOpenAiCompatibleBaseUrl(backendArgs, io, deps); + if (baseUrl.status === 'back' && backendChoice.prompted) { + attemptArgs = buildInteractiveRetryArgs(args); + continue; + } + if (baseUrl.status !== 'ready') { + return { status: baseUrl.status, projectDir: args.projectDir }; + } + + const credential = await chooseOpenAiCompatibleCredentialRef(backendArgs, io, deps); + if (credential.status === 'back' && backendChoice.prompted) { + attemptArgs = buildInteractiveRetryArgs(args, backendChoice.backend); + continue; + } + if (credential.status !== 'ready') { + return { status: credential.status, projectDir: args.projectDir }; + } + + const modelChoice = await chooseOpenAiCompatibleModel(backendArgs, io, deps); + if (modelChoice.status === 'back' && backendChoice.prompted) { + attemptArgs = buildInteractiveRetryArgs(args, backendChoice.backend); + continue; + } + if (modelChoice.status !== 'ready') { + return { status: modelChoice.status, projectDir: args.projectDir }; + } + + const preset = singleModelPreset(modelChoice.model); + const validation = await validatePresetModels( + preset, + (model) => + validateModelWithProgress('OpenAI-compatible', model, deps, () => + healthCheck(buildOpenAiCompatibleHealthConfig(baseUrl.value, credential.value, model)), + ), + io, + ); + if (validation.status !== 'ready') { + io.stderr.write(`OpenAI-compatible model health check failed: ${validation.message}\n`); + if (args.inputMode === 'disabled') { + return { status: 'failed', projectDir: args.projectDir }; + } + io.stderr.write('Choose a different base URL, credential, or model, or Back.\n'); + attemptArgs = buildInteractiveRetryArgs(args, backendChoice.backend); + continue; + } + + await persistLlmConfig( + args.projectDir, + { + backend: 'openai-compatible', + baseUrl: baseUrl.ref, + ...(credential.ref ? { credentialRef: credential.ref } : {}), + }, + validation.models, + ); + io.stdout.write(`│ LLM ready: yes (openai-compatible, ${validation.models.default})\n`); + return { status: 'ready', projectDir: args.projectDir }; + } + const credential = await chooseCredentialRef(backendArgs, io, deps); if (credential.status === 'back' && backendChoice.prompted) { attemptArgs = buildInteractiveRetryArgs(args); diff --git a/packages/cli/src/setup.ts b/packages/cli/src/setup.ts index ff1d6846c..1d2d36fb6 100644 --- a/packages/cli/src/setup.ts +++ b/packages/cli/src/setup.ts @@ -91,6 +91,10 @@ export type KtxSetupArgs = anthropicApiKeyFile?: string; vertexProject?: string; vertexLocation?: string; + llmBaseUrl?: string; + llmModel?: string; + llmApiKeyEnv?: string; + llmApiKeyFile?: string; skipLlm: boolean; embeddingBackend?: 'openai' | 'sentence-transformers'; embeddingApiKeyEnv?: string; @@ -791,6 +795,10 @@ async function runKtxSetupInner(args: KtxSetupArgs, io: KtxCliIo, deps: KtxSetup ...(args.anthropicApiKeyFile ? { anthropicApiKeyFile: args.anthropicApiKeyFile } : {}), ...(args.vertexProject ? { vertexProject: args.vertexProject } : {}), ...(args.vertexLocation ? { vertexLocation: args.vertexLocation } : {}), + ...(args.llmBaseUrl ? { llmBaseUrl: args.llmBaseUrl } : {}), + ...(args.llmModel ? { llmModel: args.llmModel } : {}), + ...(args.llmApiKeyEnv ? { llmApiKeyEnv: args.llmApiKeyEnv } : {}), + ...(args.llmApiKeyFile ? { llmApiKeyFile: args.llmApiKeyFile } : {}), forcePrompt: forcePromptSteps.has('models') || runOnly === 'models', showPromptInstructions, skipLlm: args.skipLlm || !shouldRunModels, diff --git a/packages/cli/src/status-project.ts b/packages/cli/src/status-project.ts index 2acd151a5..44e94d618 100644 --- a/packages/cli/src/status-project.ts +++ b/packages/cli/src/status-project.ts @@ -273,6 +273,20 @@ async function buildLlmStatus( fix: hint ? `Set ${hint}` : 'Set the gateway api_key or rerun `ktx setup`', }; } + if (backend === 'openai-compatible') { + const baseUrl = config.provider.openaiCompatible?.base_url; + if (!baseUrl || baseUrl.trim().length === 0) { + return { + backend, + model, + status: 'warn', + detail: 'base URL not configured', + fix: 'Set llm.provider.openaiCompatible.base_url or rerun `ktx setup`', + }; + } + const resolved = resolveRef(config.provider.openaiCompatible?.api_key, env); + return { backend, model, status: 'ok', detail: resolved.resolved.length > 0 ? 'base URL set, key set' : 'base URL set' }; + } if (backend === 'claude-code') { const modelName = model; if (options.fast === true) { diff --git a/packages/cli/test/index.test.ts b/packages/cli/test/index.test.ts index 939a1dd36..63f547ea2 100644 --- a/packages/cli/test/index.test.ts +++ b/packages/cli/test/index.test.ts @@ -1221,7 +1221,7 @@ describe('runKtxCli', () => { ); }); - it('rejects the removed --llm-model setup flag', async () => { + it('rejects the --llm-model setup flag for non-openai-compatible backends', async () => { const setup = vi.fn(async () => 0); const setupIo = makeIo(); @@ -1243,7 +1243,78 @@ describe('runKtxCli', () => { ).resolves.toBe(1); expect(setup).not.toHaveBeenCalled(); - expect(setupIo.stderr()).toContain("unknown option '--llm-model'"); + expect(setupIo.stderr()).toContain('only valid with --llm-backend openai-compatible'); + }); + + it('dispatches openai-compatible 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', + 'openai-compatible', + '--llm-base-url', + 'https://maas.example/compatible-mode/v1', + '--llm-model', + 'qwen3.7-max', + '--llm-api-key-env', + 'KTX_MAAS_KEY', + ], + setupIo.io, + { setup }, + ), + ).resolves.toBe(0); + + expect(setup).toHaveBeenCalledWith( + expect.objectContaining({ + command: 'run', + projectDir: tempDir, + inputMode: 'disabled', + cliVersion, + llmBackend: 'openai-compatible', + llmBaseUrl: 'https://maas.example/compatible-mode/v1', + llmModel: 'qwen3.7-max', + llmApiKeyEnv: 'KTX_MAAS_KEY', + skipLlm: false, + }), + setupIo.io, + ); + }); + + it('rejects conflicting openai-compatible credential setup flags', async () => { + const setup = vi.fn(async () => 0); + const setupIo = makeIo(); + + await expect( + runKtxCli( + [ + '--project-dir', + tempDir, + 'setup', + '--llm-backend', + 'openai-compatible', + '--llm-base-url', + 'https://maas.example/v1', + '--llm-model', + 'qwen3.7-max', + '--llm-api-key-env', + 'KTX_MAAS_KEY', + '--llm-api-key-file', + '/tmp/openai-compatible-key', + ], + setupIo.io, + { setup }, + ), + ).resolves.toBe(1); + + expect(setup).not.toHaveBeenCalled(); + expect(setupIo.stderr()).toContain('Choose only one OpenAI-compatible credential source'); }); it('rejects conflicting Anthropic credential setup flags', async () => { diff --git a/packages/cli/test/setup-models.test.ts b/packages/cli/test/setup-models.test.ts index 4dfe66b75..72fcee93a 100644 --- a/packages/cli/test/setup-models.test.ts +++ b/packages/cli/test/setup-models.test.ts @@ -150,6 +150,7 @@ describe('setup Anthropic model step', () => { options: [ { value: 'claude-code', label: 'Claude subscription (Pro/Max)' }, { value: 'codex', label: 'Codex subscription' }, + { value: 'openai-compatible', label: 'OpenAI-compatible endpoint (base URL + API key)' }, { value: 'anthropic', label: 'Anthropic API key' }, { value: 'vertex', label: 'Google Vertex AI for Anthropic Claude' }, { value: 'back', label: 'Back' }, @@ -158,6 +159,186 @@ describe('setup Anthropic model step', () => { ); }); + it('configures the openai-compatible backend from flags and persists the provider block', async () => { + const io = makeIo(); + const healthCheck = vi.fn(async () => ({ ok: true as const })); + const { spinner } = makeSpinnerEvents(); + + const result = await runKtxSetupAnthropicModelStep( + { + projectDir: tempDir, + inputMode: 'disabled', + llmBackend: 'openai-compatible', + llmBaseUrl: 'https://maas.example/compatible-mode/v1', + llmModel: 'qwen3.7-max', + llmApiKeyEnv: 'KTX_MAAS_KEY', + skipLlm: false, + }, + io.io, + { env: { KTX_MAAS_KEY: 'sk-maas' }, healthCheck, spinner }, // pragma: allowlist secret + ); + + expect(result.status).toBe('ready'); + expect(healthCheck).toHaveBeenCalledTimes(1); + expect(healthCheck).toHaveBeenCalledWith( + expect.objectContaining({ + backend: 'openai-compatible', + openaiCompatible: { baseURL: 'https://maas.example/compatible-mode/v1', apiKey: 'sk-maas' }, // pragma: allowlist secret + modelSlots: { default: 'qwen3.7-max' }, + }), + ); + const config = parseKtxProjectConfig(await readFile(join(tempDir, 'ktx.yaml'), 'utf-8')); + expect(config.llm.provider).toMatchObject({ + backend: 'openai-compatible', + openaiCompatible: { api_key: 'env:KTX_MAAS_KEY', base_url: 'https://maas.example/compatible-mode/v1' }, // pragma: allowlist secret + }); + expect(config.llm.models.default).toBe('qwen3.7-max'); + expect(config.llm.models.curator).toBe('qwen3.7-max'); + expect(config.scan.enrichment.mode).toBe('llm'); + }); + + it('requires the openai-compatible base URL in non-interactive mode', async () => { + const io = makeIo(); + const result = await runKtxSetupAnthropicModelStep( + { projectDir: tempDir, inputMode: 'disabled', llmBackend: 'openai-compatible', llmModel: 'qwen3.7-max', skipLlm: false }, + io.io, + { env: {} }, + ); + expect(result.status).toBe('missing-input'); + expect(io.stderr()).toContain('Missing OpenAI-compatible base URL'); + }); + + it('requires the openai-compatible model in non-interactive mode', async () => { + const io = makeIo(); + const result = await runKtxSetupAnthropicModelStep( + { + projectDir: tempDir, + inputMode: 'disabled', + llmBackend: 'openai-compatible', + llmBaseUrl: 'https://maas.example/v1', + skipLlm: false, + }, + io.io, + { env: {} }, + ); + expect(result.status).toBe('missing-input'); + expect(io.stderr()).toContain('Missing OpenAI-compatible model'); + }); + + it('reports missing-input when the openai-compatible api key env var is unset', async () => { + const io = makeIo(); + const result = await runKtxSetupAnthropicModelStep( + { + projectDir: tempDir, + inputMode: 'disabled', + llmBackend: 'openai-compatible', + llmBaseUrl: 'https://maas.example/v1', + llmModel: 'qwen3.7-max', + llmApiKeyEnv: 'KTX_MAAS_KEY', + skipLlm: false, + }, + io.io, + { env: {} }, + ); + expect(result.status).toBe('missing-input'); + expect(io.stderr()).toContain('KTX_MAAS_KEY is not set'); + }); + + it('allows a keyless openai-compatible endpoint and omits api_key', async () => { + const io = makeIo(); + const healthCheck = vi.fn(async () => ({ ok: true as const })); + const { spinner } = makeSpinnerEvents(); + const result = await runKtxSetupAnthropicModelStep( + { + projectDir: tempDir, + inputMode: 'disabled', + llmBackend: 'openai-compatible', + llmBaseUrl: 'http://localhost:8000/v1', + llmModel: 'local-model', + skipLlm: false, + }, + io.io, + { env: {}, healthCheck, spinner }, + ); + expect(result.status).toBe('ready'); + expect(healthCheck).toHaveBeenCalledWith( + expect.objectContaining({ + backend: 'openai-compatible', + openaiCompatible: { baseURL: 'http://localhost:8000/v1' }, + modelSlots: { default: 'local-model' }, + }), + ); + const config = parseKtxProjectConfig(await readFile(join(tempDir, 'ktx.yaml'), 'utf-8')); + expect(config.llm.provider).toMatchObject({ + backend: 'openai-compatible', + openaiCompatible: { base_url: 'http://localhost:8000/v1' }, + }); + expect((config.llm.provider as { openaiCompatible?: { api_key?: string } }).openaiCompatible?.api_key).toBeUndefined(); + }); + + it('fails when the openai-compatible health check fails in non-interactive mode', async () => { + const io = makeIo(); + const healthCheck = vi.fn(async () => ({ ok: false as const, message: 'model not found' })); + const { spinner } = makeSpinnerEvents(); + const result = await runKtxSetupAnthropicModelStep( + { + projectDir: tempDir, + inputMode: 'disabled', + llmBackend: 'openai-compatible', + llmBaseUrl: 'https://maas.example/v1', + llmModel: 'bad-model', + llmApiKeyEnv: 'KTX_MAAS_KEY', + skipLlm: false, + }, + io.io, + { env: { KTX_MAAS_KEY: 'sk-maas' }, healthCheck, spinner }, // pragma: allowlist secret + ); + expect(result.status).toBe('failed'); + expect(io.stderr()).toContain('OpenAI-compatible model health check failed'); + }); + + it('walks the interactive openai-compatible flow and persists provider config', async () => { + const io = makeIo(); + const healthCheck = vi.fn(async () => ({ ok: true as const })); + const { spinner } = makeSpinnerEvents(); + const textValues = ['https://maas.example/v1', 'DASHSCOPE_API_KEY', 'qwen3.7-max']; + const prompts: KtxSetupModelPromptAdapter = { + select: vi.fn(async ({ message }: { message: string }) => { + if (message.includes('Which LLM provider')) { + return 'openai-compatible'; + } + if (message.includes('authenticate to the OpenAI-compatible endpoint')) { + return 'env'; + } + return 'back'; + }), + autocomplete: vi.fn(async () => 'back'), + text: vi.fn(async () => textValues.shift() ?? ''), + password: vi.fn(async () => undefined), + cancel: vi.fn(), + }; + + const result = await runKtxSetupAnthropicModelStep( + { projectDir: tempDir, inputMode: 'auto', skipLlm: false }, + io.io, + { prompts, healthCheck, spinner, env: { DASHSCOPE_API_KEY: 'sk-dash' } }, // pragma: allowlist secret + ); + + expect(result.status).toBe('ready'); + expect(healthCheck).toHaveBeenCalledWith( + expect.objectContaining({ + backend: 'openai-compatible', + openaiCompatible: { baseURL: 'https://maas.example/v1', apiKey: 'sk-dash' }, // pragma: allowlist secret + modelSlots: { default: 'qwen3.7-max' }, + }), + ); + const config = parseKtxProjectConfig(await readFile(join(tempDir, 'ktx.yaml'), 'utf-8')); + expect(config.llm.provider).toMatchObject({ + backend: 'openai-compatible', + openaiCompatible: { api_key: 'env:DASHSCOPE_API_KEY', base_url: 'https://maas.example/v1' }, // pragma: allowlist secret + }); + }); + it('configures Claude Code backend and validates local auth', async () => { const io = makeIo(); const authProbe = vi.fn(async () => ({ ok: true as const })); From 66e3705f8e8669effb76c391a933c77c69ceca97 Mon Sep 17 00:00:00 2001 From: TimothyDing Date: Tue, 14 Jul 2026 15:04:01 +0800 Subject: [PATCH 03/14] feat(setup): expose openai embedding base_url/model/dimensions The openai embedding provider already accepts a custom base URL and dimensions, but the setup wizard hardcoded text-embedding-3-small at 1536 dimensions. Prompt for (and accept flags for) the base URL, model, and dimensions so users can point embeddings at any OpenAI-compatible endpoint, including ones that require a specific vector size. --- packages/cli/src/commands/setup-commands.ts | 18 +++ packages/cli/src/setup-embeddings.ts | 116 +++++++++++++++++- packages/cli/src/setup.ts | 6 + packages/cli/test/index.test.ts | 67 ++++++++++ .../cli/test/llm/embedding-provider.test.ts | 43 +++++++ packages/cli/test/setup-embeddings.test.ts | 89 ++++++++++++++ 6 files changed, 337 insertions(+), 2 deletions(-) diff --git a/packages/cli/src/commands/setup-commands.ts b/packages/cli/src/commands/setup-commands.ts index 8fe661e01..e423464e7 100644 --- a/packages/cli/src/commands/setup-commands.ts +++ b/packages/cli/src/commands/setup-commands.ts @@ -109,6 +109,9 @@ function shouldShowSetupEntryMenu( embeddingBackend?: string; embeddingApiKeyEnv?: string; embeddingApiKeyFile?: string; + embeddingBaseUrl?: string; + embeddingModel?: string; + embeddingDimensions?: number; skipEmbeddings?: boolean; database?: KtxSetupDatabaseDriver[]; databaseConnectionId?: string[]; @@ -187,6 +190,9 @@ function shouldShowSetupEntryMenu( 'embeddingBackend', 'embeddingApiKeyEnv', 'embeddingApiKeyFile', + 'embeddingBaseUrl', + 'embeddingModel', + 'embeddingDimensions', 'skipEmbeddings', 'databaseUrl', 'enableQueryHistory', @@ -273,6 +279,15 @@ export function registerSetupCommands(program: Command, context: KtxCliCommandCo .addOption( new Option('--embedding-api-key-file ', 'File containing the embedding provider API key').hideHelp(), ) + .addOption( + new Option('--embedding-base-url ', 'OpenAI-compatible embeddings base URL, env:NAME, or file:/path').hideHelp(), + ) + .addOption(new Option('--embedding-model ', 'Embedding model id').hideHelp()) + .addOption( + new Option('--embedding-dimensions ', 'Embedding vector dimensionality') + .argParser(positiveInteger) + .hideHelp(), + ) .addOption(new Option('--skip-embeddings', 'Leave embedding setup incomplete for now').hideHelp().default(false)) .addOption( new Option('--database ', 'Database driver to configure; repeatable') @@ -498,6 +513,9 @@ export function registerSetupCommands(program: Command, context: KtxCliCommandCo ...(options.embeddingBackend ? { embeddingBackend: options.embeddingBackend } : {}), ...(options.embeddingApiKeyEnv ? { embeddingApiKeyEnv: options.embeddingApiKeyEnv } : {}), ...(options.embeddingApiKeyFile ? { embeddingApiKeyFile: options.embeddingApiKeyFile } : {}), + ...(options.embeddingBaseUrl ? { embeddingBaseUrl: options.embeddingBaseUrl } : {}), + ...(options.embeddingModel ? { embeddingModel: options.embeddingModel } : {}), + ...(options.embeddingDimensions !== undefined ? { embeddingDimensions: options.embeddingDimensions } : {}), skipEmbeddings: options.skipEmbeddings === true, ...(options.database.length > 0 ? { databaseDrivers: options.database } : {}), ...(options.databaseConnectionId.length > 0 && creatingDatabaseConnection diff --git a/packages/cli/src/setup-embeddings.ts b/packages/cli/src/setup-embeddings.ts index 4ca6594d0..f55328510 100644 --- a/packages/cli/src/setup-embeddings.ts +++ b/packages/cli/src/setup-embeddings.ts @@ -31,6 +31,9 @@ export interface KtxSetupEmbeddingsArgs { embeddingBackend?: KtxSetupEmbeddingBackend; embeddingApiKeyEnv?: string; embeddingApiKeyFile?: string; + embeddingBaseUrl?: string; + embeddingModel?: string; + embeddingDimensions?: number; forcePrompt?: boolean; showPromptInstructions?: boolean; skipEmbeddings: boolean; @@ -46,6 +49,7 @@ export type KtxSetupEmbeddingsResult = /** @internal */ export interface KtxSetupEmbeddingsPromptAdapter { select(options: { message: string; options: KtxSetupPromptOption[] }): Promise; + text(options: { message: string; placeholder?: string }): Promise; password(options: { message: string }): Promise; cancel(message: string): void; } @@ -104,6 +108,7 @@ function buildProjectEmbeddingConfig(input: { model: string; dimensions: number; credentialRef?: string; + baseUrlRef?: string; }): KtxProjectEmbeddingConfig { if (input.backend === 'openai') { return { @@ -112,6 +117,7 @@ function buildProjectEmbeddingConfig(input: { dimensions: input.dimensions, openai: { ...(input.credentialRef ? { api_key: input.credentialRef } : {}), + ...(input.baseUrlRef ? { base_url: input.baseUrlRef } : {}), }, }; } @@ -132,6 +138,7 @@ function buildHealthConfig(input: { model: string; dimensions: number; credentialValue?: string; + baseUrl?: string; }): KtxEmbeddingConfig { if (input.backend === 'openai') { return { @@ -140,6 +147,7 @@ function buildHealthConfig(input: { dimensions: input.dimensions, openai: { ...(input.credentialValue ? { apiKey: input.credentialValue } : {}), + ...(input.baseUrl ? { baseURL: input.baseUrl } : {}), }, }; } @@ -353,6 +361,93 @@ function startHealthCheckProgress( }; } +async function resolveOpenAiEmbeddingParams( + args: KtxSetupEmbeddingsArgs, + io: KtxCliIo, + deps: KtxSetupEmbeddingsDeps, +): Promise< + | { status: 'ready'; model: string; dimensions: number; baseUrlRef?: string; baseUrlValue?: string } + | { status: 'back' | 'missing-input' } +> { + const env = deps.env ?? process.env; + const defaults = DEFAULTS.openai; + + let baseUrlRef: string | undefined; + let baseUrlValue: string | undefined; + if (args.embeddingBaseUrl) { + let value: string | undefined; + try { + value = resolveKtxConfigReference(args.embeddingBaseUrl, env); + } catch { + value = undefined; + } + if (!value) { + io.stderr.write(`Missing embedding base URL: ${args.embeddingBaseUrl} could not be resolved.\n`); + return { status: 'missing-input' }; + } + baseUrlRef = args.embeddingBaseUrl; + baseUrlValue = value; + } + + let model = args.embeddingModel?.trim() || undefined; + let dimensions = args.embeddingDimensions; + + if (args.inputMode !== 'disabled') { + const prompts = deps.prompts ?? createPromptAdapter(); + if (!args.embeddingBaseUrl) { + const url = await prompts.text({ + message: withTextInputNavigation( + 'OpenAI-compatible embeddings base URL (leave empty for the default OpenAI endpoint)', + ), + }); + if (url === undefined) { + return { status: 'back' }; + } + const trimmed = url.trim(); + if (trimmed) { + baseUrlRef = trimmed; + baseUrlValue = trimmed; + } + } + if (model === undefined) { + const value = await prompts.text({ + message: withTextInputNavigation(`Embedding model (default ${defaults.model})`), + }); + if (value === undefined) { + return { status: 'back' }; + } + model = value.trim() || defaults.model; + } + if (dimensions === undefined) { + const value = await prompts.text({ + message: withTextInputNavigation(`Embedding dimensions (default ${defaults.dimensions})`), + }); + if (value === undefined) { + return { status: 'back' }; + } + const trimmed = value.trim(); + if (!trimmed) { + dimensions = defaults.dimensions; + } else { + const parsed = Number.parseInt(trimmed, 10); + if (!Number.isInteger(parsed) || parsed <= 0) { + io.stderr.write(`Invalid embedding dimensions: ${trimmed} must be a positive integer.\n`); + return { status: 'missing-input' }; + } + dimensions = parsed; + } + } + } + + return { + status: 'ready', + model: model ?? defaults.model, + dimensions: dimensions ?? defaults.dimensions, + ...(baseUrlRef ? { baseUrlRef } : {}), + ...(baseUrlValue ? { baseUrlValue } : {}), + }; +} + export async function runKtxSetupEmbeddingsStep( args: KtxSetupEmbeddingsArgs, io: KtxCliIo, @@ -391,10 +486,12 @@ export async function runKtxSetupEmbeddingsStep( } const defaults = DEFAULTS[selectedBackend]; - const model = defaults.model; - const dimensions = defaults.dimensions; + let model = defaults.model; + let dimensions = defaults.dimensions; let credentialRef: string | undefined; let credentialValue: string | undefined; + let baseUrlRef: string | undefined; + let baseUrlValue: string | undefined; if (selectedBackend === 'openai') { const credential = await chooseCredentialRef(selectedBackend, args, io, deps); @@ -407,6 +504,19 @@ export async function runKtxSetupEmbeddingsStep( } credentialRef = credential.ref; credentialValue = credential.value; + + const params = await resolveOpenAiEmbeddingParams(args, io, deps); + if (params.status === 'back' && !args.embeddingBackend && args.inputMode !== 'disabled') { + selectedBackend = undefined; + continue; + } + if (params.status !== 'ready') { + return { status: params.status, projectDir: args.projectDir }; + } + model = params.model; + dimensions = params.dimensions; + baseUrlRef = params.baseUrlRef; + baseUrlValue = params.baseUrlValue; } let managedLocalEmbeddings: ManagedLocalEmbeddingsDaemon | undefined; @@ -443,6 +553,7 @@ export async function runKtxSetupEmbeddingsStep( model, dimensions, credentialValue, + baseUrl: baseUrlValue, }); const healthSpinner = (deps.spinner ?? (() => createCliSpinner(io)))(); const progress = startHealthCheckProgress(healthSpinner, healthCheckStartText(selectedBackend, model, dimensions)); @@ -468,6 +579,7 @@ export async function runKtxSetupEmbeddingsStep( model, dimensions, credentialRef, + baseUrlRef, }), ); io.stdout.write(`│ Embeddings ready: yes (${model}, ${dimensions} dimensions)\n`); diff --git a/packages/cli/src/setup.ts b/packages/cli/src/setup.ts index 1d2d36fb6..8fe352c67 100644 --- a/packages/cli/src/setup.ts +++ b/packages/cli/src/setup.ts @@ -99,6 +99,9 @@ export type KtxSetupArgs = embeddingBackend?: 'openai' | 'sentence-transformers'; embeddingApiKeyEnv?: string; embeddingApiKeyFile?: string; + embeddingBaseUrl?: string; + embeddingModel?: string; + embeddingDimensions?: number; skipEmbeddings: boolean; databaseDrivers?: KtxSetupDatabaseDriver[]; databaseConnectionIds?: string[]; @@ -818,6 +821,9 @@ async function runKtxSetupInner(args: KtxSetupArgs, io: KtxCliIo, deps: KtxSetup ...(args.embeddingBackend ? { embeddingBackend: args.embeddingBackend } : {}), ...(args.embeddingApiKeyEnv ? { embeddingApiKeyEnv: args.embeddingApiKeyEnv } : {}), ...(args.embeddingApiKeyFile ? { embeddingApiKeyFile: args.embeddingApiKeyFile } : {}), + ...(args.embeddingBaseUrl ? { embeddingBaseUrl: args.embeddingBaseUrl } : {}), + ...(args.embeddingModel ? { embeddingModel: args.embeddingModel } : {}), + ...(args.embeddingDimensions !== undefined ? { embeddingDimensions: args.embeddingDimensions } : {}), forcePrompt: forcePromptSteps.has('embeddings') || runOnly === 'embeddings', showPromptInstructions, skipEmbeddings: args.skipEmbeddings || !shouldRunEmbeddings, diff --git a/packages/cli/test/index.test.ts b/packages/cli/test/index.test.ts index 63f547ea2..6ff274ea5 100644 --- a/packages/cli/test/index.test.ts +++ b/packages/cli/test/index.test.ts @@ -1377,6 +1377,73 @@ describe('runKtxCli', () => { ); }); + it('dispatches OpenAI-compatible embedding customization flags to the setup runner', async () => { + const setup = vi.fn(async () => 0); + const setupIo = makeIo(); + + await expect( + runKtxCli( + [ + '--project-dir', + tempDir, + 'setup', + '--no-input', + '--skip-llm', + '--embedding-backend', + 'openai', + '--embedding-api-key-env', + 'KTX_MAAS_KEY', + '--embedding-base-url', + 'https://maas.example/compatible-mode/v1', + '--embedding-model', + 'text-embedding-v4', + '--embedding-dimensions', + '1024', + ], + setupIo.io, + { setup }, + ), + ).resolves.toBe(0); + + expect(setup).toHaveBeenCalledWith( + expect.objectContaining({ + command: 'run', + embeddingBackend: 'openai', + embeddingApiKeyEnv: 'KTX_MAAS_KEY', + embeddingBaseUrl: 'https://maas.example/compatible-mode/v1', + embeddingModel: 'text-embedding-v4', + embeddingDimensions: 1024, + skipEmbeddings: false, + }), + setupIo.io, + ); + }); + + it('rejects a non-positive --embedding-dimensions value', async () => { + const setup = vi.fn(async () => 0); + const setupIo = makeIo(); + + await expect( + runKtxCli( + [ + '--project-dir', + tempDir, + 'setup', + '--no-input', + '--embedding-backend', + 'openai', + '--embedding-dimensions', + '0', + ], + setupIo.io, + { setup }, + ), + ).resolves.toBe(1); + + expect(setup).not.toHaveBeenCalled(); + expect(setupIo.stderr()).toContain('Expected a positive integer'); + }); + it('dispatches database setup flags to the setup runner', async () => { const setup = vi.fn(async () => 0); const setupIo = makeIo(); diff --git a/packages/cli/test/llm/embedding-provider.test.ts b/packages/cli/test/llm/embedding-provider.test.ts index a46cb16f8..dd64fa698 100644 --- a/packages/cli/test/llm/embedding-provider.test.ts +++ b/packages/cli/test/llm/embedding-provider.test.ts @@ -55,6 +55,49 @@ describe('createKtxEmbeddingProvider', () => { }); }); + it('sends the configured model and dimensions to the OpenAI embeddings API', async () => { + const create = vi.fn().mockResolvedValue({ data: [{ index: 0, embedding: [0.1, 0.2, 0.3] }] }); + const createOpenAIClient = vi.fn(() => ({ embeddings: { create } })); + + const provider = createKtxEmbeddingProvider( + { + backend: 'openai', + model: 'text-embedding-v4', + dimensions: 3, + openai: { apiKey: 'sk-maas', baseURL: 'https://maas.example/compatible-mode/v1' }, // pragma: allowlist secret + }, + { createOpenAIClient }, + ); + + await expect(provider.embed('hello')).resolves.toEqual([0.1, 0.2, 0.3]); + expect(create).toHaveBeenCalledWith({ + model: 'text-embedding-v4', + input: 'hello', + dimensions: 3, + encoding_format: 'float', + }); + }); + + it('caps OpenAI embedding batches at the configured batch size', async () => { + const create = vi.fn(); + const createOpenAIClient = vi.fn(() => ({ embeddings: { create } })); + + const provider = createKtxEmbeddingProvider( + { + backend: 'openai', + model: 'text-embedding-v4', + dimensions: 3, + openai: { apiKey: 'sk-maas' }, // pragma: allowlist secret + batchSize: 2, + }, + { createOpenAIClient }, + ); + + expect(provider.maxBatchSize).toBe(2); + await expect(provider.embedMany(['a', 'b', 'c'])).rejects.toThrow('Embedding batch size 3 exceeds maximum 2'); + expect(create).not.toHaveBeenCalled(); + }); + it('supports sentence-transformers pathPrefix defaults and explicit empty prefix', async () => { const fetch = vi .fn() diff --git a/packages/cli/test/setup-embeddings.test.ts b/packages/cli/test/setup-embeddings.test.ts index 62ebda974..04303d145 100644 --- a/packages/cli/test/setup-embeddings.test.ts +++ b/packages/cli/test/setup-embeddings.test.ts @@ -40,10 +40,13 @@ function makeIo() { function makePromptAdapter(options: { selectValues?: string[]; passwordValue?: string; + textValues?: string[]; }): KtxSetupEmbeddingsPromptAdapter { const selectValues = [...(options.selectValues ?? [])]; + const textValues = [...(options.textValues ?? [])]; return { select: vi.fn(async () => selectValues.shift() ?? 'retry'), + text: vi.fn(async () => textValues.shift() ?? ''), password: vi.fn(async () => options.passwordValue ?? 'embedding-secret'), cancel: vi.fn(), }; @@ -466,6 +469,92 @@ describe('setup embeddings step', () => { expect(io.stdout()).not.toContain('sk-openai-test'); }); + it('configures an OpenAI-compatible endpoint with custom base URL, model, and dimensions from flags', async () => { + const io = makeIo(); + const healthCheck = vi.fn(async () => ({ ok: true as const })); + + const result = await runKtxSetupEmbeddingsStep( + { + projectDir: tempDir, + inputMode: 'disabled', + embeddingBackend: 'openai', + embeddingApiKeyEnv: 'KTX_MAAS_KEY', + embeddingBaseUrl: 'https://maas.example/compatible-mode/v1', + embeddingModel: 'text-embedding-v4', + embeddingDimensions: 1024, + cliVersion: '0.2.0', + runtimeInstallPolicy: 'auto', + skipEmbeddings: false, + }, + io.io, + { env: { KTX_MAAS_KEY: 'sk-maas' }, healthCheck }, // pragma: allowlist secret + ); + + expect(result.status).toBe('ready'); + expect(healthCheck).toHaveBeenCalledWith({ + backend: 'openai', + model: 'text-embedding-v4', + dimensions: 1024, + openai: { apiKey: 'sk-maas', baseURL: 'https://maas.example/compatible-mode/v1' }, // pragma: allowlist secret + }); + const config = parseKtxProjectConfig(await readFile(join(tempDir, 'ktx.yaml'), 'utf-8')); + expect(config.ingest.embeddings).toMatchObject({ + backend: 'openai', + model: 'text-embedding-v4', + dimensions: 1024, + openai: { api_key: 'env:KTX_MAAS_KEY', base_url: 'https://maas.example/compatible-mode/v1' }, // pragma: allowlist secret + }); + }); + + it('prompts for base URL, model, and dimensions in the interactive OpenAI flow', async () => { + const io = makeIo(); + const healthCheck = vi.fn(async () => ({ ok: true as const })); + const prompts = makePromptAdapter({ + selectValues: ['openai', 'env'], + textValues: ['https://maas.example/v1', 'text-embedding-v4', '1024'], + }); + + const result = await runKtxSetupEmbeddingsStep( + { projectDir: tempDir, inputMode: 'auto', cliVersion: '0.2.0', runtimeInstallPolicy: 'auto', skipEmbeddings: false }, + io.io, + { env: { OPENAI_API_KEY: 'sk-openai' }, healthCheck, prompts }, // pragma: allowlist secret + ); + + expect(result.status).toBe('ready'); + expect(healthCheck).toHaveBeenCalledWith({ + backend: 'openai', + model: 'text-embedding-v4', + dimensions: 1024, + openai: { apiKey: 'sk-openai', baseURL: 'https://maas.example/v1' }, // pragma: allowlist secret + }); + const config = parseKtxProjectConfig(await readFile(join(tempDir, 'ktx.yaml'), 'utf-8')); + expect(config.ingest.embeddings).toMatchObject({ + backend: 'openai', + model: 'text-embedding-v4', + dimensions: 1024, + openai: { api_key: 'env:OPENAI_API_KEY', base_url: 'https://maas.example/v1' }, // pragma: allowlist secret + }); + }); + + it('rejects a non-positive-integer dimensions value entered interactively', async () => { + const io = makeIo(); + const healthCheck = vi.fn(async () => ({ ok: true as const })); + const prompts = makePromptAdapter({ + selectValues: ['openai', 'env'], + textValues: ['', 'text-embedding-v4', 'not-a-number'], + }); + + const result = await runKtxSetupEmbeddingsStep( + { projectDir: tempDir, inputMode: 'auto', cliVersion: '0.2.0', runtimeInstallPolicy: 'auto', skipEmbeddings: false }, + io.io, + { env: { OPENAI_API_KEY: 'sk-openai' }, healthCheck, prompts }, // pragma: allowlist secret + ); + + expect(result.status).toBe('missing-input'); + expect(io.stderr()).toContain('must be a positive integer'); + expect(healthCheck).not.toHaveBeenCalled(); + }); + it('can fall back to OpenAI after the default local daemon is unavailable', async () => { const io = makeIo(); const prompts = makePromptAdapter({ selectValues: ['sentence-transformers', 'openai', 'env'] }); From 94fe6e941ac7217d9cc32dde269293730223ac5c Mon Sep 17 00:00:00 2001 From: TimothyDing Date: Tue, 14 Jul 2026 15:07:16 +0800 Subject: [PATCH 04/14] docs: document openai-compatible LLM backend and embedding options Document the new openai-compatible LLM backend (provider block, setup flags, single-model preset) and the openai embedding base_url/model/ dimensions options across the LLM configuration guide, the ktx setup CLI reference, and the ktx.yaml reference. --- .../content/docs/cli-reference/ktx-setup.mdx | 22 ++++++++-- .../content/docs/configuration/ktx-yaml.mdx | 4 +- .../content/docs/guides/llm-configuration.mdx | 42 +++++++++++++++++++ 3 files changed, 63 insertions(+), 5 deletions(-) diff --git a/docs-site/content/docs/cli-reference/ktx-setup.mdx b/docs-site/content/docs/cli-reference/ktx-setup.mdx index d127525a6..2e36b7634 100644 --- a/docs-site/content/docs/cli-reference/ktx-setup.mdx +++ b/docs-site/content/docs/cli-reference/ktx-setup.mdx @@ -60,18 +60,26 @@ prompts. | Flag | Description | |------|-------------| -| `--llm-backend ` | LLM backend: `anthropic`, `vertex`, `claude-code`, or `codex` | +| `--llm-backend ` | LLM backend: `anthropic`, `vertex`, `gateway`, `openai-compatible`, `claude-code`, or `codex` | | `--llm-backend claude-code` | Use the local Claude Code session for **ktx** LLM calls | | `--llm-backend codex` | Use local Codex authentication for **ktx** LLM calls | | `--anthropic-api-key-env ` | Environment variable containing the Anthropic API key | | `--anthropic-api-key-file ` | File containing the Anthropic API key | | `--vertex-project ` | Vertex AI project ID, `env:NAME`, or `file:/path` reference | | `--vertex-location ` | Vertex AI location, `env:NAME`, or `file:/path` reference | +| `--llm-base-url ` | OpenAI-compatible base URL, `env:NAME`, or `file:/path` reference | +| `--llm-model ` | OpenAI-compatible model id applied to every ktx role | +| `--llm-api-key-env ` | Environment variable containing the OpenAI-compatible API key | +| `--llm-api-key-file ` | File containing the OpenAI-compatible API key | | `--skip-llm` | Leave LLM setup incomplete | Choose only one Anthropic credential source. Anthropic credential flags are only valid with the Anthropic backend; Vertex flags are only valid with the Vertex -backend. The `claude-code` and `codex` backends use local authentication instead +backend; the `--llm-base-url`, `--llm-model`, and `--llm-api-key-*` flags are only +valid with `--llm-backend openai-compatible` (and `--llm-api-key-*` are mutually +exclusive). The `openai-compatible` backend needs `--llm-base-url` and +`--llm-model`; the API key flags are optional for keyless endpoints. The +`claude-code` and `codex` backends use local authentication instead of Anthropic API key or Vertex flags. After you choose a backend, `ktx setup` writes that backend's per-role model preset to `ktx.yaml`. To change a model, edit the matching `llm.models.` value in `ktx.yaml`. @@ -90,10 +98,16 @@ calls. | `--embedding-backend ` | Embedding backend: `openai` or `sentence-transformers` | | `--embedding-api-key-env ` | Environment variable containing the embedding provider API key | | `--embedding-api-key-file ` | File containing the embedding provider API key | +| `--embedding-base-url ` | OpenAI-compatible embeddings base URL, `env:NAME`, or `file:/path` reference | +| `--embedding-model ` | Embedding model id (default `text-embedding-3-small`) | +| `--embedding-dimensions ` | Embedding vector dimensionality (default `1536`) | | `--skip-embeddings` | Leave embedding setup incomplete | `sentence-transformers` uses the **ktx**-managed Python runtime. Choose only one -embedding credential source. +embedding credential source. The `openai` backend can target any +OpenAI-compatible embeddings endpoint: set `--embedding-base-url` and pick the +`--embedding-model` and `--embedding-dimensions` your endpoint expects (for +example `text-embedding-v4` at `1024` dimensions on Alibaba Bailian/DashScope). ### Runtime @@ -314,7 +328,7 @@ Use `ktx status` for repeatable readiness checks after setup exits. |-------|-------|----------| | Setup resumes an unexpected project | `KTX_PROJECT_DIR` or nearest `ktx.yaml` points to another directory | Pass `--project-dir ` explicitly | | Setup cannot run in CI | Required values are missing and `--no-input` disables prompts | Provide the relevant automation flags or create a fixture `ktx.yaml` | -| `Missing LLM backend: pass --llm-backend …` | `--no-input` setup ran without an LLM backend; `--target` does not select one | Pass `--llm-backend claude-code`, `codex`, `anthropic`, or `vertex` (with that backend's credential flags) | +| `Missing LLM backend: pass --llm-backend …` | `--no-input` setup ran without an LLM backend; `--target` does not select one | Pass `--llm-backend claude-code`, `codex`, `anthropic`, `vertex`, or `openai-compatible` (with that backend's credential flags; `openai-compatible` also needs `--llm-base-url` and `--llm-model`) | | Provider health check fails | Provider key, model id, Vertex project, or Vertex location is invalid | Fix the `env:` or `file:` reference and rerun setup | | Python runtime is missing | The selected setup needs runtime-backed agent, query-history, Looker, or local embedding features | Accept the interactive prompt, rerun with `--yes`, or run the suggested `ktx admin runtime install` command | | `--enable-query-history` is rejected | The selected database driver does not support query history | Use Postgres, BigQuery, or Snowflake, or rerun without query-history flags | diff --git a/docs-site/content/docs/configuration/ktx-yaml.mdx b/docs-site/content/docs/configuration/ktx-yaml.mdx index 05ebea3ee..d9e5e264d 100644 --- a/docs-site/content/docs/configuration/ktx-yaml.mdx +++ b/docs-site/content/docs/configuration/ktx-yaml.mdx @@ -473,10 +473,12 @@ llm: | Field | Type | Default | Purpose | |-------|------|---------|---------| -| `provider.backend` | `none` \| `anthropic` \| `vertex` \| `gateway` \| `claude-code` \| `codex` | `none` | Selected backend. `none` disables LLM features. `claude-code` uses the local Claude Code session and needs no API key. `codex` uses local Codex authentication and needs no API key. | +| `provider.backend` | `none` \| `anthropic` \| `vertex` \| `gateway` \| `openai-compatible` \| `claude-code` \| `codex` | `none` | Selected backend. `none` disables LLM features. `claude-code` uses the local Claude Code session and needs no API key. `codex` uses local Codex authentication and needs no API key. | | `provider.anthropic.api_key` | `string` | - | Anthropic API key. Required when `backend: anthropic`. Accepts `env:` or `file:` references. | | `provider.anthropic.base_url` | `string` | - | Override the Anthropic API base URL (proxy, self-hosted gateway). | | `provider.gateway.api_key` / `base_url` | `string` | - | Credentials for an AI Gateway provider. Required when `backend: gateway`. | +| `provider.openaiCompatible.base_url` | `string` | - | Base URL of an OpenAI-compatible chat completions endpoint. Required when `backend: openai-compatible`. Accepts `env:` or `file:` references. | +| `provider.openaiCompatible.api_key` | `string` | - | API key for the OpenAI-compatible endpoint. Optional for keyless endpoints. Accepts `env:` or `file:` references. | | `provider.vertex.project` | `string` | - | Google Cloud project ID hosting the Vertex AI endpoint. | | `provider.vertex.location` | `string` | - | Vertex AI region (for example `us-east5`). Required when the `vertex` block is present. | diff --git a/docs-site/content/docs/guides/llm-configuration.mdx b/docs-site/content/docs/guides/llm-configuration.mdx index 776cb275e..b9a961391 100644 --- a/docs-site/content/docs/guides/llm-configuration.mdx +++ b/docs-site/content/docs/guides/llm-configuration.mdx @@ -14,10 +14,52 @@ Set `llm.provider.backend` to one of these values: configured `api_key` reference. - `vertex`: Use Vertex AI Anthropic models through Google Cloud credentials. - `gateway`: Use AI Gateway-compatible Anthropic model ids. +- `openai-compatible`: Use any OpenAI-compatible chat completions endpoint + (Azure OpenAI, Alibaba Bailian/DashScope, vLLM, LiteLLM, Ollama, and similar) + through a base URL and an optional API key. - `claude-code`: Use your local Claude Code session through the Claude Agent SDK. **ktx** strips provider-routing environment variables from child processes. - `codex`: Use your local Codex authentication through the Codex SDK. +## OpenAI-compatible endpoint + +Use `openai-compatible` to route ktx through any endpoint that speaks the OpenAI +chat completions protocol. Set `base_url`; `api_key` is optional for endpoints +that need no key (for example a local Ollama or vLLM server). Because +OpenAI-compatible endpoints expose arbitrary model ids, set one model and reuse +it for every role — override per role by hand if you need to. + +```yaml +llm: + provider: + backend: openai-compatible + openaiCompatible: + api_key: env:DASHSCOPE_API_KEY + base_url: https://dashscope.aliyuncs.com/compatible-mode/v1 + models: + default: qwen-max + triage: qwen-max + candidateExtraction: qwen-max + curator: qwen-max + reconcile: qwen-max + repair: qwen-max +``` + +Configure it non-interactively: + +```bash +ktx setup \ + --llm-backend openai-compatible \ + --llm-base-url https://dashscope.aliyuncs.com/compatible-mode/v1 \ + --llm-model qwen-max \ + --llm-api-key-env DASHSCOPE_API_KEY \ + --no-input +``` + +Drop `--llm-api-key-env` (and `--llm-api-key-file`) for endpoints that require no +key. Anthropic-style `llm.promptCaching` does not apply to OpenAI-compatible +models and is ignored for this backend. + ## Claude Code Use aliases or full Claude model IDs in `llm.models`: From 5aad79274a419a23e1c5d543eb25af525ddff71d Mon Sep 17 00:00:00 2001 From: TimothyDing Date: Tue, 14 Jul 2026 16:50:35 +0800 Subject: [PATCH 05/14] fix(setup): add --embedding-batch-size for batch-limited endpoints OpenAI-compatible embedding endpoints can cap the number of inputs per request (Alibaba Bailian text-embedding-v4 rejects batches larger than 10), but ktx defaulted to 100, so schema-scan enrichment failed with an HTTP 400 during ingest. Expose the embedding batch size in the setup wizard and as a CLI flag, persisting it to ingest.embeddings.batchSize (mirrored into scan.enrichment.embeddings) so scan and ingest chunk embedding requests within the endpoint's limit. --- .../content/docs/cli-reference/ktx-setup.mdx | 4 +++ packages/cli/src/commands/setup-commands.ts | 8 +++++ packages/cli/src/setup-embeddings.ts | 29 ++++++++++++++++++- packages/cli/src/setup.ts | 2 ++ packages/cli/test/index.test.ts | 3 ++ packages/cli/test/setup-embeddings.test.ts | 10 +++++-- 6 files changed, 53 insertions(+), 3 deletions(-) diff --git a/docs-site/content/docs/cli-reference/ktx-setup.mdx b/docs-site/content/docs/cli-reference/ktx-setup.mdx index 2e36b7634..e431ac586 100644 --- a/docs-site/content/docs/cli-reference/ktx-setup.mdx +++ b/docs-site/content/docs/cli-reference/ktx-setup.mdx @@ -101,6 +101,7 @@ calls. | `--embedding-base-url ` | OpenAI-compatible embeddings base URL, `env:NAME`, or `file:/path` reference | | `--embedding-model ` | Embedding model id (default `text-embedding-3-small`) | | `--embedding-dimensions ` | Embedding vector dimensionality (default `1536`) | +| `--embedding-batch-size ` | Max texts per embedding request; set this when the endpoint caps batch size | | `--skip-embeddings` | Leave embedding setup incomplete | `sentence-transformers` uses the **ktx**-managed Python runtime. Choose only one @@ -108,6 +109,9 @@ embedding credential source. The `openai` backend can target any OpenAI-compatible embeddings endpoint: set `--embedding-base-url` and pick the `--embedding-model` and `--embedding-dimensions` your endpoint expects (for example `text-embedding-v4` at `1024` dimensions on Alibaba Bailian/DashScope). +Some endpoints reject large embedding batches (DashScope `text-embedding-v4` +caps a request at 10 inputs); set `--embedding-batch-size 10` so schema-scan and +ingest embedding requests stay within the limit. ### Runtime diff --git a/packages/cli/src/commands/setup-commands.ts b/packages/cli/src/commands/setup-commands.ts index e423464e7..55e4935e0 100644 --- a/packages/cli/src/commands/setup-commands.ts +++ b/packages/cli/src/commands/setup-commands.ts @@ -112,6 +112,7 @@ function shouldShowSetupEntryMenu( embeddingBaseUrl?: string; embeddingModel?: string; embeddingDimensions?: number; + embeddingBatchSize?: number; skipEmbeddings?: boolean; database?: KtxSetupDatabaseDriver[]; databaseConnectionId?: string[]; @@ -193,6 +194,7 @@ function shouldShowSetupEntryMenu( 'embeddingBaseUrl', 'embeddingModel', 'embeddingDimensions', + 'embeddingBatchSize', 'skipEmbeddings', 'databaseUrl', 'enableQueryHistory', @@ -288,6 +290,11 @@ export function registerSetupCommands(program: Command, context: KtxCliCommandCo .argParser(positiveInteger) .hideHelp(), ) + .addOption( + new Option('--embedding-batch-size ', 'Max texts per embedding request (some endpoints cap this, e.g. 10)') + .argParser(positiveInteger) + .hideHelp(), + ) .addOption(new Option('--skip-embeddings', 'Leave embedding setup incomplete for now').hideHelp().default(false)) .addOption( new Option('--database ', 'Database driver to configure; repeatable') @@ -516,6 +523,7 @@ export function registerSetupCommands(program: Command, context: KtxCliCommandCo ...(options.embeddingBaseUrl ? { embeddingBaseUrl: options.embeddingBaseUrl } : {}), ...(options.embeddingModel ? { embeddingModel: options.embeddingModel } : {}), ...(options.embeddingDimensions !== undefined ? { embeddingDimensions: options.embeddingDimensions } : {}), + ...(options.embeddingBatchSize !== undefined ? { embeddingBatchSize: options.embeddingBatchSize } : {}), skipEmbeddings: options.skipEmbeddings === true, ...(options.database.length > 0 ? { databaseDrivers: options.database } : {}), ...(options.databaseConnectionId.length > 0 && creatingDatabaseConnection diff --git a/packages/cli/src/setup-embeddings.ts b/packages/cli/src/setup-embeddings.ts index f55328510..84f0c5986 100644 --- a/packages/cli/src/setup-embeddings.ts +++ b/packages/cli/src/setup-embeddings.ts @@ -34,6 +34,7 @@ export interface KtxSetupEmbeddingsArgs { embeddingBaseUrl?: string; embeddingModel?: string; embeddingDimensions?: number; + embeddingBatchSize?: number; forcePrompt?: boolean; showPromptInstructions?: boolean; skipEmbeddings: boolean; @@ -109,6 +110,7 @@ function buildProjectEmbeddingConfig(input: { dimensions: number; credentialRef?: string; baseUrlRef?: string; + batchSize?: number; }): KtxProjectEmbeddingConfig { if (input.backend === 'openai') { return { @@ -119,6 +121,7 @@ function buildProjectEmbeddingConfig(input: { ...(input.credentialRef ? { api_key: input.credentialRef } : {}), ...(input.baseUrlRef ? { base_url: input.baseUrlRef } : {}), }, + ...(input.batchSize !== undefined ? { batchSize: input.batchSize } : {}), }; } const defaults = DEFAULTS[input.backend]; @@ -366,7 +369,7 @@ async function resolveOpenAiEmbeddingParams( io: KtxCliIo, deps: KtxSetupEmbeddingsDeps, ): Promise< - | { status: 'ready'; model: string; dimensions: number; baseUrlRef?: string; baseUrlValue?: string } + | { status: 'ready'; model: string; dimensions: number; baseUrlRef?: string; baseUrlValue?: string; batchSize?: number } | { status: 'back' | 'missing-input' } > { const env = deps.env ?? process.env; @@ -391,6 +394,7 @@ async function resolveOpenAiEmbeddingParams( let model = args.embeddingModel?.trim() || undefined; let dimensions = args.embeddingDimensions; + let batchSize = args.embeddingBatchSize; if (args.inputMode !== 'disabled') { const prompts = deps.prompts ?? createPromptAdapter(); @@ -437,6 +441,25 @@ async function resolveOpenAiEmbeddingParams( dimensions = parsed; } } + if (batchSize === undefined) { + const value = await prompts.text({ + message: withTextInputNavigation( + 'Embedding batch size (max texts per request; leave empty for the provider default; some endpoints cap this, e.g. 10)', + ), + }); + if (value === undefined) { + return { status: 'back' }; + } + const trimmed = value.trim(); + if (trimmed) { + const parsed = Number.parseInt(trimmed, 10); + if (!Number.isInteger(parsed) || parsed <= 0) { + io.stderr.write(`Invalid embedding batch size: ${trimmed} must be a positive integer.\n`); + return { status: 'missing-input' }; + } + batchSize = parsed; + } + } } return { @@ -445,6 +468,7 @@ async function resolveOpenAiEmbeddingParams( dimensions: dimensions ?? defaults.dimensions, ...(baseUrlRef ? { baseUrlRef } : {}), ...(baseUrlValue ? { baseUrlValue } : {}), + ...(batchSize !== undefined ? { batchSize } : {}), }; } @@ -492,6 +516,7 @@ export async function runKtxSetupEmbeddingsStep( let credentialValue: string | undefined; let baseUrlRef: string | undefined; let baseUrlValue: string | undefined; + let embeddingBatchSize: number | undefined; if (selectedBackend === 'openai') { const credential = await chooseCredentialRef(selectedBackend, args, io, deps); @@ -517,6 +542,7 @@ export async function runKtxSetupEmbeddingsStep( dimensions = params.dimensions; baseUrlRef = params.baseUrlRef; baseUrlValue = params.baseUrlValue; + embeddingBatchSize = params.batchSize; } let managedLocalEmbeddings: ManagedLocalEmbeddingsDaemon | undefined; @@ -580,6 +606,7 @@ export async function runKtxSetupEmbeddingsStep( dimensions, credentialRef, baseUrlRef, + batchSize: embeddingBatchSize, }), ); io.stdout.write(`│ Embeddings ready: yes (${model}, ${dimensions} dimensions)\n`); diff --git a/packages/cli/src/setup.ts b/packages/cli/src/setup.ts index 8fe352c67..7b5538c03 100644 --- a/packages/cli/src/setup.ts +++ b/packages/cli/src/setup.ts @@ -102,6 +102,7 @@ export type KtxSetupArgs = embeddingBaseUrl?: string; embeddingModel?: string; embeddingDimensions?: number; + embeddingBatchSize?: number; skipEmbeddings: boolean; databaseDrivers?: KtxSetupDatabaseDriver[]; databaseConnectionIds?: string[]; @@ -824,6 +825,7 @@ async function runKtxSetupInner(args: KtxSetupArgs, io: KtxCliIo, deps: KtxSetup ...(args.embeddingBaseUrl ? { embeddingBaseUrl: args.embeddingBaseUrl } : {}), ...(args.embeddingModel ? { embeddingModel: args.embeddingModel } : {}), ...(args.embeddingDimensions !== undefined ? { embeddingDimensions: args.embeddingDimensions } : {}), + ...(args.embeddingBatchSize !== undefined ? { embeddingBatchSize: args.embeddingBatchSize } : {}), forcePrompt: forcePromptSteps.has('embeddings') || runOnly === 'embeddings', showPromptInstructions, skipEmbeddings: args.skipEmbeddings || !shouldRunEmbeddings, diff --git a/packages/cli/test/index.test.ts b/packages/cli/test/index.test.ts index 6ff274ea5..e9914344c 100644 --- a/packages/cli/test/index.test.ts +++ b/packages/cli/test/index.test.ts @@ -1399,6 +1399,8 @@ describe('runKtxCli', () => { 'text-embedding-v4', '--embedding-dimensions', '1024', + '--embedding-batch-size', + '10', ], setupIo.io, { setup }, @@ -1413,6 +1415,7 @@ describe('runKtxCli', () => { embeddingBaseUrl: 'https://maas.example/compatible-mode/v1', embeddingModel: 'text-embedding-v4', embeddingDimensions: 1024, + embeddingBatchSize: 10, skipEmbeddings: false, }), setupIo.io, diff --git a/packages/cli/test/setup-embeddings.test.ts b/packages/cli/test/setup-embeddings.test.ts index 04303d145..1932fcc41 100644 --- a/packages/cli/test/setup-embeddings.test.ts +++ b/packages/cli/test/setup-embeddings.test.ts @@ -482,6 +482,7 @@ describe('setup embeddings step', () => { embeddingBaseUrl: 'https://maas.example/compatible-mode/v1', embeddingModel: 'text-embedding-v4', embeddingDimensions: 1024, + embeddingBatchSize: 10, cliVersion: '0.2.0', runtimeInstallPolicy: 'auto', skipEmbeddings: false, @@ -502,16 +503,20 @@ describe('setup embeddings step', () => { backend: 'openai', model: 'text-embedding-v4', dimensions: 1024, + batchSize: 10, openai: { api_key: 'env:KTX_MAAS_KEY', base_url: 'https://maas.example/compatible-mode/v1' }, // pragma: allowlist secret }); + // persistEmbeddingConfig mirrors the same config into the scan enrichment block, + // so the schema scan's embedding requests honor the endpoint's batch cap too. + expect(config.scan.enrichment.embeddings?.batchSize).toBe(10); }); - it('prompts for base URL, model, and dimensions in the interactive OpenAI flow', async () => { + it('prompts for base URL, model, dimensions, and batch size in the interactive OpenAI flow', async () => { const io = makeIo(); const healthCheck = vi.fn(async () => ({ ok: true as const })); const prompts = makePromptAdapter({ selectValues: ['openai', 'env'], - textValues: ['https://maas.example/v1', 'text-embedding-v4', '1024'], + textValues: ['https://maas.example/v1', 'text-embedding-v4', '1024', '10'], }); const result = await runKtxSetupEmbeddingsStep( @@ -532,6 +537,7 @@ describe('setup embeddings step', () => { backend: 'openai', model: 'text-embedding-v4', dimensions: 1024, + batchSize: 10, openai: { api_key: 'env:OPENAI_API_KEY', base_url: 'https://maas.example/v1' }, // pragma: allowlist secret }); }); From cca5c6ecdcb027dade7276383f47b165af028c0c Mon Sep 17 00:00:00 2001 From: TimothyDing Date: Wed, 15 Jul 2026 08:35:29 +0800 Subject: [PATCH 06/14] chore(git): add .qoder to .gitignore - Exclude .qoder folder from .gitignore - Remove trailing empty lines at the end of the file --- .gitignore | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/.gitignore b/.gitignore index 0ad460670..3839484d4 100644 --- a/.gitignore +++ b/.gitignore @@ -58,6 +58,7 @@ yarn-error.log* .claude .superpowers docs/superpowers +.qoder # Editors and OS files .idea/ @@ -69,4 +70,4 @@ docs/superpowers *.swo *~ .vercel -.devtools +.devtools \ No newline at end of file From 07b6f9f9f650a11a776a11a74ff00bf719bb1710 Mon Sep 17 00:00:00 2001 From: TimothyDing Date: Wed, 15 Jul 2026 15:23:00 +0800 Subject: [PATCH 07/14] feat(hologres): add Hologres scan connector and driver registration Register hologres as a first-class PostgreSQL-wire-compatible driver. KtxHologresScanConnector extends the postgres connector, reusing its transaction-free connection, introspection, sampling, and pg_stats-based statistics, and only overrides listSchemas to drop Hologres engine-internal schemas (matched exactly so user schemas like hologres_dataset_* are kept). The postgres connector now derives its driver from the connection config so snapshots stamp the correct driver. Hologres reuses the postgres SQL dialect and sqlglot dialect; CLI live-database dispatch routes it to the native Node path so it never hits the daemon (which uses transactions Hologres forbids). --- packages/cli/src/connection-drivers.ts | 1 + .../cli/src/connectors/hologres/connector.ts | 37 +++++ .../hologres/live-database-introspection.ts | 47 +++++++ .../cli/src/connectors/postgres/connector.ts | 19 ++- .../cli/src/context/connections/dialects.ts | 1 + .../cli/src/context/connections/drivers.ts | 21 +++ .../connections/local-warehouse-descriptor.ts | 1 + .../cli/src/context/project/driver-schemas.ts | 2 + packages/cli/src/context/scan/types.ts | 1 + .../cli/src/context/sql-analysis/dialect.ts | 1 + packages/cli/src/local-adapters.ts | 8 ++ .../connectors/hologres/connector.test.ts | 128 ++++++++++++++++++ .../hologres/live.integration.test.ts | 52 +++++++ .../test/context/connections/dialects.test.ts | 2 +- .../test/context/connections/drivers.test.ts | 6 + .../context/project/driver-schemas.test.ts | 1 + 16 files changed, 323 insertions(+), 5 deletions(-) create mode 100644 packages/cli/src/connectors/hologres/connector.ts create mode 100644 packages/cli/src/connectors/hologres/live-database-introspection.ts create mode 100644 packages/cli/test/connectors/hologres/connector.test.ts create mode 100644 packages/cli/test/connectors/hologres/live.integration.test.ts diff --git a/packages/cli/src/connection-drivers.ts b/packages/cli/src/connection-drivers.ts index b7b1857d0..62a114c81 100644 --- a/packages/cli/src/connection-drivers.ts +++ b/packages/cli/src/connection-drivers.ts @@ -5,6 +5,7 @@ export const KTX_DATABASE_DRIVER_IDS = [ 'sqlite', 'duckdb', 'postgres', + 'hologres', 'mysql', 'clickhouse', 'sqlserver', diff --git a/packages/cli/src/connectors/hologres/connector.ts b/packages/cli/src/connectors/hologres/connector.ts new file mode 100644 index 000000000..7e9680f99 --- /dev/null +++ b/packages/cli/src/connectors/hologres/connector.ts @@ -0,0 +1,37 @@ +import { KtxPostgresScanConnector, type KtxPostgresConnectionConfig } from '../postgres/connector.js'; + +// Hologres engine-internal schemas that never hold user data. Dropped from schema +// discovery the same way pg_catalog/information_schema are, since a user cannot +// place a table here. Matched exactly rather than by a `hologres_%` prefix because +// user schemas such as `hologres_dataset_*` share that prefix. +const HOLOGRES_SYSTEM_SCHEMAS: ReadonlySet = new Set([ + 'hologres', + 'hologres_streaming_mv', + 'hologres_statistic', + 'hologres_sample', + 'hologres_object_table', + 'hg_recyclebin', + 'hg_internal', +]); + +export type KtxHologresConnectionConfig = KtxPostgresConnectionConfig; + +export function isKtxHologresConnectionConfig( + connection: KtxHologresConnectionConfig | undefined, +): connection is KtxHologresConnectionConfig { + return String(connection?.driver ?? '').toLowerCase() === 'hologres'; +} + +/** + * Hologres scan connector. Hologres is PostgreSQL-wire-compatible, so it reuses + * the PostgreSQL connector's connection, introspection, sampling, and read-only + * single-statement execution (which never opens a transaction — Hologres forbids + * multi-statement DML/DDL transactions). The only override drops Hologres + * system schemas from discovery. + */ +export class KtxHologresScanConnector extends KtxPostgresScanConnector { + override async listSchemas(): Promise { + const schemas = await super.listSchemas(); + return schemas.filter((schema) => !HOLOGRES_SYSTEM_SCHEMAS.has(schema)); + } +} diff --git a/packages/cli/src/connectors/hologres/live-database-introspection.ts b/packages/cli/src/connectors/hologres/live-database-introspection.ts new file mode 100644 index 000000000..69ca8672b --- /dev/null +++ b/packages/cli/src/connectors/hologres/live-database-introspection.ts @@ -0,0 +1,47 @@ +import type { + LiveDatabaseIntrospectionOptions, + LiveDatabaseIntrospectionPort, +} from '../../context/ingest/adapters/live-database/types.js'; +import type { KtxProjectConnectionConfig } from '../../context/project/config.js'; +import type { + KtxPostgresConnectionConfig, + KtxPostgresEndpointResolver, + KtxPostgresPoolFactory, +} from '../postgres/connector.js'; +import { KtxHologresScanConnector } from './connector.js'; + +interface CreateHologresLiveDatabaseIntrospectionOptions { + connections: Record; + poolFactory?: KtxPostgresPoolFactory; + endpointResolver?: KtxPostgresEndpointResolver; + now?: () => Date; +} + +export function createHologresLiveDatabaseIntrospection( + options: CreateHologresLiveDatabaseIntrospectionOptions, +): LiveDatabaseIntrospectionPort { + return { + async extractSchema(connectionId: string, introspectionOptions?: LiveDatabaseIntrospectionOptions) { + const connection = options.connections[connectionId] as KtxPostgresConnectionConfig | undefined; + const connector = new KtxHologresScanConnector({ + connectionId, + connection, + poolFactory: options.poolFactory, + endpointResolver: options.endpointResolver, + now: options.now, + }); + try { + return await connector.introspect( + { + connectionId, + driver: 'hologres', + ...(introspectionOptions?.tableScope ? { tableScope: introspectionOptions.tableScope } : {}), + }, + { runId: `hologres-${connectionId}` }, + ); + } finally { + await connector.cleanup(); + } + }, + }; +} diff --git a/packages/cli/src/connectors/postgres/connector.ts b/packages/cli/src/connectors/postgres/connector.ts index c2c0f6db2..44b83ac42 100644 --- a/packages/cli/src/connectors/postgres/connector.ts +++ b/packages/cli/src/connectors/postgres/connector.ts @@ -7,6 +7,7 @@ import { scopedTableNames } from '../../context/scan/table-ref.js'; import { connectorTestFailure, createKtxConnectorCapabilities, + type KtxConnectionDriver, type KtxConnectorTestResult, type KtxColumnSampleInput, type KtxColumnSampleResult, @@ -344,6 +345,15 @@ export function isKtxPostgresConnectionConfig( return driver === 'postgres'; } +// Hologres speaks the PostgreSQL wire protocol and exposes the same +// pg_catalog/information_schema, so it shares this pool builder. Dispatch +// predicates (isKtxPostgresConnectionConfig) stay per-driver so Hologres routes +// to its own connector rather than the postgres one. +function isPostgresWireDriver(driver: unknown): boolean { + const normalized = String(driver ?? '').toLowerCase(); + return normalized === 'postgres' || normalized === 'hologres'; +} + /** @internal */ export function postgresPoolConfigFromConfig(input: { connectionId: string; @@ -351,7 +361,7 @@ export function postgresPoolConfigFromConfig(input: { env?: NodeJS.ProcessEnv; }): KtxPostgresPoolConfig { const inputDriver = input.connection?.driver ?? 'unknown'; - if (!isKtxPostgresConnectionConfig(input.connection)) { + if (!isPostgresWireDriver(input.connection?.driver)) { throw new Error(`Native PostgreSQL connector cannot run driver "${inputDriver}"`); } @@ -405,7 +415,7 @@ export function postgresPoolConfigFromConfig(input: { export class KtxPostgresScanConnector implements KtxScanConnector { readonly id: string; - readonly driver = 'postgres' as const; + readonly driver: KtxConnectionDriver; readonly capabilities = createKtxConnectorCapabilities({ tableSampling: true, columnSampling: true, @@ -431,6 +441,7 @@ export class KtxPostgresScanConnector implements KtxScanConnector { constructor(options: KtxPostgresScanConnectorOptions) { this.connectionId = options.connectionId; this.connection = options.connection ?? {}; + this.driver = String(this.connection.driver ?? '').toLowerCase() === 'hologres' ? 'hologres' : 'postgres'; this.poolConfig = postgresPoolConfigFromConfig({ connectionId: options.connectionId, connection: options.connection, @@ -440,7 +451,7 @@ export class KtxPostgresScanConnector implements KtxScanConnector { this.endpointResolver = options.endpointResolver; this.now = options.now ?? (() => new Date()); this.deadlineMs = resolveQueryDeadlineMs(this.connection); - this.id = `postgres:${options.connectionId}`; + this.id = `${this.driver}:${options.connectionId}`; } async testConnection(): Promise { @@ -465,7 +476,7 @@ export class KtxPostgresScanConnector implements KtxScanConnector { } return { connectionId: this.connectionId, - driver: 'postgres', + driver: this.driver, extractedAt: this.now().toISOString(), scope: { schemas }, metadata: { diff --git a/packages/cli/src/context/connections/dialects.ts b/packages/cli/src/context/connections/dialects.ts index 8512ed577..9ea140980 100644 --- a/packages/cli/src/context/connections/dialects.ts +++ b/packages/cli/src/context/connections/dialects.ts @@ -61,6 +61,7 @@ const sqlDialectFactories: Record KtxSqlDialect> = { duckdb: () => new KtxDuckDbDialect(), mysql: () => new KtxMysqlDialect(), postgres: () => new KtxPostgresDialect(), + hologres: () => new KtxPostgresDialect(), sqlite: () => new KtxSqliteDialect(), snowflake: () => new KtxSnowflakeDialect(), sqlserver: () => new KtxSqlServerDialect(), diff --git a/packages/cli/src/context/connections/drivers.ts b/packages/cli/src/context/connections/drivers.ts index 6a2d45890..2a472d4f4 100644 --- a/packages/cli/src/context/connections/drivers.ts +++ b/packages/cli/src/context/connections/drivers.ts @@ -169,6 +169,27 @@ export const driverRegistrations: Record { + const m = await import('../../connectors/hologres/connector.js'); + return { + isConnectionConfig: (connection) => { + const typedConnection = connection as Parameters[0]; + return m.isKtxHologresConnectionConfig(typedConnection); + }, + createScanConnector: ({ connectionId, connection }) => { + const typedConnection = connection as Parameters[0]; + if (!m.isKtxHologresConnectionConfig(typedConnection)) { + throw invalidConnectionConfig('hologres'); + } + return new m.KtxHologresScanConnector({ connectionId, connection: typedConnection }); + }, + }; + }, + }, sqlite: { driver: 'sqlite', scopeConfigKey: null, diff --git a/packages/cli/src/context/connections/local-warehouse-descriptor.ts b/packages/cli/src/context/connections/local-warehouse-descriptor.ts index 39cdddbf5..97b56cdc7 100644 --- a/packages/cli/src/context/connections/local-warehouse-descriptor.ts +++ b/packages/cli/src/context/connections/local-warehouse-descriptor.ts @@ -24,6 +24,7 @@ export interface LocalConnectionInfo { const DRIVER_TO_CONNECTION_TYPE: Record = { postgres: 'POSTGRESQL', + hologres: 'POSTGRESQL', sqlite: 'SQLITE', duckdb: 'DUCKDB', sqlserver: 'SQLSERVER', diff --git a/packages/cli/src/context/project/driver-schemas.ts b/packages/cli/src/context/project/driver-schemas.ts index 72d1a1335..edd1d9aa6 100644 --- a/packages/cli/src/context/project/driver-schemas.ts +++ b/packages/cli/src/context/project/driver-schemas.ts @@ -7,6 +7,7 @@ import { const warehouseDrivers = [ 'postgres', + 'hologres', 'mysql', 'snowflake', 'bigquery', @@ -56,6 +57,7 @@ function warehouseConnectionSchema(driver: const warehouseConnectionSchemas = [ warehouseConnectionSchema('postgres'), + warehouseConnectionSchema('hologres'), warehouseConnectionSchema('mysql'), warehouseConnectionSchema('snowflake'), warehouseConnectionSchema('bigquery'), diff --git a/packages/cli/src/context/scan/types.ts b/packages/cli/src/context/scan/types.ts index bf72558c4..cf3cfdb95 100644 --- a/packages/cli/src/context/scan/types.ts +++ b/packages/cli/src/context/scan/types.ts @@ -4,6 +4,7 @@ export type KtxConnectionDriver = | 'sqlite' | 'duckdb' | 'postgres' + | 'hologres' | 'sqlserver' | 'bigquery' | 'snowflake' diff --git a/packages/cli/src/context/sql-analysis/dialect.ts b/packages/cli/src/context/sql-analysis/dialect.ts index 830922333..dc6ed3eb5 100644 --- a/packages/cli/src/context/sql-analysis/dialect.ts +++ b/packages/cli/src/context/sql-analysis/dialect.ts @@ -8,6 +8,7 @@ import type { SqlAnalysisDialect } from './ports.js'; const SQLGLOT_DIALECTS: Record = { postgres: 'postgres', postgresql: 'postgres', + hologres: 'postgres', bigquery: 'bigquery', snowflake: 'snowflake', mysql: 'mysql', diff --git a/packages/cli/src/local-adapters.ts b/packages/cli/src/local-adapters.ts index 9f9ac60d5..5deb9582d 100644 --- a/packages/cli/src/local-adapters.ts +++ b/packages/cli/src/local-adapters.ts @@ -9,6 +9,8 @@ import { isKtxMysqlConnectionConfig } from './connectors/mysql/connector.js'; import { createPostgresLiveDatabaseIntrospection } from './connectors/postgres/live-database-introspection.js'; import { isKtxPostgresConnectionConfig, type KtxPostgresConnectionConfig } from './connectors/postgres/connector.js'; import { KtxPostgresHistoricSqlQueryClient } from './connectors/postgres/historic-sql-query-client.js'; +import { createHologresLiveDatabaseIntrospection } from './connectors/hologres/live-database-introspection.js'; +import { isKtxHologresConnectionConfig } from './connectors/hologres/connector.js'; import { createSqliteLiveDatabaseIntrospection } from './connectors/sqlite/live-database-introspection.js'; import { isKtxSqliteConnectionConfig } from './connectors/sqlite/connector.js'; import { createDuckDbLiveDatabaseIntrospection } from './connectors/duckdb/live-database-introspection.js'; @@ -118,6 +120,9 @@ function createKtxCliLiveDatabaseIntrospection( const postgres = createPostgresLiveDatabaseIntrospection({ connections: project.config.connections, }); + const hologres = createHologresLiveDatabaseIntrospection({ + connections: project.config.connections, + }); const clickhouse = createClickHouseLiveDatabaseIntrospection({ connections: project.config.connections, }); @@ -147,6 +152,9 @@ function createKtxCliLiveDatabaseIntrospection( if (isKtxPostgresConnectionConfig(connection)) { return postgres.extractSchema(connectionId, options); } + if (isKtxHologresConnectionConfig(connection)) { + return hologres.extractSchema(connectionId, options); + } if (isKtxSqliteConnectionConfig(connection)) { return sqlite.extractSchema(connectionId, options); } diff --git a/packages/cli/test/connectors/hologres/connector.test.ts b/packages/cli/test/connectors/hologres/connector.test.ts new file mode 100644 index 000000000..0a7e14047 --- /dev/null +++ b/packages/cli/test/connectors/hologres/connector.test.ts @@ -0,0 +1,128 @@ +import { describe, expect, it, vi } from 'vitest'; +import { + KtxHologresScanConnector, + isKtxHologresConnectionConfig, + type KtxHologresConnectionConfig, +} from '../../../src/connectors/hologres/connector.js'; +import type { KtxPostgresPoolConfig, KtxPostgresPoolFactory } from '../../../src/connectors/postgres/connector.js'; + +interface FakeQueryResult { + rows: Record[]; + fields?: Array<{ name: string; dataTypeID: number }>; +} + +type FakeQueryResponse = FakeQueryResult | Error; + +function fakePoolFactory(results: Map): KtxPostgresPoolFactory { + const query = vi.fn(async (sql: string, _params?: unknown[]) => { + const normalized = sql.replace(/\s+/g, ' ').trim(); + for (const [key, value] of results.entries()) { + if (normalized.includes(key)) { + if (value instanceof Error) { + throw value; + } + return value; + } + } + throw new Error(`Unexpected SQL: ${normalized}`); + }); + return { + createPool(_config: KtxPostgresPoolConfig) { + return { + async connect() { + return { query, release: vi.fn() }; + }, + end: vi.fn(async () => undefined), + }; + }, + }; +} + +const connection: KtxHologresConnectionConfig = { + driver: 'hologres', + host: 'holo.example.test', + database: 'analytics', + username: 'reader', + password: 'test-password', // pragma: allowlist secret + schema: 'public', +}; + +describe('isKtxHologresConnectionConfig', () => { + it('accepts hologres connections and rejects everything else', () => { + expect(isKtxHologresConnectionConfig(connection)).toBe(true); + expect(isKtxHologresConnectionConfig({ driver: 'postgres' })).toBe(false); + expect(isKtxHologresConnectionConfig(undefined)).toBe(false); + }); +}); + +describe('KtxHologresScanConnector', () => { + it('reports the hologres driver and id', () => { + const connector = new KtxHologresScanConnector({ + connectionId: 'wh', + connection, + poolFactory: fakePoolFactory(new Map()), + }); + expect(connector.driver).toBe('hologres'); + expect(connector.id).toBe('hologres:wh'); + }); + + it('drops Hologres system schemas from listSchemas while keeping user schemas that share the hologres_ prefix', async () => { + const results = new Map([ + [ + 'FROM information_schema.schemata', + { + rows: [ + { schema_name: 'analytics' }, + { schema_name: 'hg_internal' }, + { schema_name: 'hg_recyclebin' }, + { schema_name: 'hologres' }, + { schema_name: 'hologres_dataset_tpch_10g' }, + { schema_name: 'hologres_object_table' }, + { schema_name: 'hologres_sample' }, + { schema_name: 'hologres_statistic' }, + { schema_name: 'hologres_streaming_mv' }, + { schema_name: 'public' }, + ], + }, + ], + ]); + const connector = new KtxHologresScanConnector({ + connectionId: 'wh', + connection, + poolFactory: fakePoolFactory(results), + }); + await expect(connector.listSchemas()).resolves.toEqual(['analytics', 'hologres_dataset_tpch_10g', 'public']); + }); + + it('introspects a schema and stamps the snapshot driver as hologres', async () => { + const results = new Map([ + [ + 'c.reltuples::bigint AS row_count', + { rows: [{ table_name: 'orders', table_kind: 'r', row_count: '5', table_comment: 'order facts' }] }, + ], + [ + 'format_type(a.atttypid, a.atttypmod)', + { + rows: [ + { table_name: 'orders', column_name: 'id', data_type: 'bigint', is_nullable: false, column_comment: null }, + ], + }, + ], + ["constraint_type = 'PRIMARY KEY'", { rows: [{ table_name: 'orders', column_name: 'id' }] }], + ["constraint_type = 'FOREIGN KEY'", { rows: [] }], + ]); + const connector = new KtxHologresScanConnector({ + connectionId: 'wh', + connection, + poolFactory: fakePoolFactory(results), + now: () => new Date('2026-07-15T00:00:00.000Z'), + }); + const snapshot = await connector.introspect({ connectionId: 'wh', driver: 'hologres' }, { runId: 'run-1' }); + expect(snapshot.driver).toBe('hologres'); + expect(snapshot.scope).toEqual({ schemas: ['public'] }); + expect(snapshot.tables.map((table) => [table.db, table.name, table.kind, table.estimatedRows])).toEqual([ + ['public', 'orders', 'table', 5], + ]); + expect(snapshot.tables[0]?.columns.map((column) => [column.name, column.primaryKey])).toEqual([['id', true]]); + }); +}); diff --git a/packages/cli/test/connectors/hologres/live.integration.test.ts b/packages/cli/test/connectors/hologres/live.integration.test.ts new file mode 100644 index 000000000..4322e06ac --- /dev/null +++ b/packages/cli/test/connectors/hologres/live.integration.test.ts @@ -0,0 +1,52 @@ +import { describe, expect, it } from 'vitest'; +import { createHologresLiveDatabaseIntrospection } from '../../../src/connectors/hologres/live-database-introspection.js'; +import { KtxHologresScanConnector } from '../../../src/connectors/hologres/connector.js'; + +// Live Hologres integration. Gated on KTX_TEST_HOLOGRES_URL (a PostgreSQL-style +// DSN; URL-encode reserved characters, e.g. `$`->%24 and `#`->%23). The +// extractSchema case additionally needs KTX_TEST_HOLOGRES_SCHEMA pointed at a +// populated schema so table discovery is asserted against real data. +const url = process.env.KTX_TEST_HOLOGRES_URL; +const schema = process.env.KTX_TEST_HOLOGRES_SCHEMA; + +const HOLOGRES_SYSTEM_SCHEMAS = [ + 'hologres', + 'hologres_streaming_mv', + 'hologres_statistic', + 'hologres_sample', + 'hologres_object_table', + 'hg_recyclebin', + 'hg_internal', +]; + +describe.skipIf(!url)('createHologresLiveDatabaseIntrospection (live Hologres)', () => { + it('lists schemas with Hologres system schemas excluded', async () => { + const connector = new KtxHologresScanConnector({ + connectionId: 'wh', + connection: { driver: 'hologres', url }, + }); + try { + const schemas = await connector.listSchemas(); + expect(schemas.length).toBeGreaterThan(0); + for (const systemSchema of HOLOGRES_SYSTEM_SCHEMAS) { + expect(schemas).not.toContain(systemSchema); + } + } finally { + await connector.cleanup(); + } + }); + + it.skipIf(!schema)('extracts tables for a populated schema and stamps driver hologres', async () => { + const introspection = createHologresLiveDatabaseIntrospection({ + connections: { wh: { driver: 'hologres', url, schemas: schema ? [schema] : [] } }, + }); + const snapshot = await introspection.extractSchema('wh'); + expect(snapshot.driver).toBe('hologres'); + expect(snapshot.tables.length).toBeGreaterThan(0); + for (const table of snapshot.tables) { + expect(table.db).toBe(schema); + expect(table.name.length).toBeGreaterThan(0); + expect(table.columns.length).toBeGreaterThan(0); + } + }); +}); diff --git a/packages/cli/test/context/connections/dialects.test.ts b/packages/cli/test/context/connections/dialects.test.ts index fad193ed4..59da724fe 100644 --- a/packages/cli/test/context/connections/dialects.test.ts +++ b/packages/cli/test/context/connections/dialects.test.ts @@ -305,7 +305,7 @@ describe('getDialectForDriver', () => { it('throws with a supported-driver list for unknown drivers', () => { expect(() => getDialectForDriver('oracle')).toThrow( - 'Unsupported driver "oracle". Supported drivers: athena, bigquery, clickhouse, duckdb, mongodb, mysql, postgres, snowflake, sqlite, sqlserver', + 'Unsupported driver "oracle". Supported drivers: athena, bigquery, clickhouse, duckdb, hologres, mongodb, mysql, postgres, snowflake, sqlite, sqlserver', ); }); diff --git a/packages/cli/test/context/connections/drivers.test.ts b/packages/cli/test/context/connections/drivers.test.ts index 69ac8464b..9bb303aac 100644 --- a/packages/cli/test/context/connections/drivers.test.ts +++ b/packages/cli/test/context/connections/drivers.test.ts @@ -21,6 +21,11 @@ const connectionFixtures: Record = { url: 'postgresql://reader:secret@localhost:5432/analytics', // pragma: allowlist secret schemas: ['public'], }), + hologres: () => ({ + driver: 'hologres', + url: 'postgresql://reader:secret@localhost:80/analytics', // pragma: allowlist secret + schemas: ['public'], + }), sqlite: () => ({ driver: 'sqlite', path: 'warehouse.db' }), duckdb: (projectDir) => ({ driver: 'duckdb', path: join(projectDir, 'warehouse.duckdb') }), mongodb: () => ({ @@ -109,6 +114,7 @@ describe('driverRegistrations', () => { 'bigquery', 'clickhouse', 'duckdb', + 'hologres', 'mongodb', 'mysql', 'postgres', diff --git a/packages/cli/test/context/project/driver-schemas.test.ts b/packages/cli/test/context/project/driver-schemas.test.ts index 1c5891346..70bccdff5 100644 --- a/packages/cli/test/context/project/driver-schemas.test.ts +++ b/packages/cli/test/context/project/driver-schemas.test.ts @@ -4,6 +4,7 @@ import { connectionConfigSchema } from '../../../src/context/project/driver-sche describe('connectionConfigSchema (driver discriminated union)', () => { it.each([ ['postgres', 'postgres://user:pass@host:5432/db'], // pragma: allowlist secret + ['hologres', 'postgres://user:pass@host:80/db'], // pragma: allowlist secret ['mysql', 'mysql://user:pass@host:3306/db'], // pragma: allowlist secret ['snowflake', 'snowflake://account/db'], ['bigquery', 'bigquery://project/dataset'], From 1f1cb29b651b683f1dd8d709133a5cd99ac6f36c Mon Sep 17 00:00:00 2001 From: TimothyDing Date: Wed, 15 Jul 2026 15:35:20 +0800 Subject: [PATCH 08/14] feat(hologres): support Hologres in the setup wizard Add Hologres to the ktx setup driver picker, the --database flag parser, scope discovery (schema-scoped like postgres, default port 80), default connection id, and connection-test support. Hologres reuses the shared URL/host connection flow. Query-history wizard wiring is deferred until the hg_query_log reader lands. --- packages/cli/src/commands/setup-commands.ts | 1 + packages/cli/src/connection.ts | 1 + packages/cli/src/setup-databases.ts | 16 ++++++++++++++-- packages/cli/test/setup-databases.test.ts | 2 ++ 4 files changed, 18 insertions(+), 2 deletions(-) diff --git a/packages/cli/src/commands/setup-commands.ts b/packages/cli/src/commands/setup-commands.ts index d838e4723..acf3ef17f 100644 --- a/packages/cli/src/commands/setup-commands.ts +++ b/packages/cli/src/commands/setup-commands.ts @@ -40,6 +40,7 @@ function databaseDriver(value: string): KtxSetupDatabaseDriver { value === 'sqlite' || value === 'duckdb' || value === 'postgres' || + value === 'hologres' || value === 'mysql' || value === 'clickhouse' || value === 'sqlserver' || diff --git a/packages/cli/src/connection.ts b/packages/cli/src/connection.ts index 809eea96f..65266ad53 100644 --- a/packages/cli/src/connection.ts +++ b/packages/cli/src/connection.ts @@ -53,6 +53,7 @@ const SUPPORTED_TEST_DRIVERS = [ 'sqlite', 'duckdb', 'postgres', + 'hologres', 'mysql', 'clickhouse', 'sqlserver', diff --git a/packages/cli/src/setup-databases.ts b/packages/cli/src/setup-databases.ts index efeba2092..72867387f 100644 --- a/packages/cli/src/setup-databases.ts +++ b/packages/cli/src/setup-databases.ts @@ -66,6 +66,7 @@ export type KtxSetupDatabaseDriver = | 'sqlite' | 'duckdb' | 'postgres' + | 'hologres' | 'mysql' | 'clickhouse' | 'sqlserver' @@ -154,6 +155,7 @@ export interface KtxSetupDatabasesDeps { const DRIVER_OPTIONS: Array<{ value: KtxSetupDatabaseDriver; label: string }> = [ { value: 'postgres', label: 'PostgreSQL' }, + { value: 'hologres', label: 'Hologres' }, { value: 'bigquery', label: 'BigQuery' }, { value: 'snowflake', label: 'Snowflake' }, { value: 'mysql', label: 'MySQL' }, @@ -180,6 +182,7 @@ const DEFAULT_CONNECTION_IDS: Record = { sqlite: 'sqlite-local', duckdb: 'duckdb-local', postgres: 'postgres-warehouse', + hologres: 'hologres-warehouse', mysql: 'mysql-warehouse', clickhouse: 'clickhouse-warehouse', sqlserver: 'sqlserver-warehouse', @@ -225,6 +228,14 @@ const SCOPE_DISCOVERY_SPECS: Partial; +type UrlDriverType = Extract; const DRIVER_CONNECTION_DEFAULTS: Record = { postgres: { port: '5432' }, + hologres: { port: '80' }, mysql: { port: '3306' }, clickhouse: { port: '8123' }, sqlserver: { port: '1433' }, @@ -838,7 +850,7 @@ async function buildConnectionConfig(input: { if (path === undefined) return 'back'; return path ? { driver: 'duckdb', path } : null; } - if (driver === 'postgres' || driver === 'mysql' || driver === 'clickhouse' || driver === 'sqlserver') { + if (driver === 'postgres' || driver === 'hologres' || driver === 'mysql' || driver === 'clickhouse' || driver === 'sqlserver') { return await buildUrlConnectionConfig({ driver, connectionId: input.connectionId, diff --git a/packages/cli/test/setup-databases.test.ts b/packages/cli/test/setup-databases.test.ts index feff6613b..a808f2a3f 100644 --- a/packages/cli/test/setup-databases.test.ts +++ b/packages/cli/test/setup-databases.test.ts @@ -238,6 +238,7 @@ describe('setup databases step', () => { 'Up/Down to move, Tab to select or unselect, Enter to confirm, Escape to go back, Ctrl+C to exit.', options: [ { value: 'postgres', label: 'PostgreSQL' }, + { value: 'hologres', label: 'Hologres' }, { value: 'bigquery', label: 'BigQuery' }, { value: 'snowflake', label: 'Snowflake' }, { value: 'mysql', label: 'MySQL' }, @@ -283,6 +284,7 @@ describe('setup databases step', () => { it('offers connection URL paste first for URL-capable databases', async () => { const cases: Array<{ driver: KtxSetupDatabaseDriver; label: string }> = [ { driver: 'postgres', label: 'PostgreSQL' }, + { driver: 'hologres', label: 'Hologres' }, { driver: 'mysql', label: 'MySQL' }, { driver: 'clickhouse', label: 'ClickHouse' }, { driver: 'sqlserver', label: 'SQL Server' }, From 4305fa915c05c3d2eee2c2e96995187da1a82389 Mon Sep 17 00:00:00 2001 From: TimothyDing Date: Wed, 15 Jul 2026 15:50:51 +0800 Subject: [PATCH 09/14] feat(hologres): ingest query history from hologres.hg_query_log Add a windowed HologresQueryLogReader that aggregates hologres.hg_query_log by its digest fingerprint (executions, distinct users, p50/p95 duration, error rate from status, rows, top users) over a query_start window, plus a readiness probe runner. Missing pg_read_all_stats degrades to a warning rather than failing, since Hologres still exposes the current user's own log. Wire the hologres dialect through the historic-SQL config union, dialect maps, probe factory registry, and CLI reader selection (explicit branch so it never falls through to Snowflake), and flip the driver's hasHistoricSqlReader to true. The reader runs a single read-only statement, honoring Hologres's no-transaction constraint. --- .../hologres/historic-sql-query-client.ts | 38 +++ .../cli/src/context/connections/drivers.ts | 2 +- .../historic-sql/connection-dialect.ts | 2 +- .../historic-sql/hologres-query-log-reader.ts | 264 ++++++++++++++++++ .../ingest/adapters/historic-sql/types.ts | 4 +- .../src/context/ingest/historic-sql-probes.ts | 9 + .../historic-sql-probes/hologres-runner.ts | 95 +++++++ .../cli/src/context/ingest/local-adapters.ts | 3 +- packages/cli/src/local-adapters.ts | 34 ++- packages/cli/src/public-ingest.ts | 3 +- .../test/context/connections/drivers.test.ts | 2 +- ...-query-log-reader.live.integration.test.ts | 44 +++ .../hologres-query-log-reader.test.ts | 139 +++++++++ .../ingest/historic-sql-probes.test.ts | 5 + 14 files changed, 636 insertions(+), 8 deletions(-) create mode 100644 packages/cli/src/connectors/hologres/historic-sql-query-client.ts create mode 100644 packages/cli/src/context/ingest/adapters/historic-sql/hologres-query-log-reader.ts create mode 100644 packages/cli/src/context/ingest/historic-sql-probes/hologres-runner.ts create mode 100644 packages/cli/test/context/ingest/adapters/historic-sql/hologres-query-log-reader.live.integration.test.ts create mode 100644 packages/cli/test/context/ingest/adapters/historic-sql/hologres-query-log-reader.test.ts diff --git a/packages/cli/src/connectors/hologres/historic-sql-query-client.ts b/packages/cli/src/connectors/hologres/historic-sql-query-client.ts new file mode 100644 index 000000000..9dfe12696 --- /dev/null +++ b/packages/cli/src/connectors/hologres/historic-sql-query-client.ts @@ -0,0 +1,38 @@ +import type { KtxPostgresQueryClient } from '../../context/ingest/adapters/historic-sql/types.js'; +import type { KtxPostgresScanConnectorOptions } from '../postgres/connector.js'; +import { KtxHologresScanConnector } from './connector.js'; + +export type KtxHologresHistoricSqlQueryClientOptions = KtxPostgresScanConnectorOptions; + +export class KtxHologresHistoricSqlQueryClient implements KtxPostgresQueryClient { + private readonly connectionId: string; + private readonly connector: KtxHologresScanConnector; + + constructor(options: KtxHologresHistoricSqlQueryClientOptions) { + this.connectionId = options.connectionId; + this.connector = new KtxHologresScanConnector(options); + } + + async executeQuery( + sql: string, + params?: unknown[], + ): Promise<{ headers: string[]; rows: unknown[][]; totalRows: number }> { + const result = await this.connector.executeReadOnly( + { + connectionId: this.connectionId, + sql, + params, + }, + {} as never, + ); + return { + headers: result.headers, + rows: result.rows, + totalRows: result.totalRows, + }; + } + + async cleanup(): Promise { + await this.connector.cleanup(); + } +} diff --git a/packages/cli/src/context/connections/drivers.ts b/packages/cli/src/context/connections/drivers.ts index 2a472d4f4..cd3290035 100644 --- a/packages/cli/src/context/connections/drivers.ts +++ b/packages/cli/src/context/connections/drivers.ts @@ -172,7 +172,7 @@ export const driverRegistrations: Record { const m = await import('../../connectors/hologres/connector.js'); return { diff --git a/packages/cli/src/context/ingest/adapters/historic-sql/connection-dialect.ts b/packages/cli/src/context/ingest/adapters/historic-sql/connection-dialect.ts index 7845cbbc0..275f89e56 100644 --- a/packages/cli/src/context/ingest/adapters/historic-sql/connection-dialect.ts +++ b/packages/cli/src/context/ingest/adapters/historic-sql/connection-dialect.ts @@ -2,7 +2,7 @@ import { getDriverRegistration } from '../../../connections/drivers.js'; import type { KtxConnectionDriver } from '../../../scan/types.js'; import type { HistoricSqlDialect } from './types.js'; -const historicSqlDialects: readonly HistoricSqlDialect[] = ['postgres', 'bigquery', 'snowflake']; +const historicSqlDialects: readonly HistoricSqlDialect[] = ['postgres', 'bigquery', 'snowflake', 'hologres']; function recordOrNull(value: unknown): Record | null { return value && typeof value === 'object' && !Array.isArray(value) ? (value as Record) : null; diff --git a/packages/cli/src/context/ingest/adapters/historic-sql/hologres-query-log-reader.ts b/packages/cli/src/context/ingest/adapters/historic-sql/hologres-query-log-reader.ts new file mode 100644 index 000000000..44409d33f --- /dev/null +++ b/packages/cli/src/context/ingest/adapters/historic-sql/hologres-query-log-reader.ts @@ -0,0 +1,264 @@ +import { HistoricSqlExtensionMissingError } from './errors.js'; +import { + aggregatedTemplateSchema, + type AggregatedTemplate, + type HistoricSqlTimeWindow, + type HistoricSqlUnifiedPullConfig, +} from './types.js'; + +interface QueryResultLike { + headers: string[]; + rows: unknown[][]; + totalRows?: number; + error?: string; +} + +interface QueryClientLike { + executeQuery(sql: string, params?: unknown[]): Promise; +} + +export interface HologresQueryLogProbeResult { + warnings: string[]; + info: string[]; +} + +const HOLOGRES_QUERY_LOG_RELATION = 'hologres.hg_query_log'; +const PROBE_SQL = `SELECT 1 FROM ${HOLOGRES_QUERY_LOG_RELATION} LIMIT 1`; +const GRANTS_PROBE_SQL = "SELECT pg_has_role(current_user, 'pg_read_all_stats', 'USAGE') AS has_role"; + +const HOLOGRES_QUERY_LOG_REMEDIATION = + 'Ensure the connection role can read hologres.hg_query_log. A superuser or a member of pg_read_all_stats sees every database; a db_admin (SPM/SLPM) sees the current database.'; +const HOLOGRES_GRANTS_INFO = + 'connection role lacks pg_read_all_stats; only the current user\'s query log is visible, so hg_query_log coverage is partial'; + +// hg_query_log is a per-execution slow-query log keyed by digest (a SQL fingerprint +// Hologres computes for SELECT/INSERT/UPDATE/DELETE). Aggregating by digest over a +// query_start window mirrors the windowed Snowflake/BigQuery readers. Single +// read-only statement — no transaction (Hologres forbids multi-statement DML/DDL). +const AGGREGATE_SQL = ` +WITH filtered AS ( + SELECT digest, query, usename, status, duration, query_start, result_rows + FROM ${HOLOGRES_QUERY_LOG_RELATION} + WHERE digest IS NOT NULL + AND query_start >= $1::timestamptz + AND query_start < $2::timestamptz +), +template_stats AS ( + SELECT + digest AS template_id, + MIN(query) AS canonical_sql, + COUNT(*) AS executions, + COUNT(DISTINCT usename) AS distinct_users, + MIN(query_start) AS first_seen, + MAX(query_start) AS last_seen, + PERCENTILE_CONT(0.50) WITHIN GROUP (ORDER BY duration) AS p50_ms, + PERCENTILE_CONT(0.95) WITHIN GROUP (ORDER BY duration) AS p95_ms, + (COUNT(*) FILTER (WHERE status = 'FAILED'))::float8 / NULLIF(COUNT(*), 0) AS error_rate, + SUM(result_rows)::bigint AS rows_produced + FROM filtered + GROUP BY digest + HAVING COUNT(*) >= $3 +), +template_users AS ( + SELECT digest AS template_id, usename AS user_name, COUNT(*) AS executions, MAX(query_start) AS last_seen + FROM filtered + GROUP BY digest, usename +) +SELECT + stats.template_id, + stats.canonical_sql, + stats.executions, + stats.distinct_users, + stats.first_seen, + stats.last_seen, + stats.p50_ms, + stats.p95_ms, + stats.error_rate, + stats.rows_produced, + COALESCE( + json_agg(json_build_object('user', users.user_name, 'executions', users.executions) + ORDER BY users.executions DESC, users.last_seen DESC) + FILTER (WHERE users.user_name IS NOT NULL), + '[]'::json + )::text AS top_users +FROM template_stats AS stats +LEFT JOIN template_users AS users ON users.template_id = stats.template_id +GROUP BY + stats.template_id, + stats.canonical_sql, + stats.executions, + stats.distinct_users, + stats.first_seen, + stats.last_seen, + stats.p50_ms, + stats.p95_ms, + stats.error_rate, + stats.rows_produced +ORDER BY stats.executions DESC +`.trim(); + +function queryClient(client: unknown): QueryClientLike { + if ( + client && + typeof client === 'object' && + 'executeQuery' in client && + typeof (client as { executeQuery?: unknown }).executeQuery === 'function' + ) { + return client as QueryClientLike; + } + throw new Error('Historic SQL Hologres reader requires a query client with executeQuery(sql, params?)'); +} + +async function execute(client: QueryClientLike, sql: string, params?: unknown[]): Promise { + const result = await client.executeQuery(sql, params); + if ('error' in result && typeof result.error === 'string' && result.error.length > 0) { + throw new Error(result.error); + } + return result; +} + +function indexByHeader(headers: string[]): Map { + const out = new Map(); + headers.forEach((header, index) => out.set(header.toLowerCase(), index)); + return out; +} + +function value(row: unknown[], indexes: Map, name: string): unknown { + const index = indexes.get(name.toLowerCase()); + return index === undefined ? null : row[index]; +} + +function nullableString(raw: unknown): string | null { + if (raw === null || raw === undefined) { + return null; + } + const text = String(raw); + return text.length > 0 ? text : null; +} + +function requiredString(raw: unknown, field: string): string { + const text = nullableString(raw); + if (!text) { + throw new Error(`Hologres hg_query_log row is missing ${field}`); + } + return text; +} + +function nullableNumber(raw: unknown): number | null { + if (raw === null || raw === undefined || raw === '') { + return null; + } + const number = typeof raw === 'number' ? raw : Number(raw); + return Number.isFinite(number) ? number : null; +} + +function requiredNumber(raw: unknown, field: string): number { + const number = nullableNumber(raw); + if (number === null) { + throw new Error(`Hologres hg_query_log row has invalid ${field}: ${String(raw)}`); + } + return number; +} + +function requiredInteger(raw: unknown, field: string): number { + return Math.trunc(requiredNumber(raw, field)); +} + +function nullableInteger(raw: unknown): number | null { + const number = nullableNumber(raw); + return number === null ? null : Math.trunc(number); +} + +function isoTimestamp(raw: unknown, field: string): string { + if (raw instanceof Date) { + return raw.toISOString(); + } + const text = requiredString(raw, field); + const date = new Date(text); + if (Number.isNaN(date.getTime())) { + throw new Error(`Hologres hg_query_log row has invalid ${field}: ${text}`); + } + return date.toISOString(); +} + +function parseTopUsers(raw: unknown): Array<{ user: string | null; executions: number }> { + const text = nullableString(raw); + if (!text) { + return []; + } + try { + const parsed = JSON.parse(text) as unknown; + if (!Array.isArray(parsed)) { + return []; + } + return parsed.flatMap((entry) => { + if (!entry || typeof entry !== 'object') { + return []; + } + const user = nullableString((entry as { user?: unknown }).user); + const executions = nullableInteger((entry as { executions?: unknown }).executions); + return executions === null ? [] : [{ user, executions }]; + }); + } catch { + return []; + } +} + +function mapAggregatedRow(row: unknown[], indexes: Map, lastSeenFallback: string): AggregatedTemplate { + return aggregatedTemplateSchema.parse({ + templateId: requiredString(value(row, indexes, 'template_id'), 'template_id'), + canonicalSql: requiredString(value(row, indexes, 'canonical_sql'), 'canonical_sql'), + dialect: 'hologres', + stats: { + executions: requiredInteger(value(row, indexes, 'executions'), 'executions'), + distinctUsers: requiredInteger(value(row, indexes, 'distinct_users'), 'distinct_users'), + firstSeen: isoTimestamp(value(row, indexes, 'first_seen'), 'first_seen'), + lastSeen: isoTimestamp(value(row, indexes, 'last_seen') ?? lastSeenFallback, 'last_seen'), + p50RuntimeMs: nullableNumber(value(row, indexes, 'p50_ms')), + p95RuntimeMs: nullableNumber(value(row, indexes, 'p95_ms')), + errorRate: requiredNumber(value(row, indexes, 'error_rate'), 'error_rate'), + rowsProduced: nullableInteger(value(row, indexes, 'rows_produced')), + }, + topUsers: parseTopUsers(value(row, indexes, 'top_users')), + }); +} + +export class HologresQueryLogReader { + async probe(client: unknown): Promise { + const pgClient = queryClient(client); + try { + await execute(pgClient, PROBE_SQL); + } catch (error) { + throw new HistoricSqlExtensionMissingError({ + dialect: 'hologres', + message: 'hologres.hg_query_log is not accessible for historic-SQL ingest.', + remediation: HOLOGRES_QUERY_LOG_REMEDIATION, + cause: error, + }); + } + const warnings: string[] = []; + try { + const grants = await execute(pgClient, GRANTS_PROBE_SQL); + if (value(grants.rows[0] ?? [], indexByHeader(grants.headers), 'has_role') !== true) { + warnings.push(HOLOGRES_GRANTS_INFO); + } + } catch { + warnings.push(HOLOGRES_GRANTS_INFO); + } + return { warnings, info: [] }; + } + + async *fetchAggregated( + client: unknown, + window: HistoricSqlTimeWindow, + config: HistoricSqlUnifiedPullConfig, + ): AsyncIterable { + const pgClient = queryClient(client); + const result = await execute(pgClient, AGGREGATE_SQL, [window.start, window.end, config.minExecutions]); + const indexes = indexByHeader(result.headers); + const lastSeenFallback = window.end.toISOString(); + for (const row of result.rows) { + yield mapAggregatedRow(row, indexes, lastSeenFallback); + } + } +} diff --git a/packages/cli/src/context/ingest/adapters/historic-sql/types.ts b/packages/cli/src/context/ingest/adapters/historic-sql/types.ts index aca50c4ed..4bb2e037f 100644 --- a/packages/cli/src/context/ingest/adapters/historic-sql/types.ts +++ b/packages/cli/src/context/ingest/adapters/historic-sql/types.ts @@ -3,7 +3,7 @@ import type { SqlAnalysisPort } from '../../../../context/sql-analysis/ports.js' export const HISTORIC_SQL_SOURCE_KEY = 'historic-sql' as const; -const historicSqlDialectSchema = z.enum(['snowflake', 'bigquery', 'postgres']); +const historicSqlDialectSchema = z.enum(['snowflake', 'bigquery', 'postgres', 'hologres']); export type HistoricSqlDialect = z.infer; const filterModeSchema = z.enum(['exclude', 'include', 'mark-only']); @@ -43,7 +43,7 @@ const historicSqlCommonPullConfigSchema = z.object({ }); const historicSqlWindowedPullConfigSchema = historicSqlCommonPullConfigSchema.extend({ - dialect: z.enum(['snowflake', 'bigquery']), + dialect: z.enum(['snowflake', 'bigquery', 'hologres']), windowDays: z.number().int().positive().default(90), }); diff --git a/packages/cli/src/context/ingest/historic-sql-probes.ts b/packages/cli/src/context/ingest/historic-sql-probes.ts index 07204f3a3..161e395b2 100644 --- a/packages/cli/src/context/ingest/historic-sql-probes.ts +++ b/packages/cli/src/context/ingest/historic-sql-probes.ts @@ -87,6 +87,15 @@ const defaultHistoricSqlProbeRunnerFactories: Record< return new BigQueryJobsByProjectProbeRunner(); }, }, + hologres: { + catalogName: 'hologres.hg_query_log', + load: async () => { + const { HologresQueryLogProbeRunner } = await import( + './historic-sql-probes/hologres-runner.js' + ); + return new HologresQueryLogProbeRunner(); + }, + }, }; const DEFAULT_RUNNER_CACHE = new Map(); diff --git a/packages/cli/src/context/ingest/historic-sql-probes/hologres-runner.ts b/packages/cli/src/context/ingest/historic-sql-probes/hologres-runner.ts new file mode 100644 index 000000000..79c71e6cc --- /dev/null +++ b/packages/cli/src/context/ingest/historic-sql-probes/hologres-runner.ts @@ -0,0 +1,95 @@ +import { + HistoricSqlExtensionMissingError, + HistoricSqlGrantsMissingError, +} from '../adapters/historic-sql/errors.js'; +import { + HologresQueryLogReader, + type HologresQueryLogProbeResult, +} from '../adapters/historic-sql/hologres-query-log-reader.js'; +import { + type HistoricSqlFixAdvice, + type HistoricSqlProbeInput, + type HistoricSqlProbeRunner, + type HistoricSqlSuccessDetail, +} from '../historic-sql-probes.js'; +import { + isKtxHologresConnectionConfig, + type KtxHologresConnectionConfig, +} from '../../../connectors/hologres/connector.js'; +import { KtxHologresHistoricSqlQueryClient } from '../../../connectors/hologres/historic-sql-query-client.js'; + +interface ClientHandle { + client: unknown; + cleanup(): Promise; +} + +interface HologresQueryLogProbeRunnerOptions { + reader?: { probe(client: unknown): Promise }; + createClient?: ( + input: HistoricSqlProbeInput & { connection: KtxHologresConnectionConfig }, + ) => ClientHandle; +} + +export class HologresQueryLogProbeRunner implements HistoricSqlProbeRunner { + readonly dialect = 'hologres' as const; + readonly catalogName = 'hologres.hg_query_log'; + + private readonly reader: { probe(client: unknown): Promise }; + private readonly createClient: ( + input: HistoricSqlProbeInput & { connection: KtxHologresConnectionConfig }, + ) => ClientHandle; + + constructor(options: HologresQueryLogProbeRunnerOptions = {}) { + this.reader = options.reader ?? new HologresQueryLogReader(); + this.createClient = + options.createClient ?? + ((input) => { + const client = new KtxHologresHistoricSqlQueryClient({ + connectionId: input.connectionId, + connection: input.connection, + env: input.env, + }); + return { client, cleanup: () => client.cleanup() }; + }); + } + + async run(input: HistoricSqlProbeInput): Promise { + const inputDriver = input.connection.driver ?? 'unknown'; + if (!isKtxHologresConnectionConfig(input.connection)) { + throw new Error(`Hologres query-log probe requires a Hologres connection, got "${String(inputDriver)}"`); + } + const handle = this.createClient({ ...input, connection: input.connection }); + try { + return await this.reader.probe(handle.client); + } finally { + await handle.cleanup(); + } + } + + formatSuccessDetail(result: unknown): HistoricSqlSuccessDetail { + const hologresResult = result as HologresQueryLogProbeResult; + return { + detail: 'hologres.hg_query_log ready', + warnings: hologresResult.warnings, + }; + } + + fixAdvice(error: unknown): HistoricSqlFixAdvice { + if (error instanceof HistoricSqlExtensionMissingError) { + return { + failHeadline: 'hologres.hg_query_log is not accessible', + remediation: error.remediation, + }; + } + if (error instanceof HistoricSqlGrantsMissingError) { + return { + failHeadline: 'Hologres connection role lacks query-log access', + remediation: error.remediation, + }; + } + return { + failHeadline: `${this.catalogName} readiness check failed`, + remediation: error instanceof Error ? error.message : String(error), + }; + } +} diff --git a/packages/cli/src/context/ingest/local-adapters.ts b/packages/cli/src/context/ingest/local-adapters.ts index 75e54a330..42329e954 100644 --- a/packages/cli/src/context/ingest/local-adapters.ts +++ b/packages/cli/src/context/ingest/local-adapters.ts @@ -176,8 +176,9 @@ function isRecord(value: unknown): value is Record { return typeof value === 'object' && value !== null && !Array.isArray(value); } -const historicSqlDialectByDriver = new Map([ +const historicSqlDialectByDriver = new Map([ ['postgres', 'postgres'], + ['hologres', 'hologres'], ['bigquery', 'bigquery'], ['snowflake', 'snowflake'], ]); diff --git a/packages/cli/src/local-adapters.ts b/packages/cli/src/local-adapters.ts index 5deb9582d..6b0abccb9 100644 --- a/packages/cli/src/local-adapters.ts +++ b/packages/cli/src/local-adapters.ts @@ -10,7 +10,8 @@ import { createPostgresLiveDatabaseIntrospection } from './connectors/postgres/l import { isKtxPostgresConnectionConfig, type KtxPostgresConnectionConfig } from './connectors/postgres/connector.js'; import { KtxPostgresHistoricSqlQueryClient } from './connectors/postgres/historic-sql-query-client.js'; import { createHologresLiveDatabaseIntrospection } from './connectors/hologres/live-database-introspection.js'; -import { isKtxHologresConnectionConfig } from './connectors/hologres/connector.js'; +import { isKtxHologresConnectionConfig, type KtxHologresConnectionConfig } from './connectors/hologres/connector.js'; +import { KtxHologresHistoricSqlQueryClient } from './connectors/hologres/historic-sql-query-client.js'; import { createSqliteLiveDatabaseIntrospection } from './connectors/sqlite/live-database-introspection.js'; import { isKtxSqliteConnectionConfig } from './connectors/sqlite/connector.js'; import { createDuckDbLiveDatabaseIntrospection } from './connectors/duckdb/live-database-introspection.js'; @@ -28,6 +29,7 @@ import type { } from './context/ingest/adapters/live-database/types.js'; import { LiveDatabaseSourceAdapter } from './context/ingest/adapters/live-database/live-database.adapter.js'; import { PostgresPgssReader } from './context/ingest/adapters/historic-sql/postgres-pgss-reader.js'; +import { HologresQueryLogReader } from './context/ingest/adapters/historic-sql/hologres-query-log-reader.js'; import { SnowflakeHistoricSqlQueryHistoryReader } from './context/ingest/adapters/historic-sql/snowflake-query-history-reader.js'; import type { SourceAdapter } from './context/ingest/types.js'; import type { KtxLocalProject } from './context/project/project.js'; @@ -229,6 +231,27 @@ function createEphemeralPostgresHistoricSqlClient(project: KtxLocalProject, conn }; } +function createEphemeralHologresHistoricSqlClient(project: KtxLocalProject, connectionId: string) { + const connection = project.config.connections[connectionId] as KtxHologresConnectionConfig | undefined; + const inputDriver = connection?.driver ?? 'unknown'; + if (!isKtxHologresConnectionConfig(connection)) { + throw new Error(`Query history ingest requires a Hologres connection, got ${String(inputDriver)}`); + } + return { + async executeQuery(sql: string, params?: unknown[]) { + const client = new KtxHologresHistoricSqlQueryClient({ + connectionId, + connection, + }); + try { + return await client.executeQuery(sql, params); + } finally { + await client.cleanup(); + } + }, + }; +} + function createEphemeralBigQueryHistoricSqlClient(project: KtxLocalProject, connectionId: string) { const connection = project.config.connections[connectionId] as KtxBigQueryConnectionConfig | undefined; const inputDriver = connection?.driver ?? 'unknown'; @@ -337,6 +360,15 @@ function historicSqlOptionsForLocalRun( }; } + if (dialect === 'hologres') { + return { + ...base, + dialect, + reader: new HologresQueryLogReader() satisfies HistoricSqlReader, + queryClient: createEphemeralHologresHistoricSqlClient(project, connectionId), + }; + } + if (dialect === 'bigquery') { const inputDriver = connection?.driver ?? 'unknown'; if (!isKtxBigQueryConnectionConfig(connection)) { diff --git a/packages/cli/src/public-ingest.ts b/packages/cli/src/public-ingest.ts index beab6dca2..2087d9b07 100644 --- a/packages/cli/src/public-ingest.ts +++ b/packages/cli/src/public-ingest.ts @@ -33,7 +33,7 @@ type KtxPublicIngestStepName = 'database-schema' | 'query-history' | 'source-ing type KtxPublicIngestStepStatus = 'done' | 'skipped' | 'failed' | 'not-run'; type KtxPublicIngestInputMode = 'auto' | 'disabled'; type KtxPublicIngestQueryHistoryFlag = 'default' | 'enabled' | 'disabled'; -type HistoricSqlDialect = 'postgres' | 'bigquery' | 'snowflake'; +type HistoricSqlDialect = 'postgres' | 'bigquery' | 'snowflake' | 'hologres'; export type KtxPublicIngestArgs = { @@ -155,6 +155,7 @@ export function publicProgressMessage(message: string, target: KtxPublicIngestPl const queryHistoryDialectByDriver = new Map([ ['postgres', 'postgres'], + ['hologres', 'hologres'], ['bigquery', 'bigquery'], ['snowflake', 'snowflake'], ]); diff --git a/packages/cli/test/context/connections/drivers.test.ts b/packages/cli/test/context/connections/drivers.test.ts index 9bb303aac..2001cd3d7 100644 --- a/packages/cli/test/context/connections/drivers.test.ts +++ b/packages/cli/test/context/connections/drivers.test.ts @@ -83,7 +83,7 @@ const connectionFixtures: Record = { }; const allowedScopeKeys = new Set(['dataset_ids', 'databases', 'schemas', 'schema_names']); -const historicSqlReaderDrivers = new Set(['postgres', 'bigquery', 'snowflake']); +const historicSqlReaderDrivers = new Set(['postgres', 'hologres', 'bigquery', 'snowflake']); function assertExportedRegistryBoundaryTypes(input: { scopeConfigKey: KtxScopeConfigKey; diff --git a/packages/cli/test/context/ingest/adapters/historic-sql/hologres-query-log-reader.live.integration.test.ts b/packages/cli/test/context/ingest/adapters/historic-sql/hologres-query-log-reader.live.integration.test.ts new file mode 100644 index 000000000..11f53fbb3 --- /dev/null +++ b/packages/cli/test/context/ingest/adapters/historic-sql/hologres-query-log-reader.live.integration.test.ts @@ -0,0 +1,44 @@ +import { describe, expect, it } from 'vitest'; +import { KtxHologresHistoricSqlQueryClient } from '../../../../../src/connectors/hologres/historic-sql-query-client.js'; +import { HologresQueryLogReader } from '../../../../../src/context/ingest/adapters/historic-sql/hologres-query-log-reader.js'; +import { + historicSqlUnifiedPullConfigSchema, + type AggregatedTemplate, +} from '../../../../../src/context/ingest/adapters/historic-sql/types.js'; + +// Live Hologres query-log integration. Gated on KTX_TEST_HOLOGRES_URL (a +// PostgreSQL-style DSN; URL-encode reserved characters, e.g. `$`->%24, `#`->%23). +const url = process.env.KTX_TEST_HOLOGRES_URL; + +describe.skipIf(!url)('HologresQueryLogReader (live Hologres)', () => { + it('probes hg_query_log and aggregates recent templates by digest', async () => { + const client = new KtxHologresHistoricSqlQueryClient({ + connectionId: 'wh', + connection: { driver: 'hologres', url }, + }); + try { + const reader = new HologresQueryLogReader(); + const probe = await reader.probe(client); + expect(Array.isArray(probe.warnings)).toBe(true); + + const end = new Date(); + const start = new Date(end.getTime() - 30 * 24 * 60 * 60 * 1000); + const config = historicSqlUnifiedPullConfigSchema.parse({ dialect: 'hologres', minExecutions: 1 }); + const templates: AggregatedTemplate[] = []; + for await (const template of reader.fetchAggregated(client, { start, end }, config)) { + templates.push(template); + } + expect(templates.length).toBeGreaterThan(0); + for (const template of templates) { + expect(template.dialect).toBe('hologres'); + expect(template.templateId.length).toBeGreaterThan(0); + expect(template.canonicalSql.length).toBeGreaterThan(0); + expect(template.stats.executions).toBeGreaterThanOrEqual(1); + expect(template.stats.errorRate).toBeGreaterThanOrEqual(0); + expect(template.stats.errorRate).toBeLessThanOrEqual(1); + } + } finally { + await client.cleanup(); + } + }); +}); diff --git a/packages/cli/test/context/ingest/adapters/historic-sql/hologres-query-log-reader.test.ts b/packages/cli/test/context/ingest/adapters/historic-sql/hologres-query-log-reader.test.ts new file mode 100644 index 000000000..b5ec3cf7d --- /dev/null +++ b/packages/cli/test/context/ingest/adapters/historic-sql/hologres-query-log-reader.test.ts @@ -0,0 +1,139 @@ +import { describe, expect, it, vi } from 'vitest'; +import { HistoricSqlExtensionMissingError } from '../../../../../src/context/ingest/adapters/historic-sql/errors.js'; +import { HologresQueryLogReader } from '../../../../../src/context/ingest/adapters/historic-sql/hologres-query-log-reader.js'; +import { + historicSqlUnifiedPullConfigSchema, + type AggregatedTemplate, +} from '../../../../../src/context/ingest/adapters/historic-sql/types.js'; + +interface FakeQueryResult { + headers: string[]; + rows: unknown[][]; + totalRows?: number; + error?: string; +} + +function queryClient(results: Array) { + const executeQuery = vi.fn(async (_sql: string, _params?: unknown[]) => { + const next = results.shift(); + if (!next) { + throw new Error('unexpected query'); + } + if (next instanceof Error) { + throw next; + } + return next; + }); + return { executeQuery }; +} + +function executedSql(client: ReturnType, index: number): string { + const call = client.executeQuery.mock.calls[index]; + if (!call) { + throw new Error(`expected query client call ${index}`); + } + return call[0]; +} + +const window = { start: new Date('2026-06-15T00:00:00.000Z'), end: new Date('2026-07-15T00:00:00.000Z') }; +const config = historicSqlUnifiedPullConfigSchema.parse({ dialect: 'hologres', minExecutions: 5 }); + +describe('HologresQueryLogReader probe', () => { + it('passes with no warnings when hg_query_log is accessible and pg_read_all_stats is granted', async () => { + const client = queryClient([ + { headers: ['?column?'], rows: [[1]] }, + { headers: ['has_role'], rows: [[true]] }, + ]); + await expect(new HologresQueryLogReader().probe(client)).resolves.toEqual({ warnings: [], info: [] }); + expect(executedSql(client, 0)).toBe('SELECT 1 FROM hologres.hg_query_log LIMIT 1'); + expect(executedSql(client, 1)).toContain("pg_has_role(current_user, 'pg_read_all_stats', 'USAGE')"); + }); + + it('soft-warns instead of throwing when pg_read_all_stats is not granted', async () => { + const client = queryClient([ + { headers: ['?column?'], rows: [[1]] }, + { headers: ['has_role'], rows: [[false]] }, + ]); + const result = await new HologresQueryLogReader().probe(client); + expect(result.warnings).toHaveLength(1); + expect(result.warnings[0]).toContain('pg_read_all_stats'); + }); + + it('throws HistoricSqlExtensionMissingError when hg_query_log is not accessible', async () => { + const client = queryClient([new Error('relation "hologres.hg_query_log" does not exist')]); + await expect(new HologresQueryLogReader().probe(client)).rejects.toBeInstanceOf(HistoricSqlExtensionMissingError); + }); +}); + +describe('HologresQueryLogReader fetchAggregated', () => { + it('aggregates hg_query_log rows by digest into AggregatedTemplate shape', async () => { + const client = queryClient([ + { + headers: [ + 'template_id', + 'canonical_sql', + 'executions', + 'distinct_users', + 'first_seen', + 'last_seen', + 'p50_ms', + 'p95_ms', + 'error_rate', + 'rows_produced', + 'top_users', + ], + rows: [ + [ + 'digest-abc', + 'SELECT * FROM orders WHERE id = $1', + '10', + '2', + new Date('2026-07-01T00:00:00.000Z'), + new Date('2026-07-10T00:00:00.000Z'), + 12.5, + 340, + 0.1, + '1000', + '[{"user":"alice","executions":7},{"user":"bob","executions":3}]', + ], + ], + }, + ]); + const templates: AggregatedTemplate[] = []; + for await (const template of new HologresQueryLogReader().fetchAggregated(client, window, config)) { + templates.push(template); + } + expect(templates).toHaveLength(1); + expect(templates[0]).toMatchObject({ + templateId: 'digest-abc', + canonicalSql: 'SELECT * FROM orders WHERE id = $1', + dialect: 'hologres', + stats: { + executions: 10, + distinctUsers: 2, + firstSeen: '2026-07-01T00:00:00.000Z', + lastSeen: '2026-07-10T00:00:00.000Z', + p50RuntimeMs: 12.5, + p95RuntimeMs: 340, + errorRate: 0.1, + rowsProduced: 1000, + }, + topUsers: [ + { user: 'alice', executions: 7 }, + { user: 'bob', executions: 3 }, + ], + }); + }); + + it('filters to fingerprinted queries and binds the window and minExecutions as parameters', async () => { + const client = queryClient([{ headers: [], rows: [] }]); + const results: AggregatedTemplate[] = []; + for await (const template of new HologresQueryLogReader().fetchAggregated(client, window, config)) { + results.push(template); + } + expect(results).toHaveLength(0); + expect(executedSql(client, 0)).toContain('digest IS NOT NULL'); + expect(executedSql(client, 0)).toContain('FROM hologres.hg_query_log'); + expect(client.executeQuery.mock.calls[0]?.[1]).toEqual([window.start, window.end, 5]); + }); +}); diff --git a/packages/cli/test/context/ingest/historic-sql-probes.test.ts b/packages/cli/test/context/ingest/historic-sql-probes.test.ts index 542c4fded..c751177b8 100644 --- a/packages/cli/test/context/ingest/historic-sql-probes.test.ts +++ b/packages/cli/test/context/ingest/historic-sql-probes.test.ts @@ -45,6 +45,7 @@ function factories( fakeRunner('snowflake', 'SNOWFLAKE.ACCOUNT_USAGE.QUERY_HISTORY'); const bigquery = overrides.bigquery ?? fakeRunner('bigquery', 'INFORMATION_SCHEMA.JOBS_BY_PROJECT'); + const hologres = overrides.hologres ?? fakeRunner('hologres', 'hologres.hg_query_log'); return { postgres: { @@ -59,6 +60,10 @@ function factories( catalogName: 'INFORMATION_SCHEMA.JOBS_BY_PROJECT', load: vi.fn(async () => bigquery), }, + hologres: { + catalogName: 'hologres.hg_query_log', + load: vi.fn(async () => hologres), + }, }; } From 252b7c445f4d571f8fd73f78984b6d0c11944b1b Mon Sep 17 00:00:00 2001 From: TimothyDing Date: Wed, 15 Jul 2026 15:53:03 +0800 Subject: [PATCH 10/14] docs(hologres): document the Hologres primary source Add a Hologres section to the primary-sources integration page covering the PostgreSQL-compatible connection config (default port 80), the auto-excluded Hologres system schemas, hg_query_log query history (digest aggregation, slow-log and 30-day retention caveats, pg_read_all_stats soft-degrade), and dialect notes, and reference Hologres in the description and field table. --- .../docs/integrations/primary-sources.mdx | 105 +++++++++++++++++- 1 file changed, 101 insertions(+), 4 deletions(-) diff --git a/docs-site/content/docs/integrations/primary-sources.mdx b/docs-site/content/docs/integrations/primary-sources.mdx index b553a3a7e..e393d227e 100644 --- a/docs-site/content/docs/integrations/primary-sources.mdx +++ b/docs-site/content/docs/integrations/primary-sources.mdx @@ -1,6 +1,6 @@ --- title: Primary Sources -description: Connect ktx to PostgreSQL, Snowflake, BigQuery, MySQL, ClickHouse, SQL Server, SQLite, DuckDB, MongoDB, or Amazon Athena. +description: Connect ktx to PostgreSQL, Hologres, Snowflake, BigQuery, MySQL, ClickHouse, SQL Server, SQLite, DuckDB, MongoDB, or Amazon Athena. --- **ktx** connects to your data warehouse or database to build schema context, @@ -26,13 +26,13 @@ Agents should prefer environment or file references over literal secrets. | Field | Required | Applies to | Description | |-------|----------|------------|-------------| -| `driver` | Yes | all connections | Connector driver such as `postgres`, `snowflake`, `bigquery`, `mysql`, `clickhouse`, `sqlserver`, `sqlite`, `duckdb`, `mongodb`, or `athena` | +| `driver` | Yes | all connections | Connector driver such as `postgres`, `hologres`, `snowflake`, `bigquery`, `mysql`, `clickhouse`, `sqlserver`, `sqlite`, `duckdb`, `mongodb`, or `athena` | | `url` | One of the connection methods | URL-style connectors | Database URL, `env:NAME`, or `file:/path/to/secret` | -| `host`, `port`, `database`, `username`, `password` | One of the connection methods | PostgreSQL, MySQL, SQL Server | Field-by-field connection values | +| `host`, `port`, `database`, `username`, `password` | One of the connection methods | PostgreSQL, Hologres, MySQL, SQL Server | Field-by-field connection values | | `schema` or `schemas` | No | schema-aware warehouses | Single schema or list of schemas to scan | | `databases` | No | ClickHouse, MongoDB, Athena | List of databases to scan | | `sample_size`, `order_by` | No | MongoDB | Schema-inference sampling controls (recent documents, sort field) | -| `context.queryHistory` | No | PostgreSQL, Snowflake, BigQuery | Enables query-history ingestion when the warehouse supports it | +| `context.queryHistory` | No | PostgreSQL, Hologres, Snowflake, BigQuery | Enables query-history ingestion when the warehouse supports it | | `path` | Yes for path-style SQLite/DuckDB | SQLite, DuckDB | Local SQLite or DuckDB database path or `env:NAME` reference | | `max_bytes_billed` | No | BigQuery | Maximum bytes billed per query job | | `query_timeout_ms` | No | all warehouses | Maximum execution time for a single read-only query, in milliseconds (default 30000). A query exceeding it is cancelled server-side (or, for SQLite, by terminating the off-process executor) and returns a `query exceeded Ns` error so the agent can revise. | @@ -124,6 +124,103 @@ This helps **ktx** understand how your team actually queries the data. --- +## Hologres + +Alibaba Cloud Hologres is PostgreSQL-wire-compatible, so **ktx** reuses the +PostgreSQL connector for connection, schema introspection, foreign key +detection, sampling, and column statistics. Query history comes from Hologres's +own `hologres.hg_query_log` system table rather than `pg_stat_statements`. + +### Connection config + +```yaml title="ktx.yaml" +connections: + my-hologres: + driver: hologres + url: env:HOLOGRES_URL + schemas: + - my_schema +``` + +Or with individual fields: + +```yaml title="ktx.yaml" +connections: + my-hologres: + driver: hologres + host: hgpostcn-xxxx-cn-hangzhou.hologres.aliyuncs.com + port: 80 + database: my_db + username: env:HOLOGRES_ACCESS_ID + password: env:HOLOGRES_ACCESS_KEY + schemas: + - my_schema +``` + +Hologres endpoints commonly use port `80`, which `ktx setup` offers as the +default. Authentication mirrors PostgreSQL (`password`, `url`, and `ssl` +references). + +### System schemas + +**ktx** automatically excludes the Hologres engine-internal schemas that never +hold user data — `hologres`, `hologres_streaming_mv`, `hologres_statistic`, +`hologres_sample`, `hologres_object_table`, `hg_recyclebin`, and `hg_internal` — +the same way it drops `pg_catalog`. Exclusion is by exact name, so user schemas +that share the `hologres_` prefix (such as a `hologres_dataset_*` sample dataset) +are still scanned. + +### Features + +| Feature | Supported | Notes | +|---------|-----------|-------| +| Tables & views | Yes | Via `pg_catalog` | +| Primary keys | Yes | Via `information_schema.table_constraints` | +| Foreign keys | Yes | When declared | +| Row count estimates | Yes | Via `pg_class.reltuples` | +| Column statistics | Yes | Via `pg_stats` (`n_distinct`; the `correlation` column is not used) | +| Query history | Yes | Via `hologres.hg_query_log` | +| Table sampling | Yes | `SELECT ... LIMIT n` | + +### Query history + +Hologres query history aggregates `hologres.hg_query_log` by its `digest` +column — a SQL fingerprint Hologres computes for `SELECT`, `INSERT`, `UPDATE`, +and `DELETE` — over a `query_start` time window, feeding the same unified staged +artifact shape as Postgres, Snowflake, and BigQuery. + +```yaml + context: + queryHistory: + enabled: true + windowDays: 90 + minExecutions: 5 + filters: + dropTrivialProbes: true +``` + +**Requirements and caveats:** + +- A superuser or a member of `pg_read_all_stats` sees every database's log; a + `db_admin` (via SPM/SLPM) sees the current database. Without `pg_read_all_stats` + ingest still runs but sees only the current user's queries, and **ktx** records + a warning noting the partial coverage. +- `hg_query_log` is a slow-query log: by default it records only queries slower + than `log_min_duration_statement` (100ms), so query history is biased toward + the expensive query patterns worth optimizing. +- Hologres retains slow query logs for roughly 30 days, so the effective window + is capped regardless of a larger `windowDays`. + +### Dialect notes + +- SQL compilation reuses the PostgreSQL dialect (`LIMIT/OFFSET`, `$1`/`$2` + positional parameters, `COUNT(*) FILTER (WHERE ...)`). +- All ingest and query-history reads run as single read-only statements; **ktx** + never wraps them in a transaction, honoring Hologres's restriction against + multi-statement DML and mixed DDL/DML transactions. + +--- + ## Snowflake Connects via the Snowflake SDK. Supports multi-schema scanning, RSA key authentication, and query-history configuration for Snowflake query history. From 547435348ab09d62813e408bdd249af0b8378ed4 Mon Sep 17 00:00:00 2001 From: TimothyDing Date: Wed, 15 Jul 2026 16:07:08 +0800 Subject: [PATCH 11/14] feat(hologres): require Hologres 4.0+ via hg_version() MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Detect the Hologres version with select hg_version() (which returns e.g. "Hologres 4.2.10 ...", distinct from the compatible PostgreSQL 11.3 version) and refuse connections older than 4.0 — the first release that supports ktx — in both testConnection and introspect with a clear expected error. Parses the three-part Hologres version after the "Hologres " prefix. connectionId is now protected on the shared Postgres connector so the subclass can run the check via executeReadOnly. --- .../cli/src/connectors/hologres/connector.ts | 76 ++++++++++++++++++- .../cli/src/connectors/postgres/connector.ts | 2 +- .../connectors/hologres/connector.test.ts | 51 +++++++++++++ 3 files changed, 126 insertions(+), 3 deletions(-) diff --git a/packages/cli/src/connectors/hologres/connector.ts b/packages/cli/src/connectors/hologres/connector.ts index 7e9680f99..230e9fad7 100644 --- a/packages/cli/src/connectors/hologres/connector.ts +++ b/packages/cli/src/connectors/hologres/connector.ts @@ -1,3 +1,11 @@ +import { KtxExpectedError } from '../../errors.js'; +import { + connectorTestFailure, + type KtxConnectorTestResult, + type KtxScanContext, + type KtxScanInput, + type KtxSchemaSnapshot, +} from '../../context/scan/types.js'; import { KtxPostgresScanConnector, type KtxPostgresConnectionConfig } from '../postgres/connector.js'; // Hologres engine-internal schemas that never hold user data. Dropped from schema @@ -14,6 +22,10 @@ const HOLOGRES_SYSTEM_SCHEMAS: ReadonlySet = new Set([ 'hg_internal', ]); +// ktx supports Hologres from major version 4.0 onward. +const HOLOGRES_MIN_MAJOR_VERSION = 4; +const HOLOGRES_VERSION_SQL = 'select hg_version()'; + export type KtxHologresConnectionConfig = KtxPostgresConnectionConfig; export function isKtxHologresConnectionConfig( @@ -22,16 +34,76 @@ export function isKtxHologresConnectionConfig( return String(connection?.driver ?? '').toLowerCase() === 'hologres'; } +// hg_version() returns e.g. "Hologres 4.2.10 (tag: ...), compatible with +// PostgreSQL 11.3 ...". The three-part version follows the "Hologres " prefix, so +// anchor there to avoid matching the compatible-PostgreSQL version. +function parseHologresMajorVersion(raw: string): { version: string; major: number } | null { + const match = /Hologres\s+v?(\d+)(?:\.(\d+))?(?:\.(\d+))?/i.exec(raw); + if (!match) { + return null; + } + const major = Number.parseInt(match[1] ?? '', 10); + if (!Number.isFinite(major)) { + return null; + } + const version = [match[1], match[2], match[3]].filter((part) => part !== undefined).join('.'); + return { version, major }; +} + /** * Hologres scan connector. Hologres is PostgreSQL-wire-compatible, so it reuses * the PostgreSQL connector's connection, introspection, sampling, and read-only * single-statement execution (which never opens a transaction — Hologres forbids - * multi-statement DML/DDL transactions). The only override drops Hologres - * system schemas from discovery. + * multi-statement DML/DDL transactions). It drops Hologres system schemas from + * discovery and requires Hologres 4.0+, the first version that supports ktx. */ export class KtxHologresScanConnector extends KtxPostgresScanConnector { + override async testConnection(): Promise { + const base = await super.testConnection(); + if (!base.success) { + return base; + } + try { + await this.assertSupportedHologresVersion(this.connectionId); + return { success: true }; + } catch (error) { + return connectorTestFailure(error); + } + } + + override async introspect(input: KtxScanInput, ctx: KtxScanContext): Promise { + await this.assertSupportedHologresVersion(input.connectionId); + return super.introspect(input, ctx); + } + override async listSchemas(): Promise { const schemas = await super.listSchemas(); return schemas.filter((schema) => !HOLOGRES_SYSTEM_SCHEMAS.has(schema)); } + + private async assertSupportedHologresVersion(connectionId: string): Promise { + let raw: string; + try { + const result = await this.executeReadOnly({ connectionId, sql: HOLOGRES_VERSION_SQL }, {} as never); + const cell = result.rows[0]?.[0]; + raw = typeof cell === 'string' ? cell : String(cell ?? ''); + } catch (error) { + throw new KtxExpectedError( + 'ktx requires Hologres 4.0 or newer. Could not read the Hologres version via hg_version(); ' + + 'confirm this is a Hologres instance running 4.0 or later.', + { cause: error }, + ); + } + const parsed = parseHologresMajorVersion(raw); + if (!parsed) { + throw new KtxExpectedError( + `ktx requires Hologres 4.0 or newer, but hg_version() returned an unrecognized version string: ${raw}`, + ); + } + if (parsed.major < HOLOGRES_MIN_MAJOR_VERSION) { + throw new KtxExpectedError( + `ktx requires Hologres 4.0 or newer; this connection reports Hologres ${parsed.version}.`, + ); + } + } } diff --git a/packages/cli/src/connectors/postgres/connector.ts b/packages/cli/src/connectors/postgres/connector.ts index 44b83ac42..280acf74f 100644 --- a/packages/cli/src/connectors/postgres/connector.ts +++ b/packages/cli/src/connectors/postgres/connector.ts @@ -426,7 +426,7 @@ export class KtxPostgresScanConnector implements KtxScanConnector { estimatedRowCounts: true, }); - private readonly connectionId: string; + protected readonly connectionId: string; private readonly connection: KtxPostgresConnectionConfig; private readonly poolConfig: KtxPostgresPoolConfig; private readonly poolFactory: KtxPostgresPoolFactory; diff --git a/packages/cli/test/connectors/hologres/connector.test.ts b/packages/cli/test/connectors/hologres/connector.test.ts index 0a7e14047..80eda6df6 100644 --- a/packages/cli/test/connectors/hologres/connector.test.ts +++ b/packages/cli/test/connectors/hologres/connector.test.ts @@ -1,4 +1,5 @@ import { describe, expect, it, vi } from 'vitest'; +import { KtxExpectedError } from '../../../src/errors.js'; import { KtxHologresScanConnector, isKtxHologresConnectionConfig, @@ -47,6 +48,13 @@ const connection: KtxHologresConnectionConfig = { schema: 'public', }; +function hologresVersionResult(version: string): FakeQueryResult { + return { + fields: [{ name: 'hg_version', dataTypeID: 25 }], + rows: [{ hg_version: `Hologres ${version} (tag: -abc build: Release), compatible with PostgreSQL 11.3 on x86_64-linux` }], + }; +} + describe('isKtxHologresConnectionConfig', () => { it('accepts hologres connections and rejects everything else', () => { expect(isKtxHologresConnectionConfig(connection)).toBe(true); @@ -96,6 +104,7 @@ describe('KtxHologresScanConnector', () => { it('introspects a schema and stamps the snapshot driver as hologres', async () => { const results = new Map([ + ['hg_version()', hologresVersionResult('4.2.10')], [ 'c.reltuples::bigint AS row_count', { rows: [{ table_name: 'orders', table_kind: 'r', row_count: '5', table_comment: 'order facts' }] }, @@ -125,4 +134,46 @@ describe('KtxHologresScanConnector', () => { ]); expect(snapshot.tables[0]?.columns.map((column) => [column.name, column.primaryKey])).toEqual([['id', true]]); }); + + it('passes the connection test on Hologres 4.0 or newer', async () => { + const connector = new KtxHologresScanConnector({ + connectionId: 'wh', + connection, + poolFactory: fakePoolFactory( + new Map([ + ['SELECT 1', { rows: [] }], + ['hg_version()', hologresVersionResult('4.0.1')], + ]), + ), + }); + await expect(connector.testConnection()).resolves.toEqual({ success: true }); + }); + + it('rejects introspection against Hologres older than 4.0', async () => { + const connector = new KtxHologresScanConnector({ + connectionId: 'wh', + connection, + poolFactory: fakePoolFactory(new Map([['hg_version()', hologresVersionResult('3.9.5')]])), + }); + await expect( + connector.introspect({ connectionId: 'wh', driver: 'hologres' }, { runId: 'r' }), + ).rejects.toBeInstanceOf(KtxExpectedError); + await expect( + connector.introspect({ connectionId: 'wh', driver: 'hologres' }, { runId: 'r' }), + ).rejects.toThrow(/Hologres 4\.0 or newer/); + }); + + it('fails the connection test on Hologres older than 4.0', async () => { + const connector = new KtxHologresScanConnector({ + connectionId: 'wh', + connection, + poolFactory: fakePoolFactory( + new Map([ + ['SELECT 1', { rows: [] }], + ['hg_version()', hologresVersionResult('3.9.5')], + ]), + ), + }); + expect((await connector.testConnection()).success).toBe(false); + }); }); From 8ced73ba6c42b64de65768314efbfef02410182a Mon Sep 17 00:00:00 2001 From: TimothyDing Date: Wed, 15 Jul 2026 16:07:49 +0800 Subject: [PATCH 12/14] docs(hologres): note the Hologres 4.0+ requirement Document that ktx requires Hologres 4.0 or newer and checks it via hg_version() when testing a connection or scanning. --- docs-site/content/docs/integrations/primary-sources.mdx | 3 +++ 1 file changed, 3 insertions(+) diff --git a/docs-site/content/docs/integrations/primary-sources.mdx b/docs-site/content/docs/integrations/primary-sources.mdx index e393d227e..c7e97db0a 100644 --- a/docs-site/content/docs/integrations/primary-sources.mdx +++ b/docs-site/content/docs/integrations/primary-sources.mdx @@ -131,6 +131,9 @@ PostgreSQL connector for connection, schema introspection, foreign key detection, sampling, and column statistics. Query history comes from Hologres's own `hologres.hg_query_log` system table rather than `pg_stat_statements`. +**ktx** requires Hologres 4.0 or newer — the first release to support ktx — and +verifies the version with `hg_version()` when testing a connection or scanning. + ### Connection config ```yaml title="ktx.yaml" From 9bf80da2e87c2ff103e055156db0786618cfb557 Mon Sep 17 00:00:00 2001 From: TimothyDing Date: Wed, 15 Jul 2026 17:24:10 +0800 Subject: [PATCH 13/14] fix(scan): accept hologres via registry-derived driver allowlist Standalone scan's normalizeDriver kept a hardcoded driver allowlist that drifted from the canonical driverRegistrations registry, so a hologres connection passed the connection test but failed the scan phase. Derive the allowlist (and error message) from the registry so every registered driver is accepted and future connectors cannot regress. --- packages/cli/src/context/scan/local-scan.ts | 20 +++-------- .../cli/test/context/scan/local-scan.test.ts | 33 +++++++++++++++++++ 2 files changed, 38 insertions(+), 15 deletions(-) diff --git a/packages/cli/src/context/scan/local-scan.ts b/packages/cli/src/context/scan/local-scan.ts index ac0a90e5d..98bdd1dc9 100644 --- a/packages/cli/src/context/scan/local-scan.ts +++ b/packages/cli/src/context/scan/local-scan.ts @@ -2,6 +2,7 @@ import type { createKtxEmbeddingProvider } from '../../llm/embedding-provider.js import type { createKtxLlmProvider } from '../../llm/model-provider.js'; import type { KtxEmbeddingProvider } from '../../llm/types.js'; import { createDefaultLocalIngestAdapters } from '../../context/ingest/local-adapters.js'; +import { getDriverRegistration, listSupportedDrivers } from '../connections/drivers.js'; import { getLocalStageOnlyIngestStatus, type LocalIngestRunRecord, runLocalStageOnlyIngest } from '../../context/ingest/local-stage-ingest.js'; import type { SourceAdapter } from '../../context/ingest/types.js'; import { createLocalKtxLlmRuntimeFromConfig } from '../../context/llm/local-config.js'; @@ -137,23 +138,12 @@ const LOCAL_AUTHOR = 'ktx'; const LOCAL_AUTHOR_EMAIL = 'ktx@example.com'; function normalizeDriver(driver: string | undefined): KtxConnectionDriver { - const normalized = (driver ?? '').toLowerCase(); - if ( - normalized === 'postgres' || - normalized === 'sqlite' || - normalized === 'duckdb' || - normalized === 'mysql' || - normalized === 'clickhouse' || - normalized === 'sqlserver' || - normalized === 'bigquery' || - normalized === 'snowflake' || - normalized === 'athena' || - normalized === 'mongodb' - ) { - return normalized; + const registration = getDriverRegistration(driver ?? ''); + if (registration) { + return registration.driver; } throw new Error( - `Standalone ktx scan supports postgres/sqlite/duckdb/mysql/clickhouse/sqlserver/bigquery/snowflake/athena/mongodb in this phase, received "${driver ?? 'unknown'}"`, + `Standalone ktx scan supports ${listSupportedDrivers().join('/')} in this phase, received "${driver ?? 'unknown'}"`, ); } diff --git a/packages/cli/test/context/scan/local-scan.test.ts b/packages/cli/test/context/scan/local-scan.test.ts index 75f929d70..de2fb9aba 100644 --- a/packages/cli/test/context/scan/local-scan.test.ts +++ b/packages/cli/test/context/scan/local-scan.test.ts @@ -2209,6 +2209,39 @@ describe('local scan', () => { 'raw-sources/warehouse/live-database/2026-04-29-170000-scan-run-athena/scan-report.json', ); }); + + it('accepts hologres as a native standalone scan driver when the host supplies a live-database adapter', async () => { + await writeFile( + join(project.projectDir, 'ktx.yaml'), + [ + 'connections:', + ' warehouse:', + ' driver: hologres', + ' url: env:HOLOGRES_URL', + ' schemas:', + ' - public', + 'ingest:', + ' adapters:', + ' - live-database', + '', + ].join('\n'), + 'utf-8', + ); + project = await loadKtxProject({ projectDir: project.projectDir }); + + const result = await runLocalScan({ + project, + adapters: [fetchOnlyAdapter()], + connectionId: 'warehouse', + jobId: 'scan-run-hologres', + now: () => new Date('2026-04-29T18:00:00.000Z'), + }); + + expect(result.report.driver).toBe('hologres'); + expect(result.report.artifactPaths.reportPath).toBe( + 'raw-sources/warehouse/live-database/2026-04-29-180000-scan-run-hologres/scan-report.json', + ); + }); }); describe('resolveEnabledTables', () => { From 7a490bfced0c0a2fcdf103f0f81f9222d37b5456 Mon Sep 17 00:00:00 2001 From: TimothyDing Date: Thu, 16 Jul 2026 10:34:18 +0800 Subject: [PATCH 14/14] fix(hologres): normalize dialect to postgres for daemon SQL analysis The daemon's /sql/analyze-batch and read-only validation passed the connection dialect straight to sqlglot, which has no "hologres" dialect, so Hologres query-history ingest failed with 500 "Unknown dialect 'hologres'". Hologres is Postgres wire/SQL compatible, so normalize hologres -> postgres at both entry points before any sqlglot call. Adds regression tests for the analyze and validate paths. --- .../ktx-daemon/src/ktx_daemon/sql_analysis.py | 16 ++++++-- python/ktx-daemon/tests/test_sql_analysis.py | 38 +++++++++++++++++++ 2 files changed, 51 insertions(+), 3 deletions(-) diff --git a/python/ktx-daemon/src/ktx_daemon/sql_analysis.py b/python/ktx-daemon/src/ktx_daemon/sql_analysis.py index 54f3d0e27..d8cb9a78d 100644 --- a/python/ktx-daemon/src/ktx_daemon/sql_analysis.py +++ b/python/ktx-daemon/src/ktx_daemon/sql_analysis.py @@ -90,6 +90,15 @@ class ValidateReadOnlySqlResponse(BaseModel): ) +# sqlglot has no "hologres" dialect; Hologres is Postgres wire/SQL compatible, +# so parse and normalize it as postgres. +_SQLGLOT_DIALECT_ALIASES = {"hologres": "postgres"} + + +def _sqlglot_dialect(dialect: str) -> str: + return _SQLGLOT_DIALECT_ALIASES.get(dialect.strip().lower(), dialect) + + def _ordered_unique(values: list[str]) -> list[str]: seen: set[str] = set() result: list[str] = [] @@ -271,7 +280,7 @@ def validate_read_only_sql_response( request: ValidateReadOnlySqlRequest, ) -> ValidateReadOnlySqlResponse: try: - statements = sqlglot.parse(request.sql, read=request.dialect) + statements = sqlglot.parse(request.sql, read=_sqlglot_dialect(request.dialect)) except sqlglot.errors.SqlglotError as exc: return ValidateReadOnlySqlResponse(ok=False, error=f"Invalid expression: {exc}") @@ -314,8 +323,9 @@ def _worker_count(request: AnalyzeSqlBatchRequest) -> int: def analyze_sql_batch_response( request: AnalyzeSqlBatchRequest, ) -> AnalyzeSqlBatchResponse: - catalog = _catalog_index(request.catalog, request.dialect) - payloads = [(item.id, item.sql, request.dialect, catalog) for item in request.items] + dialect = _sqlglot_dialect(request.dialect) + catalog = _catalog_index(request.catalog, dialect) + payloads = [(item.id, item.sql, dialect, catalog) for item in request.items] if _worker_count(request) == 1: analyzed = [_analyze_payload(payload) for payload in payloads] else: diff --git a/python/ktx-daemon/tests/test_sql_analysis.py b/python/ktx-daemon/tests/test_sql_analysis.py index 2fb3970a5..aa18b936d 100644 --- a/python/ktx-daemon/tests/test_sql_analysis.py +++ b/python/ktx-daemon/tests/test_sql_analysis.py @@ -167,6 +167,32 @@ def test_analyze_sql_batch_returns_bigquery_project_dataset_table_refs() -> None ] +def test_analyze_sql_batch_accepts_hologres_dialect_as_postgres() -> None: + # Hologres has no sqlglot dialect; it must analyze as postgres, not error. + response = analyze_sql_batch_response( + AnalyzeSqlBatchRequest( + dialect="hologres", + items=[ + AnalyzeSqlBatchItem( + id="holo", + sql="select l_orderkey from public.lineitem where l_quantity > 10", + ) + ], + max_workers=1, + ) + ) + + result = response.results["holo"] + assert result.error is None + assert [item.model_dump() for item in result.tables_touched] == [ + {"catalog": None, "db": "public", "name": "lineitem"} + ] + assert result.columns_by_clause == { + "select": ["l_orderkey"], + "where": ["l_quantity"], + } + + def test_columns_from_nodes_ignores_non_expression_clause_values() -> None: assert _columns_from_nodes([True, False, None]) == [] @@ -240,3 +266,15 @@ def test_validate_read_only_sql_reports_parse_errors() -> None: assert response.ok is False assert response.error is not None assert "Invalid expression" in response.error + + +def test_validate_read_only_sql_accepts_hologres_dialect() -> None: + response = validate_read_only_sql_response( + ValidateReadOnlySqlRequest( + dialect="hologres", + sql="select l_orderkey from public.lineitem", + ) + ) + + assert response.ok is True + assert response.error is None