OmniRoute/open-sse/handlers/embeddings.ts

/**
 * Embedding Handler
 *
 * Handles POST /v1/embeddings requests.
 * Proxies to upstream embedding providers using OpenAI-compatible format.
 *
 * Request format (OpenAI-compatible):
 * {
 *   "model": "nebius/Qwen/Qwen3-Embedding-8B",
 *   "input": "text" | ["text1", "text2"],
 *   "dimensions": 4096,       // optional
 *   "encoding_format": "float" // optional
 * }
 */

import {
  getEmbeddingProvider,
  getEmbeddingModelDefaultParams,
  getEmbeddingModelModalities,
  parseEmbeddingModel,
  type EmbeddingModality,
  type EmbeddingProvider,
} from "../config/embeddingRegistry.ts";
import { saveCallLog } from "@/lib/usageDb";
import { createRequestLogger } from "../utils/requestLogger.ts";
import { isDetailedLoggingEnabled } from "@/lib/db/detailedLogs";
import { getCallLogPipelineCaptureStreamChunks } from "@/lib/logEnv";
import { toJsonErrorPayload } from "@/shared/utils/upstreamError";
import { stripStaleEncodingHeaders } from "../utils/upstreamResponseHeaders.ts";
import { sanitizeErrorMessage } from "../utils/error.ts";
import { stripTrailingSlashes } from "../utils/urlSanitize.ts";
import { fetchRemoteImage } from "@/shared/network/remoteImageFetch";
import {
  hasStructuredEmbeddingInput,
  prepareJinaMixedEmbeddingInput,
  prepareStructuredEmbeddingRequest,
} from "./embeddingStructuredInput.ts";
import { MAX_EMBEDDING_INLINE_ITEM_BYTES } from "@/shared/validation/schemas/apiV1";
import { markAccountUnavailable } from "../../src/sse/services/auth.ts";
import {
  collectJinaNativeModalities,
  isJinaNativeEmbeddingInput,
} from "@/shared/validation/jinaNativeEmbeddingInput";
import {
  collectGeminiNativeModalities,
  isGeminiEmbedding2Family,
  isGeminiNativeEmbeddingInput,
} from "@/shared/validation/geminiNativeEmbeddingInput";

interface ClientRawRequest {
  endpoint: string;
  body: Record<string, unknown>;
  headers: Record<string, string>;
}

/**
 * Flatten a single embedding item's vector to the OpenAI-spec `number[]` shape.
 *
 * Some OpenAI-compatible embedding backends — notably a llama.cpp
 * `llama-server --embedding --pooling ...` instance — return each vector wrapped in one
 * extra array level: `[[...floats]]` instead of `[...floats]` for a single input. That
 * extra level is silently spec-breaking, since a standard OpenAI-SDK consumer reading
 * `response.data[i].embedding` gets a length-1 array holding the real vector instead of
 * the vector itself. Unwrap only that single redundant level; vectors that are already
 * flat (or genuinely multi-row) are left untouched. See issue #9089.
 */
function flattenSingleRowEmbedding(item: unknown): void {
  if (!item || typeof item !== "object" || !("embedding" in item)) return;
  const record = item as { embedding: unknown };
  const embedding = record.embedding;
  if (
    Array.isArray(embedding) &&
    embedding.length === 1 &&
    Array.isArray(embedding[0]) &&
    typeof embedding[0][0] === "number"
  ) {
    record.embedding = embedding[0];
  }
}

/**
 * Handle embedding request.
 * Supports both hardcoded cloud providers and dynamic local provider_nodes.
 * When resolvedProvider is passed, uses it directly (injection pattern from route handler).
 * Falls back to hardcoded registry lookup for backward compatibility.
 */
