@strav/rag 0.4.31 → 1.0.0-alpha.20

package/src/rag_manager.ts

@@ -1,98 +1,321 @@
-import { inject, Configuration, ConfigurationError } from '@strav/kernel'
-import type { VectorStore } from './vector_store.ts'
-import type { RagConfig, StoreConfig, EmbeddingConfig, ChunkingConfig } from './types.ts'
-import { NullDriver } from './drivers/null_driver.ts'
+/**
+ * `RagManager` — the facade apps use for RAG workflows.
+ *
+ * Three concept clusters:
+ *
+ *   - **Stores.** Apps register vector stores in
+ *     `config.rag.stores`; the manager constructs them lazily on
+ *     first `store(name)` call. Custom drivers register via
+ *     `manager.extend(name, factory)`.
+ *
+ *   - **Ingest.** `manager.ingest(collection, content, opts?)`
+ *     chunks → embeds → upserts. Returns the vector ids it
+ *     wrote. Apps that already have chunks bypass the chunker
+ *     by passing `IngestOptions.chunks`.
+ *
+ *   - **Retrieve.** `manager.retrieve(query, opts?)` embeds the
+ *     query, runs the store's similarity search, and returns
+ *     ranked matches. `topK` / `threshold` / `filter` pass
+ *     through.
+ *
+ * Multitenancy: invisible. The pgvector driver relies on RLS via
+ * `app.tenant_id` session settings, so apps that wrap calls in
+ * `tenants.withTenant(...)` get per-tenant isolation for free.
+ */
+
+// biome-ignore lint/style/useImportType: PostgresDatabase value import for the container path that wires PgvectorDriver.
+import { PostgresDatabase } from '@strav/database'
+// biome-ignore lint/style/useImportType: BrainManager value import for @inject() param-type metadata.
+import { BrainManager } from '@strav/brain'
+// biome-ignore lint/style/useImportType: Application value import for the container handle.
+import { Application, inject, ulid } from '@strav/kernel'
+import { createChunker } from './chunking/chunker.ts'
 import { MemoryDriver } from './drivers/memory_driver.ts'
 import { PgvectorDriver } from './drivers/pgvector_driver.ts'
+import { EmbeddingError, RagError } from './rag_error.ts'
+import type {
+  ChunkingConfig,
+  Chunk,
+  Chunker,
+  RagConfig,
+  RetrieveOptions,
+  RetrieveResult,
+  RetrievedDocument,
+  StoreConfig,
+  VectorDocument,
+} from './types.ts'
+import type { VectorStore } from './vector_store.ts'
 
