unrag 0.2.4 → 0.2.6

package/registry/embedding/ollama.ts

@@ -0,0 +1,75 @@
+import { embed, embedMany } from "ai";
+import type { EmbeddingProvider } from "../core/types";
+import { requireOptional } from "./_shared";
+
+export type OllamaEmbeddingConfig = {
+  model?: string;
+  timeoutMs?: number;
+  baseURL?: string;
+  headers?: Record<string, string>;
+};
+
+const DEFAULT_TEXT_MODEL = "nomic-embed-text";
+
+const resolveProvider = (config: OllamaEmbeddingConfig) => {
+  const { createOllama, ollama } = requireOptional<any>({
+    id: "ollama-ai-provider-v2",
+    installHint: "bun add ollama-ai-provider-v2",
+    providerName: "ollama",
+  });
+  if (config.baseURL || config.headers) {
+    return createOllama({
+      ...(config.baseURL ? { baseURL: config.baseURL } : {}),
+      ...(config.headers ? { headers: config.headers } : {}),
+    });
+  }
+  return ollama;
+};
+
+export const createOllamaEmbeddingProvider = (
+  config: OllamaEmbeddingConfig = {}
+): EmbeddingProvider => {
+  const model =
+    config.model ?? process.env.OLLAMA_EMBEDDING_MODEL ?? DEFAULT_TEXT_MODEL;
+  const timeoutMs = config.timeoutMs;
+  const provider = resolveProvider(config);
+  const embeddingModel = provider.textEmbeddingModel(model);
+
+  return {
+    name: `ollama:${model}`,
+    dimensions: undefined,
+    embed: async ({ text }) => {
+      const abortSignal = timeoutMs
+        ? AbortSignal.timeout(timeoutMs)
+        : undefined;
+
+      const result = await embed({
+        model: embeddingModel,
+        value: text,
+        ...(abortSignal ? { abortSignal } : {}),
+      });
+
+      if (!result.embedding) {
+        throw new Error("Embedding missing from Ollama 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,
+        ...(abortSignal ? { abortSignal } : {}),
+      });
+
+      const { embeddings } = result;
+      if (!Array.isArray(embeddings)) {
+        throw new Error("Embeddings missing from Ollama embedMany response");
+      }
+      return embeddings;
+    },
+  };
+};

package/registry/embedding/openai.ts

