unrag 0.2.5 → 0.2.7

package/registry/core/ingest.ts

@@ -9,37 +9,13 @@ import type {
   IngestInput,
   IngestResult,
   IngestWarning,
+  Metadata,
   ResolvedContextEngineConfig,
 } from "./types";
+import { mergeDeep } from "./deep-merge";
 
 const now = () => performance.now();
 
-const mergeDeep = <T extends Record<string, any>>(
-  base: T,
-  overrides: any | undefined
-): T => {
-  if (!overrides) return base;
-  const out: any = Array.isArray(base) ? [...base] : { ...base };
-  for (const key of Object.keys(overrides)) {
-    const nextVal = overrides[key];
-    if (nextVal === undefined) continue;
-    const baseVal = (base as any)[key];
-    if (
-      baseVal &&
-      typeof baseVal === "object" &&
-      !Array.isArray(baseVal) &&
-      nextVal &&
-      typeof nextVal === "object" &&
-      !Array.isArray(nextVal)
-    ) {
-      out[key] = mergeDeep(baseVal, nextVal);
-    } else {
-      out[key] = nextVal;
-    }
-  }
-  return out as T;
-};
-
 const asMessage = (err: unknown) => {
   if (err instanceof Error) return err.message;
   try {
@@ -123,7 +99,7 @@ export const ingest = async (
 
   const assets: AssetInput[] = Array.isArray(input.assets) ? input.assets : [];
   type PreparedChunkSpec = Omit<Chunk, "id" | "index"> & {
-    metadata: Record<string, any>;
+    metadata: Metadata;
     embed:
       | { kind: "text"; text: string }
       | { kind: "image"; data: Uint8Array | string; mediaType?: string; assetId?: string };
@@ -140,7 +116,7 @@ export const ingest = async (
 
   const runExtractors = async (args: {
     asset: AssetInput;
-    assetMeta: Record<string, any>;
+    assetMeta: Metadata;
     assetUri?: string;
     assetMediaType?: string;
     extractors: AssetExtractor[];
@@ -528,14 +504,41 @@ export const ingest = async (
   const chunkingMs = now() - chunkingStart;
   const embeddingStart = now();
 
-  const embeddedChunks = await Promise.all(
-    prepared.map(async ({ chunk, embed }) => {
-      if (embed.kind === "image") {
-        const embedImage = config.embedding.embedImage;
-        if (!embedImage) {
-          throw new Error("Image embedding requested but provider does not support embedImage()");
-        }
-        const embedding = await embedImage({
+  const embeddedChunks: Chunk[] = new Array(prepared.length);
+
+  const textSpecs: Array<{
+    idx: number;
+    chunk: Chunk;
+    input: {
+      text: string;
+      metadata: Metadata;
+      position: number;
+      sourceId: string;
+      documentId: string;
+    };
+  }> = [];
+
+  const imageSpecs: Array<{
+    idx: number;
+    chunk: Chunk;
+    input: {
+      data: Uint8Array | string;
+      mediaType?: string;
+      metadata: Metadata;
+      position: number;
+      sourceId: string;
+      documentId: string;
+      assetId?: string;
+    };
+  }> = [];
+
+  for (let i = 0; i < prepared.length; i++) {
+    const { chunk, embed } = prepared[i]!;
+    if (embed.kind === "image") {
+      imageSpecs.push({
+        idx: i,
+        chunk,
+        input: {
           data: embed.data,
           mediaType: embed.mediaType,
           metadata: chunk.metadata,
@@ -543,21 +546,91 @@ export const ingest = async (
           sourceId: chunk.sourceId,
           documentId: chunk.documentId,
           assetId: embed.assetId,
-        });
-        return { ...chunk, embedding };
-      }
+        },
+      });
+      continue;
+    }
 
-      const embedding = await config.embedding.embed({
+    textSpecs.push({
+      idx: i,
+      chunk,
+      input: {
         text: embed.text,
         metadata: chunk.metadata,
         position: chunk.index,
         sourceId: chunk.sourceId,
         documentId: chunk.documentId,
-      });
+      },
+    });
+  }
 
-      return { ...chunk, embedding };
-    })
-  );
+  const concurrency = config.embeddingProcessing.concurrency;
+
+  // Text embeddings (prefer batch when supported).
+  if (textSpecs.length > 0) {
+    const embedMany = config.embedding.embedMany;
+    if (embedMany) {
+      const batchSize = Math.max(1, Math.floor(config.embeddingProcessing.batchSize || 1));
+      const batches: Array<typeof textSpecs> = [];
+      for (let i = 0; i < textSpecs.length; i += batchSize) {
+        batches.push(textSpecs.slice(i, i + batchSize));
+      }
+
+      const batchEmbeddings = await mapWithConcurrency(
+        batches,
+        concurrency,
+        async (batch) => {
+          const embeddings = await embedMany(batch.map((b) => b.input));
+          if (!Array.isArray(embeddings) || embeddings.length !== batch.length) {
+            throw new Error(
+              `embedMany() returned ${Array.isArray(embeddings) ? embeddings.length : "non-array"} embeddings for a batch of ${batch.length}`
+            );
+          }
+          return embeddings;
+        }
+      );
+
+      let batchIdx = 0;
+      for (const batch of batches) {
+        const embeddings = batchEmbeddings[batchIdx++]!;
+        for (let i = 0; i < batch.length; i++) {
+          const spec = batch[i]!;
+          embeddedChunks[spec.idx] = { ...spec.chunk, embedding: embeddings[i]! };
+        }
+      }
+    } else {
+      const embeddings = await mapWithConcurrency(textSpecs, concurrency, async (spec) =>
+        config.embedding.embed(spec.input)
+      );
+      for (let i = 0; i < textSpecs.length; i++) {
+        const spec = textSpecs[i]!;
+        embeddedChunks[spec.idx] = { ...spec.chunk, embedding: embeddings[i]! };
+      }
+    }
+  }
+
+  // Image embeddings (bounded concurrency).
+  if (imageSpecs.length > 0) {
+    const embedImage = config.embedding.embedImage;
+    if (!embedImage) {
+      throw new Error("Image embedding requested but provider does not support embedImage()");
+    }
+
+    const embeddings = await mapWithConcurrency(imageSpecs, concurrency, async (spec) =>
+      embedImage(spec.input)
+    );
+    for (let i = 0; i < imageSpecs.length; i++) {
+      const spec = imageSpecs[i]!;
+      embeddedChunks[spec.idx] = { ...spec.chunk, embedding: embeddings[i]! };
+    }
+  }
+
+  // Safety check: ensure all chunks got an embedding.
+  for (let i = 0; i < embeddedChunks.length; i++) {
+    if (!embeddedChunks[i]) {
+      throw new Error("Internal error: missing embedding for one or more chunks");
+    }
+  }
 
   const embeddingMs = now() - embeddingStart;
   const storageStart = now();

package/registry/core/types.ts

@@ -5,6 +5,30 @@ export type Metadata = Record<
   MetadataValue | MetadataValue[] | undefined
 >;
 
+/**
+ * Standard fields for asset-related metadata.
+ * These are added to chunk metadata when chunks are derived from assets.
+ */
+export interface AssetMetadataFields {
+  assetKind?: "image" | "pdf" | "audio" | "video" | "file";
+  assetId?: string;
+  assetUri?: string;
+  assetMediaType?: string;
+  extractor?: string;
+}
+
+/**
+ * Type guard for checking if metadata contains required asset fields.
+ */
+export function hasAssetMetadata(
+  metadata: Metadata
+): metadata is Metadata & Required<Pick<AssetMetadataFields, "assetKind" | "assetId">> {
+  return (
+    typeof metadata.assetKind === "string" &&
+    typeof metadata.assetId === "string"
+  );
+}
+
 export type Chunk = {
   id: string;
   documentId: string;
@@ -31,6 +55,24 @@ export type ContentStorageConfig = {
   storeDocumentContent: boolean;
 };
 
+/**
+ * Controls performance characteristics of embedding during ingest.
+ *
+ * These defaults are intentionally conservative to reduce rate-limit risk.
+ */
+export type EmbeddingProcessingConfig = {
+  /**
+   * Maximum number of concurrent embedding requests.
+   * This applies to both text embedding (embed/embedMany) and image embedding (embedImage).
+   */
+  concurrency: number;
+  /**
+   * Max number of text chunks per embedMany batch (when embedMany is supported).
+   * Ignored when the provider does not implement embedMany().
+   */
+  batchSize: number;
+};
+
 export type ChunkText = {
   index: number;
   content: string;
@@ -655,6 +697,11 @@ export type RetrieveResult = {
  */
 export type UnragDefaultsConfig = {
   chunking?: Partial<ChunkingOptions>;
+  /**
+   * Embedding performance defaults (batching + concurrency).
+   * These map to the engine's `embeddingProcessing` config.
+   */
+  embedding?: Partial<EmbeddingProcessingConfig>;
   retrieval?: {
     topK?: number;
   };
@@ -670,6 +717,50 @@ export type UnragEmbeddingConfig =
       provider: "ai";
       config?: import("../embedding/ai").AiEmbeddingConfig;
     }
+  | {
+      provider: "openai";
+      config?: import("../embedding/openai").OpenAiEmbeddingConfig;
+    }
+  | {
+      provider: "google";
+      config?: import("../embedding/google").GoogleEmbeddingConfig;
+    }
+  | {
+      provider: "openrouter";
+      config?: import("../embedding/openrouter").OpenRouterEmbeddingConfig;
+    }
+  | {
+      provider: "azure";
+      config?: import("../embedding/azure").AzureEmbeddingConfig;
+    }
+  | {
+      provider: "vertex";
+      config?: import("../embedding/vertex").VertexEmbeddingConfig;
+    }
+  | {
+      provider: "bedrock";
+      config?: import("../embedding/bedrock").BedrockEmbeddingConfig;
+    }
+  | {
+      provider: "cohere";
+      config?: import("../embedding/cohere").CohereEmbeddingConfig;
+    }
+  | {
+      provider: "mistral";
+      config?: import("../embedding/mistral").MistralEmbeddingConfig;
+    }
+  | {
+      provider: "together";
+      config?: import("../embedding/together").TogetherEmbeddingConfig;
+    }
+  | {
+      provider: "ollama";
+      config?: import("../embedding/ollama").OllamaEmbeddingConfig;
+    }
+  | {
+      provider: "voyage";
+      config?: import("../embedding/voyage").VoyageEmbeddingConfig;
+    }
   | {
       provider: "custom";
       /**
@@ -724,6 +815,10 @@ export type ContextEngineConfig = {
    * captions, which can still be ingested via `assets[].text` if you choose).
    */
   assetProcessing?: DeepPartial<AssetProcessingConfig>;
+  /**
+   * Embedding performance defaults for ingest (batching + concurrency).
+   */
+  embeddingProcessing?: DeepPartial<EmbeddingProcessingConfig>;
 };
 
 export type ResolvedContextEngineConfig = {
@@ -735,6 +830,5 @@ export type ResolvedContextEngineConfig = {
   extractors: AssetExtractor[];
   storage: ContentStorageConfig;
   assetProcessing: AssetProcessingConfig;
+  embeddingProcessing: EmbeddingProcessingConfig;
 };
-
-

package/registry/docs/unrag.md

@@ -8,9 +8,14 @@ Unrag installs a small RAG module into your codebase with:
 
 Add these to your environment:
 - `DATABASE_URL` (Postgres connection string)
-- `AI_GATEWAY_API_KEY` (required by the `ai` SDK when using Vercel AI Gateway)
+- (Embedding) set the environment variables required by your selected provider.
+
+If you used the default provider (Vercel AI Gateway):
+- `AI_GATEWAY_API_KEY`
 - Optional: `AI_GATEWAY_MODEL` (defaults to `openai/text-embedding-3-small`)
 
+If you picked a different provider (OpenAI / Google / Voyage / etc.), see the installed provider docs under your Unrag docs site (`/docs/providers/*`).
+
 ## Database requirements
 
 Enable pgvector:

package/registry/embedding/_shared.ts

@@ -0,0 +1,25 @@
+import { createRequire } from "node:module";
+
+const require = createRequire(import.meta.url);
+
+/**
+ * Dynamically require an optional dependency with type-safe return.
+ *
+ * @template T - The expected module type (callers must define this)
+ */
+export function requireOptional<T>(args: {
+  id: string;
+  installHint: string;
+  providerName: string;
+}): T {
+  try {
+    return require(args.id) as T;
+  } catch {
+    throw new Error(
+      `Unrag embedding provider "${args.providerName}" requires "${args.id}" to be installed.\n` +
+        `Install it with: ${args.installHint}`
+    );
+  }
+}
+
+

package/registry/embedding/ai.ts

@@ -1,5 +1,5 @@
 import { embed, embedMany } from "ai";
-import type { EmbeddingProvider, ImageEmbeddingInput } from "../core/types";
+import type { EmbeddingProvider } from "../core/types";
 
 type BaseConfig = {
   /**
@@ -9,52 +9,18 @@ type BaseConfig = {
   timeoutMs?: number;
 };
 
-export type AiEmbeddingConfig =
-  | (BaseConfig & {
-      /**
-       * Defaults to "text" for backwards compatibility.
-       * - "text": only supports embedding strings
-       * - "multimodal": additionally enables `embedImage()` for image assets (when the model supports it)
-       */
-      type?: "text";
-    })
-  | (BaseConfig & {
-      type: "multimodal";
-      /**
-       * Controls how images are translated into AI SDK `embed()` values.
-       * Different providers use different shapes; this is the escape hatch.
-       */
-      image?: {
-        value?: (input: ImageEmbeddingInput) => unknown;
-      };
-    });
+/**
+ * Text-only embedding config for the AI SDK provider.
+ */
+export type AiEmbeddingConfig = BaseConfig;
 
 const DEFAULT_TEXT_MODEL = "openai/text-embedding-3-small";
-const DEFAULT_MULTIMODAL_MODEL = "cohere/embed-v4.0";
-
-const bytesToDataUrl = (bytes: Uint8Array, mediaType: string) => {
-  const base64 = Buffer.from(bytes).toString("base64");
-  return `data:${mediaType};base64,${base64}`;
-};
-
-const defaultImageValue = (input: ImageEmbeddingInput) => {
-  const v =
-    typeof input.data === "string"
-      ? input.data
-      : bytesToDataUrl(input.data, input.mediaType ?? "image/jpeg");
-  // This matches common AI Gateway multimodal embedding inputs where
-  // the embedding value is an object containing one or more images.
-  return { image: [v] };
-};
 
 export const createAiEmbeddingProvider = (
   config: AiEmbeddingConfig = {}
 ): EmbeddingProvider => {
-  const type = (config as any).type ?? "text";
   const model =
-    config.model ??
-    process.env.AI_GATEWAY_MODEL ??
-    (type === "multimodal" ? DEFAULT_MULTIMODAL_MODEL : DEFAULT_TEXT_MODEL);
+    config.model ?? process.env.AI_GATEWAY_MODEL ?? DEFAULT_TEXT_MODEL;
   const timeoutMs = config.timeoutMs;
 
   return {
@@ -87,37 +53,11 @@ export const createAiEmbeddingProvider = (
         ...(abortSignal ? { abortSignal } : {}),
       });
 
-      const embeddings = (result as any)?.embeddings as number[][] | undefined;
-      if (!embeddings) {
+      if (!result.embeddings || result.embeddings.length === 0) {
         throw new Error("Embeddings missing from AI SDK embedMany response");
       }
-      return embeddings;
+      return result.embeddings;
     },
-    ...(type === "multimodal"
-      ? {
-          embedImage: async (input: ImageEmbeddingInput) => {
-            const abortSignal = timeoutMs
-              ? AbortSignal.timeout(timeoutMs)
-              : undefined;
-
-            const imageValue =
-              (config as any)?.image?.value?.(input) ?? defaultImageValue(input);
-
-            const result = await embed({
-              model,
-              value: imageValue,
-              ...(abortSignal ? { abortSignal } : {}),
-            });
-
-            if (!result.embedding) {
-              throw new Error("Embedding missing from AI SDK response");
-            }
-
-            return result.embedding;
-          },
-        }
-      : {}),
   };
 };
 
-

package/registry/embedding/azure.ts

@@ -0,0 +1,88 @@
+import { embed, embedMany, type EmbeddingModel } from "ai";
+import type { EmbeddingProvider } from "../core/types";
+import { requireOptional } from "./_shared";
+
+/**
+ * Azure OpenAI provider module interface.
+ */
+interface AzureModule {
+  azure: {
+    embedding: (model: string) => EmbeddingModel<string>;
+  };
+}
+
+export type AzureEmbeddingConfig = {
+  model?: string;
+  timeoutMs?: number;
+  dimensions?: number;
+  user?: string;
+};
+
+const DEFAULT_TEXT_MODEL = "text-embedding-3-small";
+
+const buildProviderOptions = (config: AzureEmbeddingConfig) => {
+  if (config.dimensions === undefined && config.user === undefined) {
+    return undefined;
+  }
+  return {
+    openai: {
+      ...(config.dimensions !== undefined ? { dimensions: config.dimensions } : {}),
+      ...(config.user ? { user: config.user } : {}),
+    },
+  };
+};
+
+export const createAzureEmbeddingProvider = (
+  config: AzureEmbeddingConfig = {}
+): EmbeddingProvider => {
+  const { azure } = requireOptional<AzureModule>({
+    id: "@ai-sdk/azure",
+    installHint: "bun add @ai-sdk/azure",
+    providerName: "azure",
+  });
+  const model =
+    config.model ?? process.env.AZURE_EMBEDDING_MODEL ?? DEFAULT_TEXT_MODEL;
+  const timeoutMs = config.timeoutMs;
+  const providerOptions = buildProviderOptions(config);
+  const embeddingModel = azure.embedding(model);
+
+  return {
+    name: `azure:${model}`,
+    dimensions: config.dimensions,
+    embed: async ({ text }) => {
+      const abortSignal = timeoutMs
+        ? AbortSignal.timeout(timeoutMs)
+        : undefined;
+
+      const result = await embed({
+        model: embeddingModel,
+        value: text,
+        ...(providerOptions ? { providerOptions } : {}),
+        ...(abortSignal ? { abortSignal } : {}),
+      });
+
+      if (!result.embedding) {
+        throw new Error("Embedding missing from Azure OpenAI response");
+      }
+
+      return result.embedding;
+    },
+    embedMany: async (inputs) => {
+      const values = inputs.map((i) => i.text);
+      const abortSignal = timeoutMs ? AbortSignal.timeout(timeoutMs) : undefined;
+
+      const result = await embedMany({
+        model: embeddingModel,
+        values,
+        ...(providerOptions ? { providerOptions } : {}),
+        ...(abortSignal ? { abortSignal } : {}),
+      });
+
+      const { embeddings } = result;
+      if (!Array.isArray(embeddings)) {
+        throw new Error("Embeddings missing from Azure OpenAI embedMany response");
+      }
+      return embeddings;
+    },
+  };
+};

package/registry/embedding/bedrock.ts

@@ -0,0 +1,88 @@
+import { embed, embedMany, type EmbeddingModel } from "ai";
+import type { EmbeddingProvider } from "../core/types";
+import { requireOptional } from "./_shared";
+
+/**
+ * Amazon Bedrock provider module interface.
+ */
+interface BedrockModule {
+  bedrock: {
+    embedding: (model: string) => EmbeddingModel<string>;
+  };
+}
+
+export type BedrockEmbeddingConfig = {
+  model?: string;
+  timeoutMs?: number;
+  dimensions?: number;
+  normalize?: boolean;
+};
+
+const DEFAULT_TEXT_MODEL = "amazon.titan-embed-text-v2:0";
+
+const buildProviderOptions = (config: BedrockEmbeddingConfig) => {
+  if (config.dimensions === undefined && config.normalize === undefined) {
+    return undefined;
+  }
+  return {
+    bedrock: {
+      ...(config.dimensions !== undefined ? { dimensions: config.dimensions } : {}),
+      ...(config.normalize !== undefined ? { normalize: config.normalize } : {}),
+    },
+  };
+};
+
+export const createBedrockEmbeddingProvider = (
+  config: BedrockEmbeddingConfig = {}
+): EmbeddingProvider => {
+  const { bedrock } = requireOptional<BedrockModule>({
+    id: "@ai-sdk/amazon-bedrock",
+    installHint: "bun add @ai-sdk/amazon-bedrock",
+    providerName: "bedrock",
+  });
+  const model =
+    config.model ?? process.env.BEDROCK_EMBEDDING_MODEL ?? DEFAULT_TEXT_MODEL;
+  const timeoutMs = config.timeoutMs;
+  const providerOptions = buildProviderOptions(config);
+  const embeddingModel = bedrock.embedding(model);
+
+  return {
+    name: `bedrock:${model}`,
+    dimensions: config.dimensions,
+    embed: async ({ text }) => {
+      const abortSignal = timeoutMs
+        ? AbortSignal.timeout(timeoutMs)
+        : undefined;
+
+      const result = await embed({
+        model: embeddingModel,
+        value: text,
+        ...(providerOptions ? { providerOptions } : {}),
+        ...(abortSignal ? { abortSignal } : {}),
+      });
+
+      if (!result.embedding) {
+        throw new Error("Embedding missing from Bedrock response");
+      }
+
+      return result.embedding;
+    },
+    embedMany: async (inputs) => {
+      const values = inputs.map((i) => i.text);
+      const abortSignal = timeoutMs ? AbortSignal.timeout(timeoutMs) : undefined;
+
+      const result = await embedMany({
+        model: embeddingModel,
+        values,
+        ...(providerOptions ? { providerOptions } : {}),
+        ...(abortSignal ? { abortSignal } : {}),
+      });
+
+      const { embeddings } = result;
+      if (!Array.isArray(embeddings)) {
+        throw new Error("Embeddings missing from Bedrock embedMany response");
+      }
+      return embeddings;
+    },
+  };
+};