feat: add AI catalog description generation
This commit is contained in:
@@ -0,0 +1,456 @@
|
||||
import { expect, test, vi } from "vitest";
|
||||
import { DescriptionGenerationWorker } from "../src/catalog/description-generation-worker.js";
|
||||
import type { DescriptionSourceSampler } from "../src/catalog/description-source-sampler.js";
|
||||
import { MemoryCatalogRepository } from "../src/catalog/memory-repository.js";
|
||||
import type { MetadataGenerationModels } from "../src/catalog/metadata-generation-models.js";
|
||||
import { ModelCompletionCancelledError } from "../src/catalog/model-completer.js";
|
||||
import type { ModelCompleter, ModelCompletionRequest } from "../src/catalog/model-completer.js";
|
||||
import { CatalogOperationCoordinator } from "../src/catalog/operation-coordinator.js";
|
||||
import type { WorkspaceRegistry } from "../src/workspaces/registry.js";
|
||||
|
||||
test("serializes Unlock with Start so stale recovery cannot release a new reservation", async () => {
|
||||
let lookupStarted!: () => void;
|
||||
const started = new Promise<void>((resolve) => { lookupStarted = resolve; });
|
||||
let releaseLookup!: () => void;
|
||||
const gate = new Promise<void>((resolve) => { releaseLookup = resolve; });
|
||||
const repository = {
|
||||
getActiveDescriptionGenerationRun: vi.fn(async () => {
|
||||
lookupStarted();
|
||||
await gate;
|
||||
return undefined;
|
||||
}),
|
||||
} as unknown as MemoryCatalogRepository;
|
||||
const resolveModel = vi.fn();
|
||||
const worker = new DescriptionGenerationWorker(
|
||||
repository,
|
||||
{} as WorkspaceRegistry,
|
||||
{
|
||||
catalog: () => ({ models: [], default: "" }),
|
||||
resolve: resolveModel,
|
||||
} as MetadataGenerationModels,
|
||||
{} as ModelCompleter,
|
||||
new CatalogOperationCoordinator(),
|
||||
{ sample: vi.fn(async () => []) },
|
||||
);
|
||||
|
||||
const unlocking = worker.unlock();
|
||||
await started;
|
||||
await expect(worker.start(
|
||||
"11111111-1111-4111-8111-111111111111",
|
||||
"openai-mini",
|
||||
"missing",
|
||||
[],
|
||||
)).rejects.toThrow("already active");
|
||||
expect(resolveModel).not.toHaveBeenCalled();
|
||||
|
||||
releaseLookup();
|
||||
await expect(unlocking).resolves.toBeUndefined();
|
||||
});
|
||||
|
||||
test("exposes an awaitable background job and absorbs provider promise rejection", async () => {
|
||||
const repository = new MemoryCatalogRepository();
|
||||
const database = await repository.create({
|
||||
workspaceId: "psd-clinical",
|
||||
engine: "postgres",
|
||||
databaseName: "warehouse",
|
||||
schema: "datawarehouse",
|
||||
binding: { transport: "postgres_direct", host: "db.internal", port: 5432, username: "reader" },
|
||||
});
|
||||
await repository.applySchemaSync(database.id, database.version, "all", [], {
|
||||
schemaVersion: 1,
|
||||
capabilities: { tables: "available", columns: "available", relationships: "available" },
|
||||
tables: [{ name: "patients", sourceComment: null }],
|
||||
columns: [{
|
||||
tableName: "patients",
|
||||
name: "birth_date",
|
||||
ordinalPosition: 1,
|
||||
dataType: "date",
|
||||
isNullable: true,
|
||||
defaultExpression: null,
|
||||
primaryKeyPosition: null,
|
||||
sourceComment: null,
|
||||
}],
|
||||
relationships: [],
|
||||
});
|
||||
const table = (await repository.listTables(database.id))[0]!;
|
||||
const column = (await repository.listColumns(database.id, table.id))[0]!;
|
||||
let rejectCompletion!: (error: Error) => void;
|
||||
const pendingCompletion = new Promise<string>((_resolve, reject) => { rejectCompletion = reject; });
|
||||
const completer: ModelCompleter = {
|
||||
complete: vi.fn(async () => await pendingCompletion),
|
||||
};
|
||||
const models: MetadataGenerationModels = {
|
||||
catalog: () => ({ models: [{ id: "openai-mini", label: "OpenAI Mini" }], default: "openai-mini" }),
|
||||
resolve: () => ({
|
||||
id: "openai-mini",
|
||||
provider: "openai",
|
||||
model: "gpt-4.1-mini",
|
||||
apiKeyEnv: "OPENAI_API_KEY",
|
||||
apiKey: "test-provider-secret",
|
||||
}),
|
||||
};
|
||||
const operations = new CatalogOperationCoordinator();
|
||||
const sourceSampler: DescriptionSourceSampler = {
|
||||
sample: vi.fn(async () => []),
|
||||
};
|
||||
const worker = new DescriptionGenerationWorker(
|
||||
repository,
|
||||
{
|
||||
read: vi.fn(async () => ({
|
||||
workspace: { workspace: { language: "it" } },
|
||||
revision: {},
|
||||
})),
|
||||
} as unknown as WorkspaceRegistry,
|
||||
models,
|
||||
completer,
|
||||
operations,
|
||||
sourceSampler,
|
||||
);
|
||||
|
||||
const run = await worker.start(database.id, "openai-mini", "selected_columns", [column.id]);
|
||||
let settled = false;
|
||||
const waiting = worker.waitForRun(run.id).then(() => { settled = true; });
|
||||
await new Promise((resolve) => setTimeout(resolve, 0));
|
||||
expect(settled).toBe(false);
|
||||
|
||||
rejectCompletion(new Error("test-provider-secret private prompt raw response"));
|
||||
await expect(waiting).resolves.toBeUndefined();
|
||||
expect(await repository.getDescriptionGenerationRun(run.id)).toMatchObject({
|
||||
status: "completed_with_errors",
|
||||
failed: 1,
|
||||
errorSummary: "Description generation completed with errors.",
|
||||
});
|
||||
const events = await repository.listDescriptionGenerationEvents(run.id);
|
||||
expect(JSON.stringify(events)).not.toMatch(/test-provider-secret|private prompt|raw response/);
|
||||
|
||||
const release = operations.reserve(database.id);
|
||||
release();
|
||||
|
||||
await expect(worker.start(
|
||||
database.id,
|
||||
"openai-mini",
|
||||
"selected_columns",
|
||||
[],
|
||||
)).rejects.toThrow("at least one target ID is required");
|
||||
});
|
||||
|
||||
test("marks an active run interrupted when the backend worker stops", async () => {
|
||||
const repository = new MemoryCatalogRepository();
|
||||
const database = await repository.create({
|
||||
workspaceId: "psd-clinical",
|
||||
engine: "postgres",
|
||||
databaseName: "warehouse",
|
||||
schema: "datawarehouse",
|
||||
binding: { transport: "postgres_direct", host: "db.internal", port: 5432, username: "reader" },
|
||||
});
|
||||
await repository.applySchemaSync(database.id, database.version, "all", [], {
|
||||
schemaVersion: 1,
|
||||
capabilities: { tables: "available", columns: "available", relationships: "available" },
|
||||
tables: [{ name: "patients", sourceComment: null }],
|
||||
columns: [{
|
||||
tableName: "patients",
|
||||
name: "status",
|
||||
ordinalPosition: 1,
|
||||
dataType: "text",
|
||||
isNullable: true,
|
||||
defaultExpression: null,
|
||||
primaryKeyPosition: null,
|
||||
sourceComment: null,
|
||||
}],
|
||||
relationships: [],
|
||||
});
|
||||
const table = (await repository.listTables(database.id))[0]!;
|
||||
const column = (await repository.listColumns(database.id, table.id))[0]!;
|
||||
const completer: ModelCompleter = {
|
||||
complete: vi.fn(async (request) => await new Promise<string>((_resolve, reject) => {
|
||||
const cancel = () => reject(new ModelCompletionCancelledError());
|
||||
if (request.signal.aborted) cancel();
|
||||
else request.signal.addEventListener("abort", cancel, { once: true });
|
||||
})),
|
||||
};
|
||||
const worker = new DescriptionGenerationWorker(
|
||||
repository,
|
||||
{
|
||||
read: vi.fn(async () => ({
|
||||
workspace: { workspace: { language: "it" } },
|
||||
revision: {},
|
||||
})),
|
||||
} as unknown as WorkspaceRegistry,
|
||||
{
|
||||
catalog: () => ({ models: [{ id: "openai-mini", label: "OpenAI Mini" }], default: "openai-mini" }),
|
||||
resolve: () => ({
|
||||
id: "openai-mini",
|
||||
provider: "openai",
|
||||
model: "gpt-4.1-mini",
|
||||
apiKeyEnv: "OPENAI_API_KEY",
|
||||
apiKey: "test-provider-secret",
|
||||
}),
|
||||
},
|
||||
completer,
|
||||
new CatalogOperationCoordinator(),
|
||||
{ sample: vi.fn(async () => []) },
|
||||
);
|
||||
|
||||
const run = await worker.start(database.id, "openai-mini", "selected_columns", [column.id]);
|
||||
await vi.waitFor(() => expect(completer.complete).toHaveBeenCalledOnce());
|
||||
await worker.stop();
|
||||
|
||||
expect(await repository.getDescriptionGenerationRun(run.id)).toMatchObject({
|
||||
status: "interrupted",
|
||||
errorSummary: "Description generation was interrupted by backend shutdown.",
|
||||
});
|
||||
expect(await repository.listDescriptionGenerationEvents(run.id)).toContainEqual(
|
||||
expect.objectContaining({
|
||||
level: "warning",
|
||||
message: "Description generation was interrupted by backend shutdown.",
|
||||
}),
|
||||
);
|
||||
});
|
||||
|
||||
test("adds only bounded transient source samples to the model request", async () => {
|
||||
const repository = new MemoryCatalogRepository();
|
||||
const database = await repository.create({
|
||||
workspaceId: "psd-clinical",
|
||||
engine: "postgres",
|
||||
databaseName: "warehouse",
|
||||
schema: "datawarehouse",
|
||||
binding: { transport: "postgres_direct", host: "db.internal", port: 5432, username: "reader" },
|
||||
});
|
||||
await repository.applySchemaSync(database.id, database.version, "all", [], {
|
||||
schemaVersion: 1,
|
||||
capabilities: { tables: "available", columns: "available", relationships: "available" },
|
||||
tables: [{ name: "patients", sourceComment: null }],
|
||||
columns: [{
|
||||
tableName: "patients",
|
||||
name: "status",
|
||||
ordinalPosition: 1,
|
||||
dataType: "text",
|
||||
isNullable: true,
|
||||
defaultExpression: null,
|
||||
primaryKeyPosition: null,
|
||||
sourceComment: null,
|
||||
}, {
|
||||
tableName: "patients",
|
||||
name: "ward",
|
||||
ordinalPosition: 2,
|
||||
dataType: "text",
|
||||
isNullable: true,
|
||||
defaultExpression: null,
|
||||
primaryKeyPosition: null,
|
||||
sourceComment: null,
|
||||
}],
|
||||
relationships: [],
|
||||
});
|
||||
const table = (await repository.listTables(database.id))[0]!;
|
||||
const columns = await repository.listColumns(database.id, table.id);
|
||||
const column = columns.find((candidate) => candidate.name === "status")!;
|
||||
const ward = columns.find((candidate) => candidate.name === "ward")!;
|
||||
const sampleSecret = "ONLY_IN_TRANSIENT_SAMPLE_7f29c8";
|
||||
const sourceSampler: DescriptionSourceSampler = {
|
||||
sample: vi.fn(async () => [{
|
||||
targetId: column.id,
|
||||
tableName: table.name,
|
||||
rows: [
|
||||
{ fields: [{ name: column.name, value: sampleSecret }] },
|
||||
{ fields: [{ name: column.name, value: "row-2" }] },
|
||||
{ fields: [{ name: column.name, value: "row-3" }] },
|
||||
],
|
||||
representativeValues: [{
|
||||
column: column.name,
|
||||
values: [sampleSecret, sampleSecret, "two", "three"],
|
||||
}],
|
||||
}, {
|
||||
targetId: ward.id,
|
||||
tableName: table.name,
|
||||
rows: [
|
||||
{ fields: [{ name: ward.name, value: "row-4" }] },
|
||||
{ fields: [{ name: ward.name, value: "row-5" }] },
|
||||
{ fields: [{ name: ward.name, value: "row-6-must-be-omitted" }] },
|
||||
],
|
||||
representativeValues: [{
|
||||
column: ward.name,
|
||||
values: ["ward-1", "ward-2", "ward-3-must-be-omitted"],
|
||||
}],
|
||||
}]),
|
||||
};
|
||||
const completer: ModelCompleter = {
|
||||
complete: vi.fn(async () => JSON.stringify({
|
||||
results: [{
|
||||
targetId: column.id,
|
||||
outcome: "generated",
|
||||
description: "Stato amministrativo del paziente.",
|
||||
}, {
|
||||
targetId: ward.id,
|
||||
outcome: "generated",
|
||||
description: "Reparto associato al paziente.",
|
||||
}],
|
||||
})),
|
||||
};
|
||||
const models: MetadataGenerationModels = {
|
||||
catalog: () => ({ models: [{ id: "openai-mini", label: "OpenAI Mini" }], default: "openai-mini" }),
|
||||
resolve: () => ({
|
||||
id: "openai-mini",
|
||||
provider: "openai",
|
||||
model: "gpt-4.1-mini",
|
||||
apiKeyEnv: "OPENAI_API_KEY",
|
||||
apiKey: "test-provider-secret",
|
||||
}),
|
||||
};
|
||||
const worker = new DescriptionGenerationWorker(
|
||||
repository,
|
||||
{
|
||||
read: vi.fn(async () => ({
|
||||
workspace: { workspace: { language: "it" } },
|
||||
revision: {},
|
||||
})),
|
||||
} as unknown as WorkspaceRegistry,
|
||||
models,
|
||||
completer,
|
||||
new CatalogOperationCoordinator(),
|
||||
sourceSampler,
|
||||
);
|
||||
|
||||
const run = await worker.start(
|
||||
database.id,
|
||||
"openai-mini",
|
||||
"selected_columns",
|
||||
[column.id, ward.id],
|
||||
);
|
||||
await worker.waitForRun(run.id);
|
||||
|
||||
expect(sourceSampler.sample).toHaveBeenCalledWith(
|
||||
expect.objectContaining({ id: database.id, binding: database.binding }),
|
||||
[
|
||||
{ targetId: column.id, tableName: table.name, columnNames: [column.name] },
|
||||
{ targetId: ward.id, tableName: table.name, columnNames: [ward.name] },
|
||||
],
|
||||
expect.any(AbortSignal),
|
||||
);
|
||||
const request = vi.mocked(completer.complete).mock.calls[0]![0] as ModelCompletionRequest;
|
||||
expect(request.messages[0]?.content).toContain("untrusted");
|
||||
const userMessage = request.messages[1]!.content;
|
||||
const context = JSON.parse(userMessage.slice(userMessage.indexOf("\n") + 1));
|
||||
const sampledRows = context.targets.flatMap(
|
||||
(targetContext: { sourceSample?: { rows: unknown[] } }) => targetContext.sourceSample?.rows ?? [],
|
||||
);
|
||||
const representativeValues = context.targets.flatMap(
|
||||
(targetContext: { sourceSample?: { representativeValues: Array<{ values: unknown[] }> } }) => (
|
||||
targetContext.sourceSample?.representativeValues.flatMap((entry) => entry.values) ?? []
|
||||
),
|
||||
);
|
||||
expect(sampledRows).toHaveLength(5);
|
||||
expect(representativeValues).toHaveLength(5);
|
||||
expect(context.targets[0].sourceSample.rows).toHaveLength(3);
|
||||
expect(context.targets[1].sourceSample.rows).toHaveLength(2);
|
||||
expect(context.targets[0].sourceSample.representativeValues).toEqual([{
|
||||
column: column.name,
|
||||
values: [sampleSecret, "two", "three"],
|
||||
}]);
|
||||
expect(context.targets[1].sourceSample.representativeValues).toEqual([{
|
||||
column: ward.name,
|
||||
values: ["ward-1", "ward-2"],
|
||||
}]);
|
||||
expect(userMessage).toContain(sampleSecret);
|
||||
expect(userMessage).not.toMatch(
|
||||
/row-6-must-be-omitted|ward-3-must-be-omitted/,
|
||||
);
|
||||
|
||||
const persisted = JSON.stringify({
|
||||
run: await repository.getDescriptionGenerationRun(run.id),
|
||||
events: await repository.listDescriptionGenerationEvents(run.id),
|
||||
database: await repository.get(database.id),
|
||||
table: await repository.getTable(database.id, table.id),
|
||||
column: await repository.getColumn(database.id, table.id, column.id),
|
||||
ward: await repository.getColumn(database.id, table.id, ward.id),
|
||||
});
|
||||
expect(persisted).not.toContain(sampleSecret);
|
||||
});
|
||||
|
||||
test("continues metadata-only with one safe warning when source sampling is unavailable", async () => {
|
||||
const repository = new MemoryCatalogRepository();
|
||||
const database = await repository.create({
|
||||
workspaceId: "psd-clinical",
|
||||
engine: "postgres",
|
||||
databaseName: "warehouse",
|
||||
schema: "datawarehouse",
|
||||
binding: {
|
||||
transport: "rest_api",
|
||||
baseUrl: "https://dwh.example.test",
|
||||
restPath: "/rpc/run_query",
|
||||
restAuth: "none",
|
||||
},
|
||||
});
|
||||
await repository.applySchemaSync(database.id, database.version, "all", [], {
|
||||
schemaVersion: 1,
|
||||
capabilities: { tables: "available", columns: "available", relationships: "available" },
|
||||
tables: [{ name: "patients", sourceComment: null }],
|
||||
columns: [{
|
||||
tableName: "patients",
|
||||
name: "status",
|
||||
ordinalPosition: 1,
|
||||
dataType: "text",
|
||||
isNullable: true,
|
||||
defaultExpression: null,
|
||||
primaryKeyPosition: null,
|
||||
sourceComment: null,
|
||||
}],
|
||||
relationships: [],
|
||||
});
|
||||
const table = (await repository.listTables(database.id))[0]!;
|
||||
const column = (await repository.listColumns(database.id, table.id))[0]!;
|
||||
const samplingFailureSecret = "UNAVAILABLE_SAMPLE_DETAIL_48b1f1";
|
||||
const sourceSampler: DescriptionSourceSampler = {
|
||||
sample: vi.fn(async () => { throw new Error(samplingFailureSecret); }),
|
||||
};
|
||||
const completer: ModelCompleter = {
|
||||
complete: vi.fn(async () => JSON.stringify({
|
||||
results: [{
|
||||
targetId: column.id,
|
||||
outcome: "generated",
|
||||
description: "Stato del paziente.",
|
||||
}],
|
||||
})),
|
||||
};
|
||||
const models: MetadataGenerationModels = {
|
||||
catalog: () => ({ models: [{ id: "openai-mini", label: "OpenAI Mini" }], default: "openai-mini" }),
|
||||
resolve: () => ({
|
||||
id: "openai-mini",
|
||||
provider: "openai",
|
||||
model: "gpt-4.1-mini",
|
||||
apiKeyEnv: "OPENAI_API_KEY",
|
||||
apiKey: "test-provider-secret",
|
||||
}),
|
||||
};
|
||||
const worker = new DescriptionGenerationWorker(
|
||||
repository,
|
||||
{
|
||||
read: vi.fn(async () => ({
|
||||
workspace: { workspace: { language: "it" } },
|
||||
revision: {},
|
||||
})),
|
||||
} as unknown as WorkspaceRegistry,
|
||||
models,
|
||||
completer,
|
||||
new CatalogOperationCoordinator(),
|
||||
sourceSampler,
|
||||
);
|
||||
|
||||
const run = await worker.start(database.id, "openai-mini", "selected_columns", [column.id]);
|
||||
await worker.waitForRun(run.id);
|
||||
|
||||
expect(await repository.getDescriptionGenerationRun(run.id)).toMatchObject({
|
||||
status: "completed",
|
||||
processed: 1,
|
||||
generated: 1,
|
||||
failed: 0,
|
||||
});
|
||||
const request = vi.mocked(completer.complete).mock.calls[0]![0] as ModelCompletionRequest;
|
||||
expect(request.messages[1]?.content).not.toMatch(/sourceSample|UNAVAILABLE_SAMPLE_DETAIL/);
|
||||
const events = await repository.listDescriptionGenerationEvents(run.id);
|
||||
expect(events.filter((event) => event.level === "warning")).toEqual([
|
||||
expect.objectContaining({
|
||||
message: "Source samples unavailable for this batch; generation continued with catalog metadata only.",
|
||||
}),
|
||||
]);
|
||||
expect(JSON.stringify(events)).not.toContain(samplingFailureSecret);
|
||||
});
|
||||
Reference in New Issue
Block a user