unrag 0.1.1

package/registry/core/chunking.ts

@@ -0,0 +1,62 @@
+import type { Chunker, ChunkingOptions, ChunkText } from "./types";
+
+const DEFAULT_CHUNK_SIZE = 200;
+const DEFAULT_CHUNK_OVERLAP = 40;
+
+export const defaultChunkingOptions: ChunkingOptions = {
+  chunkSize: DEFAULT_CHUNK_SIZE,
+  chunkOverlap: DEFAULT_CHUNK_OVERLAP,
+};
+
+const splitWords = (content: string) =>
+  content
+    .trim()
+    .split(/\s+/)
+    .filter(Boolean);
+
+export const defaultChunker: Chunker = (
+  content: string,
+  options: ChunkingOptions
+): ChunkText[] => {
+  const { chunkSize, chunkOverlap } = options;
+  const words = splitWords(content);
+  const chunks: ChunkText[] = [];
+
+  if (words.length === 0) {
+    return chunks;
+  }
+
+  let cursor = 0;
+  let index = 0;
+
+  const stride = Math.max(1, chunkSize - chunkOverlap);
+
+  while (cursor < words.length) {
+    const slice = words.slice(cursor, cursor + chunkSize);
+    const chunkContent = slice.join(" ").trim();
+
+    if (chunkContent.length === 0) {
+      break;
+    }
+
+    chunks.push({
+      index,
+      content: chunkContent,
+      tokenCount: slice.length,
+    });
+
+    cursor += stride;
+    index += 1;
+  }
+
+  return chunks;
+};
+
+export const resolveChunkingOptions = (
+  overrides?: Partial<ChunkingOptions>
+): ChunkingOptions => ({
+  ...defaultChunkingOptions,
+  ...overrides,
+});
+
+

package/registry/core/config.ts

@@ -0,0 +1,27 @@
+import type {
+  Chunker,
+  ContextEngineConfig,
+  ResolvedContextEngineConfig,
+} from "./types";
+import { defaultChunker, resolveChunkingOptions } from "./chunking";
+
+export const defineConfig = (config: ContextEngineConfig): ContextEngineConfig =>
+  config;
+
+const defaultIdGenerator = () => crypto.randomUUID();
+
+export const resolveConfig = (
+  config: ContextEngineConfig
+): ResolvedContextEngineConfig => {
+  const chunker: Chunker = config.chunker ?? defaultChunker;
+
+  return {
+    embedding: config.embedding,
+    store: config.store,
+    defaults: resolveChunkingOptions(config.defaults),
+    chunker,
+    idGenerator: config.idGenerator ?? defaultIdGenerator,
+  };
+};
+
+

package/registry/core/context-engine.ts

@@ -0,0 +1,34 @@
+import { ingest } from "./ingest";
+import { retrieve } from "./retrieve";
+import { defineConfig, resolveConfig } from "./config";
+import type {
+  ContextEngineConfig,
+  IngestInput,
+  IngestResult,
+  ResolvedContextEngineConfig,
+  RetrieveInput,
+  RetrieveResult,
+} from "./types";
+
+export class ContextEngine {
+  private readonly config: ResolvedContextEngineConfig;
+
+  constructor(config: ContextEngineConfig) {
+    this.config = resolveConfig(config);
+  }
+
+  async ingest(input: IngestInput): Promise<IngestResult> {
+    return ingest(this.config, input);
+  }
+
+  async retrieve(input: RetrieveInput): Promise<RetrieveResult> {
+    return retrieve(this.config, input);
+  }
+}
+
+export const createContextEngine = (config: ContextEngineConfig) =>
+  new ContextEngine(config);
+
+export { defineConfig };
+
+

package/registry/core/index.ts

@@ -0,0 +1,7 @@
+export { ContextEngine, createContextEngine, defineConfig } from "./context-engine";
+export { ingest } from "./ingest";
+export { retrieve } from "./retrieve";
+export { defaultChunker, resolveChunkingOptions } from "./chunking";
+export * from "./types";
+
+

package/registry/core/ingest.ts

