import { describe, expect, it, vi } from 'vitest'; vi.mock('ai', async (importOriginal) => { const actual = await importOriginal(); return { ...actual, generateText: vi.fn() }; }); import { generateText } from 'ai'; import { buildKtxColumnDescriptionPrompt, buildKtxDataSourceDescriptionPrompt, buildKtxTableDescriptionPrompt, type KtxDescriptionCachePort, KtxDescriptionGenerator, } from '../../../src/context/scan/description-generation.js'; import { createKtxConnectorCapabilities, type KtxScanConnector } from '../../../src/context/scan/types.js'; function createCache(initial: Record = {}): KtxDescriptionCachePort { const data = new Map(Object.entries(initial)); return { buildTableKey: (table) => [table.catalog, table.db, table.name].filter(Boolean).join('.'), buildColumnKey: (table, columnName) => [table.catalog, table.db, table.name, columnName].filter(Boolean).join('.'), buildConnectionKey: (connectionName) => `__connection:${connectionName}`, get: vi.fn(async (key: string) => data.get(key) ?? null), set: vi.fn(async (key: string, value: string) => { data.set(key, value); }), }; } function createLlmProvider(text = 'generated description') { vi.mocked(generateText).mockResolvedValue({ text } as never); return { generateText: vi.fn(async (input) => { const result = await generateText({ system: input.system ? { role: 'system', content: input.system } : undefined, messages: [{ role: 'user', content: input.prompt }], temperature: input.temperature, } as never); return result.text; }), generateObject: vi.fn(), runAgentLoop: vi.fn(), } as any; } function createFailingLlmProvider(message = 'timeout exceeded when trying to connect') { vi.mocked(generateText).mockRejectedValue(new Error(message) as never); return { generateText: vi.fn(async (input) => { const result = await generateText({ system: input.system ? { role: 'system', content: input.system } : undefined, messages: [{ role: 'user', content: input.prompt }], temperature: input.temperature, } as never); return result.text; }), generateObject: vi.fn(), runAgentLoop: vi.fn(), } as any; } function createConnector(): KtxScanConnector { return { id: 'test-connector', driver: 'postgres', capabilities: createKtxConnectorCapabilities({ tableSampling: true, columnSampling: true, nestedAnalysis: true, }), introspect: vi.fn(async () => { throw new Error('introspection is not used by description generation'); }), listSchemas: vi.fn(async () => []), listTables: vi.fn(async () => []), sampleColumn: vi.fn(async () => ({ values: ['paid', 'refunded', null], nullCount: 1, distinctCount: 2, })), sampleTable: vi.fn(async () => ({ headers: ['id', 'status', 'amount'], rows: [ [1, 'paid', 20], [2, 'refunded', 10], ], totalRows: 2, })), }; } describe('KTX description prompt builders', () => { it('builds column prompts with sample values, source descriptions, and nested BigQuery guidance', () => { const { system, user } = buildKtxColumnDescriptionPrompt({ columnName: 'payload', columnValues: [{ nested: true }, '[1,2]'], tableContext: 'Table: events | Columns: payload | Data source: BIGQUERY', dataSourceType: 'BIGQUERY', supportsNestedAnalysis: true, rawDescriptions: { db: 'Raw event payload', ai: 'Old AI text', user: 'User text' }, maxWords: 12, }); expect(user).toContain( ' Table: events | Columns: payload | Data source: BIGQUERY ', ); expect(user).toContain(' payload '); expect(user).toContain(' [object Object], [1,2] '); expect(user).toContain(' Raw event payload '); expect(user).not.toContain('Old AI text'); expect(user).not.toContain('User text'); expect(system).toContain('nested/structured data'); expect(system).toContain('12 words or less'); expect(user).not.toContain('12 words or less'); }); it('builds table and data-source prompts from sampled rows', () => { const sample = { headers: ['id', 'status'], rows: [ [1, 'paid'], [2, 'refunded'], ], totalRows: 2, }; const table = buildKtxTableDescriptionPrompt({ tableName: 'orders', sampleData: sample, dataSourceType: 'POSTGRESQL', rawDescriptions: { dbt: 'Fact table for commerce orders' }, }); expect(table.user).toContain('status: paid, refunded'); expect(table.system).toContain('Analyze database tables'); const datasource = buildKtxDataSourceDescriptionPrompt({ tableSamples: [['orders', sample]], dataSourceType: 'POSTGRESQL', }); expect(datasource.user).toContain('orders (2 columns, 2 sample rows)'); expect(datasource.system).toContain('Analyze databases'); }); }); describe('KtxDescriptionGenerator', () => { it('generates column descriptions with pre-fetched values, cache hits, and word-limit metadata', async () => { const cache = createCache({ 'warehouse.public.orders.cached_status': 'Cached status description' }); const llmRuntime = createLlmProvider('Payment state'); const connector = createConnector(); const generator = new KtxDescriptionGenerator({ llmRuntime, cache, settings: { columnMaxWords: 12, tableMaxWords: 18, dataSourceMaxWords: 24, temperature: 0.2, concurrencyLimit: 2, }, }); const result = await generator.generateColumnDescriptions({ connectionId: 'conn-1', connector, context: { runId: 'run-1' }, dataSourceType: 'POSTGRESQL', supportsNestedAnalysis: false, table: { catalog: 'warehouse', db: 'public', name: 'orders', columns: [ { name: 'status', sampleValues: ['paid', 'refunded'], rawDescriptions: { db: 'Payment lifecycle' } }, { name: 'cached_status', sampleValues: ['open'] }, ], }, skipExisting: false, existingDescriptions: {}, }); expect(result).toEqual({ columnDescriptions: [ ['status', 'Payment state'], ['cached_status', 'Cached status description'], ], processedColumns: ['status'], skippedColumns: ['cached_status'], }); expect(connector.sampleColumn).not.toHaveBeenCalled(); expect(generateText).toHaveBeenCalledWith( expect.objectContaining({ temperature: 0.2, system: expect.objectContaining({ role: 'system', content: expect.stringContaining('Please provide a concise description in 12 words or less.'), }), messages: expect.arrayContaining([ expect.objectContaining({ role: 'user', content: expect.stringContaining(' status '), }), ]), }), ); const lastCall = vi.mocked(generateText).mock.calls.at(-1)?.[0]; expect(lastCall?.messages?.some((message) => message.role === 'system')).toBe(false); }); it('samples through the connector when column values are not pre-fetched', async () => { const connector = createConnector(); const generator = new KtxDescriptionGenerator({ llmRuntime: createLlmProvider('Current order state'), settings: { columnMaxWords: 12, tableMaxWords: 18, dataSourceMaxWords: 24, }, }); const result = await generator.generateColumnDescriptions({ connectionId: 'conn-1', connector, context: { runId: 'run-1' }, dataSourceType: 'POSTGRESQL', supportsNestedAnalysis: false, table: { catalog: null, db: 'public', name: 'orders', columns: [{ name: 'status' }], }, }); expect(connector.sampleColumn).toHaveBeenCalledWith( { connectionId: 'conn-1', table: { catalog: null, db: 'public', name: 'orders' }, column: 'status', limit: 50, }, { runId: 'run-1' }, ); expect(result.columnDescriptions).toEqual([['status', 'Current order state']]); }); it('samples through a description sampling port without requiring structural introspection', async () => { const sampler = { id: 'description-sampler:conn-1', sampleColumn: vi.fn(async () => ({ values: ['paid', 'refunded'], nullCount: null, distinctCount: null, })), sampleTable: vi.fn(async () => ({ headers: ['id', 'status'], rows: [[1, 'paid']], totalRows: 1, })), }; const generator = new KtxDescriptionGenerator({ llmRuntime: createLlmProvider('Generated through sampler'), settings: { columnMaxWords: 12, tableMaxWords: 18, dataSourceMaxWords: 24, }, }); const result = await generator.generateColumnDescriptions({ connectionId: 'conn-1', connector: sampler, context: { runId: 'run-1' }, dataSourceType: 'POSTGRESQL', supportsNestedAnalysis: false, table: { catalog: null, db: 'public', name: 'orders', columns: [{ name: 'status' }], }, }); expect(result.columnDescriptions).toEqual([['status', 'Generated through sampler']]); expect(sampler.sampleColumn).toHaveBeenCalledWith( { connectionId: 'conn-1', table: { catalog: null, db: 'public', name: 'orders' }, column: 'status', limit: 50, }, { runId: 'run-1' }, ); expect('introspect' in sampler).toBe(false); }); it('does not turn LLM failures into generated descriptions', async () => { const cache = createCache(); const connector = createConnector(); const generator = new KtxDescriptionGenerator({ llmRuntime: createFailingLlmProvider(), cache, settings: { columnMaxWords: 12, tableMaxWords: 18, dataSourceMaxWords: 24, }, }); const columnResult = await generator.generateColumnDescriptions({ connectionId: 'conn-1', connector, context: { runId: 'run-1' }, dataSourceType: 'POSTGRESQL', supportsNestedAnalysis: false, table: { catalog: null, db: 'public', name: 'orders', columns: [{ name: 'status' }], }, }); await expect( generator.generateTableDescription({ connectionId: 'conn-1', connector, context: { runId: 'run-1' }, dataSourceType: 'POSTGRESQL', table: { catalog: null, db: 'public', name: 'orders' }, }), ).resolves.toBeNull(); expect(columnResult).toEqual({ columnDescriptions: [['status', null]], processedColumns: [], skippedColumns: [], }); expect(cache.set).not.toHaveBeenCalled(); }); it('generates and caches table and data-source descriptions', async () => { const cache = createCache(); const connector = createConnector(); const generator = new KtxDescriptionGenerator({ llmRuntime: createLlmProvider('Commerce orders'), cache, settings: { columnMaxWords: 12, tableMaxWords: 18, dataSourceMaxWords: 24, concurrencyLimit: 2, }, }); await expect( generator.generateTableDescription({ connectionId: 'conn-1', connector, context: { runId: 'run-1' }, dataSourceType: 'POSTGRESQL', table: { catalog: 'warehouse', db: 'public', name: 'orders', rawDescriptions: { db: 'Raw orders' } }, }), ).resolves.toBe('Commerce orders'); await expect( generator.generateDataSourceDescription({ connectionId: 'conn-1', connector, context: { runId: 'run-1' }, dataSourceType: 'POSTGRESQL', tables: [ { catalog: 'warehouse', db: 'public', name: 'orders' }, { catalog: 'warehouse', db: 'public', name: 'customers' }, ], connectionName: 'Warehouse', }), ).resolves.toBe('Commerce orders'); expect(cache.set).toHaveBeenCalledWith('warehouse.public.orders', 'Commerce orders'); expect(cache.set).toHaveBeenCalledWith('__connection:Warehouse', 'Commerce orders'); }); it('generates one structured table description and reuses table samples for all columns', async () => { const llmRuntime = createLlmProvider('unused'); llmRuntime.generateObject = vi.fn(async () => ({ tableDescription: 'Commerce orders', columns: [ { name: 'status', description: 'Current order state' }, { name: 'amount', description: 'Order amount in dollars' }, ], })); const connector = createConnector(); const generator = new KtxDescriptionGenerator({ llmRuntime, settings: { columnMaxWords: 12, tableMaxWords: 18, dataSourceMaxWords: 24 }, }); const result = await generator.generateBatchedTableDescriptions({ connectionId: 'conn-1', connector, context: { runId: 'run-1' }, dataSourceType: 'POSTGRESQL', supportsNestedAnalysis: false, table: { catalog: null, db: 'public', name: 'orders', rawDescriptions: { db: 'Orders fact table' }, columns: [ { name: 'status', type: 'text' }, { name: 'amount', type: 'numeric' }, ], }, }); expect(result.tableDescription).toBe('Commerce orders'); expect(Object.fromEntries(result.columnDescriptions)).toEqual({ status: 'Current order state', amount: 'Order amount in dollars', }); expect(connector.sampleTable).toHaveBeenCalledTimes(1); expect(connector.sampleColumn).not.toHaveBeenCalled(); expect(llmRuntime.generateObject).toHaveBeenCalledTimes(1); expect(llmRuntime.generateText).not.toHaveBeenCalled(); }); it('falls back to one column generateText call for each missing structured column', async () => { const llmRuntime = createLlmProvider('Fallback status'); llmRuntime.generateObject = vi.fn(async () => ({ tableDescription: 'Commerce orders', columns: [{ name: 'amount', description: 'Order amount in dollars' }], })); const connector = createConnector(); const generator = new KtxDescriptionGenerator({ llmRuntime, settings: { columnMaxWords: 12, tableMaxWords: 18, dataSourceMaxWords: 24 }, }); const result = await generator.generateBatchedTableDescriptions({ connectionId: 'conn-1', connector, context: { runId: 'run-1' }, dataSourceType: 'POSTGRESQL', supportsNestedAnalysis: false, table: { catalog: null, db: 'public', name: 'orders', columns: [ { name: 'status', type: 'text' }, { name: 'amount', type: 'numeric' }, ], }, }); expect(Object.fromEntries(result.columnDescriptions)).toEqual({ status: 'Fallback status', amount: 'Order amount in dollars', }); expect(connector.sampleColumn).not.toHaveBeenCalled(); expect(llmRuntime.generateObject).toHaveBeenCalledTimes(1); expect(llmRuntime.generateText).toHaveBeenCalledTimes(1); }); it('does not run per-column fallback when structured object generation throws', async () => { const llmRuntime = createLlmProvider('Fallback description'); llmRuntime.generateObject = vi.fn(async () => { throw new Error('object output unavailable'); }); const warnings: string[] = []; const generator = new KtxDescriptionGenerator({ llmRuntime, onWarning: (warning) => warnings.push(warning.code), settings: { columnMaxWords: 12, tableMaxWords: 18, dataSourceMaxWords: 24 }, }); const result = await generator.generateBatchedTableDescriptions({ connectionId: 'conn-1', connector: createConnector(), context: { runId: 'run-1' }, dataSourceType: 'POSTGRESQL', supportsNestedAnalysis: false, table: { catalog: null, db: 'public', name: 'orders', columns: [{ name: 'status', type: 'text' }], }, }); expect(result.tableDescription).toBeNull(); expect(Object.fromEntries(result.columnDescriptions)).toEqual({ status: null }); expect(warnings).toContain('enrichment_failed'); expect(llmRuntime.generateObject).toHaveBeenCalledTimes(1); expect(llmRuntime.generateText).not.toHaveBeenCalled(); }); }); describe('KtxDescriptionGenerator resilience', () => { function createLogger() { return { debug: vi.fn(), info: vi.fn(), warn: vi.fn(), error: vi.fn(), }; } it('retries sampleTable on transient failure and uses sampled rows when it eventually succeeds', async () => { const sampleTable = vi .fn>() .mockRejectedValueOnce(new Error('pool: transient ECONNRESET')) .mockRejectedValueOnce(new Error('pool: transient ECONNRESET')) .mockResolvedValue({ headers: ['id', 'status'], rows: [ [1, 'paid'], [2, 'refunded'], ], totalRows: 2, }); const connector: KtxScanConnector = { ...createConnector(), sampleTable, }; const logger = createLogger(); const warnings: Array<{ code: string; table?: string }> = []; const generator = new KtxDescriptionGenerator({ llmRuntime: createLlmProvider('Commerce orders'), logger, onWarning: (warning) => warnings.push({ code: warning.code, ...(warning.table ? { table: warning.table } : {}) }), settings: { columnMaxWords: 12, tableMaxWords: 18, dataSourceMaxWords: 24, concurrencyLimit: 2 }, }); const description = await generator.generateTableDescription({ connectionId: 'conn-1', connector, context: { runId: 'run-1' }, dataSourceType: 'POSTGRESQL', table: { catalog: null, db: 'public', name: 'orders' }, }); expect(description).toBe('Commerce orders'); expect(sampleTable).toHaveBeenCalledTimes(3); expect(logger.warn).toHaveBeenCalledTimes(2); expect(warnings).toEqual([]); }); it('falls back to metadata-only prompt when sampleTable retries exhaust', async () => { const sampleTable = vi .fn>() .mockRejectedValue(new Error('pool: connection refused')); const connector: KtxScanConnector = { ...createConnector(), sampleTable, }; const logger = createLogger(); const warnings: Array<{ code: string; table?: string; metadata?: Record }> = []; const generator = new KtxDescriptionGenerator({ llmRuntime: createLlmProvider('Customer reference data'), logger, onWarning: (warning) => warnings.push({ code: warning.code, ...(warning.table ? { table: warning.table } : {}), ...(warning.metadata ? { metadata: warning.metadata } : {}), }), settings: { columnMaxWords: 12, tableMaxWords: 18, dataSourceMaxWords: 24, concurrencyLimit: 2 }, }); const description = await generator.generateTableDescription({ connectionId: 'conn-1', connector, context: { runId: 'run-1' }, dataSourceType: 'POSTGRESQL', table: { catalog: null, db: 'public', name: 'customers', columns: [ { name: 'id', nativeType: 'uuid' }, { name: 'email', nativeType: 'text', comment: 'Primary contact email' }, ], }, }); expect(description).toBe('Customer reference data'); expect(sampleTable).toHaveBeenCalledTimes(3); expect(warnings.map((warning) => warning.code)).toEqual(['sampling_failed', 'description_fallback_used']); expect(warnings[1]?.metadata?.reason).toBe('sampling_failed'); const userPrompt = (vi.mocked(generateText).mock.calls.at(-1)?.[0] as { messages: Array<{ role: string; content: string }> }) .messages.find((message) => message.role === 'user')?.content; expect(userPrompt).toContain('Columns (metadata only, no sample rows)'); expect(userPrompt).toContain('email (text)'); expect(userPrompt).toContain('Primary contact email'); }); it('emits enrichment_failed and returns null when both sampling and metadata-only LLM fail', async () => { const sampleTable = vi .fn>() .mockRejectedValue(new Error('pool: connection refused')); const connector: KtxScanConnector = { ...createConnector(), sampleTable, }; const warnings: string[] = []; const generator = new KtxDescriptionGenerator({ llmRuntime: createFailingLlmProvider(), onWarning: (warning) => warnings.push(warning.code), settings: { columnMaxWords: 12, tableMaxWords: 18, dataSourceMaxWords: 24 }, }); const description = await generator.generateTableDescription({ connectionId: 'conn-1', connector, context: { runId: 'run-1' }, dataSourceType: 'POSTGRESQL', table: { catalog: null, db: 'public', name: 'orphan', columns: [{ name: 'id' }] }, }); expect(description).toBeNull(); expect(warnings).toEqual(['sampling_failed', 'enrichment_failed']); }); it('uses metadata-only fallback when connector has no sampleTable', async () => { const connector = createConnector(); const samplerWithoutTable: KtxScanConnector = { ...connector, sampleTable: undefined, }; const warnings: string[] = []; const generator = new KtxDescriptionGenerator({ llmRuntime: createLlmProvider('Orders mart'), onWarning: (warning) => warnings.push(warning.code), settings: { columnMaxWords: 12, tableMaxWords: 18, dataSourceMaxWords: 24 }, }); const description = await generator.generateTableDescription({ connectionId: 'conn-1', connector: samplerWithoutTable, context: { runId: 'run-1' }, dataSourceType: 'POSTGRESQL', table: { catalog: null, db: 'public', name: 'mart_orders', columns: [{ name: 'order_id', nativeType: 'uuid' }], }, }); expect(description).toBe('Orders mart'); expect(warnings).toEqual(['connector_capability_missing', 'description_fallback_used']); }); it('aborts retry loop when the scan context signal fires', async () => { const controller = new AbortController(); const sampleTable = vi.fn>().mockImplementation(async () => { controller.abort(); throw new Error('first attempt blew up'); }); const connector: KtxScanConnector = { ...createConnector(), sampleTable, }; const warnings: string[] = []; const generator = new KtxDescriptionGenerator({ llmRuntime: createLlmProvider('should not be called'), onWarning: (warning) => warnings.push(warning.code), settings: { columnMaxWords: 12, tableMaxWords: 18, dataSourceMaxWords: 24 }, }); await expect( generator.generateTableDescription({ connectionId: 'conn-1', connector, context: { runId: 'run-1', signal: controller.signal }, dataSourceType: 'POSTGRESQL', table: { catalog: null, db: 'public', name: 'orders' }, }), ).rejects.toThrow('aborted'); expect(sampleTable).toHaveBeenCalledTimes(1); expect(warnings).toEqual([]); }); it('generates column descriptions from rawDescriptions when sampleColumn is unavailable', async () => { const samplerWithoutColumn: KtxScanConnector = { ...createConnector(), sampleColumn: undefined, }; const logger = createLogger(); const generator = new KtxDescriptionGenerator({ llmRuntime: createLlmProvider('Payment lifecycle state'), logger, settings: { columnMaxWords: 12, tableMaxWords: 18, dataSourceMaxWords: 24 }, }); const result = await generator.generateColumnDescriptions({ connectionId: 'conn-1', connector: samplerWithoutColumn, context: { runId: 'run-1' }, dataSourceType: 'POSTGRESQL', supportsNestedAnalysis: false, table: { catalog: null, db: 'public', name: 'orders', columns: [{ name: 'status', rawDescriptions: { db: 'order lifecycle state' } }], }, }); expect(result.columnDescriptions).toEqual([['status', 'Payment lifecycle state']]); expect(logger.warn).toHaveBeenCalled(); const userPrompt = ( vi.mocked(generateText).mock.calls.at(-1)?.[0] as { messages: Array<{ role: string; content: string }> } ).messages.find((message) => message.role === 'user')?.content; expect(userPrompt).toContain(' unavailable '); expect(userPrompt).toContain(' order lifecycle state '); }); it('generates column descriptions from rawDescriptions when sampleColumn retries exhaust', async () => { const sampleColumn = vi .fn>() .mockRejectedValue(new Error('pool: connection refused')); const flakyConnector: KtxScanConnector = { ...createConnector(), sampleColumn, }; const generator = new KtxDescriptionGenerator({ llmRuntime: createLlmProvider('Customer reference identifier'), settings: { columnMaxWords: 12, tableMaxWords: 18, dataSourceMaxWords: 24 }, }); const result = await generator.generateColumnDescriptions({ connectionId: 'conn-1', connector: flakyConnector, context: { runId: 'run-1' }, dataSourceType: 'POSTGRESQL', supportsNestedAnalysis: false, table: { catalog: null, db: 'public', name: 'orders', columns: [{ name: 'customer_id', rawDescriptions: { db: 'FK to customers.id' } }], }, }); expect(sampleColumn).toHaveBeenCalledTimes(3); expect(result.columnDescriptions).toEqual([['customer_id', 'Customer reference identifier']]); }); it('skips column LLM call only when neither samples nor rawDescriptions are available', async () => { const sampleColumn = vi .fn>() .mockResolvedValue({ values: [null, null], nullCount: 2, distinctCount: 0 }); const connector: KtxScanConnector = { ...createConnector(), sampleColumn, }; vi.mocked(generateText).mockClear(); const generator = new KtxDescriptionGenerator({ llmRuntime: createLlmProvider('should not be called'), settings: { columnMaxWords: 12, tableMaxWords: 18, dataSourceMaxWords: 24 }, }); const result = await generator.generateColumnDescriptions({ connectionId: 'conn-1', connector, context: { runId: 'run-1' }, dataSourceType: 'POSTGRESQL', supportsNestedAnalysis: false, table: { catalog: null, db: 'public', name: 'orders', columns: [{ name: 'opaque_blob' }], }, }); expect(result.columnDescriptions).toEqual([['opaque_blob', null]]); expect(generateText).not.toHaveBeenCalled(); }); });