@askdb/studio 0.2.0-beta.36 → 0.2.0-beta.37

package/dist/server.js

@@ -5,21 +5,20 @@ import { extname, relative, join, resolve, dirname, sep } from "node:path";
 import { spawn } from "node:child_process";
 import { fileURLToPath } from "node:url";
 import { generateText as defaultGenerateText } from "ai";
-import { bootstrapAskDbEnv, getAskDbRuntimeConfig } from "@askdb/config";
+import { ASKDB_AI_PROVIDERS, bootstrapAskDbEnv, getAskDbRuntimeConfig, } from "@askdb/config";
 import { createAiRegistry, resolveReasoningEffort, } from "@askdb/ai";
-import { anthropicProvider } from "@askdb/ai-anthropic";
-import { azureProvider } from "@askdb/ai-azure";
-import { googleProvider } from "@askdb/ai-google";
-import { openaiProvider } from "@askdb/ai-openai";
 import { ask, formatSensitiveReference, isBuiltInDialectId, loadSchema, schemaHasSensitiveIdentifiers, SqlValidationError, validateSensitiveReferences, suggestEnrichment, tenantScopeSchema, normalizeTenantPolicy, parseTenantPolicyMarkdown, SchemaParseError, writeTenantPolicyMarkdown, tenantPolicyFrontmatterSchema, } from "@askdb/core";
-import { buildSchemaIndex, chunkContentHash, chunkSchema, createAiSdkEmbedder, createFileStore, createMemoryStore, createPgvectorStore, createRetriever, loadChunkerSourcesFromDir, } from "@askdb/rag";
+import { buildSchemaIndex, chunkContentHash, chunkSchema, aiSdkEmbedderId, createAiSdkEmbedder, createFileStore, createMemoryStore, createPgvectorStore, createRetriever, detectEmbeddingDimensions, loadChunkerSourcesFromDir, PgvectorDimensionMismatchError, } from "@askdb/rag";
 import { buildDefaultTableBody, buildFrontmatter, buildSuggestionContext, buildSuggestionTarget, buildTableDraft, loadWorkspace, replaceH2Section, replaceTableDescription, saveConcepts, saveTable, } from "@askdb/enrich";
 import { EXECUTE_DRIVER_REGISTRY, executeDialectFor, isDriverInstalled, isStudioExecuteProvider, validateExecuteSql, } from "./execute-registry.js";
 import { lockfilePackageManager, packageManagerAddArgs, packageManagerSpawnSpec, } from "./package-manager.js";
 import { resolveStudioIntrospectionPlan, runStudioIntrospection } from "./introspection.js";
 import { probeSetupState, writeSetupConfig, SetupError } from "./setup.js";
 import { checkApiRequest, checkHost, createStudioSessionToken, injectSessionToken, } from "./request-guard.js";
-const ai = createAiRegistry([openaiProvider, azureProvider, googleProvider, anthropicProvider]);
+// Batteries-included surface: every built-in provider is registered, and each
+// loads its @ai-sdk/* package only when first used, so env config alone
+// selects the provider.
+const ai = createAiRegistry();
 const DEFAULT_CLIENT_DIR = fileURLToPath(new URL("./client/", import.meta.url));
 let clientDirForTests;
 /** @internal Tests only — serve client assets from another directory. */
@@ -31,7 +30,6 @@ function clientDir() {
 }
 const STUDIO_RAG_MOCK_DIMENSIONS = 64;
 const STUDIO_RAG_MOCK_EMBEDDER_ID = `studio:mock-lexical-${STUDIO_RAG_MOCK_DIMENSIONS}`;
