feat: classify sensitive columns locally

backend/src/app.ts

@@ -3,6 +3,7 @@ import cors from "@fastify/cors";
 import cookie from "@fastify/cookie";
 import rateLimit from "@fastify/rate-limit";
 import { dirname, isAbsolute, join } from "node:path";
+import { fileURLToPath } from "node:url";
 import { tmpdir } from "node:os";
 import type { AppConfig } from "./config.js";
 import { ThtRunner } from "./tht/tht-runner.js";
@@ -57,8 +58,11 @@ import { metadataGenerationModelRoutes } from "./routes/metadata-generation-mode
 import { catalogDescriptionConsolidationRoutes } from "./routes/catalog-description-consolidation.js";
 import { PythonModelCompleter, type ModelCompleter } from "./catalog/model-completer.js";
 import { DescriptionGenerationWorker } from "./catalog/description-generation-worker.js";
-import { SensitiveDataSuggester } from "./catalog/sensitive-data-suggester.js";
-import { SensitiveDataSuggestionRunner } from "./catalog/sensitive-data-suggestion-runner.js";
+import { SensitivityAnalysisService } from "./catalog/sensitivity-analysis-service.js";
+import { SensitivityAnalysisRunner } from "./catalog/sensitivity-analysis-runner.js";
+import { SensitivityClassifier, type LocalNerDetector, type SensitivityValueSource } from "./catalog/sensitivity-classifier.js";
+import { ConcreteSensitivityValueSource } from "./catalog/sensitivity-value-source.js";
+import { PythonLocalNerDetector } from "./catalog/local-ner-detector.js";
 import {
   ConcreteDescriptionSourceSampler,
   type DescriptionSourceSampler,
@@ -93,6 +97,8 @@ export interface BuildAppDeps {
   runtimeModelCatalog?: RuntimeModelCatalog;
   modelCompleter?: ModelCompleter;
   descriptionSourceSampler?: DescriptionSourceSampler;
+  sensitivityValueSource?: SensitivityValueSource;
+  localNerDetector?: LocalNerDetector;
   workspaceRuntimeSupport?: (workspace: WorkspaceDescriptor) => boolean;
   maintenanceBarrier?: MaintenanceBarrier;
   piManagement?: PiManagementService;
@@ -194,12 +200,24 @@ export function buildApp(config: AppConfig, deps?: BuildAppDeps): FastifyInstanc
     catalogOperationCoordinator,
     descriptionSourceSampler,
   );
-  const sensitiveDataSuggester = new SensitiveDataSuggester(
+  const sensitivityValueSource = deps?.sensitivityValueSource
+    ?? new ConcreteSensitivityValueSource(catalogPostgresAccess, workspaceSecretStore);
+  const configuredNerWorker = config.sensitivityNer?.workerScript
+    ?? fileURLToPath(new URL("../python/sensitivity_ner_worker.py", import.meta.url));
+  const localNerDetector = deps?.localNerDetector ?? (config.sensitivityNer
+    ? new PythonLocalNerDetector({
+      pythonExecutable: config.sensitivityNer.pythonExecutable,
+      workerScript: configuredNerWorker,
+      modelPath: config.sensitivityNer.modelPath,
+      cwd: dirname(configuredNerWorker),
+      threads: config.sensitivityNer.threads,
+    })
+    : undefined);
+  const sensitiveDataSuggester = new SensitivityAnalysisService(
     catalogRepository,
-    metadataGenerationModels,
-    modelCompleter,
+    new SensitivityClassifier(sensitivityValueSource, localNerDetector),
   );
-  const sensitiveDataSuggestionRunner = new SensitiveDataSuggestionRunner(
+  const sensitivityAnalysisRunner = new SensitivityAnalysisRunner(
     catalogRepository,
     sensitiveDataSuggester,
   );
@@ -235,12 +253,20 @@ export function buildApp(config: AppConfig, deps?: BuildAppDeps): FastifyInstanc
   );
   app.addHook("onReady", async () => { await catalogSyncWorker.initialize(); });
   app.addHook("onReady", async () => { await descriptionGenerationWorker.initialize(); });
-  app.addHook("onReady", async () => { await sensitiveDataSuggestionRunner.initialize(); });
+  app.addHook("onReady", async () => { await sensitivityAnalysisRunner.initialize(); });
+  if (localNerDetector?.warmup) {
+    app.addHook("onReady", async () => {
+      void localNerDetector.warmup?.().catch(() => undefined);
+    });
+  }
   if (!deps?.catalogRepository && catalogRepository.close) {
     app.addHook("onClose", async () => { await catalogRepository.close?.(); });
   }
   app.addHook("onClose", async () => { await catalogSyncWorker.stop(); });
   app.addHook("onClose", async () => { await descriptionGenerationWorker.stop(); });
+  if (localNerDetector?.close) {
+    app.addHook("onClose", async () => { await localNerDetector.close?.(); });
+  }
   const workspaceDiagnoser = deps?.workspaceDiagnoser
     ?? createProductionWorkspaceDiagnoser(config.workspaceDiagnosticTimeoutMs, undefined, {
       internalQdrantUrl: config.internalQdrantUrl,
@@ -483,7 +509,7 @@ export function buildApp(config: AppConfig, deps?: BuildAppDeps): FastifyInstanc
   catalogDescriptionGenerationRoutes(app, {
     repository: catalogRepository,
     worker: descriptionGenerationWorker,
-    sensitiveDataSuggestionRunner,
+    sensitivityAnalysisRunner,
   });
   settingsRoutes(app, { cfg: config, getSettings });
   piManagementRoutes(app, { service: piManagement });