unrag 0.2.6 → 0.2.7

package/registry/connectors/google-drive/types.ts

@@ -1,5 +1,78 @@
 import type { ContextEngine, AssetInput, IngestInput } from "../../core";
 
+/**
+ * Service account credentials structure.
+ */
+export interface ServiceAccountCredentials {
+  client_email: string;
+  private_key: string;
+  [key: string]: unknown;
+}
+
+/**
+ * OAuth auth with an existing client instance.
+ */
+export type GoogleDriveOAuthClientAuth = {
+  /** Use an existing OAuth2 client instance (recommended if your app already has one). */
+  kind: "oauth";
+  oauthClient: unknown;
+  clientId?: never;
+  clientSecret?: never;
+  redirectUri?: never;
+  refreshToken?: never;
+  accessToken?: never;
+};
+
+/**
+ * OAuth auth with credentials for building a client.
+ */
+export type GoogleDriveOAuthConfigAuth = {
+  /**
+   * Convenience form for OAuth2: the connector will construct an OAuth2 client
+   * and set credentials including the refresh token.
+   */
+  kind: "oauth";
+  clientId: string;
+  clientSecret: string;
+  redirectUri: string;
+  refreshToken: string;
+  /** Optional access token if you already have one. */
+  accessToken?: string;
+  oauthClient?: never;
+};
+
+/**
+ * OAuth auth (either form).
+ */
+export type GoogleDriveOAuthAuth = GoogleDriveOAuthClientAuth | GoogleDriveOAuthConfigAuth;
+
+/**
+ * Service account auth.
+ */
+export type GoogleDriveServiceAccountAuth = {
+  /**
+   * Service account credentials. This supports both:
+   * - direct service-account access (files must be shared to the service account)
+   * - Workspace domain-wide delegation (DWD) when `subject` is provided
+   */
+  kind: "service_account";
+  credentialsJson: string | ServiceAccountCredentials;
+  /**
+   * DWD impersonation subject email (Workspace only).
+   * When provided, the service account will impersonate this user.
+   */
+  subject?: string;
+};
+
+/**
+ * Google Auth escape hatch.
+ */
+export type GoogleDriveGoogleAuthAuth = {
+  /** Escape hatch: provide a pre-configured GoogleAuth (or equivalent) instance. */
+  kind: "google_auth";
+  auth: unknown;
+};
+
 /**
  * A plug-and-play auth input for Google Drive.
  *
@@ -8,43 +81,9 @@ import type { ContextEngine, AssetInput, IngestInput } from "../../core";
  * by the CLI (`unrag add google-drive`).
  */
 export type GoogleDriveAuth =
-  | {
-      /** Use an existing OAuth2 client instance (recommended if your app already has one). */
-      kind: "oauth";
-      oauthClient: unknown;
-    }
-  | {
-      /**
-       * Convenience form for OAuth2: the connector will construct an OAuth2 client
-       * and set credentials including the refresh token.
-       */
-      kind: "oauth";
-      clientId: string;
-      clientSecret: string;
-      redirectUri: string;
-      refreshToken: string;
-      /** Optional access token if you already have one. */
-      accessToken?: string;
-    }
-  | {
-      /**
-       * Service account credentials. This supports both:
-       * - direct service-account access (files must be shared to the service account)
-       * - Workspace domain-wide delegation (DWD) when `subject` is provided
-       */
-      kind: "service_account";
-      credentialsJson: string | Record<string, unknown>;
-      /**
-       * DWD impersonation subject email (Workspace only).
-       * When provided, the service account will impersonate this user.
-       */
-      subject?: string;
-    }
-  | {
-      /** Escape hatch: provide a pre-configured GoogleAuth (or equivalent) instance. */
-      kind: "google_auth";
-      auth: unknown;
-    };
+  | GoogleDriveOAuthAuth
+  | GoogleDriveServiceAccountAuth
+  | GoogleDriveGoogleAuthAuth;
 
 export type GoogleDriveSyncProgressEvent =
   | { type: "file:start"; fileId: string; sourceId: string }

package/registry/connectors/notion/client.ts

@@ -1,4 +1,4 @@
-import { Client } from "@notionhq/client";
+import { Client, type ClientOptions } from "@notionhq/client";
 
 export type NotionClient = Client;
 
@@ -7,16 +7,25 @@ export type CreateNotionClientInput = {
   timeoutMs?: number;
 };
 