@@ -0,0 +1,79 @@
+import type {
+  Chunk,
+  IngestInput,
+  IngestResult,
+  ResolvedContextEngineConfig,
+} from "./types";
+
+const now = () => performance.now();
+
+export const ingest = async (
+  config: ResolvedContextEngineConfig,
+  input: IngestInput
+): Promise<IngestResult> => {
+  const totalStart = now();
+  const chunkingStart = now();
+
+  const chunkingOptions = {
+    ...config.defaults,
+    ...input.chunking,
+  };
+
+  const metadata = input.metadata ?? {};
+  const documentId = config.idGenerator();
+
+  const chunks = config.chunker(input.content, chunkingOptions).map<Chunk>(
+    (chunk) => ({
+      id: config.idGenerator(),
+      documentId,
+      sourceId: input.sourceId,
+      index: chunk.index,
+      content: chunk.content,
+      tokenCount: chunk.tokenCount,
+      metadata,
+      documentContent: input.content,
+    })
+  );
+
+  const chunkingMs = now() - chunkingStart;
+  const embeddingStart = now();
+
+  const embeddedChunks = await Promise.all(
+    chunks.map(async (chunk) => {
+      const embedding = await config.embedding.embed({
+        text: chunk.content,
+        metadata,
+        position: chunk.index,
+        sourceId: chunk.sourceId,
+        documentId: chunk.documentId,
+      });
+
+      return {
+        ...chunk,
+        embedding,
+      };
+    })
+  );
+
+  const embeddingMs = now() - embeddingStart;
+  const storageStart = now();
+
+  await config.store.upsert(embeddedChunks);
+
+  const storageMs = now() - storageStart;
+  const totalMs = now() - totalStart;
+
+  return {
+    documentId,
+    chunkCount: embeddedChunks.length,
+    embeddingModel: config.embedding.name,
+    durations: {
+      totalMs,
+      chunkingMs,
+      embeddingMs,
+      storageMs,
+    },
+  };
+};
+
+

package/registry/core/retrieve.ts

@@ -0,0 +1,48 @@
+import type {
+  RetrieveInput,
+  RetrieveResult,
+  ResolvedContextEngineConfig,
+} from "./types";
+
+const now = () => performance.now();
+
+const DEFAULT_TOP_K = 8;
+
+export const retrieve = async (
+  config: ResolvedContextEngineConfig,
+  input: RetrieveInput
+): Promise<RetrieveResult> => {
+  const totalStart = now();
+
+  const embeddingStart = now();
+  const queryEmbedding = await config.embedding.embed({
+    text: input.query,
+    metadata: {},
+    position: 0,
+    sourceId: "query",
+    documentId: "query",
+  });
+  const embeddingMs = now() - embeddingStart;
+
+  const retrievalStart = now();
+  const chunks = await config.store.query({
+    embedding: queryEmbedding,
+    topK: input.topK ?? DEFAULT_TOP_K,
+    scope: input.scope,
+  });
+  const retrievalMs = now() - retrievalStart;
+
+  const totalMs = now() - totalStart;
+
+  return {
+    chunks,
+    embeddingModel: config.embedding.name,
+    durations: {
+      totalMs,
+      embeddingMs,
+      retrievalMs,
+    },
+  };
+};
+
+

package/registry/core/types.ts

