ai-rag-vectors

What this skill does

# AI RAG Vectors

Vector embedding layer that chunks parsed PDF pages, generates embeddings via AI Gateway, stores them in pgvector, and provides cosine similarity search for RAG retrieval.

## Prerequisites

- Next.js app with `src/` directory and App Router
- `db` skill installed (Drizzle ORM + Postgres)
- `ai-core` skill installed (`getModel()` at `@/lib/ai`)
- `ai-rag-ingest` skill installed (`document` + `documentPage` tables in `@/db/schema/rag`)
- Docker running with PostgreSQL

## Installation

```bash
bun add ai
```

> `ai` is likely already installed from `ai-core`. No additional packages needed — pgvector is a Postgres extension, and embeddings use the AI SDK `embed`/`embedMany` functions.

## Docker Update

Update your `docker-compose.yml` to use the pgvector image instead of plain Postgres:

Find this:

```yaml
  db:
    image: postgres:17-alpine
```

Replace with:

```yaml
  db:
    image: pgvector/pgvector:pg17
```

After updating, recreate the container:

```bash
docker compose down db && docker compose up -d db
```

## What Gets Created

```
src/
├── db/
│   └── schema/
│       └── rag.ts                          # Add documentChunk table (extend existing)
├── lib/
│   └── rag/
│       ├── chunker.ts                      # Recursive text splitter
│       ├── embeddings.ts                   # embed/embedMany wrappers + indexDocument()
│       └── search.ts                       # searchChunks() semantic search
└── app/
    └── api/
        └── rag/
            ├── documents/
            │   └── [documentId]/
            │       └── index/
            │           └── route.ts        # POST trigger indexing
            └── search/
                └── route.ts                # POST semantic search
```

## Database

After applying this skill, enable the pgvector extension and push the schema:

```bash
# Enable pgvector extension (run once)
docker exec -it postgres psql -U postgres -d app -c "CREATE EXTENSION IF NOT EXISTS vector;"

# Push schema
bunx drizzle-kit push
```

Then create the HNSW index for fast cosine similarity search:

```bash
docker exec -it postgres psql -U postgres -d app -c "CREATE INDEX IF NOT EXISTS document_chunk_embedding_idx ON document_chunk USING hnsw (embedding vector_cosine_ops);"
```

## Setup Steps

### Step 1: Extend `src/db/schema/rag.ts`

Add the `documentChunk` table and the custom vector column type to the existing RAG schema file.

Add these imports at the top:

```typescript
import { customType } from "drizzle-orm/pg-core";
```

Add this after the existing `documentPage` table:

```typescript
const vector = customType<{ data: number[]; dpiType: string }>({
  dataType() {
    return "vector(768)";
  },
  toDriver(value: number[]): string {
    return `[${value.join(",")}]`;
  },
  fromDriver(value: unknown): number[] {
    if (typeof value === "string") {
      return value
        .replace(/[\[\]]/g, "")
        .split(",")
        .map(Number);
    }
    return value as number[];
  },
});

export const documentChunk = pgTable("document_chunk", {
  id: uuid("id").defaultRandom().primaryKey(),
  documentId: uuid("document_id")
    .notNull()
    .references(() => document.id, { onDelete: "cascade" }),
  pageNumber: integer("page_number").notNull(),
  chunkIndex: integer("chunk_index").notNull(),
  textContent: text("text_content").notNull(),
  embedding: vector("embedding"),
  createdAt: timestamp("created_at", { withTimezone: true }).defaultNow().notNull(),
});
```

### Step 2: Create `src/lib/rag/chunker.ts`

