unrag 0.2.5 → 0.2.7

package/registry/embedding/vertex.ts

@@ -0,0 +1,111 @@
+import { embed, embedMany, type EmbeddingModel } from "ai";
+import type { EmbeddingProvider } from "../core/types";
+import { requireOptional } from "./_shared";
+
+/**
+ * Vertex AI provider module interface.
+ */
+interface VertexModule {
+  vertex: {
+    embeddingModel: (model: string) => EmbeddingModel<string>;
+  };
+}
+
+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<VertexModule>({
+    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,169 @@
+import { embed, embedMany, type EmbeddingModel } from "ai";
+import type { EmbeddingProvider, ImageEmbeddingInput } from "../core/types";
+import { requireOptional } from "./_shared";
+
+/**
+ * Voyage AI provider module interface.
+ */
+interface VoyageModule {
+  voyage: {
+    embeddingModel?: (model: string) => EmbeddingModel<string>;
+    textEmbeddingModel?: (model: string) => EmbeddingModel<string>;
+    multimodalEmbeddingModel?: (model: string) => EmbeddingModel<unknown>;
+  };
+}
+
+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<VoyageModule>({
+    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 textEmbeddingModel =
+    type === "multimodal"
+      ? undefined
+      : typeof voyage.embeddingModel === "function"
+        ? voyage.embeddingModel(model)
+        : voyage.textEmbeddingModel?.(model);
+  const multimodalEmbeddingModel =
+    type === "multimodal" ? voyage.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;
+          },
+        }
+      : {}),
+  };
+};

package/registry/extractors/audio-transcribe/index.ts

@@ -1,7 +1,21 @@
-import { experimental_transcribe as transcribe } from "ai";
-import type { AssetExtractor } from "../../core/types";
+import { experimental_transcribe as transcribe, type TranscriptionModel } from "ai";
+import type { AssetExtractor, ExtractedTextItem } from "../../core/types";
 import { getAssetBytes } from "../_shared/fetch";
 
+/**
+ * Model reference type that accepts both string gateway IDs and TranscriptionModel instances.
+ */
+type TranscriptionModelRef = string | TranscriptionModel;
+
+/**
+ * Transcription segment from the AI SDK.
+ */
+interface TranscriptionSegment {
+  text?: string;
+  startSecond?: number;
+  endSecond?: number;
+}
+
 /**
  * Audio transcription via the AI SDK `transcribe()` API.
  */