-@inject
-export default class RagManager {
-  private static _config: RagConfig
-  private static _stores = new Map<string, VectorStore>()
-  private static _extensions = new Map<string, (config: StoreConfig) => VectorStore>()
-
-  constructor(config: Configuration) {
-    RagManager._config = {
-      default: config.get('rag.default', 'null') as string,
-      prefix: config.get('rag.prefix', '') as string,
-      embedding: config.get('rag.embedding', {
-        provider: 'openai',
-        model: 'text-embedding-3-small',
-        dimension: 1536,
-      }) as EmbeddingConfig,
-      chunking: config.get('rag.chunking', {
-        strategy: 'recursive',
-        chunkSize: 512,
-        overlap: 64,
-      }) as ChunkingConfig,
-      stores: config.get('rag.stores', {}) as Record<string, StoreConfig>,
-    }
-  }
+export interface IngestOptions {
+  /**
+   * Override the store. Defaults to the manager's default store.
+   */
+  store?: string
+  /** App-defined pointer back to the source row this content came from. */
+  sourceId?: string
+  /** Metadata attached to every chunk. Combined with per-chunk metadata if `chunks` is supplied. */
+  metadata?: Record<string, unknown>
+  /** Override chunking strategy + size for this call. */
+  chunking?: Partial<ChunkingConfig>
+  /**
+   * Pre-chunked content — skips the chunker entirely. The chunker
+   * helpers (`startOffset`, `endOffset`, `index`) carry through
+   * as metadata.
+   */
+  chunks?: readonly Chunk[]
+  /**
+   * Optional per-chunk sanitizer applied AFTER chunking, BEFORE
+   * embedding. Return `null` to drop the chunk; otherwise return
+   * the (possibly modified) text. Use to scrub PII / secrets /
+   * prompt-injection markers from untrusted source content.
+   */
+  sanitize?(chunk: { content: string; index: number }): string | null | Promise<string | null>
+  /** Override the brain provider used for the embedding call. */
+  embedProvider?: string
+  /** Override the embedding model. */
+  embedModel?: string
+}
+
+export interface RagManagerOptions {
+  config: RagConfig
+  brain: BrainManager
+  /** Optional — required only if pgvector stores are configured. */
+  db?: PostgresDatabase
+}
+
+/** Factory for custom drivers — apps register via `manager.extend(name, factory)`. */
+export type StoreFactory = (config: StoreConfig) => VectorStore
+
+@inject()
+export class RagManager {
+  readonly config: RagConfig
+  private readonly brain: BrainManager
+  private readonly db: PostgresDatabase | undefined
+  private readonly stores = new Map<string, VectorStore>()
+  private readonly extensions = new Map<string, StoreFactory>()
 
-  static get config(): RagConfig {
-    if (!RagManager._config) {
-      throw new ConfigurationError(
-        'RagManager not configured. Resolve it through the container first.'
+  constructor(options: RagManagerOptions) {
+    if (!options.config.stores[options.config.default]) {
+      throw new RagError(
+        `RagManager: default store "${options.config.default}" is not configured.`,
+        {
+          context: {
+            default: options.config.default,
+            available: Object.keys(options.config.stores),
+          },
+        },
       )
     }
-    return RagManager._config
+    this.config = options.config
+    this.brain = options.brain
+    this.db = options.db
   }
 
-  static store(name?: string): VectorStore {
-    const key = name ?? RagManager.config.default
+  // ─── Store management ─────────────────────────────────────────────────
 
-    let store = RagManager._stores.get(key)
-    if (store) return store
-
-    const storeConfig = RagManager.config.stores[key]
-    if (!storeConfig) {
-      throw new ConfigurationError(`RAG store "${key}" is not configured.`)
+  /**
+   * Resolve a vector store by name (or the default when omitted).
+   * Stores are constructed lazily on first use + memoized.
+   */
+  store(name?: string): VectorStore {
+    const key = name ?? this.config.default
+    const cached = this.stores.get(key)
+    if (cached) return cached
+    const cfg = this.config.stores[key]
+    if (!cfg) {
+      throw new RagError(`RagManager: store "${key}" is not configured.`, {
+        context: { requested: key, available: Object.keys(this.config.stores) },
+      })
     }
-
-    store = RagManager.createStore(key, storeConfig)
-    RagManager._stores.set(key, store)
+    const store = this.createStore(cfg)
+    this.stores.set(key, store)
     return store
   }
 
-  static get prefix(): string {
-    return RagManager._config?.prefix ?? ''
+  /** Register a custom driver. Subsequent `store(...)` calls can resolve `driver: <name>`. */
+  extend(name: string, factory: StoreFactory): void {
+    this.extensions.set(name, factory)
   }
 
-  static collectionName(name: string): string {
-    return RagManager.prefix ? `${RagManager.prefix}${name}` : name
+  /** Hand-wire a store instance under a name (tests / one-off drivers). */
+  useStore(name: string, store: VectorStore): void {
+    this.stores.set(name, store)
   }
 
-  static extend(name: string, factory: (config: StoreConfig) => VectorStore): void {
-    RagManager._extensions.set(name, factory)
+  /**
+   * Compose the configured prefix with `name` — apps that want to
+   * namespace collections (per-tenant, per-app) set `config.rag.prefix`
+   * and call `manager.collectionName(...)` to resolve at runtime.
+   */
+  collectionName(name: string): string {
+    return this.config.prefix ? `${this.config.prefix}${name}` : name
   }
 
-  static useStore(store: VectorStore): void {
-    RagManager._stores.set(store.name, store)
+  // ─── Ingest ───────────────────────────────────────────────────────────
+
+  /**
+   * Chunk → embed → upsert. Returns the vector ids it wrote. The
+   * caller-supplied `collection` is composed with the configured
+   * `prefix` before hitting the store.
+   */
+  async ingest(
+    collection: string,
+    content: string,
+    options: IngestOptions = {},
+  ): Promise<string[]> {
+    const fullCollection = this.collectionName(collection)
+    const chunkerConfig: ChunkingConfig = {
+      strategy: options.chunking?.strategy ?? this.config.chunking.strategy,
+      chunkSize: options.chunking?.chunkSize ?? this.config.chunking.chunkSize,
+      overlap: options.chunking?.overlap ?? this.config.chunking.overlap,
+      ...(options.chunking?.separators ?? this.config.chunking.separators
+        ? { separators: options.chunking?.separators ?? this.config.chunking.separators }
+        : {}),
+    }
+
+    let chunks: Chunk[] = options.chunks
+      ? [...options.chunks]
+      : this.chunker(chunkerConfig).chunk(content)
+    if (chunks.length === 0) return []
+
+    if (options.sanitize) {
+      const filtered: Chunk[] = []
+      for (const chunk of chunks) {
+        const next = await options.sanitize({ content: chunk.content, index: chunk.index })
+        if (next === null) continue
+        filtered.push({ ...chunk, content: next })
+      }
+      chunks = filtered
+      if (chunks.length === 0) return []
+    }
+
+    const texts = chunks.map((c) => c.content)
+    let embeddings: number[][]
+    try {
+      const result = await this.brain.embed(texts, {
+        provider: options.embedProvider ?? this.config.embedding.provider,
+        model: options.embedModel ?? this.config.embedding.model,
+      })
+      embeddings = result.embeddings as number[][]
+    } catch (cause) {
+      throw new EmbeddingError(
+        `RagManager.ingest: embedding ${texts.length} chunks failed.`,
+        { context: { collection: fullCollection }, cause },
+      )
+    }
+
+    const documents: VectorDocument[] = chunks.map((chunk, i) => ({
+      id: ulid(),
+      ...(options.sourceId !== undefined ? { sourceId: options.sourceId } : {}),
+      content: chunk.content,
+      embedding: embeddings[i]!,
+      metadata: {
+        ...(options.metadata ?? {}),
+        chunkIndex: chunk.index,
+        startOffset: chunk.startOffset,
+        endOffset: chunk.endOffset,
+      },
+    }))
+
+    await this.store(options.store).upsert(fullCollection, documents)
+    return documents.map((d) => d.id as string)
+  }
+
+  // ─── Retrieve ─────────────────────────────────────────────────────────
+
+  async retrieve(
+    query: string,
+    options: RetrieveOptions = {},
+  ): Promise<RetrieveResult> {
+    const fullCollection = this.collectionName(
+      options.collection ?? this.config.default,
+    )
+    const start = performance.now()
+
+    let embedding: number[]
+    try {
+      const result = await this.brain.embed([query], {
+        provider: options.embedProvider ?? this.config.embedding.provider,
+        model: options.embedModel ?? this.config.embedding.model,
+      })
+      embedding = result.embeddings[0] as number[]
+    } catch (cause) {
+      throw new EmbeddingError(
+        `RagManager.retrieve: embedding query failed.`,
+        { context: { collection: fullCollection }, cause },
+      )
+    }
+
+    const queryOpts: { topK?: number; threshold?: number; filter?: Record<string, unknown> } = {}
+    if (options.topK !== undefined) queryOpts.topK = options.topK
+    if (options.threshold !== undefined) queryOpts.threshold = options.threshold
+    if (options.filter !== undefined) queryOpts.filter = options.filter
+
+    const result = await this.store(options.store).query(
+      fullCollection,
+      embedding,
+      queryOpts,
+    )
+
+    const matches: RetrievedDocument[] = result.matches.map((m) => ({
+      id: m.id,
+      content: m.content,
+      score: m.score,
+      similarity: m.score,
+      metadata: m.metadata,
+      ...(m.sourceId !== undefined ? { sourceId: m.sourceId } : {}),
+    }))
+
+    return {
+      query,
+      matches,
+      processingTimeMs: performance.now() - start,
+    }
   }
 
-  static reset(): void {
-    RagManager._stores.clear()
-    RagManager._extensions.clear()
-    RagManager._config = undefined as any
+  /**
+   * Create a collection on the active (or named) store. For
+   * pgvector this is a no-op — every collection lives in the
+   * same table. For MemoryDriver this allocates the bucket.
+   * Apps call it once at boot or before the first ingest of a
+   * new collection name.
+   */
+  async createCollection(
+    collection: string,
+    options: { store?: string; dimension?: number } = {},
+  ): Promise<void> {
+    const fullCollection = this.collectionName(collection)
+    const dimension = options.dimension ?? this.config.embedding.dimension
+    await this.store(options.store).createCollection(fullCollection, dimension)
   }
 
-  private static createStore(name: string, config: StoreConfig): VectorStore {
-    const driverName = config.driver ?? name
+  // ─── Internals ────────────────────────────────────────────────────────
 
-    const extension = RagManager._extensions.get(driverName)
-    if (extension) return extension(config)
+  private chunker(config: ChunkingConfig): Chunker {
+    return createChunker(config)
+  }
 
-    switch (driverName) {
-      case 'pgvector':
-        return new PgvectorDriver(config)
+  private createStore(config: StoreConfig): VectorStore {
+    const ext = this.extensions.get(config.driver)
+    if (ext) return ext(config)
+    switch (config.driver) {
       case 'memory':
         return new MemoryDriver()
-      case 'null':
-        return new NullDriver()
+      case 'pgvector':
+        if (!this.db) {
+          throw new RagError(
+            'RagManager: pgvector driver requires a PostgresDatabase. Register DatabaseProvider before RagProvider, or pass `db` to the manager constructor.',
+          )
+        }
+        return PgvectorDriver.fromConfig(this.db, config)
       default:
-        throw new ConfigurationError(
-          `Unknown RAG driver "${driverName}". Register it with RagManager.extend().`
+        throw new RagError(
+          `RagManager: unknown driver "${config.driver}". Register it via \`manager.extend(...)\`.`,
+          { context: { driver: config.driver } },
         )
     }
   }
 }
+
+/** Public alias for the container-resolution helper apps occasionally pass around. */
+export type RagManagerResolver = (app: Application) => RagManager

package/src/rag_provider.ts

@@ -1,16 +1,94 @@
-import { ServiceProvider } from '@strav/kernel'
-import type { Application } from '@strav/kernel'
-import RagManager from './rag_manager.ts'
+/**
+ * `RagProvider` — `ServiceProvider` that wires `RagManager` into
+ * the container from `config.rag`.
+ *
+ * Eager construction at boot — a malformed config or a missing
+ * pgvector dependency should fail before the first call hits.
+ * Apps register `BrainProvider` and `DatabaseProvider` (when
+ * pgvector is in the store list) before this one; the
+ * `dependencies` array makes the order explicit.
+ *
+ * Config defaults: if `config.rag` is absent entirely, the
+ * provider boots a sensible in-memory setup so apps can try
+ * `rag.ingest()` / `rag.retrieve()` in dev without configuration
+ * — the memory driver is registered as the default store and a
+ * `recursive` chunker is configured. Production apps override
+ * via a real `config/rag.ts`.
+ */
 
-export default class RagProvider extends ServiceProvider {
-  readonly name = 'rag'
-  override readonly dependencies = ['config']
+// biome-ignore lint/style/useImportType: PostgresDatabase value import — required when any pgvector store is configured. Loaded conditionally below.
+import { PostgresDatabase } from '@strav/database'
+// biome-ignore lint/style/useImportType: BrainManager value import for c.resolve.
+import { BrainManager } from '@strav/brain'
+import {
+  type Application,
+  ConfigError,
+  ConfigRepository,
+  ServiceProvider,
+} from '@strav/kernel'
+import { RagManager, type RagManagerOptions } from './rag_manager.ts'
+import type { RagConfig } from './types.ts'
+
+export class RagProvider extends ServiceProvider {
+  override readonly name = 'rag'
+  override readonly dependencies = ['config', 'brain']
 
   override register(app: Application): void {
-    app.singleton(RagManager)
+    app.singleton(RagManager, (c) => {
+      const raw = c.resolve(ConfigRepository).get('rag') as Partial<RagConfig> | undefined
+      const config = applyDefaults(raw)
+
+      const brain = c.resolve(BrainManager)
+      const opts: RagManagerOptions = { config, brain }
+
+      // Only resolve the database when at least one store needs it.
+      const needsDb = Object.values(config.stores).some((s) => s.driver === 'pgvector')
+      if (needsDb) {
+        try {
+          opts.db = c.resolve(PostgresDatabase)
+        } catch (cause) {
+          throw new ConfigError(
+            'RagProvider: at least one store uses `driver: "pgvector"` but PostgresDatabase is not registered. Register DatabaseProvider before RagProvider.',
+            { cause },
+          )
+        }
+      }
+      return new RagManager(opts)
+    })
   }
 
   override boot(app: Application): void {
+    // Force-resolve so config errors surface at boot, not on first call.
     app.resolve(RagManager)
   }
 }
+
+/**
+ * Fill in defaults for omitted config fields. Apps with no
+ * `config/rag.ts` at all get a working in-memory setup.
+ */
+function applyDefaults(raw: Partial<RagConfig> | undefined): RagConfig {
+  const config: Partial<RagConfig> = raw ?? {}
+  const stores = config.stores ?? { memory: { driver: 'memory' } }
+  const def = config.default ?? Object.keys(stores)[0] ?? 'memory'
+  if (!stores[def]) {
+    throw new ConfigError(
+      `RagProvider: default store "${def}" is not declared in config.rag.stores.`,
+    )
+  }
+  return {
+    default: def,
+    ...(config.prefix !== undefined ? { prefix: config.prefix } : {}),
+    embedding: config.embedding ?? {
+      provider: 'openai',
+      model: 'text-embedding-3-small',
+      dimension: 1536,
+    },
+    chunking: config.chunking ?? {
+      strategy: 'recursive',
+      chunkSize: 512,
+      overlap: 64,
+    },
+    stores,
+  }
+}

package/src/rag_vector_schema.ts

@@ -0,0 +1,56 @@
+/**
+ * `ragVectorSchema` — the framework-known shape of the vectors
+ * table. Carries every column EXCEPT `embedding`, which lives
+ * outside the framework schema system because pgvector's
+ * `vector(N)` type isn't expressible via `defineSchema`. The
+ * `applyRagVectorMigration` helper attaches the embedding column
+ * + HNSW index in the migration step.
+ *
+ * Columns the framework manages:
+ *
+ *   - `id`         ULID primary key.
+ *   - `tenant_id`  Auto-injected by `tenanted: true`. RLS policies
+ *                  scope reads + writes by
+ *                  `current_setting('app.tenant_id')` so apps that
+ *                  wrap calls in `tenants.withTenant(...)` get
+ *                  per-tenant isolation for free.
+ *   - `collection` Logical bucket — apps create one collection per
+ *                  conceptual corpus (`articles`, `support_docs`,
+ *                  per-user notebooks, ...).
+ *   - `source_id`  Optional pointer back to the source row a chunk
+ *                  came from. `deleteBySource(collection, id)`
+ *                  drops every chunk for one source in a single
+ *                  DELETE — handy when re-indexing on update.
+ *   - `content`    The chunk text. Plain `text` for full storage.
+ *   - `metadata`   Free-form JSONB. Indexable in the recommended
+ *                  migration via a GIN index when apps query by
+ *                  metadata filters.
+ *   - `created_at` Insert timestamp. Useful for audit + soft
+ *                  recency filtering.
+ *
+ * Columns attached by `applyRagVectorMigration`:
+ *
+ *   - `embedding vector(<dimension>) NOT NULL` — the vector itself.
+ *   - `HNSW idx_<table>_embedding` on `(embedding vector_cosine_ops)`.
+ *
+ * Apps register the schema with `SchemaRegistry` at boot (mirrors
+ * every other framework schema), then call
+ * `applyRagVectorMigration(db, { dimension, registry })` inside
+ * the migration's `up()`.
+ */
+
+import { Archetype, defineSchema } from '@strav/database'
+
+export const ragVectorSchema = defineSchema(
+  'rag_vector',
+  Archetype.Entity,
+  (t) => {
+    t.id()
+    t.string('collection').max(128).notNull()
+    t.string('source_id').max(128).nullable()
+    t.text('content').notNull()
+    t.json('metadata').notNull().default({})
+    t.timestamp('created_at').notNull()
+  },
+  { tenanted: true },
+)