```typescript
type Chunk = {
  text: string;
  pageNumber: number;
  chunkIndex: number;
};

const CHUNK_SIZE = 2000; // ~512 tokens
const CHUNK_OVERLAP = 200;

const SEPARATORS = ["\n\n", "\n", ". ", " ", ""];

/**
 * Recursive text splitter that respects paragraph/sentence boundaries.
 * Splits on the largest separator that produces chunks within the size limit.
 */
function splitText(text: string, separators: string[]): string[] {
  if (text.length <= CHUNK_SIZE) {
    return [text];
  }

  const separator = separators[0];
  const nextSeparators = separators.slice(1);

  if (separator === "") {
    // Last resort: hard split by character count
    const chunks: string[] = [];
    for (let i = 0; i < text.length; i += CHUNK_SIZE - CHUNK_OVERLAP) {
      chunks.push(text.slice(i, i + CHUNK_SIZE));
    }
    return chunks;
  }

  const parts = text.split(separator);
  const chunks: string[] = [];
  let current = "";

  for (const part of parts) {
    const candidate = current ? current + separator + part : part;

    if (candidate.length > CHUNK_SIZE) {
      if (current) {
        chunks.push(current);
      }

      if (part.length > CHUNK_SIZE) {
        // Part itself is too large — recurse with finer separators
        const subChunks = splitText(part, nextSeparators);
        chunks.push(...subChunks);
        current = "";
      } else {
        current = part;
      }
    } else {
      current = candidate;
    }
  }

  if (current.trim()) {
    chunks.push(current);
  }

  return chunks;
}

/**
 * Add overlap between adjacent chunks for context continuity.
 */
function addOverlap(chunks: string[]): string[] {
  if (chunks.length <= 1) return chunks;

  const result: string[] = [chunks[0]];

  for (let i = 1; i < chunks.length; i++) {
    const prevChunk = chunks[i - 1];
    const overlap = prevChunk.slice(-CHUNK_OVERLAP);
    result.push(overlap + chunks[i]);
  }

  return result;
}

/**
 * Chunk an array of page texts into overlapping chunks with metadata.
 * @param pageTexts - Array of text strings, one per page (1-indexed page numbers)
 * @returns Array of chunks with page number and chunk index
 */
export function chunkPages(pageTexts: string[]): Chunk[] {
  const allChunks: Chunk[] = [];

  for (let pageIdx = 0; pageIdx < pageTexts.length; pageIdx++) {
    const pageText = pageTexts[pageIdx].trim();
    if (!pageText) continue;

    const rawChunks = splitText(pageText, SEPARATORS);
    const overlappedChunks = addOverlap(rawChunks);

    for (let chunkIdx = 0; chunkIdx < overlappedChunks.length; chunkIdx++) {
      const text = overlappedChunks[chunkIdx].trim();
      if (!text) continue;

      allChunks.push({
        text,
        pageNumber: pageIdx + 1,
        chunkIndex: chunkIdx,
      });
    }
  }

  return allChunks;
}
```

### Step 3: Create `src/lib/rag/embeddings.ts`

```typescript
import { embed, embedMany } from "ai";
import { gateway } from "@ai-sdk/gateway";
import { db } from "@/lib/db";
import { documentPage, documentChunk, document } from "@/lib/db/schema/rag";
import { eq } from "drizzle-orm";
import { chunkPages } from "./chunker";

const EMBEDDING_MODEL = "google/text-embedding-004";
const EMBED_BATCH_SIZE = 50;

function getEmbeddingModel() {
  return gateway.textEmbeddingModel(EMBEDDING_MODEL);
}

/**
 * Generate a single embedding vector for a text string.
 */
export async function embedText(text: string): Promise<number[]> {
  const { embedding } = await embed({
    model: getEmbeddingModel(),
    value: text,
  });
  return embedding;
}

/**
 * Generate embeddings for multiple text strings in batch.
 */
export async function embedTexts(texts: string[]): Promise<number[][]> {
  const { embeddings } = await embedMany({
    model: getEmbeddingModel(),
    values: texts,
  });
  return embeddings;
}

/**
 * Index a document: chunk all pages, generate embeddings, and store in documentChunk table.
 * Call this after ai-rag-ingest has processed the document (status = "ready").
 */
export async function indexDocument(documentId: string): Promise<number> {
  // Fetch all pages for the document
  const pages = await db
    .select({
      pageNumber: documentPage.pageNumber,
      textContent: documentPage.textContent,
    })
    .from(documentPage)
    .where(eq(documentPage.documentId, documentId))
    .orderBy(documentPage.pageNumber);

  if (pages.length === 0) {
    throw new Error("No pages found for document. Process the PDF first.");
  }

  // Build pa