@strav/rag 0.4.31 → 1.0.0-alpha.19

package/src/drivers/pgvector_driver.ts

@@ -1,157 +1,222 @@
-import { Database } from '@strav/database'
+/**
+ * `PgvectorDriver` — `VectorStore` backed by Postgres + the
+ * `pgvector` extension. Single table per app (`rag_vector` by
+ * default), `collection` is a column inside it.
+ *
+ * Multitenancy: every query relies on RLS scoping by
+ * `current_setting('app.tenant_id')`. Apps wrap calls in
+ * `tenants.withTenant(tenantId, async () => { ... })` — the
+ * driver itself has no tenant awareness.
+ *
+ * Why one table instead of one-per-collection:
+ *
+ *   - `defineSchema` doesn't support runtime table creation.
+ *   - HNSW indexes work fine with `collection` as a leading
+ *     column; if a collection grows past tens of millions and
+ *     wants its own partial HNSW, that's a one-line follow-up
+ *     migration.
+ *   - One RLS policy, one set of grants, fewer surprises.
+ *
+ * Why this driver doesn't extend `Repository`:
+ *
+ *   - The framework repository hydrates rows into a `Model`, but
+ *     `embedding vector(N)` isn't expressible in the framework's
+ *     type system. The driver uses raw `db.query` / `db.execute`
+ *     on the table and returns plain objects.
+ *   - All vector ops (`<=>`, `vector_cosine_ops`) are
+ *     pgvector-specific; the framework's query builder can't
+ *     model them.
+ */
+
+import type { PostgresDatabase } from '@strav/database'
+import { VectorQueryError } from '../rag_error.ts'
+import { ragVectorSchema } from '../rag_vector_schema.ts'
+import type {
+  QueryOptions,
+  QueryResult,
+  StoreConfig,
+  VectorDocument,
+  VectorMatch,
+} from '../types.ts'
 import type { VectorStore } from '../vector_store.ts'