@@ -25,43 +39,45 @@ export function createAudioTranscribeExtractor(): AssetExtractor {
       const abortSignal = AbortSignal.timeout(cfg.timeoutMs);
 
       const result = await transcribe({
-        model: cfg.model as any,
+        model: cfg.model as TranscriptionModelRef,
         audio: bytes,
         abortSignal,
       });
 
-      const segments: any[] = Array.isArray((result as any)?.segments)
-        ? (result as any).segments
+      const segments: TranscriptionSegment[] = Array.isArray(result.segments)
+        ? result.segments
         : [];
 
       if (segments.length > 0) {
+        const textItems: ExtractedTextItem[] = segments
+          .map((s, i) => {
+            const t = String(s?.text ?? "").trim();
+            if (!t) return null;
+            const start = Number(s?.startSecond ?? NaN);
+            const end = Number(s?.endSecond ?? NaN);
+            return {
+              label: `segment-${i + 1}`,
+              content: t,
+              ...(Number.isFinite(start) && Number.isFinite(end)
+                ? { timeRangeSec: [start, end] as [number, number] }
+                : {}),
+            };
+          })
+          .filter((item): item is ExtractedTextItem => item !== null);
+
         return {
-          texts: segments
-            .map((s, i) => {
-              const t = String(s?.text ?? "").trim();
-              if (!t) return null;
-              const start = Number(s?.startSecond ?? NaN);
-              const end = Number(s?.endSecond ?? NaN);
-              return {
-                label: `segment-${i + 1}`,
-                content: t,
-                ...(Number.isFinite(start) && Number.isFinite(end)
-                  ? { timeRangeSec: [start, end] as [number, number] }
-                  : {}),
-              };
-            })
-            .filter(Boolean) as any,
+          texts: textItems,
           diagnostics: {
             model: cfg.model,
             seconds:
-              typeof (result as any)?.durationInSeconds === "number"
-                ? (result as any).durationInSeconds
+              typeof result.durationInSeconds === "number"
+                ? result.durationInSeconds
                 : undefined,
           },
         };
       }
 
-      const text = String((result as any)?.text ?? "").trim();
+      const text = (result.text ?? "").trim();
       if (!text) return { texts: [], diagnostics: { model: cfg.model } };
 
       return {

package/registry/extractors/file-docx/index.ts

@@ -3,6 +3,13 @@ import { getAssetBytes } from "../_shared/fetch";
 import { extFromFilename, normalizeMediaType } from "../_shared/media";
 import { capText } from "../_shared/text";
 
+/**
+ * Minimal mammoth module interface.
+ */
+interface MammothModule {
+  extractRawText(options: { arrayBuffer: ArrayBuffer }): Promise<{ value?: string }>;
+}
+
 const DOCX_MEDIA =
   "application/vnd.openxmlformats-officedocument.wordprocessingml.document";
 
@@ -33,7 +40,7 @@ export function createFileDocxExtractor(): AssetExtractor {
       });
 
       // Dynamic import so the core package can be used without mammoth unless this extractor is installed.
-      const mammoth: any = await import("mammoth");
+      const mammoth = (await import("mammoth")) as MammothModule;
       const arrayBuffer = bytes.buffer.slice(
         bytes.byteOffset,
         bytes.byteOffset + bytes.byteLength

package/registry/extractors/file-pptx/index.ts

@@ -3,6 +3,27 @@ import { getAssetBytes } from "../_shared/fetch";
 import { extFromFilename, normalizeMediaType } from "../_shared/media";
 import { capText } from "../_shared/text";
 
+/**
+ * Zip file entry interface.
+ */
+interface ZipFile {
+  async(type: "string"): Promise<string>;
+}
+
+/**
+ * JSZip instance interface.
+ */
+interface JSZipInstance {
+  files: Record<string, ZipFile>;
+}
+
+/**
+ * JSZip constructor interface.
+ */
+interface JSZipConstructor {
+  loadAsync(data: Uint8Array): Promise<JSZipInstance>;
+}
+
 const PPTX_MEDIA =
   "application/vnd.openxmlformats-officedocument.presentationml.presentation";
 
@@ -41,7 +62,7 @@ export function createFilePptxExtractor(): AssetExtractor {
       });
 
       // Dynamic import to avoid hard dependency unless installed.
-      const JSZip: any = (await import("jszip")).default;
+      const JSZip = (await import("jszip")).default as unknown as JSZipConstructor;
       const zip = await JSZip.loadAsync(bytes);
 
       const slidePaths = Object.keys(zip.files).filter((p) =>

package/registry/extractors/file-xlsx/index.ts

@@ -3,6 +3,29 @@ import { getAssetBytes } from "../_shared/fetch";
 import { extFromFilename, normalizeMediaType } from "../_shared/media";
 import { capText } from "../_shared/text";
 
+/**
+ * XLSX sheet interface.
+ */
+type XLSXSheet = unknown;
+
+/**
+ * XLSX workbook interface.
+ */
+interface XLSXWorkbook {
+  SheetNames?: string[];
+  Sheets?: Record<string, XLSXSheet>;
+}
+
+/**
+ * Minimal xlsx module interface.
+ */
+interface XLSXModule {
+  read(data: Buffer, options: { type: string }): XLSXWorkbook;
+  utils: {
+    sheet_to_csv(sheet: XLSXSheet): string;
+  };
+}
+
 const XLSX_MEDIA =
   "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet";
 
@@ -32,7 +55,7 @@ export function createFileXlsxExtractor(): AssetExtractor {
         defaultMediaType: XLSX_MEDIA,
       });
 
-      const xlsx: any = await import("xlsx");
+      const xlsx = (await import("xlsx")) as XLSXModule;
       const wb = xlsx.read(Buffer.from(bytes), { type: "buffer" });
 
       const parts: string[] = [];

package/registry/extractors/image-caption-llm/index.ts

@@ -1,9 +1,14 @@
-import { generateText } from "ai";
+import { generateText, type LanguageModel } from "ai";
 import type { AssetExtractor } from "../../core/types";
 import { getAssetBytes } from "../_shared/fetch";
 import { normalizeMediaType } from "../_shared/media";
 import { capText } from "../_shared/text";
 
+/**
+ * Model reference type that accepts both string gateway IDs and LanguageModel instances.
+ */
+type ModelRef = string | LanguageModel;
+
 /**
  * Caption generation for images via a vision-capable LLM.
  *
@@ -29,7 +34,7 @@ export function createImageCaptionLlmExtractor(): AssetExtractor {
       const abortSignal = AbortSignal.timeout(cfg.timeoutMs);
 
       const result = await generateText({
-        model: cfg.model as any,
+        model: cfg.model as ModelRef,
         abortSignal,
         messages: [
           {
@@ -46,7 +51,7 @@ export function createImageCaptionLlmExtractor(): AssetExtractor {
         ],
       });
 
-      const caption = String((result as any)?.text ?? "").trim();
+      const caption = (result.text ?? "").trim();
       if (!caption) return { texts: [], diagnostics: { model: cfg.model } };
 
       return {

package/registry/extractors/image-ocr/index.ts

@@ -1,13 +1,18 @@
-import { generateText } from "ai";
+import { generateText, type LanguageModel } from "ai";
 import type { AssetExtractor } from "../../core/types";
 import { getAssetBytes } from "../_shared/fetch";
 import { normalizeMediaType } from "../_shared/media";
 import { capText } from "../_shared/text";
 
+/**
+ * Model reference type that accepts both string gateway IDs and LanguageModel instances.
+ */
+type ModelRef = string | LanguageModel;
+
 /**
  * Image OCR via a vision-capable LLM.
  *
- * This extractor is intended for screenshots, charts, diagrams, and any image with embedded text.
+ * This extractor is intended for screenshots, charts, diagrams, and images with embedded text.
  */
 export function createImageOcrExtractor(): AssetExtractor {
   return {
@@ -29,7 +34,7 @@ export function createImageOcrExtractor(): AssetExtractor {
       const abortSignal = AbortSignal.timeout(cfg.timeoutMs);
 
       const result = await generateText({
-        model: cfg.model as any,
+        model: cfg.model as ModelRef,
         abortSignal,
         messages: [
           {
@@ -46,7 +51,7 @@ export function createImageOcrExtractor(): AssetExtractor {
         ],
       });
 
-      const text = String((result as any)?.text ?? "").trim();
+      const text = (result.text ?? "").trim();
       if (!text) return { texts: [], diagnostics: { model: cfg.model } };
 
       return {

package/registry/extractors/pdf-llm/index.ts

@@ -1,9 +1,14 @@
-import { generateText } from "ai";
+import { generateText, type LanguageModel } from "ai";
 import type { AssetData, AssetExtractor, AssetFetchConfig } from "../../core/types";
 import { getAssetBytes } from "../_shared/fetch";
 import { normalizeMediaType } from "../_shared/media";
 import { capText } from "../_shared/text";
 
+/**
+ * Model reference type that accepts both string gateway IDs and LanguageModel instances.
+ */
+type ModelRef = string | LanguageModel;
+
 async function getPdfBytes(args: {
   data: AssetData;
   fetchConfig: AssetFetchConfig;
@@ -49,8 +54,8 @@ export function createPdfLlmExtractor(): AssetExtractor {
       const abortSignal = AbortSignal.timeout(llm.timeoutMs);
 
       const result = await generateText({
-        // Intentionally allow string model ids for AI Gateway usage.
-        model: llm.model as any,
+        // String model IDs are supported for AI Gateway routing.
+        model: llm.model as ModelRef,
         abortSignal,
         messages: [
           {
@@ -68,7 +73,7 @@ export function createPdfLlmExtractor(): AssetExtractor {
         ],
       });
 
-      const text = String((result as any)?.text ?? "").trim();
+      const text = (result.text ?? "").trim();
       if (!text) return { texts: [], diagnostics: { model: llm.model } };
 
       const capped = capText(text, llm.maxOutputChars);

package/registry/extractors/pdf-text-layer/index.ts

@@ -2,6 +2,27 @@ import type { AssetExtractor } from "../../core/types";
 import { getAssetBytes } from "../_shared/fetch";
 import { capText } from "../_shared/text";
 
+/**
+ * Text content item from pdfjs-dist.
+ */
+interface PdfTextItem {
+  str?: string;
+}
+
+/**
+ * Minimal pdfjs-dist module interface.
+ */
+interface PdfJsModule {
+  getDocument(params: { data: Uint8Array }): {
+    promise: Promise<{
+      numPages: number;
+      getPage(pageNum: number): Promise<{
+        getTextContent(): Promise<{ items?: PdfTextItem[] }>;
+      }>;
+    }>;
+  };
+}
+
 /**
  * Fast/cheap PDF extraction using the PDF's built-in text layer.
  *
@@ -29,7 +50,7 @@ export function createPdfTextLayerExtractor(): AssetExtractor {
       });
 
       // Dynamic import so the core package can be used without pdfjs unless this extractor is installed.
-      const pdfjs: any = await import("pdfjs-dist/legacy/build/pdf.mjs");
+      const pdfjs = (await import("pdfjs-dist/legacy/build/pdf.mjs")) as PdfJsModule;
 
       const doc = await pdfjs.getDocument({ data: bytes }).promise;
       const totalPages: number = Number(doc?.numPages ?? 0);
@@ -42,7 +63,7 @@ export function createPdfTextLayerExtractor(): AssetExtractor {
       for (let pageNum = 1; pageNum <= maxPages; pageNum++) {
         const page = await doc.getPage(pageNum);
         const textContent = await page.getTextContent();
-        const items: any[] = Array.isArray(textContent?.items)
+        const items: PdfTextItem[] = Array.isArray(textContent?.items)
           ? textContent.items
           : [];
         const pageText = items

package/registry/extractors/video-frames/index.ts

@@ -1,4 +1,4 @@
-import { generateText } from "ai";
+import { generateText, type LanguageModel } from "ai";
 import { spawn } from "node:child_process";
 import { mkdir, readdir, readFile, rm, writeFile } from "node:fs/promises";
 import os from "node:os";
@@ -7,6 +7,11 @@ import type { AssetExtractor } from "../../core/types";
 import { getAssetBytes } from "../_shared/fetch";
 import { capText } from "../_shared/text";
 
+/**
+ * Model reference type that accepts both string gateway IDs and LanguageModel instances.
+ */
+type ModelRef = string | LanguageModel;
+
 const run = async (cmd: string, args: string[], opts: { cwd: string }) => {
   return await new Promise<{ stdout: string; stderr: string }>((resolve, reject) => {
     const child = spawn(cmd, args, { cwd: opts.cwd, stdio: ["ignore", "pipe", "pipe"] });
@@ -87,7 +92,7 @@ export function createVideoFramesExtractor(): AssetExtractor {
 
           const imgBytes = await readFile(path.join(tmpDir, f));
           const result = await generateText({
-            model: cfg.model as any,
+            model: cfg.model as ModelRef,
             abortSignal: abortPerFrame(cfg.timeoutMs),
             messages: [
               {
@@ -100,7 +105,7 @@ export function createVideoFramesExtractor(): AssetExtractor {
             ],
           });
 
-          const t = String((result as any)?.text ?? "").trim();
+          const t = (result.text ?? "").trim();
           if (!t) continue;
 
           const capped = capText(t, cfg.maxOutputChars - totalChars);