@@ -0,0 +1,111 @@
+export type MetadataValue = string | number | boolean | null;
+
+export type Metadata = Record<
+  string,
+  MetadataValue | MetadataValue[] | undefined
+>;
+
+export type Chunk = {
+  id: string;
+  documentId: string;
+  sourceId: string;
+  index: number;
+  content: string;
+  tokenCount: number;
+  metadata: Metadata;
+  embedding?: number[];
+  documentContent?: string;
+};
+
+export type ChunkText = {
+  index: number;
+  content: string;
+  tokenCount: number;
+};
+
+export type ChunkingOptions = {
+  chunkSize: number;
+  chunkOverlap: number;
+};
+
+export type Chunker = (content: string, options: ChunkingOptions) => ChunkText[];
+
+export type EmbeddingInput = {
+  text: string;
+  metadata: Metadata;
+  position: number;
+  sourceId: string;
+  documentId: string;
+};
+
+export type EmbeddingProvider = {
+  name: string;
+  dimensions?: number;
+  embed: (input: EmbeddingInput) => Promise<number[]>;
+};
+
+export type VectorStore = {
+  upsert: (chunks: Chunk[]) => Promise<void>;
+  query: (params: {
+    embedding: number[];
+    topK: number;
+    scope?: {
+      sourceId?: string;
+    };
+  }) => Promise<Array<Chunk & { score: number }>>;
+};
+
+export type IngestInput = {
+  sourceId: string;
+  content: string;
+  metadata?: Metadata;
+  chunking?: Partial<ChunkingOptions>;
+};
+
+export type IngestResult = {
+  documentId: string;
+  chunkCount: number;
+  embeddingModel: string;
+  durations: {
+    totalMs: number;
+    chunkingMs: number;
+    embeddingMs: number;
+    storageMs: number;
+  };
+};
+
+export type RetrieveInput = {
+  query: string;
+  topK?: number;
+  scope?: {
+    sourceId?: string;
+  };
+};
+
+export type RetrieveResult = {
+  chunks: Array<Chunk & { score: number }>;
+  embeddingModel: string;
+  durations: {
+    totalMs: number;
+    embeddingMs: number;
+    retrievalMs: number;
+  };
+};
+
+export type ContextEngineConfig = {
+  embedding: EmbeddingProvider;
+  store: VectorStore;
+  defaults?: Partial<ChunkingOptions>;
+  chunker?: Chunker;
+  idGenerator?: () => string;
+};
+
+export type ResolvedContextEngineConfig = {
+  embedding: EmbeddingProvider;
+  store: VectorStore;
+  defaults: ChunkingOptions;
+  chunker: Chunker;
+  idGenerator: () => string;
+};
+
+

package/registry/docs/unrag.md

@@ -0,0 +1,82 @@
+# Unrag setup
+
+Unrag installs a small RAG module into your codebase with:
+- chunk → embed → store on ingest
+- embed → vector similarity search on retrieve
+
+## Environment variables
+
+Add these to your environment:
+- `DATABASE_URL` (Postgres connection string)
+- `AI_GATEWAY_API_KEY` (required by the `ai` SDK when using Vercel AI Gateway)
+- Optional: `AI_GATEWAY_MODEL` (defaults to `openai/text-embedding-3-small`)
+
+## Database requirements
+
+Enable pgvector:
+
+```sql
+create extension if not exists vector;
+```
+
+## Schema (Postgres)
+
+You are responsible for migrations. Create these tables:
+
+```sql
+create table documents (
+  id uuid primary key,
+  source_id text not null,
+  metadata jsonb,
+  created_at timestamp default now()
+);
+
+create table chunks (
+  id uuid primary key,
+  document_id uuid not null references documents(id) on delete cascade,
+  source_id text not null,
+  idx integer not null,
+  content text not null,
+  token_count integer not null,
+  metadata jsonb,
+  created_at timestamp default now()
+);
+
+create table embeddings (
+  chunk_id uuid primary key references chunks(id) on delete cascade,
+  embedding vector,
+  embedding_dimension integer,
+  created_at timestamp default now()
+);
+```
+
+Recommended indexes:
+
+```sql
+create index if not exists chunks_source_id_idx on chunks(source_id);
+create index if not exists documents_source_id_idx on documents(source_id);
+```
+
+<!-- __UNRAG_ADAPTER_NOTES__ -->
+
+## Usage (Next.js)
+
+- Use the engine only on the server (Route Handlers / Server Actions).
+- Prefer a singleton DB client/pool pattern to avoid hot-reload connection storms.
+- If Unrag detected Next.js, it added:
+  - `@unrag/*` path alias to your installed module directory
+  - `@unrag/config` path alias to `./unrag.config.ts`
+
+Example route handler:
+
+```ts
+import { createUnragEngine } from "@unrag/config";
+
+export async function GET() {
+  const engine = createUnragEngine();
+  const result = await engine.retrieve({ query: "hello", topK: 5 });
+  return Response.json(result);
+}
+```
+
+