-import type { VectorDocument, QueryOptions, QueryResult, VectorMatch, StoreConfig } from '../types.ts'
-import { VectorQueryError } from '../errors.ts'
+
+export interface PgvectorDriverOptions {
+  /** PostgresDatabase instance — typically resolved from the container. */
+  db: PostgresDatabase
+  /** Override table name. Defaults to `rag_vector`. */
+  table?: string
+}
 
 export class PgvectorDriver implements VectorStore {
   readonly name = 'pgvector'
-  private initialized = false
 
-  constructor(_config: StoreConfig) {}
+  private readonly db: PostgresDatabase
+  private readonly table: string
 
-  async createCollection(collection: string, dimension: number): Promise<void> {
-    await this.ensureTable(dimension)
+  constructor(options: PgvectorDriverOptions) {
+    this.db = options.db
+    this.table = options.table ?? ragVectorSchema.name
+  }
 
-    const indexName = `idx_strav_vectors_hnsw_${collection.replace(/[^a-z0-9_]/gi, '_')}`
-    try {
-      await Database.raw.unsafe(
-        `CREATE INDEX IF NOT EXISTS "${indexName}"
-         ON _strav_vectors USING hnsw (embedding vector_cosine_ops)
-         WHERE collection = '${collection}'`
-      )
-    } catch {
-      // Index may already exist
-    }
+  /**
+   * Factory used by `RagManager.createStore` — accepts the raw
+   * `StoreConfig` from `config.rag.stores[<name>]` and resolves
+   * the `db` from the container. Apps that want explicit control
+   * `new PgvectorDriver({ db, table })` directly.
+   */
+  static fromConfig(db: PostgresDatabase, config: StoreConfig): PgvectorDriver {
+    return new PgvectorDriver({
+      db,
+      ...(typeof config.table === 'string' ? { table: config.table } : {}),
+    })
+  }
+
+  // ─── Collections ──────────────────────────────────────────────────────
+
+  async createCollection(_collection: string, _dimension: number): Promise<void> {
+    // No-op: every collection lives in the same table. The
+    // `applyRagVectorMigration` helper attached the
+    // `vector(<dimension>)` column at migration time, so the
+    // dimension is fixed per table and enforced at INSERT.
   }
 
   async deleteCollection(collection: string): Promise<void> {
-    await Database.raw.unsafe(
-      `DELETE FROM _strav_vectors WHERE collection = $1`,
-      [collection]
+    await this.db.execute(
+      `DELETE FROM "${this.table}" WHERE "collection" = $1`,
+      [collection],
     )
   }
 
-  async upsert(collection: string, documents: VectorDocument[]): Promise<void> {
-    const sql = Database.raw
+  // ─── Mutations ────────────────────────────────────────────────────────
 
+  async upsert(
+    collection: string,
+    documents: readonly VectorDocument[],
+  ): Promise<void> {
+    if (documents.length === 0) return
+    // pgvector accepts the vector as a stringified array literal —
+    // `[0.12,0.34,...]` — cast with `::vector` at the boundary.
     for (const doc of documents) {
-      const embeddingStr = `[${doc.embedding.join(',')}]`
-      const metadata = JSON.stringify(doc.metadata ?? {})
-      const id = doc.id != null ? String(doc.id) : crypto.randomUUID()
-      const sourceId = doc.sourceId != null ? String(doc.sourceId) : null
-
-      await sql.unsafe(
-        `INSERT INTO _strav_vectors (collection, source_id, content, metadata, embedding)
-         VALUES ($1, $2, $3, $4::jsonb, $5::vector)`,
-        [collection, sourceId, doc.content, metadata, embeddingStr]
+      const id = doc.id ?? crypto.randomUUID()
+      const embeddingLiteral = `[${doc.embedding.join(',')}]`
+      await this.db.execute(
+        `INSERT INTO "${this.table}"
+          ("id", "collection", "source_id", "content", "metadata", "embedding", "created_at")
+         VALUES ($1, $2, $3, $4, $5::jsonb, $6::vector, NOW())
+         ON CONFLICT ("id") DO UPDATE SET
+           "collection" = EXCLUDED."collection",
+           "source_id"  = EXCLUDED."source_id",
+           "content"    = EXCLUDED."content",
+           "metadata"   = EXCLUDED."metadata",
+           "embedding"  = EXCLUDED."embedding"`,
+        [
+          id,
+          collection,
+          doc.sourceId ?? null,
+          doc.content,
+          JSON.stringify(doc.metadata ?? {}),
+          embeddingLiteral,
+        ],
       )
     }
   }
 
-  async delete(collection: string, ids: (string | number)[]): Promise<void> {
+  async delete(collection: string, ids: readonly string[]): Promise<void> {
     if (ids.length === 0) return
     const placeholders = ids.map((_, i) => `$${i + 2}`).join(', ')
-    await Database.raw.unsafe(
-      `DELETE FROM _strav_vectors WHERE collection = $1 AND id IN (${placeholders})`,
-      [collection, ...ids]
+    await this.db.execute(
+      `DELETE FROM "${this.table}" WHERE "collection" = $1 AND "id" IN (${placeholders})`,
+      [collection, ...ids],
     )
   }
 
-  async deleteBySource(collection: string, sourceId: string | number): Promise<void> {
-    await Database.raw.unsafe(
-      `DELETE FROM _strav_vectors WHERE collection = $1 AND source_id = $2`,
-      [collection, String(sourceId)]
+  async deleteBySource(collection: string, sourceId: string): Promise<void> {
+    await this.db.execute(
+      `DELETE FROM "${this.table}" WHERE "collection" = $1 AND "source_id" = $2`,
+      [collection, sourceId],
     )
   }
 
   async flush(collection: string): Promise<void> {
-    await Database.raw.unsafe(
-      `DELETE FROM _strav_vectors WHERE collection = $1`,
-      [collection]
+    await this.db.execute(
+      `DELETE FROM "${this.table}" WHERE "collection" = $1`,
+      [collection],
     )
   }
 
+  // ─── Query ────────────────────────────────────────────────────────────
+
   async query(
     collection: string,
-    vector: number[],
-    options?: QueryOptions
+    vector: readonly number[],
+    options: QueryOptions = {},
   ): Promise<QueryResult> {
     const start = performance.now()
-    const topK = options?.topK ?? 5
-    const threshold = options?.threshold ?? 0
-    const embeddingStr = `[${vector.join(',')}]`
+    const topK = options.topK ?? 5
+    const threshold = options.threshold
 
-    let whereClause = 'collection = $1'
-    const params: unknown[] = [collection]
-    let paramIndex = 2
+    // pgvector's `<=>` is cosine distance in [0, 2]; `1 - (a <=> b)`
+    // is cosine similarity. We further map cos similarity in
+    // [-1, 1] → [0, 1] via `(s + 1) / 2` to match MemoryDriver so
+    // scores are comparable across drivers.
+    const params: unknown[] = [collection, `[${vector.join(',')}]`]
+    const where: string[] = [`"collection" = $1`]
 
-    if (options?.filter) {
+    if (options.filter) {
       for (const [key, value] of Object.entries(options.filter)) {
-        whereClause += ` AND metadata->>'${key}' = $${paramIndex}`
-        params.push(String(value))
-        paramIndex++
+        params.push(JSON.stringify(value))
+        where.push(`"metadata" @> jsonb_build_object('${escapeJsonbKey(key)}', $${params.length}::jsonb)`)
       }
     }
 
-    if (threshold > 0) {
-      whereClause += ` AND (embedding <=> $${paramIndex}::vector) <= $${paramIndex + 1}`
-      params.push(embeddingStr, 1 - threshold)
-      paramIndex += 2
+    let sql = `
+      SELECT "id", "source_id", "content", "metadata",
+             ((1 - ("embedding" <=> $2::vector)) + 1) / 2 AS score
+      FROM "${this.table}"
+      WHERE ${where.join(' AND ')}
+    `
+    if (threshold !== undefined) {
+      params.push(threshold)
+      sql += ` AND ((1 - ("embedding" <=> $2::vector)) + 1) / 2 >= $${params.length}`
     }
-
+    params.push(topK)
+    sql += ` ORDER BY "embedding" <=> $2::vector LIMIT $${params.length}`
+
+    let rows: Array<{
+      id: string
+      source_id: string | null
+      content: string
+      metadata: Record<string, unknown> | string
+      score: number | string
+    }>
     try {
-      const rows = (await Database.raw.unsafe(
-        `SELECT id, source_id, content, metadata,
-                1 - (embedding <=> $${paramIndex}::vector) AS score
-         FROM _strav_vectors
-         WHERE ${whereClause}
-         ORDER BY embedding <=> $${paramIndex}::vector
-         LIMIT $${paramIndex + 1}`,
-        [...params, embeddingStr, topK]
-      )) as any[]
-
-      const matches: VectorMatch[] = rows.map((row: any) => ({
-        id: row.source_id ?? row.id,
-        content: row.content,
-        score: parseFloat(row.score),
-        metadata: typeof row.metadata === 'string' ? JSON.parse(row.metadata) : row.metadata ?? {},
-      }))
-
-      return {
-        matches,
-        processingTimeMs: performance.now() - start,
-      }
-    } catch (err) {
-      throw new VectorQueryError(collection, err instanceof Error ? err.message : String(err))
-    }
-  }
-
-  private async ensureTable(dimension: number): Promise<void> {
-    if (this.initialized) return
-
-    const sql = Database.raw
-
-    await sql.unsafe(`CREATE EXTENSION IF NOT EXISTS vector`)
-
-    await sql.unsafe(`
-      CREATE TABLE IF NOT EXISTS _strav_vectors (
-        id BIGSERIAL PRIMARY KEY,
-        collection VARCHAR(255) NOT NULL,
-        source_id VARCHAR(255),
-        content TEXT NOT NULL,
-        metadata JSONB DEFAULT '{}',
-        embedding vector(${dimension}),
-        created_at TIMESTAMPTZ DEFAULT NOW()
+      rows = await this.db.query(sql, params)
+    } catch (cause) {
+      throw new VectorQueryError(
+        `pgvector query failed for collection "${collection}".`,
+        { context: { collection, table: this.table }, cause },
       )
-    `)
+    }
 
-    await sql.unsafe(
-      `CREATE INDEX IF NOT EXISTS idx_strav_vectors_collection ON _strav_vectors(collection)`
-    )
-    await sql.unsafe(
-      `CREATE INDEX IF NOT EXISTS idx_strav_vectors_source ON _strav_vectors(collection, source_id)`
-    )
+    const matches: VectorMatch[] = rows.map((r) => ({
+      id: r.id,
+      content: r.content,
+      score: typeof r.score === 'string' ? Number.parseFloat(r.score) : r.score,
+      metadata: typeof r.metadata === 'string' ? JSON.parse(r.metadata) : r.metadata,
+      sourceId: r.source_id,
+    }))
+    return { matches, processingTimeMs: performance.now() - start }
+  }
+}
 
-    this.initialized = true
+/**
+ * Escape a JSONB object key for embedding in an SQL string. Keys
+ * are app-supplied so we sanitize defensively — backslash-escape
+ * single quotes; refuse keys with NUL bytes.
+ */
+function escapeJsonbKey(key: string): string {
+  if (key.includes('\0')) {
+    throw new VectorQueryError(`pgvector filter key contains NUL byte: ${JSON.stringify(key)}`)
   }
+  return key.replace(/'/g, "''")
 }

package/src/index.ts

@@ -1,48 +1,52 @@
-// Manager
-export { default, default as RagManager } from './rag_manager.ts'
+// Public API of `@strav/rag`.
+//
+// V1: vector store abstraction + memory & pgvector drivers +
+// fixed-size & recursive chunkers + RagManager + RagProvider.
+// Composes with `@strav/brain` for embeddings and `@strav/database`
+// for pgvector persistence + multitenancy.
+//
+// Deferred to follow-up slices: `retrievable()` repository mixin,
+// CLI commands (`rag:reindex`, `rag:flush`), re-ranking strategies.
 
-// Provider
-export { default as RagProvider } from './rag_provider.ts'
-
-// Store interface
-export type { VectorStore } from './vector_store.ts'
-
-// Drivers
-export { NullDriver } from './drivers/null_driver.ts'
-export { MemoryDriver } from './drivers/memory_driver.ts'
-export { PgvectorDriver } from './drivers/pgvector_driver.ts'
-
-// Mixin
-export { retrievable } from './retrievable.ts'
-export type { RetrievableInstance, RetrievableModel } from './retrievable.ts'
-
-// Helper
-export { rag } from './helpers.ts'
-
-// Chunking
 export { createChunker } from './chunking/chunker.ts'
 export { FixedSizeChunker } from './chunking/fixed_size_chunker.ts'
 export { RecursiveChunker } from './chunking/recursive_chunker.ts'
-
-// Errors
-export { RagError, CollectionNotFoundError, VectorQueryError, EmbeddingError } from './errors.ts'
-
-// Types
+export { MemoryDriver } from './drivers/memory_driver.ts'
+export {
+  PgvectorDriver,
+  type PgvectorDriverOptions,
+} from './drivers/pgvector_driver.ts'
+export {
+  applyRagVectorMigration,
+  type ApplyRagVectorMigrationOptions,
+} from './migrations.ts'
+export {
+  CollectionNotFoundError,
+  EmbeddingError,
+  RagError,
+  VectorQueryError,
+} from './rag_error.ts'
+export {
+  type IngestOptions,
+  RagManager,
+  type RagManagerOptions,
+  type StoreFactory,
+} from './rag_manager.ts'
+export { RagProvider } from './rag_provider.ts'
+export { ragVectorSchema } from './rag_vector_schema.ts'
 export type {
-  RagConfig,
-  StoreConfig,
-  EmbeddingConfig,
+  Chunk,
+  Chunker,
   ChunkingConfig,
-  VectorDocument,
+  EmbeddingConfig,
   QueryOptions,
   QueryResult,
-  VectorMatch,
+  RagConfig,
   RetrieveOptions,
-  RerankOptions,
   RetrieveResult,
   RetrievedDocument,
-  Chunk,
-  Chunker,
+  StoreConfig,
+  VectorDocument,
+  VectorMatch,
 } from './types.ts'
-
-export type { IngestOptions } from './helpers.ts'
+export type { VectorStore } from './vector_store.ts'

package/src/migrations.ts

@@ -0,0 +1,116 @@
+/**
+ * Migration helpers — emit the DDL apps need to put `rag_vector`
+ * into a working state. The framework's `emitCreateTable` handles
+ * everything except the pgvector-specific bits (the `vector(N)`
+ * column type and the HNSW index). This module fills the gap.
+ *
+ * Apps drop one call into their migration:
+ *
+ * ```ts
+ * import { SchemaRegistry, emitDropTable, type Migration } from '@strav/database'
+ * import { applyRagVectorMigration, ragVectorSchema } from '@strav/rag'
+ *
+ * export const migration: Migration = {
+ *   name: '20260601000000_create_rag_vector',
+ *   async up(db) {
+ *     await applyRagVectorMigration(db, {
+ *       dimension: 1536,           // match the embedding model
+ *       registry,
+ *     })
+ *   },
+ *   async down(db) {
+ *     await db.execute(emitDropTable(ragVectorSchema.name).sql)
+ *   },
+ * }
+ * ```
+ *
+ * The helper is idempotent against `IF NOT EXISTS` clauses where
+ * Postgres supports them, but apps should still rely on the
+ * migration runner's tracking table for re-run safety rather than
+ * the helper itself.
+ */
+
+import {
+  emitCreateTable,
+  type DatabaseExecutor,
+  type SchemaRegistry,
+} from '@strav/database'
+import { ragVectorSchema } from './rag_vector_schema.ts'
+
+export interface ApplyRagVectorMigrationOptions {
+  /**
+   * Vector dimension. Must match the configured embedding model
+   * (OpenAI's `text-embedding-3-small` → 1536,
+   * `text-embedding-3-large` → 3072, Gemini's
+   * `text-embedding-004` → 768, etc.). Mismatched dimensions
+   * cause `vector` casts at INSERT to throw.
+   */
+  dimension: number
+  /**
+   * Schema registry — required for `emitCreateTable` to resolve
+   * foreign-key references (the tenant registry, in this case).
+   */
+  registry: SchemaRegistry
+  /**
+   * Optional override table name. Defaults to `rag_vector` (the
+   * `ragVectorSchema.name`). Apps that need multiple vector
+   * tables (e.g., one per dimension) override this here AND
+   * register their own schema variant under the override name.
+   */
+  table?: string
+  /**
+   * HNSW construction parameter `m`. Default Postgres-level
+   * default (16). Higher = better recall, slower builds.
+   */
+  hnswM?: number
+  /**
+   * HNSW construction parameter `ef_construction`. Default 64.
+   * Higher = better recall, slower builds.
+   */
+  hnswEfConstruction?: number
+}
+
+export async function applyRagVectorMigration(
+  db: DatabaseExecutor,
+  options: ApplyRagVectorMigrationOptions,
+): Promise<void> {
+  const table = options.table ?? ragVectorSchema.name
+  const { dimension, registry } = options
+
+  await db.execute(`CREATE EXTENSION IF NOT EXISTS vector`)
+
+  // Framework table + RLS + tenant_id column come from emitCreateTable.
+  await db.execute(emitCreateTable(ragVectorSchema, { registry }).sql)
+
+  // Vector column — pgvector-specific. NOT NULL because every
+  // ingested chunk has an embedding by construction.
+  await db.execute(
+    `ALTER TABLE "${table}" ADD COLUMN IF NOT EXISTS "embedding" vector(${dimension}) NOT NULL`,
+  )
+
+  // HNSW index on cosine ops — pgvector's default for similarity
+  // search. Partial index per collection isn't possible at
+  // CREATE INDEX time without a literal value; apps that have
+  // very large per-collection corpora add `WHERE collection = '...'`
+  // partial indexes in a separate migration.
+  const hnswOpts: string[] = []
+  if (options.hnswM !== undefined) hnswOpts.push(`m = ${options.hnswM}`)
+  if (options.hnswEfConstruction !== undefined) {
+    hnswOpts.push(`ef_construction = ${options.hnswEfConstruction}`)
+  }
+  const withClause = hnswOpts.length > 0 ? ` WITH (${hnswOpts.join(', ')})` : ''
+  await db.execute(
+    `CREATE INDEX IF NOT EXISTS "idx_${table}_embedding_hnsw"
+     ON "${table}" USING hnsw ("embedding" vector_cosine_ops)${withClause}`,
+  )
+
+  // Helpful secondary indexes for the standard access patterns.
+  await db.execute(
+    `CREATE INDEX IF NOT EXISTS "idx_${table}_collection"
+     ON "${table}" ("collection")`,
+  )
+  await db.execute(
+    `CREATE INDEX IF NOT EXISTS "idx_${table}_source_id"
+     ON "${table}" ("source_id") WHERE "source_id" IS NOT NULL`,
+  )
+}

package/src/rag_error.ts

@@ -0,0 +1,76 @@
+/**
+ * `RagError` hierarchy — typed wrappers for failures in the RAG
+ * stack. Each subclass carries a specific error code so apps can
+ * branch on the failure mode at the call site instead of parsing
+ * error messages.
+ *
+ * Three concrete subclasses ship in V1:
+ *
+ *   - `CollectionNotFoundError` — `rag.retrieve` against a
+ *     collection that doesn't exist on the active store. Apps
+ *     create the collection via `rag.createCollection(...)`
+ *     before the first ingest.
+ *
+ *   - `VectorQueryError` — the underlying store rejected the
+ *     query (bad dimension, malformed filter, etc.). Cause
+ *     carries the driver-native error.
+ *
+ *   - `EmbeddingError` — the brain provider rejected the
+ *     embedding call. Wraps the brain-side error so apps can
+ *     `error.cause instanceof BrainError` for retry logic.
+ */
+
+import { StravError } from '@strav/kernel'
+
+export class RagError extends StravError {
+  constructor(
+    message: string,
+    options: {
+      code?: string
+      status?: number
+      context?: Record<string, unknown>
+      cause?: unknown
+    } = {},
+  ) {
+    super(
+      message,
+      { code: options.code ?? 'rag.error', status: options.status ?? 500 },
+      { ...(options.context ? { context: options.context } : {}), ...(options.cause !== undefined ? { cause: options.cause } : {}) },
+    )
+  }
+}
+
+export class CollectionNotFoundError extends RagError {
+  constructor(collection: string, store: string) {
+    super(
+      `RAG collection "${collection}" does not exist on store "${store}". Call \`rag.createCollection("${collection}", dim)\` before the first ingest.`,
+      {
+        code: 'rag.collection_not_found',
+        status: 404,
+        context: { collection, store },
+      },
+    )
+  }
+}
+
+export class VectorQueryError extends RagError {
+  constructor(message: string, options: { context?: Record<string, unknown>; cause?: unknown } = {}) {
+    super(message, {
+      code: 'rag.vector_query',
+      status: 500,
+      ...(options.context ? { context: options.context } : {}),
+      ...(options.cause !== undefined ? { cause: options.cause } : {}),
+    })
+  }
+}
+
+export class EmbeddingError extends RagError {
+  constructor(message: string, options: { context?: Record<string, unknown>; cause?: unknown } = {}) {
+    super(message, {
+      code: 'rag.embedding',
+      status: 500,
+      ...(options.context ? { context: options.context } : {}),
+      ...(options.cause !== undefined ? { cause: options.cause } : {}),
+    })
+  }
+}