@kaelio/ktx 0.4.1 → 0.6.0

package/dist/context/scan/local-enrichment.js

@@ -0,0 +1,479 @@
+import pLimit from 'p-limit';
+import { getDialectForDriver } from '../connections/dialects.js';
+import { buildDefaultKtxProjectConfig } from '../project/config.js';
+import { KtxDescriptionGenerator } from './description-generation.js';
+import { buildKtxColumnEmbeddingText } from './embedding-text.js';
+import { completedKtxScanEnrichmentStateSummary, computeKtxScanEnrichmentInputHash, summarizeKtxScanEnrichmentState, } from './enrichment-state.js';
+import { skippedKtxScanEnrichmentSummary } from './enrichment-summary.js';
+import { discoverKtxRelationships } from './relationship-discovery.js';
+const DESCRIPTION_TABLE_CONCURRENCY = 4;
+function tableId(table) {
+    return [table.catalog, table.db, table.name].filter((value) => Boolean(value)).join('.');
+}
+function columnId(table, column) {
+    return `${tableId(table)}.${column.name}`;
+}
+function tableRef(table) {
+    return {
+        catalog: table.catalog,
+        db: table.db,
+        name: table.name,
+    };
+}
+function endpoint(table, column) {
+    return {
+        tableId: table.id,
+        columnIds: [column.id],
+        table: table.ref,
+        columns: [column.name],
+    };
+}
+function relationshipId(from, to) {
+    return `${from.tableId}:(${from.columnIds.join(',')})->${to.tableId}:(${to.columnIds.join(',')})`;
+}
+function targetMatchesForeignKey(table, foreignKey) {
+    return (table.ref.name === foreignKey.toTable &&
+        (foreignKey.toCatalog === null || table.ref.catalog === foreignKey.toCatalog) &&
+        (foreignKey.toDb === null || table.ref.db === foreignKey.toDb));
+}
+function assertConnectorDriverMatchesSnapshot(input) {
+    if (input.connector.driver !== input.snapshot.driver) {
+        throw new Error(`ktx scan connector driver "${input.connector.driver}" does not match snapshot driver "${input.snapshot.driver}" for connection "${input.connectionId}"`);
+    }
+}
+function formalRelationshipsFromSnapshot(snapshot, tables) {
+    const tableById = new Map(tables.map((table) => [table.id, table]));
+    const relationships = [];
+    for (const sourceTableSnapshot of snapshot.tables) {
+        const sourceTable = tableById.get(tableId(sourceTableSnapshot));
+        if (!sourceTable) {
+            continue;
+        }
+        for (const foreignKey of sourceTableSnapshot.foreignKeys) {
+            const sourceColumn = sourceTable.columns.find((column) => column.name === foreignKey.fromColumn);
+            const targetTable = tables.find((table) => targetMatchesForeignKey(table, foreignKey));
+            const targetColumn = targetTable?.columns.find((column) => column.name === foreignKey.toColumn);
+            if (!sourceColumn || !targetTable || !targetColumn) {
+                continue;
+            }
+            const from = endpoint(sourceTable, sourceColumn);
+            const to = endpoint(targetTable, targetColumn);
+            relationships.push({
+                id: relationshipId(from, to),
+                source: 'formal',
+                from,
+                to,
+                relationshipType: 'many_to_one',
+                confidence: 1,
+                isPrimaryKeyReference: true,
+            });
+        }
+    }
+    return relationships.sort((left, right) => left.id.localeCompare(right.id));
+}
+function providerlessEnrichedWarning(relationshipDetection) {
+    return {
+        code: 'scan_enrichment_backend_not_configured',
+        message: 'Skipping description and embedding enrichment because scan.enrichment.mode is not configured; relationship discovery still ran.',
+        recoverable: true,
+        metadata: {
+            skippedStages: ['descriptions', 'embeddings'],
+            relationshipDetection,
+        },
+    };
+}
+export function createDeterministicLocalScanEnrichmentProviders() {
+    return {
+        llmRuntime: deterministicLlmRuntime(),
+    };
+}
+function deterministicLlmRuntime() {
+    return {
+        async generateText(input) {
+            return `Deterministic description for ${input.prompt.slice(0, 64).trim() || 'data source'}`;
+        },
+        async generateObject(input) {
+            if (input.prompt.includes('Sample rows:')) {
+                const columns = Array.from(input.prompt.matchAll(/^- ([^\s(]+)/gm), (match) => ({
+                    name: match[1] ?? 'column',
+                    description: `Deterministic description for ${match[1] ?? 'column'}`,
+                }));
+                return {
+                    tableDescription: `Deterministic description for ${input.prompt.slice(0, 64).trim() || 'table'}`,
+                    columns,
+                };
+            }
+            return { pkCandidates: [], fkCandidates: [] };
+        },
+        async runAgentLoop() {
+            return { stopReason: 'natural' };
+        },
+    };
+}
+export function snapshotToKtxEnrichedSchema(snapshot, embeddingsByColumnId = new Map()) {
+    const tables = snapshot.tables.map((table) => {
+        const id = tableId(table);
+        const ref = tableRef(table);
+        const columns = table.columns.map((column) => {
+            const idForColumn = columnId(table, column);
+            return {
+                id: idForColumn,
+                tableId: id,
+                tableRef: ref,
+                name: column.name,
+                nativeType: column.nativeType,
+                normalizedType: column.normalizedType,
+                dimensionType: column.dimensionType,
+                nullable: column.nullable,
+                primaryKey: column.primaryKey,
+                parentColumnId: null,
+                descriptions: {
+                    ...(column.comment ? { db: column.comment } : {}),
+                },
+                embedding: embeddingsByColumnId.get(idForColumn) ?? null,
+                sampleValues: null,
+                cardinality: null,
+            };
+        });
+        return {
+            id,
+            ref,
+            enabled: true,
+            descriptions: {
+                ...(table.comment ? { db: table.comment } : {}),
+            },
+            columns,
+        };
+    });
+    return {
+        connectionId: snapshot.connectionId,
+        tables,
+        relationships: formalRelationshipsFromSnapshot(snapshot, tables),
+    };
+}
+function embeddingBatchSize(maxBatchSize) {
+    return Number.isInteger(maxBatchSize) && maxBatchSize > 0 ? maxBatchSize : 100;
+}
+async function generateDescriptions(input) {
+    const warningSink = input.warnings;
+    const generator = new KtxDescriptionGenerator({
+        llmRuntime: input.providers.llmRuntime,
+        ...(input.context.logger ? { logger: input.context.logger } : {}),
+        ...(warningSink
+            ? {
+                onWarning: (warning) => {
+                    warningSink.push(warning);
+                },
+            }
+            : {}),
+        settings: {
+            columnMaxWords: 16,
+            tableMaxWords: 24,
+            dataSourceMaxWords: 32,
+            concurrencyLimit: 4,
+        },
+    });
+    const updates = [];
+    const totalTables = input.snapshot.tables.length;
+    if (totalTables === 0) {
+        await input.progress?.update(1, 'No tables to describe');
+        return updates;
+    }
+    const limitTable = pLimit(DESCRIPTION_TABLE_CONCURRENCY);
+    const tableUpdates = await Promise.all(input.snapshot.tables.map((table, index) => limitTable(async () => {
+        await input.progress?.update((index + 1) / totalTables, `Generating descriptions ${index + 1}/${totalTables} tables`, {
+            transient: true,
+        });
+        const batched = await generator.generateBatchedTableDescriptions({
+            connectionId: input.snapshot.connectionId,
+            connector: input.connector,
+            context: input.context,
+            dataSourceType: input.snapshot.driver,
+            supportsNestedAnalysis: input.connector.capabilities.nestedAnalysis,
+            table: {
+                catalog: table.catalog,
+                db: table.db,
+                name: table.name,
+                rawDescriptions: table.comment ? { db: table.comment } : {},
+                columns: table.columns.map((column) => ({
+                    name: column.name,
+                    type: column.nativeType,
+                    ...(column.comment ? { rawDescriptions: { db: column.comment } } : {}),
+                })),
+            },
+        });
+        return {
+            table: tableRef(table),
+            tableDescription: batched.tableDescription,
+            columnDescriptions: Object.fromEntries(batched.columnDescriptions),
+        };
+    })));
+    updates.push(...tableUpdates);
+    await input.progress?.update(1, `Generated descriptions for ${totalTables} tables`);
+    return updates;
+}
+async function buildEmbeddings(input) {
+    const descriptionByTable = new Map(input.descriptions.map((item) => [item.table.name, item]));
+    const texts = [];
+    for (const table of input.snapshot.tables) {
+        const tableDescriptions = descriptionByTable.get(table.name);
+        for (const column of table.columns) {
+            const id = columnId(table, column);
+            const text = buildKtxColumnEmbeddingText({
+                tableName: table.name,
+                columnName: column.name,
+                columnType: column.nativeType,
+                resolvedDescription: tableDescriptions?.columnDescriptions[column.name] ?? column.comment,
+                resolvedTableDescription: tableDescriptions?.tableDescription ?? table.comment,
+                sampleValues: column.comment ? [column.comment] : null,
+                foreignKeys: {
+                    outgoing: (table.foreignKeys ?? [])
+                        .filter((foreignKey) => foreignKey.fromColumn === column.name)
+                        .map((foreignKey) => ({ toTable: foreignKey.toTable, toColumn: foreignKey.toColumn })),
+                    incoming: [],
+                },
+            });
+            texts.push({ columnId: id, text });
+        }
+    }
+    const embeddings = [];
+    const maxBatchSize = embeddingBatchSize(input.embedding.maxBatchSize);
+    const embeddingTexts = texts.map((item) => item.text);
+    const batchCount = Math.ceil(embeddingTexts.length / maxBatchSize);
+    if (batchCount === 0) {
+        await input.progress?.update(1, 'No embeddings to build');
+    }
+    for (let offset = 0; offset < embeddingTexts.length; offset += maxBatchSize) {
+        const batchIndex = Math.floor(offset / maxBatchSize) + 1;
+        await input.progress?.update(batchIndex / batchCount, `Building embeddings ${batchIndex}/${batchCount} batches`, {
+            transient: true,
+        });
+        const batch = embeddingTexts.slice(offset, offset + maxBatchSize);
+        const batchEmbeddings = await input.embedding.embedBatch(batch);
+        if (batchEmbeddings.length !== batch.length) {
+            throw new Error(`expected ${batch.length} embeddings, received ${batchEmbeddings.length}`);
+        }
+        embeddings.push(...batchEmbeddings);
+    }
+    const byColumnId = new Map();
+    const updates = texts.map((item, index) => {
+        const embedding = embeddings[index] ?? [];
+        byColumnId.set(item.columnId, embedding);
+        return {
+            columnId: item.columnId,
+            text: item.text,
+            embedding,
+        };
+    });
+    if (batchCount > 0) {
+        await input.progress?.update(1, `Built embeddings for ${updates.length} columns`);
+    }
+    return { updates, byColumnId };
+}
+async function runEnrichmentStage(input) {
+    const existing = await input.stateStore?.findCompletedStage({
+        runId: input.runId,
+        stage: input.stage,
+        inputHash: input.inputHash,
+    });
+    if (existing) {
+        input.resumedStages.push(input.stage);
+        input.completedStages.push(input.stage);
+        return existing.output;
+    }
+    try {
+        const output = await input.compute();
+        input.completedStages.push(input.stage);
+        await input.stateStore?.saveCompletedStage({
+            runId: input.runId,
+            connectionId: input.connectionId,
+            syncId: input.syncId,
+            mode: input.mode,
+            stage: input.stage,
+            inputHash: input.inputHash,
+            output,
+            updatedAt: input.now().toISOString(),
+        });
+        return output;
+    }
+    catch (error) {
+        input.failedStages.push(input.stage);
+        await input.stateStore?.saveFailedStage({
+            runId: input.runId,
+            connectionId: input.connectionId,
+            syncId: input.syncId,
+            mode: input.mode,
+            stage: input.stage,
+            inputHash: input.inputHash,
+            errorMessage: error instanceof Error ? error.message : String(error),
+            updatedAt: input.now().toISOString(),
+        });
+        throw error;
+    }
+}
+function embeddingsByColumnId(updates) {
+    return new Map(updates.map((update) => [update.columnId, update.embedding]));
+}
+export async function runLocalScanEnrichment(input) {
+    const progress = input.context.progress;
+    await progress?.update(0, 'Loading enrichment schema snapshot');
+    const snapshot = input.snapshot ??
+        (await input.connector.introspect({
+            connectionId: input.connectionId,
+            driver: input.connector.driver,
+            mode: input.mode,
+            detectRelationships: input.detectRelationships,
+        }, input.context));
+    await progress?.update(0.05, `Loaded schema snapshot with ${snapshot.tables.length} tables`);
+    assertConnectorDriverMatchesSnapshot({
+        connector: input.connector,
+        snapshot,
+        connectionId: input.connectionId,
+    });
+    const dialect = getDialectForDriver(snapshot.driver);
+    const now = input.now ?? (() => new Date());
+    const state = completedKtxScanEnrichmentStateSummary();
+    const syncId = input.syncId ?? input.context.runId;
+    const relationshipSettings = input.relationshipSettings ?? buildDefaultKtxProjectConfig().scan.relationships;
+    const inputHash = computeKtxScanEnrichmentInputHash({
+        snapshot,
+        mode: input.mode,
+        detectRelationships: input.detectRelationships ?? false,
+        providerIdentity: input.providerIdentity ?? {},
+        relationshipSettings,
+    });
+    const warnings = [];
+    let descriptions = [];
+    let embeddingUpdates = [];
+    let schema = snapshotToKtxEnrichedSchema(snapshot);
+    const summary = { ...skippedKtxScanEnrichmentSummary };
+    const relationshipDetectionEnabled = relationshipSettings.enabled;
+    const shouldDetectRelationships = relationshipDetectionEnabled &&
+        (input.mode === 'relationships' || input.mode === 'enriched' || (input.detectRelationships ?? false));
+    if (input.mode === 'enriched' && !input.providers) {
+        warnings.push(providerlessEnrichedWarning(shouldDetectRelationships));
+    }
+    if (input.mode === 'enriched' && input.providers) {
+        const providers = input.providers;
+        const descriptionProgress = progress?.startPhase(0.45);
+        descriptions = await runEnrichmentStage({
+            stateStore: input.stateStore,
+            runId: input.context.runId,
+            connectionId: input.connectionId,
+            syncId,
+            mode: input.mode,
+            stage: 'descriptions',
+            inputHash,
+            now,
+            resumedStages: state.resumedStages,
+            completedStages: state.completedStages,
+            failedStages: state.failedStages,
+            compute: () => generateDescriptions({
+                snapshot,
+                connector: input.connector,
+                context: input.context,
+                providers,
+                progress: descriptionProgress,
+                warnings,
+            }),
+        });
+        summary.dataDictionary = input.connector.sampleColumn ? 'completed' : 'skipped';
+        summary.tableDescriptions = 'completed';
+        summary.columnDescriptions = 'completed';
+        const embeddingProgress = progress?.startPhase(0.2);
+        const embedding = providers.embedding;
+        if (embedding) {
+            embeddingUpdates = await runEnrichmentStage({
+                stateStore: input.stateStore,
+                runId: input.context.runId,
+                connectionId: input.connectionId,
+                syncId,
+                mode: input.mode,
+                stage: 'embeddings',
+                inputHash,
+                now,
+                resumedStages: state.resumedStages,
+                completedStages: state.completedStages,
+                failedStages: state.failedStages,
+                compute: async () => {
+                    const embeddings = await buildEmbeddings({
+                        snapshot,
+                        embedding,
+                        descriptions,
+                        progress: embeddingProgress,
+                    });
+                    return embeddings.updates;
+                },
+            });
+            schema = snapshotToKtxEnrichedSchema(snapshot, embeddingsByColumnId(embeddingUpdates));
+            summary.embeddings = 'completed';
+        }
+    }
+    let relationshipUpdate = null;
+    let relationshipProfile = null;
+    let resolvedRelationships = null;
+    let compositeRelationships = null;
+    let relationships = { accepted: 0, review: 0, rejected: 0, skipped: 0 };
+    if (shouldDetectRelationships) {
+        const relationshipProgress = progress?.startPhase(0.25);
+        const relationshipStage = await runEnrichmentStage({
+            stateStore: input.stateStore,
+            runId: input.context.runId,
+            connectionId: input.connectionId,
+            syncId,
+            mode: input.mode,
+            stage: 'relationships',
+            inputHash,
+            now,
+            resumedStages: state.resumedStages,
+            completedStages: state.completedStages,
+            failedStages: state.failedStages,
+            compute: async () => {
+                await relationshipProgress?.update(0, 'Detecting relationships');
+                const detection = await discoverKtxRelationships({
+                    connectionId: input.connectionId,
+                    dialect,
+                    connector: input.connector,
+                    schema,
+                    context: input.context,
+                    settings: relationshipSettings,
+                    llmRuntime: input.providers?.llmRuntime ?? null,
+                });
+                await relationshipProgress?.update(1, `Relationship detection found ${detection.relationships.accepted} accepted, ${detection.relationships.review} review`);
+                return {
+                    relationshipUpdate: detection.relationshipUpdate,
+                    relationshipProfile: detection.profile,
+                    resolvedRelationships: detection.resolvedRelationships,
+                    compositeRelationships: detection.compositeRelationships,
+                    relationships: detection.relationships,
+                    statisticalValidation: detection.statisticalValidation,
+                    llmRelationshipValidation: detection.llmRelationshipValidation,
+                    warnings: detection.warnings,
+                };
+            },
+        });
+        summary.deterministicRelationships = 'completed';
+        summary.llmRelationshipValidation = relationshipStage.llmRelationshipValidation;
+        summary.statisticalValidation = relationshipStage.statisticalValidation;
+        relationshipUpdate = relationshipStage.relationshipUpdate;
+        relationshipProfile = relationshipStage.relationshipProfile;
+        resolvedRelationships = relationshipStage.resolvedRelationships;
+        compositeRelationships = relationshipStage.compositeRelationships;
+        relationships = relationshipStage.relationships;
+        warnings.push(...relationshipStage.warnings);
+    }
+    await progress?.update(1, 'Enrichment complete');
+    return {
+        snapshot,
+        summary,
+        relationships,
+        state: summarizeKtxScanEnrichmentState(state),
+        warnings,
+        descriptionUpdates: descriptions,
+        embeddingUpdates,
+        relationshipUpdate,
+        relationshipProfile,
+        resolvedRelationships,
+        compositeRelationships,
+    };
+}

package/dist/context/scan/local-scan.d.ts

@@ -0,0 +1,62 @@
+import type { KtxEmbeddingProvider } from '../../llm/types.js';
+import { type LocalIngestRunRecord } from '../../context/ingest/local-stage-ingest.js';
+import type { SourceAdapter } from '../../context/ingest/types.js';
+import type { KtxLocalProject } from '../../context/project/project.js';
+import { type KtxLocalScanEnrichmentProviders } from './local-enrichment.js';
+import { SqliteLocalScanEnrichmentStateStore } from './sqlite-local-enrichment-state-store.js';
+import type { KtxProgressPort, KtxScanConnector, KtxScanMode, KtxScanReport, KtxScanTrigger } from './types.js';
+export interface RunLocalScanOptions {
+    project: KtxLocalProject;
+    connectionId: string;
+    mode?: KtxScanMode;
+    detectRelationships?: boolean;
+    dryRun?: boolean;
+    trigger?: KtxScanTrigger;
+    databaseIntrospectionUrl?: string;
+    adapters?: SourceAdapter[];
+    jobId?: string;
+    now?: () => Date;
+    connector?: KtxScanConnector;
+    createConnector?: (connectionId: string) => KtxScanConnector | Promise<KtxScanConnector>;
+    enrichmentProviders?: KtxLocalScanEnrichmentProviders | null;
+    enrichmentStateStore?: SqliteLocalScanEnrichmentStateStore | null;
+    progress?: KtxProgressPort;
+    embeddingProvider?: KtxEmbeddingProvider | null;
+}
+export interface LocalScanRunResult {
+    runId: string;
+    status: 'done';
+    done: true;
+    connectionId: string;
+    mode: KtxScanMode;
+    dryRun: boolean;
+    syncId: string;
+    report: KtxScanReport;
+}
+/** @internal */
+export interface LocalScanStatusResponse {
+    runId: string;
+    status: LocalIngestRunRecord['status'];
+    done: boolean;
+    connectionId: string;
+    mode: KtxScanMode;
+    dryRun: boolean;
+    syncId: string;
+    progress: number;
+    startedAt: string;
+    completedAt: string;
+    reportPath: string | null;
+    warnings: KtxScanReport['warnings'];
+}
+export interface LocalScanMcpOptions {
+    adapters?: SourceAdapter[];
+    databaseIntrospectionUrl?: string;
+    jobIdFactory?: () => string;
+    now?: () => Date;
+    createConnector?: (connectionId: string) => KtxScanConnector | Promise<KtxScanConnector>;
+}
+export declare function runLocalScan(options: RunLocalScanOptions): Promise<LocalScanRunResult>;
+/** @internal */
+export declare function getLocalScanReport(project: KtxLocalProject, runId: string): Promise<KtxScanReport | null>;
+/** @internal */
+export declare function getLocalScanStatus(project: KtxLocalProject, runId: string): Promise<LocalScanStatusResponse | null>;