@@ -0,0 +1,79 @@
+import { embed, embedMany } from "ai";
+import type { EmbeddingProvider } from "../core/types";
+import { requireOptional } from "./_shared";
+
+export type OpenAiEmbeddingConfig = {
+  model?: string;
+  timeoutMs?: number;
+  dimensions?: number;
+  user?: string;
+};
+
+const DEFAULT_TEXT_MODEL = "text-embedding-3-small";
+
+const buildProviderOptions = (config: OpenAiEmbeddingConfig) => {
+  if (config.dimensions === undefined && config.user === undefined) {
+    return undefined;
+  }
+  return {
+    openai: {
+      ...(config.dimensions !== undefined ? { dimensions: config.dimensions } : {}),
+      ...(config.user ? { user: config.user } : {}),
+    },
+  };
+};
+
+export const createOpenAiEmbeddingProvider = (
+  config: OpenAiEmbeddingConfig = {}
+): EmbeddingProvider => {
+  const { openai } = requireOptional<any>({
+    id: "@ai-sdk/openai",
+    installHint: "bun add @ai-sdk/openai",
+    providerName: "openai",
+  });
+  const model =
+    config.model ?? process.env.OPENAI_EMBEDDING_MODEL ?? DEFAULT_TEXT_MODEL;
+  const timeoutMs = config.timeoutMs;
+  const providerOptions = buildProviderOptions(config);
+  const embeddingModel = openai.embedding(model);
+
+  return {
+    name: `openai:${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 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 OpenAI embedMany response");
+      }
+      return embeddings;
+    },
+  };
+};

package/registry/embedding/openrouter.ts

@@ -0,0 +1,85 @@
+import type { EmbeddingProvider } from "../core/types";
+import { requireOptional } from "./_shared";
+
+export type OpenRouterEmbeddingConfig = {
+  model?: string;
+  timeoutMs?: number;
+  apiKey?: string;
+  baseURL?: string;
+  headers?: Record<string, string>;
+  referer?: string;
+  title?: string;
+};
+
+const DEFAULT_TEXT_MODEL = "text-embedding-3-small";
+
+const buildHeaders = (config: OpenRouterEmbeddingConfig) => {
+  const headers: Record<string, string> = { ...(config.headers ?? {}) };
+  if (config.referer) headers["HTTP-Referer"] = config.referer;
+  if (config.title) headers["X-Title"] = config.title;
+  return headers;
+};
+
+export const createOpenRouterEmbeddingProvider = (
+  config: OpenRouterEmbeddingConfig = {}
+): EmbeddingProvider => {
+  const { OpenRouter } = requireOptional<any>({
+    id: "@openrouter/sdk",
+    installHint: "bun add @openrouter/sdk",
+    providerName: "openrouter",
+  });
+  const model =
+    config.model ?? process.env.OPENROUTER_EMBEDDING_MODEL ?? DEFAULT_TEXT_MODEL;
+  const timeoutMs = config.timeoutMs;
+  const headers = buildHeaders(config);
+
+  const client = new OpenRouter({
+    apiKey: config.apiKey ?? process.env.OPENROUTER_API_KEY ?? "",
+    ...(config.baseURL ? { baseURL: config.baseURL } : {}),
+    ...(Object.keys(headers).length ? { headers } : {}),
+  } as any);
+
+  return {
+    name: `openrouter:${model}`,
+    dimensions: undefined,
+    embed: async ({ text }) => {
+      const abortSignal = timeoutMs
+        ? AbortSignal.timeout(timeoutMs)
+        : undefined;
+
+      const result = await (client as any).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;
+      if (!embedding) {
+        throw new Error("Embedding missing from OpenRouter response");
+      }
+
+      return embedding as number[];
+    },
+    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(
+        { input: values, model },
+        abortSignal ? { fetchOptions: { signal: abortSignal } } : undefined
+      );
+
+      const embeddings = (result as any)?.data?.map(
+        (item: { embedding?: number[] }) => item.embedding
+      );
+
+      if (!Array.isArray(embeddings) || embeddings.some((e) => !Array.isArray(e))) {
+        throw new Error("Embeddings missing from OpenRouter response");
+      }
+
+      return embeddings as number[][];
+    },
+  };
+};

package/registry/embedding/together.ts

@@ -0,0 +1,67 @@
+import { embed, embedMany } from "ai";
+import type { EmbeddingProvider } from "../core/types";
+import { requireOptional } from "./_shared";
+
+export type TogetherEmbeddingConfig = {
+  model?: string;
+  timeoutMs?: number;
+};
+
+const DEFAULT_TEXT_MODEL = "togethercomputer/m2-bert-80M-2k-retrieval";
+
+export const createTogetherEmbeddingProvider = (
+  config: TogetherEmbeddingConfig = {}
+): EmbeddingProvider => {
+  const { togetherai } = requireOptional<any>({
+    id: "@ai-sdk/togetherai",
+    installHint: "bun add @ai-sdk/togetherai",
+    providerName: "together",
+  });
+  const model =
+    config.model ??
+    process.env.TOGETHER_AI_EMBEDDING_MODEL ??
+    DEFAULT_TEXT_MODEL;
+  const timeoutMs = config.timeoutMs;
+  const embeddingModel =
+    "embeddingModel" in togetherai
+      ? (togetherai as any).embeddingModel(model)
+      : (togetherai as any).textEmbeddingModel(model);
+
+  return {
+    name: `together:${model}`,
+    dimensions: undefined,
+    embed: async ({ text }) => {
+      const abortSignal = timeoutMs
+        ? AbortSignal.timeout(timeoutMs)
+        : undefined;
+
+      const result = await embed({
+        model: embeddingModel,
+        value: text,
+        ...(abortSignal ? { abortSignal } : {}),
+      });
+
+      if (!result.embedding) {
+        throw new Error("Embedding missing from Together.ai 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,
+        ...(abortSignal ? { abortSignal } : {}),
+      });
+
+      const { embeddings } = result;
+      if (!Array.isArray(embeddings)) {
+        throw new Error("Embeddings missing from Together.ai embedMany response");
+      }
+      return embeddings;
+    },
+  };
+};

package/registry/embedding/vertex.ts

@@ -0,0 +1,102 @@
+import { embed, embedMany } from "ai";
+import type { EmbeddingProvider } from "../core/types";
+import { requireOptional } from "./_shared";
+
+export type VertexEmbeddingTaskType =
+  | "SEMANTIC_SIMILARITY"
+  | "CLASSIFICATION"
+  | "CLUSTERING"
+  | "RETRIEVAL_DOCUMENT"
+  | "RETRIEVAL_QUERY"
+  | "QUESTION_ANSWERING"
+  | "FACT_VERIFICATION"
+  | "CODE_RETRIEVAL_QUERY";
+
+export type VertexEmbeddingConfig = {
+  model?: string;
+  timeoutMs?: number;
+  outputDimensionality?: number;
+  taskType?: VertexEmbeddingTaskType;
+  title?: string;
+  autoTruncate?: boolean;
+};
+
+const DEFAULT_TEXT_MODEL = "text-embedding-004";
+
+const buildProviderOptions = (config: VertexEmbeddingConfig) => {
+  if (
+    config.outputDimensionality === undefined &&
+    !config.taskType &&
+    config.autoTruncate === undefined &&
+    !config.title
+  ) {
+    return undefined;
+  }
+  return {
+    google: {
+      ...(config.outputDimensionality !== undefined
+        ? { outputDimensionality: config.outputDimensionality }
+        : {}),
+      ...(config.taskType ? { taskType: config.taskType } : {}),
+      ...(config.autoTruncate !== undefined ? { autoTruncate: config.autoTruncate } : {}),
+      ...(config.title ? { title: config.title } : {}),
+    },
+  };
+};
+
+export const createVertexEmbeddingProvider = (
+  config: VertexEmbeddingConfig = {}
+): EmbeddingProvider => {
+  const { vertex } = requireOptional<any>({
+    id: "@ai-sdk/google-vertex",
+    installHint: "bun add @ai-sdk/google-vertex",
+    providerName: "vertex",
+  });
+  const model =
+    config.model ??
+    process.env.GOOGLE_VERTEX_EMBEDDING_MODEL ??
+    DEFAULT_TEXT_MODEL;
+  const timeoutMs = config.timeoutMs;
+  const providerOptions = buildProviderOptions(config);
+  const embeddingModel = vertex.embeddingModel(model);
+
+  return {
+    name: `vertex:${model}`,
+    dimensions: config.outputDimensionality,
+    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 Vertex 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 Vertex embedMany response");
+      }
+      return embeddings;
+    },
+  };
+};

package/registry/embedding/voyage.ts

@@ -0,0 +1,159 @@
+import { embed, embedMany, type EmbeddingModel } from "ai";
+import type { EmbeddingProvider, ImageEmbeddingInput } from "../core/types";
+import { requireOptional } from "./_shared";
+
+type BaseConfig = {
+  model?: string;
+  timeoutMs?: number;
+};
+
+export type VoyageEmbeddingConfig =
+  | (BaseConfig & {
+      type?: "text";
+    })
+  | (BaseConfig & {
+      type: "multimodal";
+      text?: {
+        value?: (text: string) => unknown;
+      };
+      image?: {
+        value?: (input: ImageEmbeddingInput) => unknown;
+      };
+    });
+
+const DEFAULT_TEXT_MODEL = "voyage-3.5-lite";
+const DEFAULT_MULTIMODAL_MODEL = "voyage-multimodal-3";
+
+const bytesToDataUrl = (bytes: Uint8Array, mediaType: string) => {
+  const base64 = Buffer.from(bytes).toString("base64");
+  return `data:${mediaType};base64,${base64}`;
+};
+
+const defaultTextValue = (text: string) => ({
+  text: [text],
+});
+
+const defaultImageValue = (input: ImageEmbeddingInput) => {
+  const v =
+    typeof input.data === "string"
+      ? input.data
+      : bytesToDataUrl(input.data, input.mediaType ?? "image/jpeg");
+  return { image: [v] };
+};
+
+export const createVoyageEmbeddingProvider = (
+  config: VoyageEmbeddingConfig = {}
+): EmbeddingProvider => {
+  const { voyage } = requireOptional<any>({
+    id: "voyage-ai-provider",
+    installHint: "bun add voyage-ai-provider",
+    providerName: "voyage",
+  });
+  const type = config.type ?? "text";
+  const isMultimodal = config.type === "multimodal";
+  const model =
+    config.model ??
+    process.env.VOYAGE_MODEL ??
+    (type === "multimodal" ? DEFAULT_MULTIMODAL_MODEL : DEFAULT_TEXT_MODEL);
+  const timeoutMs = config.timeoutMs;
+
+  const voyageProvider = voyage as any;
+  const textEmbeddingModel =
+    type === "multimodal"
+      ? undefined
+      : typeof voyageProvider.embeddingModel === "function"
+        ? voyageProvider.embeddingModel(model)
+        : voyageProvider.textEmbeddingModel(model);
+  const multimodalEmbeddingModel =
+    type === "multimodal" ? (voyage as any).multimodalEmbeddingModel(model) : undefined;
+
+  // AI SDK 6 types only accept string inputs; cast multimodal models/values.
+  const multimodalModel = multimodalEmbeddingModel as unknown as EmbeddingModel;
+
+  const resolveTextValue = (text: string) => {
+    if (isMultimodal && config.text?.value) {
+      return config.text.value(text);
+    }
+    return defaultTextValue(text);
+  };
+
+  const resolveImageValue = (input: ImageEmbeddingInput) => {
+    if (isMultimodal && config.image?.value) {
+      return config.image.value(input);
+    }
+    return defaultImageValue(input);
+  };
+
+  return {
+    name: `voyage:${model}`,
+    dimensions: undefined,
+    embed: async ({ text }) => {
+      const abortSignal = timeoutMs
+        ? AbortSignal.timeout(timeoutMs)
+        : undefined;
+
+      const result =
+        type === "multimodal"
+          ? await embed({
+              model: multimodalModel,
+              value: resolveTextValue(text) as unknown as string,
+              ...(abortSignal ? { abortSignal } : {}),
+            })
+          : await embed({
+              model: textEmbeddingModel!,
+              value: text,
+              ...(abortSignal ? { abortSignal } : {}),
+            });
+
+      if (!result.embedding) {
+        throw new Error("Embedding missing from Voyage response");
+      }
+
+      return result.embedding;
+    },
+    embedMany: async (inputs) => {
+      const abortSignal = timeoutMs ? AbortSignal.timeout(timeoutMs) : undefined;
+      const result =
+        type === "multimodal"
+          ? await embedMany({
+              model: multimodalModel,
+              values: inputs.map((i) => resolveTextValue(i.text)) as unknown as string[],
+              ...(abortSignal ? { abortSignal } : {}),
+            })
+          : await embedMany({
+              model: textEmbeddingModel!,
+              values: inputs.map((i) => i.text),
+              ...(abortSignal ? { abortSignal } : {}),
+            });
+
+      const { embeddings } = result;
+      if (!Array.isArray(embeddings)) {
+        throw new Error("Embeddings missing from Voyage embedMany response");
+      }
+      return embeddings;
+    },
+    ...(type === "multimodal"
+      ? {
+          embedImage: async (input: ImageEmbeddingInput) => {
+            const abortSignal = timeoutMs
+              ? AbortSignal.timeout(timeoutMs)
+              : undefined;
+
+            const value = resolveImageValue(input);
+
+            const result = await embed({
+              model: multimodalModel,
+              value: value as unknown as string,
+              ...(abortSignal ? { abortSignal } : {}),
+            });
+
+            if (!result.embedding) {
+              throw new Error("Embedding missing from Voyage response");
+            }
+
+            return result.embedding;
+          },
+        }
+      : {}),
+  };
+};