unrag 0.2.6 → 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;
   };
@@ -768,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 = {
@@ -779,4 +830,5 @@ export type ResolvedContextEngineConfig = {
   extractors: AssetExtractor[];
   storage: ContentStorageConfig;
   assetProcessing: AssetProcessingConfig;
+  embeddingProcessing: EmbeddingProcessingConfig;
 };

package/registry/embedding/_shared.ts

@@ -2,7 +2,12 @@ import { createRequire } from "node:module";
 
 const require = createRequire(import.meta.url);
 
-export function requireOptional<T = any>(args: {
+/**
+ * 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;

package/registry/embedding/ai.ts

@@ -53,11 +53,10 @@ 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;
     },
   };
 };

package/registry/embedding/azure.ts

@@ -1,7 +1,16 @@
-import { embed, embedMany } from "ai";
+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;
@@ -26,7 +35,7 @@ const buildProviderOptions = (config: AzureEmbeddingConfig) => {
 export const createAzureEmbeddingProvider = (
   config: AzureEmbeddingConfig = {}
 ): EmbeddingProvider => {
-  const { azure } = requireOptional<any>({
+  const { azure } = requireOptional<AzureModule>({
     id: "@ai-sdk/azure",
     installHint: "bun add @ai-sdk/azure",
     providerName: "azure",

package/registry/embedding/bedrock.ts

@@ -1,7 +1,16 @@
-import { embed, embedMany } from "ai";
+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;
@@ -26,7 +35,7 @@ const buildProviderOptions = (config: BedrockEmbeddingConfig) => {
 export const createBedrockEmbeddingProvider = (
   config: BedrockEmbeddingConfig = {}
 ): EmbeddingProvider => {
-  const { bedrock } = requireOptional<any>({
+  const { bedrock } = requireOptional<BedrockModule>({
     id: "@ai-sdk/amazon-bedrock",
     installHint: "bun add @ai-sdk/amazon-bedrock",
     providerName: "bedrock",

package/registry/embedding/cohere.ts

@@ -1,7 +1,16 @@
-import { embed, embedMany } from "ai";
+import { embed, embedMany, type EmbeddingModel } from "ai";
 import type { EmbeddingProvider } from "../core/types";
 import { requireOptional } from "./_shared";
 
+/**
+ * Cohere provider module interface.
+ */
+interface CohereModule {
+  cohere: {
+    embedding: (model: string) => EmbeddingModel<string>;
+  };
+}
+
 export type CohereEmbeddingConfig = {
   model?: string;
   timeoutMs?: number;
@@ -26,7 +35,7 @@ const buildProviderOptions = (config: CohereEmbeddingConfig) => {
 export const createCohereEmbeddingProvider = (
   config: CohereEmbeddingConfig = {}
 ): EmbeddingProvider => {
-  const { cohere } = requireOptional<any>({
+  const { cohere } = requireOptional<CohereModule>({
     id: "@ai-sdk/cohere",
     installHint: "bun add @ai-sdk/cohere",
     providerName: "cohere",

package/registry/embedding/google.ts

@@ -1,7 +1,16 @@
-import { embed, embedMany } from "ai";
+import { embed, embedMany, type EmbeddingModel } from "ai";
 import type { EmbeddingProvider } from "../core/types";
 import { requireOptional } from "./_shared";
 
+/**
+ * Google AI provider module interface.
+ */
+interface GoogleModule {
+  google: {
+    embedding: (model: string) => EmbeddingModel<string>;
+  };
+}
+
 export type GoogleEmbeddingTaskType =
   | "SEMANTIC_SIMILARITY"
   | "CLASSIFICATION"
@@ -38,7 +47,7 @@ const buildProviderOptions = (config: GoogleEmbeddingConfig) => {
 export const createGoogleEmbeddingProvider = (
   config: GoogleEmbeddingConfig = {}
 ): EmbeddingProvider => {
-  const { google } = requireOptional<any>({
+  const { google } = requireOptional<GoogleModule>({
     id: "@ai-sdk/google",
     installHint: "bun add @ai-sdk/google",
     providerName: "google",

package/registry/embedding/mistral.ts

@@ -1,7 +1,16 @@
-import { embed, embedMany } from "ai";
+import { embed, embedMany, type EmbeddingModel } from "ai";
 import type { EmbeddingProvider } from "../core/types";
 import { requireOptional } from "./_shared";
 
+/**
+ * Mistral provider module interface.
+ */
+interface MistralModule {
+  mistral: {
+    embedding: (model: string) => EmbeddingModel<string>;
+  };
+}
+
 export type MistralEmbeddingConfig = {
   model?: string;
   timeoutMs?: number;
@@ -12,7 +21,7 @@ const DEFAULT_TEXT_MODEL = "mistral-embed";
 export const createMistralEmbeddingProvider = (
   config: MistralEmbeddingConfig = {}
 ): EmbeddingProvider => {
-  const { mistral } = requireOptional<any>({
+  const { mistral } = requireOptional<MistralModule>({
     id: "@ai-sdk/mistral",
     installHint: "bun add @ai-sdk/mistral",
     providerName: "mistral",

package/registry/embedding/ollama.ts

@@ -1,7 +1,22 @@
-import { embed, embedMany } from "ai";
+import { embed, embedMany, type EmbeddingModel } from "ai";
 import type { EmbeddingProvider } from "../core/types";
 import { requireOptional } from "./_shared";
 
+/**
+ * Ollama provider instance interface.
+ */
+interface OllamaProvider {
+  textEmbeddingModel: (model: string) => EmbeddingModel<string>;
+}
+
+/**
+ * Ollama provider module interface.
+ */
+interface OllamaModule {
+  createOllama: (config: { baseURL?: string; headers?: Record<string, string> }) => OllamaProvider;
+  ollama: OllamaProvider;
+}
+
 export type OllamaEmbeddingConfig = {
   model?: string;
   timeoutMs?: number;
@@ -11,8 +26,8 @@ export type OllamaEmbeddingConfig = {
 
 const DEFAULT_TEXT_MODEL = "nomic-embed-text";
 
-const resolveProvider = (config: OllamaEmbeddingConfig) => {
-  const { createOllama, ollama } = requireOptional<any>({
+const resolveProvider = (config: OllamaEmbeddingConfig): OllamaProvider => {
+  const { createOllama, ollama } = requireOptional<OllamaModule>({
     id: "ollama-ai-provider-v2",
     installHint: "bun add ollama-ai-provider-v2",
     providerName: "ollama",

package/registry/embedding/openai.ts

@@ -1,7 +1,16 @@
-import { embed, embedMany } from "ai";
+import { embed, embedMany, type EmbeddingModel } from "ai";
 import type { EmbeddingProvider } from "../core/types";
 import { requireOptional } from "./_shared";
 
+/**
+ * OpenAI provider module interface.
+ */
+interface OpenAiModule {
+  openai: {
+    embedding: (model: string) => EmbeddingModel<string>;
+  };
+}
+
 export type OpenAiEmbeddingConfig = {
   model?: string;
   timeoutMs?: number;
@@ -26,7 +35,7 @@ const buildProviderOptions = (config: OpenAiEmbeddingConfig) => {
 export const createOpenAiEmbeddingProvider = (
   config: OpenAiEmbeddingConfig = {}
 ): EmbeddingProvider => {
-  const { openai } = requireOptional<any>({
+  const { openai } = requireOptional<OpenAiModule>({
     id: "@ai-sdk/openai",
     installHint: "bun add @ai-sdk/openai",
     providerName: "openai",

package/registry/embedding/openrouter.ts

@@ -11,6 +11,49 @@ export type OpenRouterEmbeddingConfig = {
   title?: string;
 };
 
+/**
+ * OpenRouter embedding result item.
+ */
+interface EmbeddingDataItem {
+  embedding?: number[];
+}
+
+/**
+ * OpenRouter embedding response.
+ */
+interface EmbeddingResponse {
+  data?: EmbeddingDataItem[];
+  embedding?: number[];
+}
+
+/**
+ * OpenRouter client embeddings interface.
+ */
+interface EmbeddingsClient {
+  generate(
+    params: { input: string | string[]; model: string },
+    options?: { fetchOptions?: { signal?: AbortSignal } }
+  ): Promise<EmbeddingResponse>;
+}
+
+/**
+ * OpenRouter client interface.
+ */
+interface OpenRouterClient {
+  embeddings: EmbeddingsClient;
+}
+
+/**
+ * OpenRouter SDK module interface.
+ */
+interface OpenRouterModule {
+  OpenRouter: new (config: {
+    apiKey: string;
+    baseURL?: string;
+    headers?: Record<string, string>;
+  }) => OpenRouterClient;
+}
+
 const DEFAULT_TEXT_MODEL = "text-embedding-3-small";
 
 const buildHeaders = (config: OpenRouterEmbeddingConfig) => {
@@ -23,7 +66,7 @@ const buildHeaders = (config: OpenRouterEmbeddingConfig) => {
 export const createOpenRouterEmbeddingProvider = (
   config: OpenRouterEmbeddingConfig = {}
 ): EmbeddingProvider => {
-  const { OpenRouter } = requireOptional<any>({
+  const { OpenRouter } = requireOptional<OpenRouterModule>({
     id: "@openrouter/sdk",
     installHint: "bun add @openrouter/sdk",
     providerName: "openrouter",
@@ -37,7 +80,7 @@ export const createOpenRouterEmbeddingProvider = (
     apiKey: config.apiKey ?? process.env.OPENROUTER_API_KEY ?? "",
     ...(config.baseURL ? { baseURL: config.baseURL } : {}),
     ...(Object.keys(headers).length ? { headers } : {}),
-  } as any);
+  });
 
   return {
     name: `openrouter:${model}`,
@@ -47,35 +90,34 @@ export const createOpenRouterEmbeddingProvider = (
         ? AbortSignal.timeout(timeoutMs)
         : undefined;
 
-      const result = await (client as any).embeddings.generate(
+      const result = await client.embeddings.generate(
         { input: text, model },
         abortSignal ? { fetchOptions: { signal: abortSignal } } : undefined
       );
 
       const embedding =
-        (result as any)?.data?.[0]?.embedding ??
-        (result as any)?.embedding ??
-        (result as any)?.data?.embedding;
+        result.data?.[0]?.embedding ??
+        result.embedding;
       if (!embedding) {
         throw new Error("Embedding missing from OpenRouter response");
       }
 
-      return embedding as number[];
+      return embedding;
     },
     embedMany: async (inputs) => {
       const values = inputs.map((i) => i.text);
       const abortSignal = timeoutMs ? AbortSignal.timeout(timeoutMs) : undefined;
 
-      const result = await (client as any).embeddings.generate(
+      const result = await client.embeddings.generate(
         { input: values, model },
         abortSignal ? { fetchOptions: { signal: abortSignal } } : undefined
       );
 
-      const embeddings = (result as any)?.data?.map(
-        (item: { embedding?: number[] }) => item.embedding
+      const embeddings = result.data?.map(
+        (item) => item.embedding
       );
 
-      if (!Array.isArray(embeddings) || embeddings.some((e) => !Array.isArray(e))) {
+      if (!embeddings || embeddings.some((e) => !Array.isArray(e))) {
         throw new Error("Embeddings missing from OpenRouter response");
       }
 

package/registry/embedding/together.ts

@@ -1,7 +1,17 @@
-import { embed, embedMany } from "ai";
+import { embed, embedMany, type EmbeddingModel } from "ai";
 import type { EmbeddingProvider } from "../core/types";
 import { requireOptional } from "./_shared";
 
+/**
+ * Together AI provider module interface.
+ */
+interface TogetherAiModule {
+  togetherai: {
+    embeddingModel?: (model: string) => EmbeddingModel<string>;
+    textEmbeddingModel?: (model: string) => EmbeddingModel<string>;
+  };
+}
+
 export type TogetherEmbeddingConfig = {
   model?: string;
   timeoutMs?: number;
@@ -12,7 +22,7 @@ const DEFAULT_TEXT_MODEL = "togethercomputer/m2-bert-80M-2k-retrieval";
 export const createTogetherEmbeddingProvider = (
   config: TogetherEmbeddingConfig = {}
 ): EmbeddingProvider => {
-  const { togetherai } = requireOptional<any>({
+  const { togetherai } = requireOptional<TogetherAiModule>({
     id: "@ai-sdk/togetherai",
     installHint: "bun add @ai-sdk/togetherai",
     providerName: "together",
@@ -23,9 +33,9 @@ export const createTogetherEmbeddingProvider = (
     DEFAULT_TEXT_MODEL;
   const timeoutMs = config.timeoutMs;
   const embeddingModel =
-    "embeddingModel" in togetherai
-      ? (togetherai as any).embeddingModel(model)
-      : (togetherai as any).textEmbeddingModel(model);
+    typeof togetherai.embeddingModel === "function"
+      ? togetherai.embeddingModel(model)
+      : togetherai.textEmbeddingModel?.(model);
 
   return {
     name: `together:${model}`,