export async function handleEmbedding({
  body,
  credentials,
  log,
  resolvedProvider = null,
  resolvedModel = null,
  clientRawRequest = null,
  apiKeyId = null,
  apiKeyName = null,
  connectionId = null,
}: {
  body: Record<string, unknown>;
  credentials: {
    apiKey?: string | null;
    accessToken?: string | null;
    providerSpecificData?: Record<string, unknown> | null;
  } | null;
  log?: { info: (...args: unknown[]) => void; error: (...args: unknown[]) => void };
  resolvedProvider?: EmbeddingProvider | null;
  resolvedModel?: string | null;
  clientRawRequest?: ClientRawRequest | null;
  apiKeyId?: string | null;
  apiKeyName?: string | null;
  connectionId?: string | null;
}) {
  // Use pre-resolved provider/model from route handler if available (supports dynamic provider_nodes).
  let provider: string | null;
  let model: string | null;
  let providerConfig: EmbeddingProvider | null;

  if (resolvedProvider) {
    provider = resolvedProvider.id;
    model = resolvedModel;
    providerConfig = resolvedProvider;
  } else {
    const parsed = parseEmbeddingModel(body.model as string);
    provider = parsed.provider;
    model = parsed.model;
    providerConfig = provider ? getEmbeddingProvider(provider) : null;
  }

  const startTime = Date.now();

  // Set up request logger for pipeline artifact capture
  const detailedLoggingEnabled = await isDetailedLoggingEnabled();
  const captureStreamChunks = getCallLogPipelineCaptureStreamChunks();
  const reqLogger = await createRequestLogger(
    provider || "openai",
    "openai",
    body.model as string,
    {
      enabled: detailedLoggingEnabled,
      captureStreamChunks,
      connectionId: connectionId || undefined,
      model: model || (body.model as string),
      provider: provider || undefined,
    }
  );

  // Log client raw request
  if (clientRawRequest) {
    reqLogger.logClientRawRequest(
      clientRawRequest.endpoint,
      clientRawRequest.body,
      clientRawRequest.headers
    );
  }

  // Summarized request body for call log (avoid storing large embedding input arrays)
  const logRequestBody = {
    model: body.model,
    input_count: Array.isArray(body.input) ? body.input.length : 1,
    dimensions: body.dimensions || undefined,
  };

  if (!provider) {
    return {
      success: false,
      status: 400,
      error: `Invalid embedding model: ${body.model}. Use format: provider/model`,
    };
  }

  if (!providerConfig) {
    return {
      success: false,
      status: 400,
      error: `Unknown embedding provider: ${provider}`,
    };
  }

  const structuredItems = Array.isArray(body.input)
    ? body.input.filter(
        (item): item is { type: EmbeddingModality } =>
          typeof item === "object" && item !== null && "type" in item
      )
    : [];
  const nativeModalities = [
    ...(isJinaNativeEmbeddingInput(body.input) ? collectJinaNativeModalities(body.input) : []),
    ...(isGeminiNativeEmbeddingInput(body.input) ? collectGeminiNativeModalities(body.input) : []),
  ].filter((modality) => modality !== "text");
  if (structuredItems.length > 0 || nativeModalities.length > 0) {
    const supportedModalities = getEmbeddingModelModalities(providerConfig, model);
    if (!supportedModalities) {
      return {
        success: false,
        status: 400,
        error: `Embedding model ${body.model} does not advertise structured embedding input support`,
      };
    }
    const unsupportedCanonical = structuredItems.find(
      (item) => !supportedModalities.includes(item.type)
    );
    if (unsupportedCanonical) {
      return {
        success: false,
        status: 400,
        error: `Embedding model ${body.model} does not support ${unsupportedCanonical.type} input`,
      };
    }
    const unsupportedNative = nativeModalities.find(
      (modality) => !supportedModalities.includes(modality)
    );
    if (unsupportedNative) {
      return {
        success: false,
        status: 400,
        error: `Embedding model ${body.model} does not support ${unsupportedNative} input`,
      };
    }
  }

  // Build upstream request — start with standard fields, then forward extra fields
  // the client sent (e.g. input_type, user, truncate for NVIDIA NIM asymmetric models).
  const KNOWN_FIELDS = new Set(["model", "input", "dimensions", "encoding_format"]);

  let upstreamBody: Record<string, unknown> = {
    model: model,
    input: body.input,
  };

  if (body.dimensions !== undefined) upstreamBody.dimensions = body.dimensions;
  if (body.encoding_format !== undefined) upstreamBody.encoding_format = body.encoding_format;

  for (const [key, value] of Object.entries(body)) {
    if (!KNOWN_FIELDS.has(key) && value !== undefined) {
      upstreamBody[key] = value;
    }
  }

  // Gemini embedding models (gemini-embedding-001 / -2-preview / text-embedding-004)
  // default to 3072-dim vectors. Clients targeting pgvector-style schemas typically
  // request a smaller size (e.g. 1536) via OpenAI's `dimensions` field, but Google's
  // OpenAI-compatibility shim at /v1beta/openai/embeddings does not document the
  // `dimensions` → `outputDimensionality` translation. Mirror the request value into
  // the Gemini-native `outputDimensionality` field so the upstream actually returns
  // the requested vector size. Ported from upstream decolua/9router#1366.
  if (provider === "gemini" && upstreamBody.outputDimensionality === undefined) {
    const outputDimensionality = Number(body.dimensions);
    if (Number.isFinite(outputDimensionality) && outputDimensionality > 0) {
      upstreamBody.outputDimensionality = outputDimensionality;
    }
  }

  // Inject model-level default params (e.g. NVIDIA NIM asymmetric models require
  // `input_type`) only for keys the client did not already supply, so a
  // client-sent value is never overwritten. Symmetric models carry no defaults
  // and are unaffected. See issue #1378.
  const defaultParams = getEmbeddingModelDefaultParams(providerConfig, model);
  if (defaultParams) {
    for (const [key, value] of Object.entries(defaultParams)) {
      if (upstreamBody[key] === undefined) {
        upstreamBody[key] = value;
      }
    }
  }

  let upstreamUrl = providerConfig.baseUrl;
  if (provider === "ollama-local" || provider === "lmstudio") {
    // Keyless local servers (#2824 ollama-local, #11233 lmstudio): honor the
    // configured connection's baseUrl when one was hydrated, and fall back to
    // the static localhost registry default otherwise.
    const configuredBaseUrl = credentials?.providerSpecificData?.baseUrl;
    const rawBaseUrl =
      typeof configuredBaseUrl === "string" && configuredBaseUrl.trim().length > 0
        ? configuredBaseUrl
        : providerConfig.baseUrl;
    // Use the shared O(n) helper instead of `/\/+$/` — that regex is
    // vulnerable to polynomial backtracking on adversarial input
    // (CodeQL js/polynomial-redos) since baseUrl is operator-configured
    // per-connection data. See open-sse/utils/urlSanitize.ts.
    const normalizedBaseUrl = stripTrailingSlashes(rawBaseUrl.trim());
    const localServerHost = normalizedBaseUrl
      .replace(/\/v1\/(?:chat\/completions|embeddings)$/i, "")
      .replace(/\/api\/chat$/i, "")
      .replace(/\/v1$/i, "");
    upstreamUrl = `${localServerHost}/v1/embeddings`;
  }
  let normalizeProviderResponse:
    ((data: Record<string, unknown>) => Record<string, unknown>) | null = null;

  // Build headers
  const headers: Record<string, string> = {
    "Content-Type": "application/json",
  };

  // Skip credential injection for local providers (authType: "none")
  const token =
    providerConfig.authType === "none" ? null : credentials?.apiKey || credentials?.accessToken;
  if (token) {
    if (providerConfig.authHeader === "bearer") {
      headers["Authorization"] = `Bearer ${token}`;
    } else if (providerConfig.authHeader === "x-api-key") {
      headers["x-api-key"] = token;
    }
  } else if (providerConfig.authType !== "none") {
    return {
      success: false,
      status: 401,
      error: `No valid authentication token for provider ${provider}. Check provider credentials.`,
    };
  }

  // Jina v5 Omni native docs ({ text }, { image: url|base64 }, { content: [...] })
  // must reach api.jina.ai unchanged. Do not fetch those image URLs or collapse
  // to string[]. Canonical { type, source } items still go through the translator.
  const jinaNative = isJinaNativeEmbeddingInput(body.input);
  const geminiNative = isGeminiNativeEmbeddingInput(body.input);
  const canonicalStructured = hasStructuredEmbeddingInput(body.input);
  const passThroughJinaNative =
    providerConfig.structuredInputProtocol === "jina-v1" && jinaNative && !canonicalStructured;
  // gemini-embedding-2 aggregates a string[] on Google's OpenAI shim into one
  // vector. Always use embedContent / batchEmbedContents so N input items
  // become N embeddings. Native multimodal parts take the same path.
  const useGeminiNativeTransport =
    providerConfig.structuredInputProtocol === "gemini-embed-content" &&
    (isGeminiEmbedding2Family(model) || canonicalStructured || geminiNative || jinaNative);

  if (providerConfig.structuredInputProtocol === "jina-v1" && jinaNative && canonicalStructured) {
    try {
      const mixed = Array.isArray(body.input) ? body.input : [body.input];
      upstreamBody.input = await prepareJinaMixedEmbeddingInput(mixed, async (url) => {
        const result = await fetchRemoteImage(url, {
          guard: "public-only",
          maxBytes: MAX_EMBEDDING_INLINE_ITEM_BYTES,
          pinDns: true,
        });
        return { buffer: result.buffer, contentType: result.contentType || null };
      });
    } catch (error) {
      return { success: false, status: 400, error: sanitizeErrorMessage(error) };
    }
  } else if (useGeminiNativeTransport || (!passThroughJinaNative && canonicalStructured)) {
    if (!model) {
      return {
        success: false,
        status: 400,
        error: `Invalid embedding model: ${body.model}. Use format: provider/model`,
      };
    }
    try {
      const prepared = await prepareStructuredEmbeddingRequest(
        providerConfig,
        model,
        body,
        token ?? "",
        {
          fetchMedia: async (url) => {
            const result = await fetchRemoteImage(url, {
              guard: "public-only",
              maxBytes: MAX_EMBEDDING_INLINE_ITEM_BYTES,
              pinDns: true,
            });
            return { buffer: result.buffer, contentType: result.contentType || null };
          },
        }
      );
      upstreamBody = prepared.body;
      upstreamUrl = prepared.url;
      normalizeProviderResponse = prepared.normalizeResponse ?? null;
      if (prepared.authHeader) {
        delete headers.Authorization;
        delete headers["x-api-key"];
        headers[prepared.authHeader.name] = prepared.authHeader.value;
      }
    } catch (error) {
      return { success: false, status: 400, error: sanitizeErrorMessage(error) };
    }
  }

  if (log) {
    log.info(
      "EMBED",
      `${provider}/${model} | input: ${Array.isArray(body.input) ? body.input.length + " items" : "1 item"}`
    );
  }

  try {
    // Quota share enforcement (fail-open: errors allow the request through)
    if (apiKeyId && connectionId && provider) {
      try {
        const { enforceQuotaShare } = await import("@/lib/quota/enforce");
        const quotaDecision = await enforceQuotaShare({
          apiKeyId,
          connectionId,
          provider,
          // Per-(key,model) cap — resolved embedding model id (same scope used in logs/routing).
          model: model || undefined,
        });
        if (quotaDecision.kind === "block") {
          return {
            success: false,
            status: quotaDecision.httpStatus ?? 429,
            error: quotaDecision.reason || "Quota share limit reached",
          };
        }
      } catch {
        // fail-open per B16
      }
    }

    // Log provider request
    reqLogger.logTargetRequest(upstreamUrl, headers, upstreamBody);

    const response = await fetch(upstreamUrl, {
      method: "POST",
      headers,
      body: JSON.stringify(upstreamBody),
    });

    if (!response.ok) {
      const errorText = await response.text();
      if (log) {
        log.error("EMBED", `${provider} error ${response.status}: ${errorText.slice(0, 200)}`);
      }

      // Log provider response
      reqLogger.logProviderResponse(response.status, "", response.headers, errorText.slice(0, 500));

      // Build client error response
      const clientErrorBody = toJsonErrorPayload(
        errorText.slice(0, 500),
        "Embedding provider error"
      );
      reqLogger.logConvertedResponse(clientErrorBody);

      const pipelinePayloads = detailedLoggingEnabled ? reqLogger.getPipelinePayloads() : null;

      // Save error call log for Logger panel
      saveCallLog({
        method: "POST",
        path: "/v1/embeddings",
        status: response.status,
        model: `${provider}/${model}`,
        provider,
        duration: Date.now() - startTime,
        error: errorText.slice(0, 500),
        requestBody: logRequestBody,
        pipelinePayloads,
        apiKeyId,
        apiKeyName,
        connectionId,
      }).catch(() => {});

      // #10347 — persist a connection-level failure marker on a hard upstream failure so
      // the dead account is not re-selected and re-hit on the next embed request (chat
      // parity). markAccountUnavailable classifies the status via checkFallbackError: a
      // payment-required 402 becomes the TERMINAL state credits_exhausted (the terminal
      // marker excludes the account from selection until an operator resets it), benign
      // 4xx are a no-op, and terminal statuses are never overwritten. honors per-connection
      // disableCooling. The write must never break the error response path, so it is
      // best-effort.
      if (connectionId) {
        try {
          await markAccountUnavailable(connectionId, response.status, errorText, provider, model);
        } catch {
          // swallow — the upstream error response takes priority
        }
      }

      return {
        success: false,
        status: response.status,
        error: errorText,
        headers: stripStaleEncodingHeaders(response.headers),
      };
    }

    const rawData = (await response.json()) as Record<string, unknown>;
    const data = (normalizeProviderResponse ? normalizeProviderResponse(rawData) : rawData) as {
      data?: unknown[] | unknown;
      usage?: { prompt_tokens?: number; total_tokens?: number };
    };

    // Log provider response
    reqLogger.logProviderResponse(response.status, "", response.headers, data);

    // OpenAI-spec compliance (#9089): each item's `embedding` must be a flat number[].
    // Some OpenAI-compatible backends (e.g. a llama.cpp `llama-server --embedding`
    // instance) return the vector wrapped in one extra array level — `[[...floats]]`
    // instead of `[...floats]` — for a single input, which silently breaks any standard
    // OpenAI-SDK consumer doing `response.data[i].embedding`. Flatten that one redundant
    // level without touching providers that already return flat vectors.
    const responseItems = data.data || data;
    if (Array.isArray(responseItems)) {
      for (const item of responseItems) {
        flattenSingleRowEmbedding(item);
      }
    }

    // Normalize response to OpenAI format
    const normalizedResponse = {
      object: "list",
      data: data.data || data,
      model: `${provider}/${model}`,
      usage: data.usage || { prompt_tokens: 0, total_tokens: 0 },
    };

    // Log client response
    reqLogger.logConvertedResponse(normalizedResponse);

    const pipelinePayloads = detailedLoggingEnabled ? reqLogger.getPipelinePayloads() : null;

    // Save success call log for Logger panel
    // Embeddings only have input tokens (prompt_tokens + total_tokens), no output/completion tokens
    saveCallLog({
      method: "POST",
      path: "/v1/embeddings",
      status: 200,
      model: `${provider}/${model}`,
      provider,
      duration: Date.now() - startTime,
      tokens: {
        prompt_tokens: data.usage?.prompt_tokens || data.usage?.total_tokens || 0,
        completion_tokens: 0,
      },
      requestBody: logRequestBody,
      responseBody: {
        usage: data.usage || null,
        object: "list",
        data_count: Array.isArray(data.data) ? data.data.length : 0,
      },
      pipelinePayloads,
      apiKeyId,
      apiKeyName,
      connectionId,
    }).catch(() => {});

    // Record quota consumption (fire-and-forget, never blocks)
    if (apiKeyId && connectionId && provider) {
      try {
        const { scheduleRecordConsumption } = await import("@/lib/quota/spendRecorder");
        scheduleRecordConsumption({
          apiKeyId,
          connectionId,
          provider,
          // Per-(key,model) cap accounting — same resolved model id used at enforce time.
          model: model || undefined,
          cost: {
            tokens: data.usage?.prompt_tokens || data.usage?.total_tokens || 0,
            requests: 1,
          },
        });
      } catch {
        // fail-open per B29
      }
    }

    return {
      success: true,
      data: normalizedResponse,
      headers: stripStaleEncodingHeaders(response.headers),
    };
  } catch (err) {
    if (log) {
      log.error("EMBED", `${provider} fetch error: ${err.message}`);
    }

    // Log error
    reqLogger.logError(err, upstreamBody);

    const pipelinePayloads = detailedLoggingEnabled ? reqLogger.getPipelinePayloads() : null;

    // Save exception call log for Logger panel
    saveCallLog({
      method: "POST",
      path: "/v1/embeddings",
      status: 502,
      model: `${provider}/${model}`,
      provider,
      duration: Date.now() - startTime,
      error: err.message,
      requestBody: logRequestBody,
      pipelinePayloads,
      apiKeyId,
      apiKeyName,
      connectionId,
    }).catch(() => {});

    return {
      success: false,
      status: 502,
      error: `Embedding provider error: ${sanitizeErrorMessage(err.message)}`,
    };
  }
}