@axiom-lattice/gateway 8.0.0 → 8.0.2

package/dist/index.mjs

@@ -16986,6 +16986,7 @@ function registerA2ABridgeRoutes(app2) {
 
 // src/routes/stt.ts
 import { sttModelLatticeManager } from "@axiom-lattice/core";
+import { ASR_DEFAULT_KEY } from "@axiom-lattice/protocols";
 function tenantIdOf(request) {
   try {
     return resolveTenantScope(request, false).tenantId;
@@ -17046,12 +17047,19 @@ var registerSTTRoutes = (app2) => {
         return reply.status(400).send({ error: `Unsupported audio format: ${format}` });
       }
       try {
+        const requestedKey = modelKey2 || ASR_DEFAULT_KEY;
         const client = await sttModelLatticeManager.getSTTClientWithTenantAsync(
           tenantIdOf(request),
-          modelKey2
+          requestedKey
+        ) ?? (requestedKey === ASR_DEFAULT_KEY ? void 0 : await sttModelLatticeManager.getSTTClientWithTenantAsync(
+          tenantIdOf(request),
+          ASR_DEFAULT_KEY
+        )) ?? await sttModelLatticeManager.getSTTClientWithTenantAsync(
+          tenantIdOf(request),
+          "default"
         );
         if (!client) {
-          return reply.status(404).send({ error: `STT model "${modelKey2}" not found` });
+          return reply.status(404).send({ error: `STT model "${requestedKey}" not found` });
         }
         const result = await client.transcribe(buffer, format);
         return reply.send({
@@ -21671,8 +21679,12 @@ import { randomUUID as randomUUID23 } from "crypto";
 import { toModelKey } from "@axiom-lattice/protocols";
 import {
   assertSafeUpstreamUrl,
+  buildModelEmbeddingsConfig,
   buildModelLlmConfig,
   DEFAULT_MODEL_EXTRA,
+  embeddingProbeCandidates,
+  resolveEmbeddingEndpoint,
+  resolveEmbeddingShape,
   UnsupportedModelProviderError,
   UnsafeUpstreamUrlError
 } from "@axiom-lattice/core";
@@ -21844,7 +21856,6 @@ var openAiCompatibleAdapter = {
 };
 async function probeEmbeddingInference(input, signal) {
   const startedAt = Date.now();
-  const endpoint = buildOpenAiEndpoint(input.baseURL, "embeddings");
   const headers = {
     Accept: "application/json",
     "Content-Type": "application/json"
@@ -21852,17 +21863,25 @@ async function probeEmbeddingInference(input, signal) {
   if (input.apiKey) {
     headers.Authorization = `Bearer ${input.apiKey}`;
   }
+  const body = input.shape === "contents" ? JSON.stringify({
+    model: input.modelId,
+    input: { contents: [{ text: "ping" }] }
+  }) : input.shape === "texts" ? JSON.stringify({
+    model: input.modelId,
+    input: { texts: ["ping"] }
+  }) : JSON.stringify({ model: input.modelId, input: "ping" });
+  let status;
   try {
     let response;
     try {
-      response = await fetch(endpoint, {
+      response = await fetch(input.endpoint, {
         method: "POST",
         headers,
         // Never follow redirects: a redirected request could reach an address
         // that bypassed the SSRF checks.
         redirect: "manual",
         signal,
-        body: JSON.stringify({ model: input.modelId, input: "ping" })
+        body
       });
     } catch (error) {
       if (error instanceof Error && error.name === "AbortError") {
@@ -21880,11 +21899,12 @@ async function probeEmbeddingInference(input, signal) {
         true
       );
     }
+    status = response.status;
     assertUpstreamOk(response);
-    const body = await readBodyWithLimit(response, MAX_RESPONSE_BYTES);
+    const text = await readBodyWithLimit(response, MAX_RESPONSE_BYTES);
     let payload;
     try {
-      payload = JSON.parse(body);
+      payload = JSON.parse(text);
     } catch {
       throw new ModelProviderServiceError(
         "UPSTREAM_INVALID_RESPONSE",
@@ -21895,14 +21915,16 @@ async function probeEmbeddingInference(input, signal) {
     return {
       reachable: true,
       latencyMs: Date.now() - startedAt,
-      dimensions: readEmbeddingDimensions(payload)
+      dimensions: input.shape === "openai" ? readEmbeddingDimensions(payload) : readServiceDimensions(payload),
+      status
     };
   } catch (error) {
     if (error instanceof ModelProviderServiceError && isUpstreamErrorCode(error.code)) {
       return {
         reachable: false,
         latencyMs: Date.now() - startedAt,
-        error: error.code
+        error: error.code,
+        ...status === void 0 ? {} : { status }
       };
     }
     throw error;
@@ -21923,6 +21945,25 @@ function readEmbeddingDimensions(payload) {
   const embedding = first.embedding;
   return Array.isArray(embedding) ? embedding.length : void 0;
 }
+function readServiceDimensions(payload) {
+  if (typeof payload !== "object" || payload === null) {
+    return void 0;
+  }
+  const output = payload.output;
+  if (typeof output !== "object" || output === null) {
+    return void 0;
+  }
+  const embeddings = output.embeddings;
+  if (!Array.isArray(embeddings) || embeddings.length === 0) {
+    return void 0;
+  }
+  const first = embeddings[0];
+  if (typeof first !== "object" || first === null) {
+    return void 0;
+  }
+  const embedding = first.embedding;
+  return Array.isArray(embedding) ? embedding.length : void 0;
+}
 async function readBodyWithLimit(response, maxBytes) {
   if (!response.body) {
     return response.text();
@@ -22369,14 +22410,57 @@ var ModelProviderService = class {
         );
       }
       try {
-        const outcome = await probeEmbeddingInference(
-          { baseURL: input.baseURL, apiKey, modelId: input.modelId },
+        const provider = input.providerId ? await this.deps.providerStore.getProviderById(
+          tenantId,
+          input.providerId
+        ) : null;
+        const modelRow = provider && input.providerId ? await this.deps.providerStore.getModel(
+          tenantId,
+          input.providerId,
+          input.modelId
+        ) : null;
+        const extra = modelRow?.extra ?? {};
+        const explicitShape = extra.embeddingShape === "openai" || extra.embeddingShape === "contents" || extra.embeddingShape === "texts";
+        const baseURL = provider?.baseURL ?? input.baseURL;
+        const initialShape = provider && modelRow ? buildModelEmbeddingsConfig(provider, modelRow).shape : resolveEmbeddingShape(input.modelId, {});
+        const attempt = (shape) => probeEmbeddingInference(
+          {
+            endpoint: resolveEmbeddingEndpoint(baseURL, shape, extra),
+            apiKey,
+            modelId: input.modelId,
+            shape
+          },
           controller.signal
         );
+        const candidates = explicitShape ? [initialShape] : embeddingProbeCandidates(initialShape, input.modelId);
+        let outcome = await attempt(candidates[0]);
+        let usedShape = candidates[0];
+        for (const candidate of candidates.slice(1)) {
+          if (outcome.reachable || outcome.status !== 400 && outcome.status !== 404) {
+            break;
+          }
+          const next = await attempt(candidate);
+          if (!next.reachable) {
+            continue;
+          }
+          outcome = next;
+          usedShape = candidate;
+          if (provider && modelRow) {
+            await this.persistEmbeddingShape(
+              tenantId,
+              provider,
+              modelRow,
+              candidate,
+              extra
+            );
+          }
+          break;
+        }
         return {
           reachable: outcome.reachable,
           latencyMs: outcome.latencyMs,
           models: [],
+          shape: usedShape,
           ...outcome.dimensions === void 0 ? {} : { dimensions: outcome.dimensions },
           ...outcome.error ? { error: outcome.error } : {}
         };
@@ -22477,48 +22561,94 @@ var ModelProviderService = class {
     return this.toPublicEntry(provider, models.length);
   }
   /**
-   * Give a tenant's sole **chat** model the default flag.
+   * Give each kind's sole model its default flag.
    *
-   * A tenant with exactly one runnable chat model has an obvious default, and
-   * saying so is what keeps the invariant the rest of the feature leans on: a
-   * tenant either has no chat model at all — and falls back to the deployment's
-   * default — or has a default, which cannot be deleted. Without this a tenant
-   * could sit with one model and no default, and deleting that model would
-   * leave nothing behind.
+   * A tenant with exactly one runnable model **of a kind** has an obvious
+   * default for that kind, and saying so keeps the invariant the rest of the
+   * feature leans on: a kind either has no model at all — and falls back to the
+   * deployment's default — or has a default, which cannot be deleted.
    *
-   * Chat only: a bare `"default"` key resolves to a chat client, so an
-   * embedding or ASR row must never be marked as one.
+   * Per kind because each kind publishes its own alias key (`default`,
+   * `default:embedding`, `default:asr`).
    *
-   * Deliberately narrow: more than one model is the tenant's choice to make,
-   * and a disabled model is not a candidate for "the model this tenant runs".
+   * Deliberately narrow: more than one model of a kind is the tenant's choice to
+   * make, and a disabled model is not a candidate.
    *
    * @param tenantId - The tenant whose model set just changed
    */
   async ensureDefaultModel(tenantId) {
     const models = await this.deps.providerStore.listModels(tenantId);
-    const chatModels = models.filter((model) => model.kind === "chat");
-    if (chatModels.length !== 1) {
+    const toMark = ["chat", "embedding", "asr"].flatMap((kind) => {
+      const ofKind = models.filter((model) => model.kind === kind);
+      if (ofKind.length !== 1) {
+        return [];
+      }
+      const [only] = ofKind;
+      return only.enabled && !only.isDefault ? [only] : [];
+    });
+    if (toMark.length === 0) {
       return;
     }
-    const [only] = chatModels;
-    if (only.isDefault || !only.enabled) {
-      return;
+    const byProvider = /* @__PURE__ */ new Map();
+    for (const model of models) {
+      const list = byProvider.get(model.providerId) ?? [];
+      list.push(model);
+      byProvider.set(model.providerId, list);
     }
-    const siblings = models.filter(
-      (model) => model.providerId === only.providerId
+    for (const [providerId2, providerModels] of byProvider) {
+      if (!providerModels.some((model) => toMark.includes(model))) {
+        continue;
+      }
+      await this.deps.providerStore.replaceProviderModels(
+        tenantId,
+        providerId2,
+        providerModels.map((model) => ({
+          id: model.modelId,
+          enabled: model.enabled,
+          stale: model.stale,
+          extra: model.extra,
+          kind: model.kind,
+          isDefault: toMark.includes(model) ? true : model.isDefault
+        }))
+      );
+    }
+  }
+  /**
+   * Persist the embedding shape a probe detected.
+   *
+   * A probe is read-only except for this: when the initial shape fails with
+   * 400/404 and another shape works, writing it down is what makes the
+   * runtime use it. The whole provider set is written because
+   * `replaceProviderModels` replaces it — passing one row would delete its
+   * siblings.
+   */
+  async persistEmbeddingShape(tenantId, provider, model, shape, extra) {
+    const rows = await this.deps.providerStore.listModels(
+      tenantId,
+      provider.id
     );
+    const nextExtra = { ...extra, embeddingShape: shape };
+    if (shape === "openai" && typeof extra.embeddingPath !== "string") {
+      delete nextExtra.embeddingPath;
+    }
     await this.deps.providerStore.replaceProviderModels(
       tenantId,
-      only.providerId,
-      siblings.map((model) => ({
-        id: model.modelId,
-        enabled: model.enabled,
-        stale: model.stale,
-        extra: model.extra,
-        kind: model.kind,
-        isDefault: model.modelId === only.modelId
+      provider.id,
+      rows.map((row) => ({
+        id: row.modelId,
+        enabled: row.enabled,
+        stale: row.stale,
+        kind: row.kind,
+        isDefault: row.isDefault,
+        extra: row.modelId === model.modelId ? nextExtra : row.extra
       }))
     );
+    const updated = await this.deps.providerStore.listModels(
+      tenantId,
+      provider.id
+    );
+    this.syncRegistry(tenantId, provider, updated);
+    this.notifyChanged(tenantId, provider.id);
   }
   syncRegistry(tenantId, provider, models) {
     this.deps.registry.evictProviderCache?.(tenantId, provider.id);