mptyl/ThothII
1
0
Fork 0
Code Issues 25 Pull Requests Actions Packages Projects Releases 1 Wiki Activity

feat: refine metadata catalog workflows

Browse Source
This commit is contained in:
Codex
2026-09-02 11:38:47 +02:00
parent 076c9742c5
commit 4531746038
48 changed files with 1345 additions and 493 deletions
Download Patch File Download Diff File Expand all files Collapse all files
backend/src/catalog/description-generation-worker.ts
+86 -39
View File
@@ -1,7 +1,7 @@
import { z } from "zod";
import type { WorkspaceRegistry } from "../workspaces/registry.js";
import type { MetadataGenerationModels, ResolvedMetadataGenerationModel } from "./metadata-generation-models.js";
import type { ModelCompleter, ModelCompletionMessage } from "./model-completer.js";
import type { ModelCompleter, ModelCompletionMessage, ModelCompletionResult } from "./model-completer.js";
import {
ModelCompletionCancelledError,
ModelCompletionProviderError,
@@ -38,6 +38,7 @@ const MAX_SAMPLE_FIELDS_PER_ROW = 4;
const MAX_SAMPLE_COLUMNS = 4;
const MAX_REPRESENTATIVE_VALUES_PER_REQUEST = 5;
const MAX_TARGET_SAMPLE_JSON_BYTES = 8 * 1024;
const MAX_COMPLETION_ATTEMPTS_PER_BATCH = 2;
const generatedOutcomeSchema = z.object({
targetId: z.uuid(),
outcome: z.literal("generated"),
@@ -55,7 +56,9 @@ const completionResponseSchema = z.object({
results: z.array(outcomeSchema).min(1).max(MAX_TARGETS_PER_BATCH),
}).strict();
class InvalidModelOutcomeError extends Error {}
class InvalidModelJsonError extends Error {}
class InvalidModelSchemaError extends Error {}
class MissingModelTargetsError extends Error {}
class DescriptionGenerationBatchError extends Error {
constructor(
readonly failure: unknown,
@@ -143,6 +146,9 @@ interface DescriptionGenerationCounters {
generated: number;
nonGeneratable: number;
failed: number;
inputTokens: number;
cacheReadTokens: number;
outputTokens: number;
consecutiveTechnicalFailures: number;
}
@@ -152,6 +158,9 @@ function persistedCounters(counters: DescriptionGenerationCounters) {
generated: counters.generated,
nonGeneratable: counters.nonGeneratable,
failed: counters.failed,
inputTokens: counters.inputTokens,
cacheReadTokens: counters.cacheReadTokens,
outputTokens: counters.outputTokens,
};
}
@@ -625,36 +634,47 @@ function messagesFor(
}
function parseOutcomes(content: string, expectedTargetIds: readonly string[]): Map<string, ParsedOutcome> {
const trimmed = content.trim();
const fenced = /^```(?:json)?[ \t]*\r?\n([\s\S]*?)\r?\n```$/iu.exec(trimmed);
let parsed: unknown;
try {
const trimmed = content.trim();
const fenced = /^```(?:json)?[ \t]*\r?\n([\s\S]*?)\r?\n```$/iu.exec(trimmed);
const outcomes = completionResponseSchema.parse(JSON.parse(fenced?.[1] ?? trimmed)).results;
const expected = new Set(expectedTargetIds);
if (outcomes.length !== expectedTargetIds.length || expected.size !== expectedTargetIds.length) {
throw new InvalidModelOutcomeError();
}
const mapped = new Map<string, ParsedOutcome>();
for (const outcome of outcomes) {
if (!expected.has(outcome.targetId) || mapped.has(outcome.targetId)) {
throw new InvalidModelOutcomeError();
}
mapped.set(outcome.targetId, outcome.outcome === "generated"
? { ...outcome, description: outcome.description.trim() }
: outcome);
}
if (mapped.size !== expected.size) throw new InvalidModelOutcomeError();
return mapped;
} catch (error) {
if (error instanceof InvalidModelOutcomeError) throw error;
throw new InvalidModelOutcomeError();
parsed = JSON.parse(fenced?.[1] ?? trimmed);
} catch {
throw new InvalidModelJsonError();
}
const response = completionResponseSchema.safeParse(parsed);
if (!response.success) throw new InvalidModelSchemaError();
const outcomes = response.data.results;
const expected = new Set(expectedTargetIds);
if (outcomes.length !== expectedTargetIds.length || expected.size !== expectedTargetIds.length) {
throw new MissingModelTargetsError();
}
const mapped = new Map<string, ParsedOutcome>();
for (const outcome of outcomes) {
if (!expected.has(outcome.targetId) || mapped.has(outcome.targetId)) {
throw new MissingModelTargetsError();
}
mapped.set(outcome.targetId, outcome.outcome === "generated"
? { ...outcome, description: outcome.description.trim() }
: outcome);
}
if (mapped.size !== expected.size) throw new MissingModelTargetsError();
return mapped;
}
function safeFailure(error: unknown): string {
if (error instanceof DescriptionGenerationFailureStreakError) return error.message;
const failure = error instanceof DescriptionGenerationBatchError ? error.failure : error;
if (failure instanceof ModelCompletionProviderError) return "The model provider request failed.";
if (failure instanceof InvalidModelOutcomeError) return "The model response was invalid.";
if (failure instanceof InvalidModelJsonError) return "The model response was not valid JSON.";
if (failure instanceof InvalidModelSchemaError) {
return "The model response did not match the required schema.";
}
if (failure instanceof MissingModelTargetsError) {
return "The model response was missing one or more requested targets.";
}
return "Description generation failed.";
}
@@ -907,6 +927,9 @@ export class DescriptionGenerationWorker {
generated: 0,
nonGeneratable: 0,
failed: 0,
inputTokens: 0,
cacheReadTokens: 0,
outputTokens: 0,
consecutiveTechnicalFailures: 0,
};
await this.processTargets(run, database, plan.columnTargets, model, counters, signal);
@@ -973,20 +996,41 @@ export class DescriptionGenerationWorker {
);
}
throwIfCancelled(signal);
let outcomes: Map<string, ParsedOutcome>;
try {
const content = await this.completer.complete({
model,
messages: messagesFor(database, batch, run.language, sourceSamples),
signal,
});
throwIfCancelled(signal);
outcomes = parseOutcomes(content, batch.map((target) => (
target.kind === "column" ? target.column.id : target.table.id
)));
} catch (error) {
if (error instanceof ModelCompletionCancelledError) throw error;
const batchError = new DescriptionGenerationBatchError(error, batch.map(failureTarget));
const expectedTargetIds = batch.map((target) => (
target.kind === "column" ? target.column.id : target.table.id
));
const messages = messagesFor(database, batch, run.language, sourceSamples);
let outcomes: Map<string, ParsedOutcome> | undefined;
let terminalFailure: unknown;
for (let attempt = 1; attempt <= MAX_COMPLETION_ATTEMPTS_PER_BATCH; attempt += 1) {
try {
const completion = await this.completer.complete({ model, messages, signal });
const result: ModelCompletionResult = typeof completion === "string"
? { content: completion, usage: { input: 0, cacheRead: 0, output: 0 } }
: completion;
counters.inputTokens += result.usage.input;
counters.cacheReadTokens += result.usage.cacheRead;
counters.outputTokens += result.usage.output;
throwIfCancelled(signal);
outcomes = parseOutcomes(result.content, expectedTargetIds);
break;
} catch (error) {
if (error instanceof ModelCompletionCancelledError) throw error;
terminalFailure = error;
if (attempt < MAX_COMPLETION_ATTEMPTS_PER_BATCH) {
await this.appendEvent(
run.id,
"warning",
`${safeFailure(error)} Retrying batch (attempt ${attempt + 1} of ${MAX_COMPLETION_ATTEMPTS_PER_BATCH}).`,
);
}
}
}
if (!outcomes) {
const batchError = new DescriptionGenerationBatchError(
terminalFailure,
batch.map(failureTarget),
);
counters.processed += batch.length;
counters.failed += batch.length;
counters.consecutiveTechnicalFailures += 1;
@@ -1009,7 +1053,10 @@ export class DescriptionGenerationWorker {
const targetId = target.kind === "column" ? target.column.id : target.table.id;
const outcome = outcomes.get(targetId);
if (!outcome) {
throw new DescriptionGenerationBatchError(new InvalidModelOutcomeError(), [failureTarget(target)]);
throw new DescriptionGenerationBatchError(
new MissingModelTargetsError(),
[failureTarget(target)],
);
}
const generatedDescription = outcome.outcome === "generated"
? outcome.description
backend/src/catalog/description-source-sampler.ts
+132 -43
View File
@@ -1,5 +1,8 @@
import { readFile } from "node:fs/promises";
import type { WorkspaceSecretStore } from "../workspaces/secret-store.js";
import type { CatalogPostgresAccess } from "./postgres-access.js";
import type { WorkspaceDatabase } from "./types.js";
import { CATALOG_SECRET_IDS } from "./secrets.js";
import { CatalogConnectorError, type WorkspaceDatabase } from "./types.js";
const MAX_SOURCE_ROWS = 5;
const MAX_REPRESENTATIVE_VALUES = 5;
@@ -79,15 +82,82 @@ function distinctKey(value: Exclude<DescriptionSourceSampleValue, null>): string
return `${typeof value}:${String(value)}`;
}
/** PostgreSQL-wire sampler. REST bindings remain unsupported by CatalogPostgresAccess. */
export class PostgresDescriptionSourceSampler implements DescriptionSourceSampler {
constructor(private readonly access: CatalogPostgresAccess) {}
function columnsFor(target: DescriptionSourceSamplingTarget): string[] {
return [...new Set(target.columnNames)].slice(0, MAX_SOURCE_COLUMNS_PER_TARGET);
}
function normalizedSample(
target: DescriptionSourceSamplingTarget,
columnNames: readonly string[],
sourceRows: readonly Record<string, unknown>[],
): DescriptionTargetSourceSample {
const rows = sourceRows.slice(0, MAX_SOURCE_ROWS).map((row) => ({
fields: columnNames.flatMap((name) => {
const value = normalizeValue(row[name]);
return value === undefined ? [] : [{ name, value }];
}),
}));
const valuesByColumn = new Map<string, Exclude<DescriptionSourceSampleValue, null>[]>();
const seenByColumn = new Map<string, Set<string>>();
let representativeValueCount = 0;
for (const row of rows) {
for (const field of row.fields) {
if (representativeValueCount === MAX_REPRESENTATIVE_VALUES) break;
if (field.value === null) continue;
const seen = seenByColumn.get(field.name) ?? new Set<string>();
const key = distinctKey(field.value);
if (seen.has(key)) continue;
seen.add(key);
seenByColumn.set(field.name, seen);
const values = valuesByColumn.get(field.name) ?? [];
values.push(field.value);
valuesByColumn.set(field.name, values);
representativeValueCount += 1;
}
if (representativeValueCount === MAX_REPRESENTATIVE_VALUES) break;
}
return {
targetId: target.targetId,
tableName: target.tableName,
rows,
representativeValues: columnNames.flatMap((column) => {
const values = valuesByColumn.get(column);
return values && values.length > 0 ? [{ column, values }] : [];
}),
};
}
function samplingSql(
database: WorkspaceDatabase,
target: DescriptionSourceSamplingTarget,
columns: readonly string[],
): string {
const projections = columns.map((columnName) => {
const identifier = quoteIdentifier(columnName);
return `LEFT((${identifier})::text, ${MAX_SOURCE_VALUE_BYTES}) AS ${identifier}`;
});
return [
`SELECT ${projections.join(", ")}`,
`FROM ${quoteIdentifier(database.schema)}.${quoteIdentifier(target.tableName)}`,
`LIMIT ${MAX_SOURCE_ROWS}`,
].join(" ");
}
/** Bounded source sampler that follows the database's PostgreSQL-wire or REST binding. */
export class ConcreteDescriptionSourceSampler implements DescriptionSourceSampler {
constructor(
private readonly access: CatalogPostgresAccess,
private readonly secretStore?: Pick<WorkspaceSecretStore, "materialize">,
) {}
async sample(
database: WorkspaceDatabase,
targets: readonly DescriptionSourceSamplingTarget[],
signal: AbortSignal,
): Promise<readonly DescriptionTargetSourceSample[]> {
if (database.binding.transport === "rest_api") {
return await this.sampleRest(database, targets, signal);
}
const client = await this.access.connect(database, signal);
let transactionOpen = false;
try {
@@ -95,7 +165,7 @@ export class PostgresDescriptionSourceSampler implements DescriptionSourceSample
transactionOpen = true;
const samples: DescriptionTargetSourceSample[] = [];
for (const target of targets) {
const columnNames = [...new Set(target.columnNames)].slice(0, MAX_SOURCE_COLUMNS_PER_TARGET);
const columnNames = columnsFor(target);
if (columnNames.length === 0) {
samples.push({
targetId: target.targetId,
@@ -115,44 +185,7 @@ export class PostgresDescriptionSourceSampler implements DescriptionSourceSample
"LIMIT $2",
].join(" ");
const result = await client.query(sql, [MAX_SOURCE_VALUE_BYTES, MAX_SOURCE_ROWS]);
const rows = result.rows.slice(0, MAX_SOURCE_ROWS).map((row) => ({
fields: columnNames.flatMap((name) => {
const value = normalizeValue(row[name]);
return value === undefined ? [] : [{ name, value }];
}),
}));
const valuesByColumn = new Map<
string,
Exclude<DescriptionSourceSampleValue, null>[]
>();
const seenByColumn = new Map<string, Set<string>>();
let representativeValueCount = 0;
for (const row of rows) {
for (const field of row.fields) {
if (representativeValueCount === MAX_REPRESENTATIVE_VALUES) break;
if (field.value === null) continue;
const seen = seenByColumn.get(field.name) ?? new Set<string>();
const key = distinctKey(field.value);
if (seen.has(key)) continue;
seen.add(key);
seenByColumn.set(field.name, seen);
const values = valuesByColumn.get(field.name) ?? [];
values.push(field.value);
valuesByColumn.set(field.name, values);
representativeValueCount += 1;
}
if (representativeValueCount === MAX_REPRESENTATIVE_VALUES) break;
}
const representativeValues = columnNames.flatMap((column) => {
const values = valuesByColumn.get(column);
return values && values.length > 0 ? [{ column, values }] : [];
});
samples.push({
targetId: target.targetId,
tableName: target.tableName,
rows,
representativeValues,
});
samples.push(normalizedSample(target, columnNames, result.rows));
}
return samples;
} finally {
@@ -160,4 +193,60 @@ export class PostgresDescriptionSourceSampler implements DescriptionSourceSample
await client.end().catch(() => undefined);
}
}
private async sampleRest(
database: WorkspaceDatabase,
targets: readonly DescriptionSourceSamplingTarget[],
signal: AbortSignal,
): Promise<readonly DescriptionTargetSourceSample[]> {
if (!this.secretStore) throw new CatalogConnectorError("REST source sampling is not configured");
const auth = database.binding.restAuth ?? "bearer";
const materialized = this.secretStore.materialize(
database.workspaceId,
auth === "none" ? [] : [CATALOG_SECRET_IDS.apiKey],
);
try {
const headers: Record<string, string> = { "content-type": "application/json" };
if (auth !== "none") {
const credentialFile = materialized.files.get(CATALOG_SECRET_IDS.apiKey);
if (!credentialFile) throw new CatalogConnectorError("REST API key is not configured");
const credential = (await readFile(credentialFile, "utf8")).trim();
if (auth === "bearer") headers.authorization = `Bearer ${credential}`;
else headers["x-api-key"] = credential;
}
const baseUrl = database.binding.baseUrl?.replace(/\/+$/, "");
if (!baseUrl) throw new CatalogConnectorError("Database binding is incomplete");
const samples: DescriptionTargetSourceSample[] = [];
for (const target of targets) {
const columnNames = columnsFor(target);
if (columnNames.length === 0) {
samples.push(normalizedSample(target, columnNames, []));
continue;
}
const response = await fetch(`${baseUrl}/rpc/run_query`, {
method: "POST",
headers,
body: JSON.stringify({ query_text: samplingSql(database, target, columnNames) }),
signal,
});
if (!response.ok) throw new CatalogConnectorError("REST source sampling failed");
const body: unknown = await response.json();
if (!Array.isArray(body)
|| body.some((row) => !row || typeof row !== "object" || Array.isArray(row))) {
throw new CatalogConnectorError("REST source sampling response is invalid");
}
samples.push(normalizedSample(
target,
columnNames,
body as Array<Record<string, unknown>>,
));
}
return samples;
} catch (error) {
if (error instanceof CatalogConnectorError) throw error;
throw new CatalogConnectorError("REST source sampling failed");
} finally {
materialized.release();
}
}
}
backend/src/catalog/memory-repository.ts
+8 -2
View File
@@ -335,7 +335,10 @@ export class MemoryCatalogRepository implements CatalogRepository {
processed: 0,
generated: 0,
nonGeneratable: 0,
failed: 0,
failed: 0,
inputTokens: 0,
cacheReadTokens: 0,
outputTokens: 0,
createdAt: now,
startedAt: null,
updatedAt: now,
@@ -444,7 +447,10 @@ export class MemoryCatalogRepository implements CatalogRepository {
status: "running",
total: 0,
suggestedSensitive: 0,
suggestedNonSensitive: 0,
suggestedNonSensitive: 0,
inputTokens: 0,
cacheReadTokens: 0,
outputTokens: 0,
createdAt: now,
startedAt: now,
updatedAt: now,
backend/src/catalog/migrate.ts
+2
View File
@@ -11,6 +11,7 @@ import * as descriptionGenerationRunsMigration from "./migrations/005_descriptio
import * as sensitiveDataFlagMigration from "./migrations/006_sensitive_data_flag.js";
import * as sensitiveDataSuggestionRunsMigration from "./migrations/007_sensitive_data_suggestion_runs.js";
import * as catalogLogicalRelationshipsMigration from "./migrations/008_catalog_logical_relationships.js";
import * as aiTokenUsageMigration from "./migrations/009_ai_token_usage.js";
const connectionString = process.env.THT_CATALOG_MIGRATOR_DATABASE_URL;
const host = process.env.THT_CATALOG_DB_HOST;
@@ -44,6 +45,7 @@ const provider: MigrationProvider = {
"006_sensitive_data_flag": sensitiveDataFlagMigration,
"007_sensitive_data_suggestion_runs": sensitiveDataSuggestionRunsMigration,
"008_catalog_logical_relationships": catalogLogicalRelationshipsMigration,
"009_ai_token_usage": aiTokenUsageMigration,
};
},
};
backend/src/catalog/migrations/009_ai_token_usage.ts
+20
View File
@@ -0,0 +1,20 @@
import type { Kysely } from "kysely";
import type { CatalogDatabase } from "../repository.js";
export async function up(db: Kysely<CatalogDatabase>): Promise<void> {
for (const table of ["description_generation_runs", "sensitive_data_suggestion_runs"] as const) {
await db.schema.alterTable(table)
.addColumn("input_tokens", "integer", (col) => col.notNull().defaultTo(0))
.addColumn("cache_read_tokens", "integer", (col) => col.notNull().defaultTo(0))
.addColumn("output_tokens", "integer", (col) => col.notNull().defaultTo(0))
.execute();
}
}
export async function down(db: Kysely<CatalogDatabase>): Promise<void> {
for (const table of ["sensitive_data_suggestion_runs", "description_generation_runs"] as const) {
await db.schema.alterTable(table)
.dropColumn("input_tokens").dropColumn("cache_read_tokens").dropColumn("output_tokens")
.execute();
}
}
backend/src/catalog/model-completer.ts
+7 -5
View File
@@ -4,7 +4,7 @@ import type { ResolvedMetadataGenerationModel } from "./metadata-generation-mode
const MAX_HELPER_OUTPUT_BYTES = 64 * 1024;
const helperOutputSchema = z.discriminatedUnion("ok", [
z.object({ ok: z.literal(true), content: z.string() }).strict(),
z.object({ ok: z.literal(true), content: z.string(), usage: z.object({ input: z.number().int().nonnegative(), cacheRead: z.number().int().nonnegative(), output: z.number().int().nonnegative() }).strict().optional() }).strict(),
z.object({ ok: z.literal(false), error: z.literal("provider_failure") }).strict(),
]);
@@ -18,10 +18,12 @@ export interface ModelCompletionRequest {
messages: readonly ModelCompletionMessage[];
signal: AbortSignal;
}
export interface ModelCompletionUsage { input: number; cacheRead: number; output: number; }
export interface ModelCompletionResult { content: string; usage: ModelCompletionUsage; }
/** The provider boundary used by Description Generation. */
export interface ModelCompleter {
complete(request: ModelCompletionRequest): Promise<string>;
complete(request: ModelCompletionRequest): Promise<string | ModelCompletionResult>;
}
export class ModelCompletionProviderError extends Error {
@@ -47,7 +49,7 @@ export class PythonModelCompleter implements ModelCompleter {
terminationGraceMs?: number;
}) {}
async complete(request: ModelCompletionRequest): Promise<string> {
async complete(request: ModelCompletionRequest): Promise<ModelCompletionResult> {
if (request.signal.aborted) throw new ModelCompletionCancelledError();
const payload = {
model: `${request.model.provider}/${request.model.model}`,
@@ -62,7 +64,7 @@ export class PythonModelCompleter implements ModelCompleter {
...(request.model.disableThinking === true ? { disable_thinking: true } : {}),
};
return await new Promise<string>((resolve, reject) => {
return await new Promise<ModelCompletionResult>((resolve, reject) => {
const child = spawn(
this.options.pythonExecutable,
["-m", this.options.helperModule ?? "tht.internal.litellm_completion"],
@@ -135,7 +137,7 @@ export class PythonModelCompleter implements ModelCompleter {
if (!output.ok) return fail();
settled = true;
cleanup();
resolve(output.content);
resolve({ content: output.content, usage: output.usage ?? { input: 0, cacheRead: 0, output: 0 } });
} catch {
fail();
}
backend/src/catalog/repository.ts
+12
View File
@@ -171,6 +171,9 @@ interface DescriptionGenerationRunTable {
generated: number;
nonGeneratable: number;
failed: number;
inputTokens: number;
cacheReadTokens: number;
outputTokens: number;
createdAt: Timestamp;
startedAt: Timestamp | null;
updatedAt: Timestamp;
@@ -195,6 +198,9 @@ interface SensitiveDataSuggestionRunTable {
total: number;
suggestedSensitive: number;
suggestedNonSensitive: number;
inputTokens: number;
cacheReadTokens: number;
outputTokens: number;
createdAt: Timestamp;
startedAt: Timestamp;
updatedAt: Timestamp;
@@ -816,6 +822,9 @@ export class KyselyCatalogRepository implements CatalogRepository {
generated: 0,
nonGeneratable: 0,
failed: 0,
inputTokens: 0,
cacheReadTokens: 0,
outputTokens: 0,
startedAt: null,
finishedAt: null,
errorSummary: null,
@@ -943,6 +952,9 @@ export class KyselyCatalogRepository implements CatalogRepository {
total: 0,
suggestedSensitive: 0,
suggestedNonSensitive: 0,
inputTokens: 0,
cacheReadTokens: 0,
outputTokens: 0,
finishedAt: null,
errorSummary: null,
}).returningAll().executeTakeFirstOrThrow();
backend/src/catalog/sensitive-data-suggester.ts
+10 -2
View File
@@ -1,6 +1,6 @@
import { z } from "zod";
import type { MetadataGenerationModels } from "./metadata-generation-models.js";
import type { ModelCompleter, ModelCompletionMessage } from "./model-completer.js";
import type { ModelCompleter, ModelCompletionMessage, ModelCompletionResult, ModelCompletionUsage } from "./model-completer.js";
import type {
CatalogColumn,
CatalogRepository,
@@ -201,6 +201,8 @@ export class SensitiveDataSuggester {
targetIds: readonly string[],
signal: AbortSignal,
onPrepared?: (total: number) => void | Promise<void>,
onProgress?: (processed: number, suggestions: readonly SensitiveDataSuggestion[]) => void | Promise<void>,
onUsage?: (usage: ModelCompletionUsage) => void | Promise<void>,
): Promise<readonly SensitiveDataSuggestion[]> {
const database = await this.repository.get(databaseId);
if (!database) throw new SensitiveDataSuggestionTargetNotFoundError("database");
@@ -212,11 +214,16 @@ export class SensitiveDataSuggester {
for (const batch of batchesFor(database, columns)) {
let received: Map<string, { columnId: string; sensitive: boolean }> | undefined;
for (let attempt = 0; attempt < 2 && !received; attempt += 1) {
const content = await this.completer.complete({
const completion = await this.completer.complete({
model,
signal,
messages: [systemMessage, { role: "user", content: userContent(database, batch) }],
});
const result: ModelCompletionResult = typeof completion === "string"
? { content: completion, usage: { input: 0, cacheRead: 0, output: 0 } }
: completion;
await onUsage?.(result.usage);
const content = result.content;
try {
const parsed = responseSchema.parse(JSON.parse(content));
const expected = new Set(batch.map((column) => column.columnId));
@@ -240,6 +247,7 @@ export class SensitiveDataSuggester {
currentSensitive: column.currentSensitive,
sensitive: received!.get(column.columnId)!.sensitive,
})));
await onProgress?.(suggestions.length, suggestions.slice(-batch.length));
}
return suggestions;
}
backend/src/catalog/sensitive-data-suggestion-runner.ts
+26
View File
@@ -10,6 +10,7 @@ import type {
SensitiveDataSuggestionRun,
SensitiveDataSuggestionScope,
} from "./types.js";
import type { ModelCompletionUsage } from "./model-completer.js";
const interruptedMessage = "Sensitive-field suggestion generation was interrupted by backend restart.";
const failedMessage = "Sensitive-field suggestion generation failed.";
@@ -73,6 +74,31 @@ export class SensitiveDataSuggestionRunner {
});
if (!prepared) throw new Error("Sensitive Data Suggestion Run disappeared");
},
async (processed, batch) => {
const suggestedSensitive = batch.filter((suggestion) => suggestion.sensitive).length;
const suggestedNonSensitive = batch.length - suggestedSensitive;
const current = await this.repository.getSensitiveDataSuggestionRun(started.id);
if (!current) throw new Error("Sensitive Data Suggestion Run disappeared");
const progress = await this.repository.updateSensitiveDataSuggestionRun(started.id, {
suggestedSensitive: current.suggestedSensitive + suggestedSensitive,
suggestedNonSensitive: current.suggestedNonSensitive + suggestedNonSensitive,
});
if (!progress) throw new Error("Sensitive Data Suggestion Run disappeared");
await this.repository.appendSensitiveDataSuggestionEvent(
started.id,
"info",
`Classified ${processed} of ${progress.total} columns.`,
);
},
async (usage: ModelCompletionUsage) => {
const current = await this.repository.getSensitiveDataSuggestionRun(started.id);
if (!current) throw new Error("Sensitive Data Suggestion Run disappeared");
await this.repository.updateSensitiveDataSuggestionRun(started.id, {
inputTokens: current.inputTokens + usage.input,
cacheReadTokens: current.cacheReadTokens + usage.cacheRead,
outputTokens: current.outputTokens + usage.output,
});
},
);
const suggestedSensitive = suggestions.filter((suggestion) => suggestion.sensitive).length;
const suggestedNonSensitive = suggestions.length - suggestedSensitive;
backend/src/catalog/types.ts
+12
View File
@@ -234,6 +234,9 @@ export interface DescriptionGenerationRun {
generated: number;
nonGeneratable: number;
failed: number;
inputTokens: number;
cacheReadTokens: number;
outputTokens: number;
createdAt: string;
startedAt: string | null;
updatedAt: string;
@@ -250,6 +253,9 @@ export interface DescriptionGenerationRunUpdate {
startedAt?: string | null;
finishedAt?: string | null;
errorSummary?: string | null;
inputTokens?: number;
cacheReadTokens?: number;
outputTokens?: number;
}
export interface DescriptionGenerationEvent {
@@ -272,6 +278,9 @@ export interface SensitiveDataSuggestionRun {
total: number;
suggestedSensitive: number;
suggestedNonSensitive: number;
inputTokens: number;
cacheReadTokens: number;
outputTokens: number;
createdAt: string;
startedAt: string;
updatedAt: string;
@@ -286,6 +295,9 @@ export interface SensitiveDataSuggestionRunUpdate {
suggestedNonSensitive?: number;
finishedAt?: string | null;
errorSummary?: string | null;
inputTokens?: number;
cacheReadTokens?: number;
outputTokens?: number;
}
export interface SensitiveDataSuggestionEvent {