@strav/rag 0.4.31 → 1.0.0-alpha.20

package/src/retrievable.ts

@@ -1,179 +1,270 @@
-import type { BaseModel } from '@strav/database'
-import type { NormalizeConstructor } from '@strav/kernel'
-import { Emitter } from '@strav/kernel'
-import { brain } from '@strav/brain'
-import RagManager from './rag_manager.ts'
-import { createChunker } from './chunking/chunker.ts'
-import type { VectorDocument, RetrieveOptions, RetrieveResult } from './types.ts'
-
-export function retrievable<T extends NormalizeConstructor<typeof BaseModel>>(Base: T) {
-  return class Retrievable extends Base {
-    private static _retrievalBooted = false
-
-    static retrievableAs(): string {
-      return (this as unknown as typeof BaseModel).tableName
+/**
+ * `retrievable(Repository)` — class mixin that bolts vector-index
+ * methods onto a Repository so apps can re-index a row and search
+ * its collection without juggling `RagManager` calls by hand.
+ *
+ * ```ts
+ * @inject()
+ * export class ArticleRepository extends retrievable(Repository<Article>) {
+ *   static override readonly schema = articleSchema
+ *   static override readonly model = Article
+ *
+ *   constructor(db: PostgresDatabase, events: EventBus, rag: RagManager) {
+ *     super(db, events)
+ *     this.rag = rag
+ *   }
+ *
+ *   // Override the extension points as needed:
+ *   protected override toContent(a: Article): string {
+ *     return `${a.title}\n\n${a.body}`
+ *   }
+ *
+ *   protected override toMetadata(a: Article): Record<string, unknown> {
+ *     return { authorId: a.author_id, tags: a.tags }
+ *   }
+ * }
+ * ```
+ *
+ * Usage:
+ *
+ * ```ts
+ * const article = await articles.create(...)
+ * await articles.vectorize(article)              // index it
+ *
+ * const { matches } = await articles.retrieve('query')   // search
+ *
+ * await articles.delete(article)
+ * await articles.vectorRemove(article)           // drop from index
+ * ```
+ *
+ * Why not auto-vectorize on `create` / `update`?
+ *
+ *   V1 ships the explicit pattern. An auto-hook tied to repository
+ *   events would couple persistence to the embedding provider's
+ *   availability — a transient rate-limit on the embedder would
+ *   fail the create call. Apps that want auto-vectorize wire it
+ *   themselves via `events.on('article.created', m =>
+ *   articles.vectorize(m))` so they control the failure mode
+ *   (fire-and-forget vs awaited vs queued via `@strav/queue`).
+ *
+ * Extension points (all optional overrides):
+ *
+ *   - `collectionName()` — defaults to the table name from the
+ *     schema. Override when the collection should differ from the
+ *     table, or to compose a per-tenant suffix dynamically.
+ *
+ *   - `toContent(model)` — defaults to concatenating every string
+ *     field on the model with `\n`. The default works for simple
+ *     row shapes; apps with structured content override.
+ *
+ *   - `toMetadata(model)` — defaults to `{}`. Apps return fields
+ *     they want to filter on (e.g. `author_id`, `lang`, `kind`).
+ *
+ *   - `shouldRetrieve(model)` — gates indexing. Return `false` for
+ *     draft / soft-deleted / private rows. The default is `true`.
+ */
+
+import type { Repository } from '@strav/database'
+import type { RagManager } from './rag_manager.ts'
+import type {
+  RetrieveOptions,
+  RetrieveResult,
+  VectorMatch,
+} from './types.ts'
+
+/** Minimal constructor type we can mix into. Wider than `typeof Repository` so subclasses with extra ctor args still type-check. */
+// biome-ignore lint/suspicious/noExplicitAny: mixin constructor signatures intentionally accept any[]; the user-side subclass narrows.
+type RepositoryConstructor<TModel extends object> = new (...args: any[]) => Repository<TModel>
+
+/**
+ * Returns a subclass that extends `Base` with `vectorize` /
+ * `vectorRemove` / `retrieve` plus override-points
+ * (`collectionName`, `toContent`, `toMetadata`,
+ * `shouldRetrieve`). The user-side class declares an explicit
+ * constructor that calls `super(...)` and assigns `this.rag`.
+ */
+export function retrievable<TModel extends object, TBase extends RepositoryConstructor<TModel>>(
+  Base: TBase,
+) {
+  abstract class RetrievableRepository extends Base {
+    /**
+     * The framework's `RagManager`. Assigned by the user-side
+     * subclass constructor. Public on purpose — apps that want to
+     * drop down to raw `rag.store()` / `rag.ingest(...)` access
+     * have a hook.
+     */
+    rag!: RagManager
+
+    /**
+     * Collection name for vector storage. Defaults to the table
+     * name from `static schema`. Override to point at a different
+     * collection (or to compose per-tenant / per-env suffixes).
+     */
+    protected collectionName(): string {
+      const ctor = this.constructor as unknown as { schema: { name: string } }
+      return ctor.schema.name
     }
 
-    toRetrievableContent(): string {
+    /**
+     * Build the indexable text from a model row. Default
+     * concatenates every non-underscore string field with `\n`.
+     * Apps with structured content override this — typically
+     * something like `` `${a.title}\n\n${a.body}` ``.
+     */
+    protected toContent(model: TModel): string {
       const parts: string[] = []
-      for (const key of Object.keys(this)) {
+      for (const [key, value] of Object.entries(model as Record<string, unknown>)) {
         if (key.startsWith('_')) continue
-        const val = (this as any)[key]
-        if (typeof val === 'string' && val.length > 0) parts.push(val)
+        if (typeof value === 'string' && value.length > 0) parts.push(value)
       }
       return parts.join('\n')
     }
 
-    toRetrievableMetadata(): Record<string, unknown> {
+    /**
+     * Build the metadata bag attached to every chunk. Apps return
+     * fields they want to filter retrievals on. The framework
+     * automatically adds `chunkIndex`, `startOffset`, `endOffset`
+     * — overrides shouldn't try to re-add those.
+     */
+    protected toMetadata(_model: TModel): Record<string, unknown> {
       return {}
     }
 
-    shouldBeRetrievable(): boolean {
+    /**
+     * Whether the model should currently be indexed. Override to
+     * skip drafts, soft-deleted rows, private records, etc. The
+     * default `true` indexes every model — fine for the common
+     * case.
+     */
+    protected shouldRetrieve(_model: TModel): boolean {
       return true
     }
 
-    // ── Instance methods ──────────────────────────────────────────────
-
-    async vectorize(): Promise<void> {
-      if (!this.shouldBeRetrievable()) return
-
-      const ctor = this.constructor as typeof Retrievable
-      const collection = RagManager.collectionName(ctor.retrievableAs())
-      const config = RagManager.config
-      const pkProp = (ctor as unknown as typeof BaseModel).primaryKeyProperty
-      const id = (this as any)[pkProp]
-
-      const content = this.toRetrievableContent()
-      if (!content) return
-
-      // Remove existing chunks for this model instance
-      await RagManager.store().deleteBySource(collection, id)
-
-      const chunker = createChunker(config.chunking)
-      const chunks = chunker.chunk(content)
-      if (chunks.length === 0) return
-
-      const texts = chunks.map(c => c.content)
-      const embeddings = await brain.embed(texts, {
-        provider: config.embedding.provider,
-        model: config.embedding.model,
-      })
-
-      const metadata = this.toRetrievableMetadata()
-      const documents: VectorDocument[] = chunks.map((chunk, i) => ({
-        id: `${id}_${i}`,
+    /**
+     * (Re-)index a single model. Drops any existing chunks for
+     * the model's id, then ingests fresh chunks of the current
+     * content. When `shouldRetrieve(model)` returns `false`, the
+     * chunks are dropped without re-ingest — apps don't need a
+     * separate "this just became private" path.
+     *
+     * Returns the vector ids written. Empty array when content
+     * was empty or `shouldRetrieve` returned `false`.
+     */
+    async vectorize(model: TModel): Promise<string[]> {
+      const collection = this.collectionName()
+      const id = modelId(model)
+
+      // Drop existing chunks for this source first so updates
+      // replace cleanly. (RagManager.ingest writes fresh ids per
+      // call; without this step every re-vectorize would
+      // duplicate.)
+      await this.rag
+        .store()
+        .deleteBySource(this.rag.collectionName(collection), id)
+
+      if (!this.shouldRetrieve(model)) return []
+
+      const content = this.toContent(model)
+      if (!content) return []
+
+      return this.rag.ingest(collection, content, {
         sourceId: id,
-        content: chunk.content,
-        embedding: embeddings[i]!,
-        metadata: {
-          ...metadata,
-          modelId: id,
-          chunkIndex: chunk.index,
-        },
-      }))
-
-      await RagManager.store().upsert(collection, documents)
+        metadata: this.toMetadata(model),
+      })
     }
 
-    async vectorRemove(): Promise<void> {
-      const ctor = this.constructor as typeof Retrievable
-      const collection = RagManager.collectionName(ctor.retrievableAs())
-      const pkProp = (ctor as unknown as typeof BaseModel).primaryKeyProperty
-      const id = (this as any)[pkProp]
-      await RagManager.store().deleteBySource(collection, id)
+    /**
+     * Drop every chunk for one model. Apps call this after
+     * `delete(model)` in their domain code. The mixin doesn't
+     * auto-hook the delete lifecycle for the same reason it
+     * doesn't auto-hook create/update — keeps embedding-provider
+     * availability out of the persistence path.
+     */
+    async vectorRemove(model: TModel): Promise<void> {
+      const collection = this.collectionName()
+      const id = modelId(model)
+      await this.rag
+        .store()
+        .deleteBySource(this.rag.collectionName(collection), id)
     }
 
-    // ── Static methods ────────────────────────────────────────────────
-
-    static async retrieve(query: string, options?: RetrieveOptions): Promise<RetrieveResult> {
-      const { rag } = await import('./helpers.ts')
-      return rag.retrieve(query, {
+    /**
+     * Semantic search over this repository's collection. Default
+     * `collection` is the mixin's `collectionName()` — apps that
+     * want to retrieve from another collection pass it explicitly.
+     */
+    async retrieve(
+      query: string,
+      options: Omit<RetrieveOptions, 'collection'> & { collection?: string } = {},
+    ): Promise<RetrieveResult> {
+      return this.rag.retrieve(query, {
         ...options,
-        collection: options?.collection ?? this.retrievableAs(),
+        collection: options.collection ?? this.collectionName(),
       })
     }
 
-    static async importAll(chunkSize: number = 100): Promise<number> {
-      const ModelCtor = this as unknown as typeof BaseModel & typeof Retrievable
-      const collection = RagManager.collectionName(this.retrievableAs())
-      const config = RagManager.config
-      const db = ModelCtor.db
-      const table = ModelCtor.tableName
-      const pkCol = ModelCtor.primaryKeyColumn
-
-      await RagManager.store().createCollection(collection, config.embedding.dimension)
-
-      let imported = 0
+    /**
+     * Re-index every row in the repository. Walks rows in batches
+     * of `batchSize` and vectorizes each. Useful for backfilling
+     * a new collection or recovering after a schema change.
+     *
+     * The CLI's `rag:reindex <repository>` doesn't ship in V1 —
+     * apps that want one wire it as their own console command
+     * pointing at this method.
+     *
+     * Returns the total count of rows processed (NOT the chunk
+     * count — chunks per row vary with content size).
+     */
+    async reindexAll(batchSize: number = 100): Promise<number> {
+      let processed = 0
       let offset = 0
-
       while (true) {
-        const rows = (await db.sql.unsafe(
-          `SELECT * FROM "${table}" ORDER BY "${pkCol}" LIMIT $1 OFFSET $2`,
-          [chunkSize, offset]
-        )) as Record<string, unknown>[]
-
+        const rows = await this.query().orderBy('id', 'asc').limit(batchSize).offset(offset).get()
         if (rows.length === 0) break
-
-        for (const row of rows) {
-          const instance = ModelCtor.hydrate(row) as InstanceType<typeof Retrievable>
-          if (instance.shouldBeRetrievable()) {
-            try {
-              await instance.vectorize()
-              imported++
-            } catch {
-              // Vectorization is secondary — continue on failure
-            }
-          }
-        }
-
-        offset += chunkSize
-        if (rows.length < chunkSize) break
+        for (const row of rows) await this.vectorize(row)
+        processed += rows.length
+        offset += rows.length
+        if (rows.length < batchSize) break
       }
-
-      return imported
-    }
-
-    static async flushVectors(): Promise<void> {
-      const collection = RagManager.collectionName(this.retrievableAs())
-      await RagManager.store().flush(collection)
-    }
-
-    static async createVectorCollection(): Promise<void> {
-      const collection = RagManager.collectionName(this.retrievableAs())
-      await RagManager.store().createCollection(collection, RagManager.config.embedding.dimension)
+      return processed
     }
 
-    static bootRetrieval(eventPrefix: string): void {
-      if (this._retrievalBooted) return
-      this._retrievalBooted = true
-
-      const vectorizeFn = async (model: unknown) => {
-        if (model && typeof (model as any).vectorize === 'function') {
-          try {
-            await (model as any).vectorize()
-          } catch {
-            // Vectorization is secondary — failures should not break the event pipeline
-          }
-        }
-      }
-
-      const removeFn = async (model: unknown) => {
-        if (model && typeof (model as any).vectorRemove === 'function') {
-          try {
-            await (model as any).vectorRemove()
-          } catch {
-            // Vector removal is secondary
-          }
-        }
+    /**
+     * Match-to-models helper. Takes the `matches` array from
+     * `retrieve(...)` and hydrates the source rows by id, in
+     * match order. Matches whose `sourceId` doesn't resolve to a
+     * row (deleted between index time + retrieval) are dropped.
+     */
+    async resolveMatches(matches: readonly VectorMatch[]): Promise<TModel[]> {
+      const ids = [...new Set(matches.map((m) => m.sourceId).filter((s): s is string => !!s))]
+      if (ids.length === 0) return []
+      const found = await this.findMany(ids as unknown as readonly string[])
+      const byId = new Map<string, TModel>(
+        found.map((m) => [modelId(m), m]),
+      )
+      const out: TModel[] = []
+      for (const match of matches) {
+        if (!match.sourceId) continue
+        const row = byId.get(match.sourceId)
+        if (row) out.push(row)
       }
-
-      Emitter.on(`${eventPrefix}.created`, vectorizeFn)
-      Emitter.on(`${eventPrefix}.updated`, vectorizeFn)
-      Emitter.on(`${eventPrefix}.synced`, vectorizeFn)
-      Emitter.on(`${eventPrefix}.deleted`, removeFn)
+      return out
     }
   }
+  return RetrievableRepository
 }
 
-export type RetrievableInstance = InstanceType<ReturnType<typeof retrievable>>
-export type RetrievableModel = ReturnType<typeof retrievable>
+/**
+ * Coerce a model's `id` to a string. Repositories use ULID or UUID
+ * ids by default, both of which round-trip through `String(...)`
+ * cleanly; integer PKs (bigSerial) coerce the same way.
+ */
+function modelId(model: object): string {
+  const id = (model as { id?: unknown }).id
+  if (id === undefined || id === null) {
+    throw new Error(
+      `retrievable: model has no \`id\` to use as a vector sourceId. The mixin only works on models with a single-column id.`,
+    )
+  }
+  return String(id)
+}

package/src/types.ts

@@ -1,47 +1,84 @@
-// ── Vector Documents ─────────────────────────────────────────────────────
-
+/**
+ * `@strav/rag` types — the data shapes apps see when reading and
+ * writing vectors and when running retrieval.
+ *
+ * Three concept clusters:
+ *
+ *   - **Vector docs + queries** — the storage layer. A
+ *     `VectorDocument` is one indexed unit (a chunk of source
+ *     content, its embedding, and free-form metadata).
+ *     `query()` returns `VectorMatch[]` ranked by similarity.
+ *
+ *   - **Retrieval pipeline** — `RetrieveOptions` /
+ *     `RetrieveResult`. Apps call `rag.retrieve(query, ...)`,
+ *     the manager embeds the query through `@strav/brain`,
+ *     queries the active store, and returns matches with
+ *     normalized similarity scores.
+ *
+ *   - **Chunking** — `Chunk`, `Chunker`. The chunker takes raw
+ *     content and produces overlapping segments suitable for
+ *     embedding. Two strategies ship: `fixed` (mechanical N-char
+ *     windows with overlap) and `recursive` (paragraph-aware,
+ *     better for prose).
+ */
+
+// ─── Vector documents + queries ──────────────────────────────────────────
+
+/**
+ * One indexed unit. `id` is provider-assigned (ULID by default);
+ * `sourceId` is the optional app-defined pointer back to the row
+ * the chunk came from (e.g., `article_id`) — `deleteBySource`
+ * removes every chunk for one source in a single call.
+ */
 export interface VectorDocument {
-  id?: string | number
-  sourceId?: string | number
+  id?: string
+  sourceId?: string | null
   content: string
   embedding: number[]
   metadata?: Record<string, unknown>
 }
 
-// ── Query Options & Results ──────────────────────────────────────────────
-
 export interface QueryOptions {
+  /** Top-K matches to return. Default `5`. */
   topK?: number
+  /** Minimum similarity threshold (0–1). Matches below this are filtered out. */
   threshold?: number
+  /** Metadata filter — flat key/value AND. Driver-specific operators are NOT supported in V1. */
   filter?: Record<string, unknown>
 }
 
 export interface QueryResult {
   matches: VectorMatch[]
-  processingTimeMs?: number
+  /** Time the underlying store took to compute the query, in ms. */
+  processingTimeMs: number
 }
 
 export interface VectorMatch {
-  id: string | number
+  id: string
   content: string
+  /** Similarity score in [0, 1]. 1.0 = identical embeddings, 0 = orthogonal. */
   score: number
   metadata: Record<string, unknown>
+  sourceId?: string | null
 }
 
-// ── Retrieval (high-level pipeline) ──────────────────────────────────────
+// ─── Retrieval pipeline ─────────────────────────────────────────────────
 
 export interface RetrieveOptions {
+  /** Override the collection. Defaults to the manager's default. */
   collection?: string
+  /** Top-K matches. Default `5`. */
   topK?: number
+  /** Minimum similarity threshold. */
   threshold?: number
+  /** Metadata filter — flat key/value AND. */
   filter?: Record<string, unknown>
-  rerank?: RerankOptions
-}
-
-export interface RerankOptions {
-  authorityWeight?: number
-  recencyWeight?: number
-  similarityWeight?: number
+  /** Override the store. Defaults to the manager's default store. */
+  store?: string
+  /** Override the embedding model used to encode the query. */
+  embedModel?: string
+  /** Override the brain provider used for embedding. */
+  embedProvider?: string
 }
 
 export interface RetrieveResult {
@@ -51,19 +88,24 @@ export interface RetrieveResult {
 }
 
 export interface RetrievedDocument {
-  id: string | number
+  id: string
   content: string
+  /** Same as `VectorMatch.score` — kept as a separate field so future re-ranking can diverge `score` from raw `similarity`. */
   score: number
   similarity: number
   metadata: Record<string, unknown>
+  sourceId?: string | null
 }
 
-// ── Chunking ─────────────────────────────────────────────────────────────
+// ─── Chunking ────────────────────────────────────────────────────────────
 
 export interface Chunk {
   content: string
+  /** 0-based ordinal within the source. */
   index: number
+  /** Character offset of the chunk's first character in the source. */
   startOffset: number
+  /** Character offset one past the chunk's last character. */
   endOffset: number
 }
 
@@ -71,30 +113,46 @@ export interface Chunker {
   chunk(content: string): Chunk[]
 }
 
-// ── Configuration ────────────────────────────────────────────────────────
+// ─── Configuration ──────────────────────────────────────────────────────
 
+/**
+ * `config.rag` shape. Apps that don't configure rag get a sensible
+ * default (memory driver, OpenAI text-embedding-3-small, recursive
+ * chunking) — see `RagProvider.boot()` for the defaults.
+ */
 export interface RagConfig {
+  /** Default store name — must be a key in `stores`. */
   default: string
-  prefix: string
+  /** Optional collection-name prefix. Used to namespace per-app or per-tenant. */
+  prefix?: string
   embedding: EmbeddingConfig
   chunking: ChunkingConfig
   stores: Record<string, StoreConfig>
 }
 
 export interface EmbeddingConfig {
+  /** `@strav/brain` provider key (e.g., `'openai'`, `'gemini'`, `'ollama'`). */
   provider: string
+  /** Model identifier — passed to `brain.embed(..., { model })`. */
   model: string
+  /** Vector dimension. Must match the chosen model. */
   dimension: number
 }
 
 export interface ChunkingConfig {
-  strategy: string
+  /** `'fixed'` or `'recursive'`. Custom strategies aren't pluggable in V1. */
+  strategy: 'fixed' | 'recursive'
   chunkSize: number
   overlap: number
-  separators?: string[]
+  /** Custom separators for the recursive strategy. Defaults to `['\n\n', '\n', '. ', ' ']`. */
+  separators?: readonly string[]
 }
 
 export interface StoreConfig {
+  /** `'memory'` or `'pgvector'`; custom drivers register via `rag.extend(name, factory)`. */
   driver: string
+  /** Pgvector: explicit table name override. Default `'rag_vector'`. */
+  table?: string
+  /** Free-form fields driver-specific (e.g., HNSW tuning for pgvector). */
   [key: string]: unknown
 }

package/src/vector_store.ts

@@ -1,15 +1,55 @@
-import type { VectorDocument, QueryOptions, QueryResult } from './types.ts'
+/**
+ * `VectorStore` — the storage abstraction every driver
+ * (`MemoryDriver`, `PgvectorDriver`, custom drivers registered
+ * via `rag.extend(...)`) implements.
+ *
+ * Lifecycle:
+ *
+ *   - `createCollection(name, dimension)` — idempotent. For
+ *     pgvector this is mostly a no-op (the table holds every
+ *     collection); the dimension is enforced at INSERT.
+ *   - `deleteCollection(name)` — drops every vector under
+ *     `collection = name`.
+ *
+ * Reads + writes:
+ *
+ *   - `upsert(collection, docs)` — inserts (and overwrites by id
+ *     when supplied).
+ *   - `delete(collection, ids)` — removes specific vectors.
+ *   - `deleteBySource(collection, sourceId)` — removes every
+ *     vector with the matching `source_id`. Apps call this when
+ *     re-indexing a source row.
+ *   - `flush(collection)` — drops every vector in the
+ *     collection. Faster than `deleteCollection` for the common
+ *     "wipe + re-ingest" pattern because the collection's
+ *     identity stays intact.
+ *   - `query(collection, vector, opts)` — top-K similarity
+ *     search.
+ *
+ * Multitenancy lives BELOW this interface — the pgvector driver
+ * relies on `app.tenant_id` session settings (set by
+ * `tenants.withTenant`) to enforce isolation via RLS. The
+ * `MemoryDriver` is single-tenant by construction and ignores
+ * tenancy.
+ */
+
+import type { QueryOptions, QueryResult, VectorDocument } from './types.ts'
 
 export interface VectorStore {
+  /** Driver identifier — `'memory'`, `'pgvector'`, or the name passed to `rag.extend`. */
   readonly name: string
 
   createCollection(collection: string, dimension: number): Promise<void>
   deleteCollection(collection: string): Promise<void>
 
-  upsert(collection: string, documents: VectorDocument[]): Promise<void>
-  delete(collection: string, ids: (string | number)[]): Promise<void>
-  deleteBySource(collection: string, sourceId: string | number): Promise<void>
+  upsert(collection: string, documents: readonly VectorDocument[]): Promise<void>
+  delete(collection: string, ids: readonly string[]): Promise<void>
+  deleteBySource(collection: string, sourceId: string): Promise<void>
   flush(collection: string): Promise<void>
 
-  query(collection: string, vector: number[], options?: QueryOptions): Promise<QueryResult>
+  query(
+    collection: string,
+    vector: readonly number[],
+    options?: QueryOptions,
+  ): Promise<QueryResult>
 }