-const DEFAULT_EMBEDDING_MODEL = "text-embedding-3-small";
 let studioPgvectorStoreFactoryForTests;
 export function setStudioPgvectorStoreFactoryForTests(factory) {
     studioPgvectorStoreFactoryForTests = factory;
@@ -224,8 +222,8 @@ export function serializeWorkspace(workspace) {
         dialect,
         warnings: workspace.warnings,
         aiConfigured: Boolean(aiConfig),
-        model: aiConfig?.model ?? "gpt-4o-mini",
-        aiProvider: aiConfig?.provider ?? "openai",
+        model: aiConfig?.model ?? rt.ai.language.model ?? "provider default",
+        aiProvider: aiConfig?.provider ?? rt.ai.language.provider,
         tables: workspace.tables.map((table) => {
             const draft = buildTableDraft(table.physical, table.parsed);
             return {
@@ -373,7 +371,8 @@ function parseSetupConfigBody(body) {
         throw new StudioHttpError(400, "Request body must be a JSON object.");
     }
     const databases = ["postgres", "mysql", "sqlite", "sqlserver", "prisma"];
-    const aiProviders = ["openai", "anthropic", "google", "azure", "foundry"];
+    // Every provider with an askdb.config.* branch; @askdb/client's provider-config-drift test asserts this matches @askdb/ai's built-in table.
+    const aiProviders = ASKDB_AI_PROVIDERS;
     const ragStores = ["file", "memory", "pgvector"];
     const executeProviders = ["postgres", "mysql", "sqlite", "sqlserver"];
     if (typeof body.database !== "string" || !databases.includes(body.database)) {
@@ -659,14 +658,17 @@ async function getRagStatus(state) {
     const sources = loadChunkerSourcesFromDir(state.schemaDir);
     const chunkResult = chunkSchema(sources);
     const lockPath = join(state.schemaDir, "schema.lock.json");
-    const lock = readOptionalJson(lockPath);
+    const lock = resolveStudioRagStoreConfig(state).kind === "memory"
+        ? state.ragMemoryIndex
+        : readOptionalJson(lockPath);
     const currentHashes = Object.fromEntries(chunkResult.chunks.map((chunk) => [chunk.id, chunkContentHash(chunk.text)]));
     const lockHashes = lock?.hashes ?? {};
     const hashIds = Object.keys(currentHashes);
-    const store = await openStudioRagStore(state, config.dimensions);
+    const store = await openStudioRagStore(state, { dimensions: config.dimensions });
     try {
         const chunksIndexed = await countStudioRagStoreChunks(store, sources.schema.schemaId);
         const stale = !lock ||
+            lock.incomplete === true ||
             lock.embedderId !== config.embedderId ||
             Object.keys(lockHashes).length !== hashIds.length ||
             hashIds.some((id) => lockHashes[id] !== currentHashes[id]) ||
@@ -707,8 +709,9 @@ async function getRagStatus(state) {
             updatedAt: lock?.updatedAt ?? null,
             chunksTotal: chunkResult.chunks.length,
             chunksIndexed,
-            dimensions: config.dimensions,
-            expectedDimensions: config.dimensions,
+            // The width the last index build learned from the model, else the one configured.
+            dimensions: lock?.dimensions ?? config.dimensions ?? null,
+            expectedDimensions: config.dimensions ?? null,
             sensitiveExcluded: chunkResult.stats.sensitiveExcluded,
             sensitiveIncluded: chunkResult.stats.sensitiveIncluded,
             files: fileArtifacts,
@@ -721,27 +724,40 @@ async function getRagStatus(state) {
 async function indexRag(state) {
     const config = resolveStudioRagEmbedderConfig();
     if (!config.configured) {
-        throw new StudioHttpError(400, studioRagAiSdkKeyMissingMessage());
+        throw new StudioHttpError(400, studioRagAiSdkKeyMissingMessage(config));
     }
     const usage = createRequestUsageCollector();
     clearIncompatibleRagStore(state, config);
     const sources = loadChunkerSourcesFromDir(state.schemaDir);
-    const store = await openStudioRagStore(state, config.dimensions);
+    const embedder = await createStudioRagEmbedder(config, usage);
     let result;
+    let store;
     try {
+        // A new pgvector table needs its width up front: the configured one, else the model's own.
+        const dimensions = config.dimensions ??
+            (resolveStudioRagStoreConfig(state).kind === "pgvector" ? await detectEmbeddingDimensions(embedder) : undefined);
+        store = await openStudioRagStore(state, { dimensions, provision: true });
         result = await buildSchemaIndex({
             schema: sources,
-            embedder: await createStudioRagEmbedder(config, usage),
+            embedder,
             store: store.store,
             embedderId: config.embedderId,
             lockFilePath: join(state.schemaDir, "schema.lock.json"),
         });
+        if (store.kind === "memory") {
+            state.ragMemoryIndex = {
+                embedderId: config.embedderId,
+                updatedAt: new Date().toISOString(),
+                dimensions: store.store.describe?.().dimensions,
+                hashes: Object.fromEntries(result.chunks.map((chunk) => [chunk.id, chunkContentHash(chunk.text)])),
+            };
+        }
     }
     catch (error) {
         throw formatStudioRagOperationError(error, config);
     }
     finally {
-        await store.dispose();
+        await store?.dispose();
     }
     return {
         status: await getRagStatus(state),
@@ -779,7 +795,7 @@ async function queryRag(state, query) {
 async function createCurrentStudioRagIndex(state, action, usage) {
     const config = resolveStudioRagEmbedderConfig();
     if (!config.configured) {
-        throw new StudioHttpError(400, studioRagAiSdkKeyMissingMessage());
+        throw new StudioHttpError(400, studioRagAiSdkKeyMissingMessage(config));
     }
     const status = (await getRagStatus(state));
     if (!status.hasIndex) {
@@ -791,7 +807,7 @@ async function createCurrentStudioRagIndex(state, action, usage) {
     if (!status.schemaId || typeof status.chunksTotal !== "number") {
         throw new StudioHttpError(500, "Studio RAG status is missing schema metadata.");
     }
-    const store = await openStudioRagStore(state, config.dimensions);
+    const store = await openStudioRagStore(state, { dimensions: config.dimensions });
     return {
         config,
         status: {
@@ -809,22 +825,26 @@ async function createCurrentStudioRagIndex(state, action, usage) {
 }
 function resolveStudioRagStoreConfig(state) {
     const rt = getAskDbRuntimeConfig();
-    const kind = rt.structured.rag.store;
+    const kind = rt.rag.store;
     if (kind === "memory")
         return { kind };
     if (kind === "file") {
-        const basePath = rt.structured.rag.storeConfig.file?.basePath?.trim();
+        const basePath = rt.rag.storeConfig.file?.basePath?.trim();
         return { kind, basePath: basePath ? resolve(basePath) : join(state.schemaDir, "schema") };
     }
     const connectionString = pickFlat(rt.flat, "ASKDB_PGVECTOR_URL");
     return {
         kind,
         connectionString,
-        table: rt.structured.rag.storeConfig.pgvector?.table?.trim() || undefined,
+        table: rt.rag.storeConfig.pgvector?.table?.trim() || undefined,
         indexStrategy: pickFlat(rt.flat, "ASKDB_PGVECTOR_INDEX_STRATEGY"),
     };
 }
-async function openStudioRagStore(state, dimensions) {
+/**
+ * Opens the configured store. `provision` creates the pgvector table when it's missing, at
+ * `dimensions`; without it, a pgvector store only reads and writes an existing table.
+ */
+async function openStudioRagStore(state, { dimensions, provision = false }) {
     const config = resolveStudioRagStoreConfig(state);
     if (config.kind === "memory") {
         state.ragMemoryStore ??= createMemoryStore();
@@ -852,10 +872,19 @@ async function openStudioRagStore(state, dimensions) {
     const store = pgvectorFactory({
         connectionString: config.connectionString,
         table: config.table,
-        dimensions,
+        ...(dimensions !== undefined ? { dimensions } : {}),
         ...(config.indexStrategy ? { indexStrategy: config.indexStrategy } : {}),
     });
-    await store.ensureSchema();
+    if (provision) {
+        try {
+            await store.ensureSchema();
+        }
+        catch (error) {
+            // The caller never receives this store, so close its pool here (a width refusal has opened it).
+            await store.close().catch(() => { });
+            throw error;
+        }
+    }
     return {
         kind: "pgvector",
         store,
@@ -867,6 +896,10 @@ async function openStudioRagStore(state, dimensions) {
     };
 }
 async function countStudioRagStoreChunks(store, schemaId) {
+    // A pgvector table that doesn't exist yet holds nothing; the first index build creates it.
+    if (typeof store.store.tableDimensions === "function" && (await store.store.tableDimensions()) === undefined) {
+        return 0;
+    }
     if (typeof store.store.count === "function") {
         return store.store.count({ schemaId });
     }
@@ -879,33 +912,9 @@ function pickFlat(flat, key) {
     const value = flat[key];
     return typeof value === "string" && value.trim() !== "" ? value.trim() : undefined;
 }
-function pickEnv(env, key) {
-    const v = env[key];
-    return typeof v === "string" && v.trim() !== "" ? v.trim() : undefined;
-}
 function resolveStudioRagEmbedderConfig() {
-    const rt = getAskDbRuntimeConfig();
-    const base = rt.ai.aiEnv;
-    const explicitKind = pickEnv(base, "ASKDB_RAG_EMBEDDER");
-    const kind = explicitKind?.toLowerCase();
-    if (kind === "mock") {
-        return {
-            kind: "mock",
-            embedderId: STUDIO_RAG_MOCK_EMBEDDER_ID,
-            dimensions: STUDIO_RAG_MOCK_DIMENSIONS,
-            configured: true,
-            label: "Mock lexical",
-        };
-    }
-    if (kind !== undefined && kind !== "ai-sdk" && kind !== "openai") {
-        throw new StudioHttpError(400, `Unsupported Studio RAG embedder: ${kind}`);
-    }
-    const env = buildStudioRagEmbeddingEnv(kind, base);
-    const aiConfig = ai.resolveEmbeddingConfig(env, {
-        modelEnvVar: "ASKDB_RAG_EMBEDDER_MODEL",
-        modelDefault: DEFAULT_EMBEDDING_MODEL,
-    });
-    if (!aiConfig && kind === undefined) {
+    const embedding = getAskDbRuntimeConfig().ai.embedding;
+    if (!embedding) {
         return {
             kind: "mock",
             embedderId: STUDIO_RAG_MOCK_EMBEDDER_ID,
@@ -914,50 +923,38 @@ function resolveStudioRagEmbedderConfig() {
             label: "Mock lexical",
         };
     }
-    const provider = aiConfig?.provider ?? fallbackStudioRagProvider(kind, base);
-    const model = aiConfig?.model ?? DEFAULT_EMBEDDING_MODEL;
-    const dimensionOverride = readPositiveIntegerEnv(pickEnv(base, "ASKDB_RAG_EMBEDDER_DIMENSIONS"));
-    const dimensions = dimensionOverride ?? defaultEmbeddingDimensions(model);
+    // The env holds the embedding connection only, so no other provider's key can reach it (#345).
+    const aiConfig = ai.resolveEmbeddingConfig(embedding.env);
+    // The adapter's canonical name (`foundry` resolves to `azure`) keeps existing index ids stable.
+    const provider = aiConfig?.provider ?? embedding.provider;
+    const model = embedding.model;
+    const dimensions = embedding.dimensions;
     return {
         kind: "ai-sdk",
         provider,
-        embedderId: `ai-sdk:${provider}:${model}:${dimensions}`,
+        configuredProvider: embedding.provider,
+        connection: embedding.connection,
+        embedderId: aiSdkEmbedderId({ provider, model, dimensions }),
         dimensions,
         configured: Boolean(aiConfig),
         label: `AI SDK (${provider})`,
         model,
         baseUrl: aiConfig?.baseURL,
         aiConfig,
-        requestDimensions: dimensionOverride,
+        // Only an explicit ai.embedding.dimensions is sent; otherwise the model's own width applies.
+        requestDimensions: embedding.dimensions,
     };
 }
-function buildStudioRagEmbeddingEnv(kind, base) {
-    const apiKeyOverride = pickEnv(base, "ASKDB_RAG_EMBEDDER_API_KEY");
-    const baseUrlOverride = pickEnv(base, "ASKDB_RAG_EMBEDDER_BASE_URL");
-    return {
-        ...base,
-        ...(kind === "openai" ? { ASKDB_AI_PROVIDER: "openai" } : {}),
-        ...(apiKeyOverride ? { ASKDB_AI_API_KEY: apiKeyOverride } : {}),
-        ...(baseUrlOverride ? { ASKDB_AI_BASE_URL: baseUrlOverride } : {}),
-    };
-}
-function fallbackStudioRagProvider(kind, base) {
-    if (kind === "openai")
-        return "openai";
-    const raw = (pickEnv(base, "ASKDB_AI_PROVIDER") ?? "").toLowerCase();
-    return raw === "azure" || raw === "azure-openai" || raw === "foundry"
-        ? "azure"
-        : "openai";
-}
-function studioRagAiSdkKeyMissingMessage() {
-    return ("Studio RAG AI SDK embeddings require a configured AI provider key. " +
-        "Set ASKDB_AI_API_KEY or the provider-native key, or set ASKDB_RAG_EMBEDDER=mock for the local lexical embedder.");
+function studioRagAiSdkKeyMissingMessage(config) {
+    return (`Studio RAG embeddings need an API key on the ai.embedding connection ` +
+        `("${config.connection}" in ai.providerConfig.${config.configuredProvider}). ` +
+        'Set it in askdb.config.*, or set rag.embedder: "mock" for the local lexical embedder.');
 }
 async function createStudioRagEmbedder(config, usage) {
     if (config.kind === "mock")
         return createStudioMockEmbedder(config.dimensions);
     if (!config.aiConfig) {
-        throw new StudioHttpError(400, studioRagAiSdkKeyMissingMessage());
+        throw new StudioHttpError(400, studioRagAiSdkKeyMissingMessage(config));
     }
     const model = await ai.createEmbeddingModel(config.aiConfig, {
         dimensions: config.requestDimensions,
@@ -977,6 +974,9 @@ async function createStudioRagEmbedder(config, usage) {
 function formatStudioRagOperationError(error, config) {
     if (error instanceof StudioHttpError)
         return error;
+    if (error instanceof PgvectorDimensionMismatchError) {
+        return new StudioHttpError(409, studioRagWidthMismatchMessage(error, config));
+    }
     if (config.kind === "mock") {
         return new StudioHttpError(500, error instanceof Error ? error.message : String(error));
     }
@@ -985,7 +985,7 @@ function formatStudioRagOperationError(error, config) {
         return new StudioHttpError(500, error instanceof Error ? error.message : String(error));
     }
     const parts = [
-        `Studio RAG embedding request failed for provider ${config.provider}, model ${config.model}.`,
+        `Studio RAG embedding request failed for ai.embedding (provider ${config.configuredProvider}, connection ${config.connection}, model ${config.model}).`,
     ];
     if (config.baseUrl)
         parts.push(`Base URL: ${config.baseUrl}.`);
@@ -998,6 +998,20 @@ function formatStudioRagOperationError(error, config) {
         parts.push(`Error: ${error.message}`);
     return new StudioHttpError(502, parts.join(" "));
 }
+/** An existing pgvector table of another width: say where the new width comes from and how to resolve it, in config terms. */
+function studioRagWidthMismatchMessage(error, config) {
+    const mismatch = `pgvector table "${error.table}" holds ${error.tableDimensions}-dimension vectors, but`;
+    const rebuild = `Drop the table (DROP TABLE "${error.table}";) and build the index again`;
+    const newTable = "point rag.storeConfig.pgvector.table at a new table";
+    if (config.kind === "mock") {
+        return `${mismatch} the mock embedder writes ${error.dimensions}. ${rebuild}, or ${newTable}.`;
+    }
+    const source = config.dimensions !== undefined
+        ? `ai.embedding.dimensions asks for ${error.dimensions}`
+        : `embedding model ${config.model} returns ${error.dimensions}`;
+    return (`${mismatch} ${source}. ${rebuild}, ${newTable}, ` +
+        `or set ai.embedding.dimensions: ${error.tableDimensions} if model ${config.model} supports that width.`);
+}
 function findApiCallError(error) {
     const queue = [error];
     for (let i = 0; i < queue.length; i++) {
@@ -1032,30 +1046,20 @@ function clearIncompatibleRagStore(state, config) {
         return;
     const embeddingsJsonPath = `${store.basePath}.embeddings.json`;
     const embeddings = readOptionalJson(embeddingsJsonPath);
-    if (embeddings?.dimensions === undefined || embeddings.dimensions === config.dimensions)
+    if (embeddings?.dimensions === undefined)
+        return;
+    // The file store refuses mixed widths. When the width isn't known up front, vectors from a
+    // different embedder may not match it, so they go too.
+    const lockPath = join(state.schemaDir, "schema.lock.json");
+    const compatible = config.dimensions !== undefined
+        ? embeddings.dimensions === config.dimensions
+        : readOptionalJson(lockPath)?.embedderId === config.embedderId;
+    if (compatible)
         return;
-    for (const path of [
-        `${store.basePath}.embeddings.json`,
-        `${store.basePath}.embeddings.bin`,
-        join(state.schemaDir, "schema.lock.json"),
-    ]) {
+    for (const path of [`${store.basePath}.embeddings.json`, `${store.basePath}.embeddings.bin`, lockPath]) {
         rmSync(path, { force: true });
     }
 }
-function defaultEmbeddingDimensions(model) {
-    if (model === "text-embedding-3-large")
-        return 3072;
-    return 1536;
-}
-function readPositiveIntegerEnv(value) {
-    if (value === undefined || value.trim() === "")
-        return undefined;
-    const parsed = Number(value);
-    if (!Number.isInteger(parsed) || parsed <= 0) {
-        throw new StudioHttpError(400, `Invalid Studio RAG embedding dimensions: ${value}`);
-    }
-    return parsed;
-}
 function serializeRagResult(result) {
     return {
         id: result.id,