package/registry/embedding/ai.ts

@@ -0,0 +1,42 @@
+import { embed } from "ai";
+import type { EmbeddingProvider } from "../core/types";
+
+export type AiEmbeddingConfig = {
+  /**
+   * AI Gateway model id, e.g. "openai/text-embedding-3-small"
+   */
+  model?: string;
+  timeoutMs?: number;
+};
+
+const DEFAULT_MODEL = "openai/text-embedding-3-small";
+
+export const createAiEmbeddingProvider = (
+  config: AiEmbeddingConfig = {}
+): EmbeddingProvider => {
+  const model = config.model ?? process.env.AI_GATEWAY_MODEL ?? DEFAULT_MODEL;
+
+  return {
+    name: `ai-sdk:${model}`,
+    dimensions: undefined,
+    embed: async ({ text }) => {
+      const abortSignal = config.timeoutMs
+        ? AbortSignal.timeout(config.timeoutMs)
+        : undefined;
+
+      const result = await embed({
+        model,
+        value: text,
+        ...(abortSignal ? { abortSignal } : {}),
+      });
+
+      if (!result.embedding) {
+        throw new Error("Embedding missing from AI SDK response");
+      }
+
+      return result.embedding;
+    },
+  };
+};
+
+

package/registry/store/drizzle-postgres-pgvector/index.ts

@@ -0,0 +1,4 @@
+export { createDrizzleVectorStore } from "./store";
+export { documents, chunks, embeddings, schema } from "./schema";
+
+

package/registry/store/drizzle-postgres-pgvector/schema.ts

@@ -0,0 +1,63 @@
+import {
+  customType,
+  integer,
+  jsonb,
+  pgTable,
+  primaryKey,
+  text,
+  timestamp,
+  uuid,
+} from "drizzle-orm/pg-core";
+
+const vector = (name: string, dimensions?: number) =>
+  customType<{ data: number[]; driverData: string }>({
+    dataType: () => (dimensions ? `vector(${dimensions})` : "vector"),
+    toDriver: (value) => {
+      const content = Array.isArray(value) ? value : [];
+      return `[${content.join(",")}]`;
+    },
+  })(name);
+
+export const documents = pgTable("documents", {
+  id: uuid("id").primaryKey(),
+  sourceId: text("source_id").notNull(),
+  content: text("content").notNull(),
+  metadata: jsonb("metadata").$type<Record<string, unknown> | null>(),
+  createdAt: timestamp("created_at", { mode: "date", withTimezone: false }).defaultNow(),
+});
+
+export const chunks = pgTable("chunks", {
+  id: uuid("id").primaryKey(),
+  documentId: uuid("document_id")
+    .notNull()
+    .references(() => documents.id, { onDelete: "cascade" }),
+  sourceId: text("source_id").notNull(),
+  index: integer("idx").notNull(),
+  content: text("content").notNull(),
+  tokenCount: integer("token_count").notNull(),
+  metadata: jsonb("metadata").$type<Record<string, unknown> | null>(),
+  createdAt: timestamp("created_at", { mode: "date", withTimezone: false }).defaultNow(),
+});
+
+export const embeddings = pgTable(
+  "embeddings",
+  {
+    chunkId: uuid("chunk_id")
+      .notNull()
+      .references(() => chunks.id, { onDelete: "cascade" }),
+    embedding: vector("embedding"),
+    embeddingDimension: integer("embedding_dimension"),
+    createdAt: timestamp("created_at", { mode: "date", withTimezone: false }).defaultNow(),
+  },
+  (table) => ({
+    pk: primaryKey({ columns: [table.chunkId] }),
+  })
+);
+
+export const schema = {
+  documents,
+  chunks,
+  embeddings,
+};
+
+

package/registry/store/drizzle-postgres-pgvector/store.ts