+/**
+ * Extended client options that include timeoutMs (supported by @notionhq/client).
+ */
+type NotionClientOptions = ClientOptions & {
+  timeoutMs?: number;
+};
+
 export function createNotionClient(input: CreateNotionClientInput): NotionClient {
   const token = input.token?.trim();
   if (!token) throw new Error("NOTION token is required");
 
-  return new Client({
+  const options: NotionClientOptions = {
     auth: token,
     // @notionhq/client uses undici/fetch under the hood; timeout is supported.
     // If unsupported in a future version, callers can wrap requests.
     timeoutMs: input.timeoutMs ?? 30_000,
-  } as any);
+  };
+
+  return new Client(options);
 }
 
 

package/registry/connectors/notion/render.ts

@@ -15,6 +15,31 @@ export type NotionBlockNode = {
   children: NotionBlockNode[];
 };
 
+/**
+ * Notion block content payload that may contain rich_text and other properties.
+ */
+interface BlockPayload {
+  rich_text?: RichText[];
+  checked?: boolean;
+  language?: string;
+  caption?: RichText[];
+  type?: string;
+  external?: { url?: string };
+  file?: { url?: string };
+  media_type?: string;
+}
+
+/**
+ * Get block-type-specific payload from a Notion block.
+ */
+const getBlockPayload = (block: NotionBlock, type: string): BlockPayload | undefined => {
+  const payload = block[type];
+  if (typeof payload === "object" && payload !== null) {
+    return payload as BlockPayload;
+  }
+  return undefined;
+};
+
 const rt = (value: unknown): string => {
   const items = Array.isArray(value) ? (value as RichText[]) : [];
   return items.map((t) => t?.plain_text ?? "").join("");
@@ -37,25 +62,30 @@ const toAssetKind = (notionType: string): AssetKind | null => {
   return supportedAssetKinds.has(t) ? t : null;
 };
 
-const pickUrl = (payload: any): string | undefined => {
-  const type = String(payload?.type ?? "");
-  if (type === "external") return asString(payload?.external?.url);
-  if (type === "file") return asString(payload?.file?.url);
+const pickUrl = (payload: BlockPayload | undefined): string | undefined => {
+  if (!payload) return undefined;
+  const type = String(payload.type ?? "");
+  if (type === "external") return asString(payload.external?.url);
+  if (type === "file") return asString(payload.file?.url);
   return undefined;
 };
 
-const pickCaption = (payload: any): string => {
+const pickCaption = (payload: BlockPayload | undefined): string => {
   // Notion captions are typically an array of rich text items.
   return rt(payload?.caption);
 };
 
-const inferMediaType = (assetKind: AssetKind, payload: any): string | undefined => {
+const inferMediaType = (assetKind: AssetKind, payload: BlockPayload | undefined): string | undefined => {
   if (assetKind === "pdf") return "application/pdf";
   // Notion does not consistently include media types; keep it optional.
   return asString(payload?.media_type) || undefined;
 };
 
-const asMetadata = (obj: Record<string, unknown>): Metadata => obj as any;
+/**
+ * Convert a plain object to Metadata type.
+ * The Metadata type allows string, number, boolean, null values.
+ */
+const toMetadata = (obj: Record<string, string>): Metadata => obj;
 
 export function extractNotionAssets(
   nodes: NotionBlockNode[],
@@ -66,10 +96,10 @@ export function extractNotionAssets(
 
   const walk = (node: NotionBlockNode, depth: number) => {
     if (depth > maxDepth) return;
-    const b = node.block as any;
+    const b = node.block;
     const kind = toAssetKind(String(b.type ?? ""));
     if (kind) {
-      const payload = b[kind];
+      const payload = getBlockPayload(b, kind);
       const url = pickUrl(payload);
       if (url) {
         const caption = pickCaption(payload).trim();
@@ -80,7 +110,7 @@ export function extractNotionAssets(
           data: { kind: "url", url, ...(mediaType ? { mediaType } : {}) },
           uri: url,
           ...(caption ? { text: caption } : {}),
-          metadata: asMetadata({
+          metadata: toMetadata({
             connector: "notion",
             notionBlockId: String(b.id),
             notionBlockType: String(b.type),
@@ -108,40 +138,49 @@ export function renderNotionBlocksToText(
   const walk = (node: NotionBlockNode, depth: number, listDepth: number) => {
     if (depth > maxDepth) return;
     const b = node.block;
-
     const t = b.type;
 
     if (t === "paragraph") {
-      const text = rt((b as any).paragraph?.rich_text);
+      const payload = getBlockPayload(b, "paragraph");
+      const text = rt(payload?.rich_text);
       if (text.trim()) lines.push(text);
     } else if (t === "heading_1") {
-      const text = rt((b as any).heading_1?.rich_text);
+      const payload = getBlockPayload(b, "heading_1");
+      const text = rt(payload?.rich_text);
       if (text.trim()) lines.push(`# ${text}`);
     } else if (t === "heading_2") {
-      const text = rt((b as any).heading_2?.rich_text);
+      const payload = getBlockPayload(b, "heading_2");
+      const text = rt(payload?.rich_text);
       if (text.trim()) lines.push(`## ${text}`);
     } else if (t === "heading_3") {
-      const text = rt((b as any).heading_3?.rich_text);
+      const payload = getBlockPayload(b, "heading_3");
+      const text = rt(payload?.rich_text);
       if (text.trim()) lines.push(`### ${text}`);
     } else if (t === "bulleted_list_item") {
-      const text = rt((b as any).bulleted_list_item?.rich_text);
+      const payload = getBlockPayload(b, "bulleted_list_item");
+      const text = rt(payload?.rich_text);
       if (text.trim()) lines.push(`${indent(listDepth)}- ${text}`);
     } else if (t === "numbered_list_item") {
-      const text = rt((b as any).numbered_list_item?.rich_text);
+      const payload = getBlockPayload(b, "numbered_list_item");
+      const text = rt(payload?.rich_text);
       if (text.trim()) lines.push(`${indent(listDepth)}- ${text}`);
     } else if (t === "to_do") {
-      const text = rt((b as any).to_do?.rich_text);
-      const checked = Boolean((b as any).to_do?.checked);
+      const payload = getBlockPayload(b, "to_do");
+      const text = rt(payload?.rich_text);
+      const checked = Boolean(payload?.checked);
       if (text.trim()) lines.push(`${indent(listDepth)}- [${checked ? "x" : " "}] ${text}`);
     } else if (t === "quote") {
-      const text = rt((b as any).quote?.rich_text);
+      const payload = getBlockPayload(b, "quote");
+      const text = rt(payload?.rich_text);
       if (text.trim()) lines.push(`> ${text}`);
     } else if (t === "callout") {
-      const text = rt((b as any).callout?.rich_text);
+      const payload = getBlockPayload(b, "callout");
+      const text = rt(payload?.rich_text);
       if (text.trim()) lines.push(text);
     } else if (t === "code") {
-      const text = rt((b as any).code?.rich_text);
-      const lang = String((b as any).code?.language ?? "").trim();
+      const payload = getBlockPayload(b, "code");
+      const text = rt(payload?.rich_text);
+      const lang = String(payload?.language ?? "").trim();
       lines.push("```" + lang);
       if (text.trim()) lines.push(text);
       lines.push("```");

package/registry/connectors/notion/sync.ts

@@ -1,4 +1,10 @@
-import type { IngestResult } from "../../core";
+import type { IngestResult, Metadata } from "../../core";
+import { isFullPage } from "@notionhq/client";
+import type {
+  GetPageResponse,
+  ListBlockChildrenResponse,
+  RichTextItemResponse,
+} from "@notionhq/client/build/src/api-endpoints";
 import { createNotionClient, type NotionClient } from "./client";
 import { normalizeNotionPageId32, toUuidHyphenated } from "./ids";
 import {
@@ -37,17 +43,16 @@ export function buildNotionPageIngestInput(
   };
 }
 
-const richTextToText = (richText: any[] | undefined) =>
-  (Array.isArray(richText) ? richText : [])
-    .map((t) => String(t?.plain_text ?? ""))
-    .join("");
+const richTextToText = (richText: RichTextItemResponse[] | undefined): string =>
+  (richText ?? []).map((t) => t.plain_text).join("");
 
-const getNotionPageTitle = (page: any): string => {
-  const props = page?.properties ?? {};
+const getNotionPageTitle = (page: GetPageResponse): string => {
+  if (!isFullPage(page)) return "";
+  const props = page.properties;
   for (const key of Object.keys(props)) {
     const p = props[key];
-    if (p?.type === "title") {
-      return richTextToText(p?.title);
+    if (p.type === "title") {
+      return richTextToText(p.title);
     }
   }
   return "";
@@ -61,15 +66,15 @@ async function listAllBlockChildren(
   let cursor: string | undefined = undefined;
 
   while (true) {
-    const res: any = await notion.blocks.children.list({
+    const res: ListBlockChildrenResponse = await notion.blocks.children.list({
       block_id: blockId,
       start_cursor: cursor,
       page_size: 100,
     });
 
-    blocks.push(...((res?.results ?? []) as NotionBlock[]));
-    if (!res?.has_more) break;
-    cursor = res?.next_cursor ?? undefined;
+    blocks.push(...(res.results as NotionBlock[]));
+    if (!res.has_more) break;
+    cursor = res.next_cursor ?? undefined;
     if (!cursor) break;
   }
 
@@ -105,30 +110,30 @@ export async function loadNotionPageDocument(args: {
   const pageId = normalizeNotionPageId32(args.pageIdOrUrl);
   const apiId = toUuidHyphenated(pageId);
 
-  const page: any = await args.notion.pages.retrieve({ page_id: apiId });
+  const page: GetPageResponse = await args.notion.pages.retrieve({ page_id: apiId });
   const title = getNotionPageTitle(page);
-  const url = String(page?.url ?? "");
-  const lastEditedTime = String(page?.last_edited_time ?? "");
+  const url = isFullPage(page) ? page.url : "";
+  const lastEditedTime = isFullPage(page) ? page.last_edited_time : "";
 
   const tree = await buildBlockTree(args.notion, apiId, 0, args.maxDepth ?? 4);
   const body = renderNotionBlocksToText(tree);
   const content = [title.trim(), body.trim()].filter(Boolean).join("\n\n");
   const assets = extractNotionAssets(tree);
 
-  const metadata = {
+  const metadata: Metadata = {
     connector: "notion",
     kind: "page",
     pageId,
     url,
     title,
     lastEditedTime,
-  } as const;
+  };
 
   const ingest = buildNotionPageIngestInput({
     pageId,
     content,
     assets,
-    metadata: metadata as any,
+    metadata,
     sourceIdPrefix: args.sourceIdPrefix,
   });
 
@@ -140,10 +145,12 @@ export async function loadNotionPageDocument(args: {
   };
 }
 
-const isNotFound = (err: any) => {
-  const status = Number(err?.status ?? err?.statusCode ?? err?.code);
+const isNotFound = (err: unknown): boolean => {
+  if (typeof err !== "object" || err === null) return false;
+  const e = err as Record<string, unknown>;
+  const status = Number(e.status ?? e.statusCode ?? e.code ?? 0);
   if (status === 404) return true;
-  const msg = String(err?.message ?? "");
+  const msg = String(e.message ?? "");
   return msg.toLowerCase().includes("could not find");
 };
 
@@ -187,7 +194,7 @@ export async function syncNotionPages(
         sourceId: doc.sourceId,
         content: doc.content,
         assets: doc.assets,
-        metadata: doc.metadata as any,
+        metadata: doc.metadata,
       });
 
       succeeded += 1;

package/registry/core/assets.ts

@@ -1,4 +1,5 @@
 import type { AssetKind, Chunk } from "./types";
+import { hasAssetMetadata } from "./types";
 
 export type ChunkAssetRef = {
   assetId: string;
@@ -21,26 +22,26 @@ const assetKinds = new Set<AssetKind>(["image", "pdf", "audio", "video", "file"]
 export function getChunkAssetRef(
   chunk: Pick<Chunk, "metadata">
 ): ChunkAssetRef | null {
-  const meta = chunk.metadata as any;
-  const kind = meta?.assetKind;
-  const id = meta?.assetId;
+  const meta = chunk.metadata;
 
-  if (typeof kind !== "string" || !assetKinds.has(kind as AssetKind)) {
+  if (!hasAssetMetadata(meta)) {
     return null;
   }
-  if (typeof id !== "string" || !id) {
+
+  const kind = meta.assetKind;
+  if (!assetKinds.has(kind)) {
     return null;
   }
 
-  const assetUri = typeof meta?.assetUri === "string" ? meta.assetUri : undefined;
+  const assetUri = typeof meta.assetUri === "string" ? meta.assetUri : undefined;
   const assetMediaType =
-    typeof meta?.assetMediaType === "string" ? meta.assetMediaType : undefined;
+    typeof meta.assetMediaType === "string" ? meta.assetMediaType : undefined;
   const extractor =
-    typeof meta?.extractor === "string" ? meta.extractor : undefined;
+    typeof meta.extractor === "string" ? meta.extractor : undefined;
 
   return {
-    assetId: id,
-    assetKind: kind as AssetKind,
+    assetId: meta.assetId,
+    assetKind: kind,
     ...(assetUri ? { assetUri } : {}),
     ...(assetMediaType ? { assetMediaType } : {}),
     ...(extractor ? { extractor } : {}),

package/registry/core/config.ts

@@ -3,10 +3,11 @@ import type {
   ContextEngineConfig,
   ResolvedContextEngineConfig,
   AssetProcessingConfig,
-  DeepPartial,
   ContentStorageConfig,
+  EmbeddingProcessingConfig,
 } from "./types";
 import { defaultChunker, resolveChunkingOptions } from "./chunking";
+import { mergeDeep } from "./deep-merge";
 
 export const defineConfig = (config: ContextEngineConfig): ContextEngineConfig =>
   config;
@@ -123,30 +124,9 @@ export const defaultContentStorageConfig: ContentStorageConfig = {
   storeDocumentContent: true,
 };
 
-const mergeDeep = <T extends Record<string, any>>(
-  base: T,
-  overrides: DeepPartial<T> | undefined
-): T => {
-  if (!overrides) return base;
-  const out: any = Array.isArray(base) ? [...base] : { ...base };
-  for (const key of Object.keys(overrides) as Array<keyof T>) {
-    const nextVal = overrides[key];
-    if (nextVal === undefined) continue;
-    const baseVal = base[key];
-    if (
-      baseVal &&
-      typeof baseVal === "object" &&
-      !Array.isArray(baseVal) &&
-      nextVal &&
-      typeof nextVal === "object" &&
-      !Array.isArray(nextVal)
-    ) {
-      out[key] = mergeDeep(baseVal, nextVal as any);
-    } else {
-      out[key] = nextVal as any;
-    }
-  }
-  return out as T;
+export const defaultEmbeddingProcessingConfig: EmbeddingProcessingConfig = {
+  concurrency: 4,
+  batchSize: 32,
 };
 
 export const resolveAssetProcessingConfig = (
@@ -157,6 +137,10 @@ export const resolveContentStorageConfig = (
   overrides?: DeepPartial<ContentStorageConfig>
 ): ContentStorageConfig => mergeDeep(defaultContentStorageConfig, overrides);
 
+export const resolveEmbeddingProcessingConfig = (
+  overrides?: DeepPartial<EmbeddingProcessingConfig>
+): EmbeddingProcessingConfig => mergeDeep(defaultEmbeddingProcessingConfig, overrides);
+
 export const resolveConfig = (
   config: ContextEngineConfig
 ): ResolvedContextEngineConfig => {
@@ -171,6 +155,7 @@ export const resolveConfig = (
     extractors: config.extractors ?? [],
     storage: resolveContentStorageConfig(config.storage),
     assetProcessing: resolveAssetProcessingConfig(config.assetProcessing),
+    embeddingProcessing: resolveEmbeddingProcessingConfig(config.embeddingProcessing),
   };
 };
 

package/registry/core/context-engine.ts

@@ -142,6 +142,7 @@ export const defineUnragConfig = <T extends DefineUnragConfigInput>(config: T) =
 
   const defaults = {
     chunking: config.defaults?.chunking ?? {},
+    embedding: config.defaults?.embedding ?? {},
     retrieval: {
       topK: config.defaults?.retrieval?.topK ?? 8,
     },
@@ -157,6 +158,10 @@ export const defineUnragConfig = <T extends DefineUnragConfigInput>(config: T) =
     return defineConfig({
       ...(config.engine ?? {}),
       defaults: defaults.chunking,
+      embeddingProcessing: {
+        ...(defaults.embedding ?? {}),
+        ...(config.engine?.embeddingProcessing ?? {}),
+      },
       embedding: getEmbeddingProvider(),
       store: runtime.store,
       extractors,

package/registry/core/deep-merge.ts

@@ -0,0 +1,45 @@
+import type { DeepPartial } from "./types";
+
+/**
+ * Type guard to check if a value is a plain object (not array, not null).
+ */
+export function isRecord(value: unknown): value is Record<string, unknown> {
+  return value !== null && typeof value === "object" && !Array.isArray(value);
+}
+
+/**
+ * Deep merge utility that recursively merges overrides into base.
+ * Arrays are replaced, not merged. Undefined values in overrides are skipped.
+ */
+export function mergeDeep<T extends Record<string, unknown>>(
+  base: T,
+  overrides: DeepPartial<T> | undefined
+): T {
+  if (!overrides) return base;
+
+  const out = (Array.isArray(base) ? [...base] : { ...base }) as T;
+
+  for (const key of Object.keys(overrides) as Array<keyof T>) {
+    const nextVal = overrides[key as keyof typeof overrides];
+    if (nextVal === undefined) continue;
+
+    const baseVal = base[key];
+
+    if (
+      isRecord(baseVal) &&
+      isRecord(nextVal) &&
+      !Array.isArray(baseVal) &&
+      !Array.isArray(nextVal)
+    ) {
+      out[key] = mergeDeep(
+        baseVal,
+        nextVal as DeepPartial<typeof baseVal>
+      ) as T[keyof T];
+    } else {
+      out[key] = nextVal as T[keyof T];
+    }
+  }
+
+  return out;
+}
+