@@ -0,0 +1,145 @@
+import { documents, chunks, embeddings } from "./schema";
+import type { Chunk, VectorStore } from "../../core/types";
+import { sql, type SQL } from "drizzle-orm";
+import type { PgDatabase } from "drizzle-orm/pg-core";
+
+type DrizzleDb = PgDatabase<any, any, any>;
+
+const sanitizeMetadata = (metadata: unknown) => {
+  if (metadata === undefined) {
+    return null;
+  }
+
+  try {
+    return JSON.parse(JSON.stringify(metadata));
+  } catch {
+    return null;
+  }
+};
+
+const toDocumentRow = (chunk: Chunk) => ({
+  id: chunk.documentId,
+  sourceId: chunk.sourceId,
+  content: chunk.documentContent ?? "",
+  metadata: sanitizeMetadata(chunk.metadata) as Record<string, unknown> | null,
+});
+
+const toChunkRow = (chunk: Chunk) => ({
+  id: chunk.id,
+  documentId: chunk.documentId,
+  sourceId: chunk.sourceId,
+  index: chunk.index,
+  content: chunk.content,
+  tokenCount: chunk.tokenCount,
+  metadata: sanitizeMetadata(chunk.metadata) as Record<string, unknown> | null,
+});
+
+export const createDrizzleVectorStore = (db: DrizzleDb): VectorStore => ({
+  upsert: async (chunkItems) => {
+    if (chunkItems.length === 0) {
+      return;
+    }
+
+    await db.transaction(async (tx) => {
+      const head = chunkItems[0]!;
+      const documentRow = toDocumentRow(head);
+
+      await tx
+        .insert(documents)
+        .values(documentRow)
+        .onConflictDoUpdate({
+          target: documents.id,
+          set: {
+            sourceId: documentRow.sourceId,
+            content: documentRow.content,
+            metadata: documentRow.metadata,
+          },
+        });
+
+      for (const chunk of chunkItems) {
+        const chunkRow = toChunkRow(chunk);
+
+        await tx
+          .insert(chunks)
+          .values(chunkRow)
+          .onConflictDoUpdate({
+            target: chunks.id,
+            set: {
+              content: chunkRow.content,
+              tokenCount: chunkRow.tokenCount,
+              metadata: chunkRow.metadata,
+              index: chunkRow.index,
+              sourceId: chunkRow.sourceId,
+            },
+          });
+
+        if (!chunk.embedding) {
+          continue;
+        }
+
+        await tx
+          .insert(embeddings)
+          .values({
+            chunkId: chunk.id,
+            embedding: chunk.embedding,
+            embeddingDimension: chunk.embedding.length,
+          })
+          .onConflictDoUpdate({
+            target: embeddings.chunkId,
+            set: {
+              embedding: chunk.embedding,
+              embeddingDimension: chunk.embedding.length,
+            },
+          });
+      }
+    });
+  },
+
+  query: async ({ embedding, topK, scope = {} }) => {
+    const filters: SQL[] = [];
+
+    if (scope.sourceId) {
+      // Interpret scope.sourceId as a prefix so callers can namespace content
+      // (e.g. `tenant:acme:`) without needing separate tables.
+      filters.push(sql`c.source_id like ${scope.sourceId + "%"}`);
+    }
+
+    const whereClause =
+      filters.length > 0 ? sql`where ${sql.join(filters, sql` and `)}` : sql``;
+
+    const vectorLiteral = `[${embedding.join(",")}]`;
+
+    const rows = await db.execute(
+      sql`
+        select
+          c.id,
+          c.document_id,
+          c.source_id,
+          c.idx,
+          c.content,
+          c.token_count,
+          c.metadata,
+          (e.embedding <=> ${vectorLiteral}) as score
+        from ${chunks} as c
+        join ${embeddings} as e on e.chunk_id = c.id
+        join ${documents} as d on d.id = c.document_id
+        ${whereClause}
+        order by score asc
+        limit ${topK}
+      `
+    );
+
+    return (rows as Array<Record<string, unknown>>).map((row) => ({
+      id: String(row.id),
+      documentId: String(row.document_id),
+      sourceId: String(row.source_id),
+      index: Number(row.idx),
+      content: String(row.content),
+      tokenCount: Number(row.token_count),
+      metadata: (row.metadata ?? {}) as Chunk["metadata"],
+      score: Number(row.score),
+    }));
+  },
+});
+
+