@strav/rag 1.0.0-alpha.19 → 1.0.0-alpha.21

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
package/package.json CHANGED
@@ -1,6 +1,6 @@
1
1
  {
2
2
  "name": "@strav/rag",
3
- "version": "1.0.0-alpha.19",
3
+ "version": "1.0.0-alpha.21",
4
4
  "description": "Strav RAG module — vector store abstraction, pgvector + in-memory drivers, chunking strategies. Composes with @strav/brain for embeddings and @strav/database for persistence.",
5
5
  "type": "module",
6
6
  "main": "./src/index.ts",
@@ -19,9 +19,10 @@
19
19
  "access": "public"
20
20
  },
21
21
  "dependencies": {
22
- "@strav/brain": "1.0.0-alpha.19",
23
- "@strav/database": "1.0.0-alpha.19",
24
- "@strav/kernel": "1.0.0-alpha.19"
22
+ "@strav/brain": "1.0.0-alpha.21",
23
+ "@strav/cli": "1.0.0-alpha.21",
24
+ "@strav/database": "1.0.0-alpha.21",
25
+ "@strav/kernel": "1.0.0-alpha.21"
25
26
  },
26
27
  "peerDependencies": {
27
28
  "@types/bun": ">=1.3.14"
@@ -0,0 +1,3 @@
1
+ export { RagConsoleProvider } from './rag_console_provider.ts'
2
+ export { RagFlush } from './rag_flush.ts'
3
+ export { RagList } from './rag_list.ts'
@@ -0,0 +1,17 @@
1
+ /**
2
+ * `RagConsoleProvider` — declares the rag console commands.
3
+ *
4
+ * Apps add it to `bootstrap/providers.ts` alongside `RagProvider`.
5
+ * Separate provider (mirrors `QueueConsoleProvider`) so apps
6
+ * that don't use the CLI don't pay the cost of resolving the
7
+ * commands at boot.
8
+ */
9
+
10
+ import { ConsoleProvider } from '@strav/cli'
11
+ import { RagFlush } from './rag_flush.ts'
12
+ import { RagList } from './rag_list.ts'
13
+
14
+ export class RagConsoleProvider extends ConsoleProvider {
15
+ override readonly name = 'console.rag'
16
+ override readonly commands = [RagFlush, RagList] as const
17
+ }
@@ -0,0 +1,51 @@
1
+ /**
2
+ * `bun strav rag:flush <collection> [--store=name] [--force]` —
3
+ * drop every vector in a collection on the active (or named)
4
+ * store.
5
+ *
6
+ * Use cases:
7
+ *
8
+ * - Wiping a corrupted index before re-ingest.
9
+ * - Cleaning up a dev / staging environment.
10
+ * - Recovering after a dimension / model change.
11
+ *
12
+ * The command confirms before running unless `--force` is set.
13
+ * Doesn't touch the source data — apps run their own re-ingest
14
+ * afterward, typically via `retrievable` repo's `reindexAll()`.
15
+ */
16
+
17
+ import { Command, type ExecuteArgs, ExitCode } from '@strav/cli'
18
+ import { RagManager } from '../rag_manager.ts'
19
+
20
+ export class RagFlush extends Command {
21
+ static signature = 'rag:flush {collection} {--store=} {--force}'
22
+ static description = 'Delete every vector in a collection (on the active or --store= named store).'
23
+ static providers = ['config', 'logger', 'brain', 'rag']
24
+
25
+ override async execute({ args, flags }: ExecuteArgs): Promise<number> {
26
+ const collection = args.collection as string
27
+ const storeName = typeof flags.store === 'string' && flags.store.length > 0
28
+ ? flags.store
29
+ : undefined
30
+
31
+ const manager = this.app.resolve(RagManager)
32
+ const fullCollection = manager.collectionName(collection)
33
+ const storeLabel = storeName ?? manager.config.default
34
+
35
+ if (flags.force !== true) {
36
+ const ok = await this.confirm(
37
+ `Delete every vector in collection "${fullCollection}" on store "${storeLabel}"? This is irreversible.`,
38
+ )
39
+ if (!ok) {
40
+ this.info('Aborted.')
41
+ return ExitCode.Success
42
+ }
43
+ }
44
+
45
+ await manager.store(storeName).flush(fullCollection)
46
+ this.success(
47
+ `Flushed collection "${fullCollection}" on store "${storeLabel}".`,
48
+ )
49
+ return ExitCode.Success
50
+ }
51
+ }
@@ -0,0 +1,48 @@
1
+ /**
2
+ * `bun strav rag:list` — print the configured RAG stores +
3
+ * chunker + embedding setup.
4
+ *
5
+ * Diagnostic only — no mutations. Useful for verifying that
6
+ * `config/rag.ts` parses correctly and that the registered
7
+ * driver names match what's expected.
8
+ */
9
+
10
+ import { Command, type ExecuteArgs, ExitCode } from '@strav/cli'
11
+ import { RagManager } from '../rag_manager.ts'
12
+
13
+ export class RagList extends Command {
14
+ static signature = 'rag:list'
15
+ static description = 'List configured RAG stores + embedding + chunking settings.'
16
+ static providers = ['config', 'logger', 'brain', 'rag']
17
+
18
+ override async execute(_args: ExecuteArgs): Promise<number> {
19
+ const manager = this.app.resolve(RagManager)
20
+ const config = manager.config
21
+
22
+ this.info(`Default store: ${config.default}`)
23
+ if (config.prefix) this.info(`Collection prefix: ${config.prefix}`)
24
+
25
+ this.info('')
26
+ this.info('Stores:')
27
+ for (const [name, store] of Object.entries(config.stores)) {
28
+ const flag = name === config.default ? ' (default)' : ''
29
+ this.info(` ${name}${flag}: driver=${store.driver}`)
30
+ }
31
+
32
+ this.info('')
33
+ this.info('Embedding:')
34
+ this.info(` provider: ${config.embedding.provider}`)
35
+ this.info(` model: ${config.embedding.model}`)
36
+ this.info(` dim: ${config.embedding.dimension}`)
37
+
38
+ this.info('')
39
+ this.info('Chunking:')
40
+ this.info(` strategy: ${config.chunking.strategy}`)
41
+ this.info(` chunkSize: ${config.chunking.chunkSize}`)
42
+ this.info(` overlap: ${config.chunking.overlap}`)
43
+ if (config.chunking.separators) {
44
+ this.info(` separators: ${JSON.stringify(config.chunking.separators)}`)
45
+ }
46
+ return ExitCode.Success
47
+ }
48
+ }
@@ -28,7 +28,11 @@
28
28
  * model them.
29
29
  */
30
30
 
31
- import type { PostgresDatabase } from '@strav/database'
31
+ import {
32
+ currentTransactionalContext,
33
+ type DatabaseExecutor,
34
+ type PostgresDatabase,
35
+ } from '@strav/database'
32
36
  import { VectorQueryError } from '../rag_error.ts'
33
37
  import { ragVectorSchema } from '../rag_vector_schema.ts'
34
38
  import type {
@@ -71,6 +75,20 @@ export class PgvectorDriver implements VectorStore {
71
75
  })
72
76
  }
73
77
 
78
+ /**
79
+ * Route reads + writes through the ambient `UnitOfWork`
80
+ * transaction when one is active (e.g., inside
81
+ * `tenants.withTenant(...)`); fall back to the raw pool
82
+ * otherwise. Mirrors how `Repository.executor(opts)` works in
83
+ * `@strav/database`, so RLS scoping + transactional event
84
+ * flushing apply uniformly across framework + driver code.
85
+ */
86
+ private exec(): DatabaseExecutor {
87
+ const ambient = currentTransactionalContext()
88
+ if (ambient) return ambient.tx
89
+ return this.db as unknown as DatabaseExecutor
90
+ }
91
+
74
92
  // ─── Collections ──────────────────────────────────────────────────────
75
93
 
76
94
  async createCollection(_collection: string, _dimension: number): Promise<void> {
@@ -81,7 +99,7 @@ export class PgvectorDriver implements VectorStore {
81
99
  }
82
100
 
83
101
  async deleteCollection(collection: string): Promise<void> {
84
- await this.db.execute(
102
+ await this.exec().execute(
85
103
  `DELETE FROM "${this.table}" WHERE "collection" = $1`,
86
104
  [collection],
87
105
  )
@@ -96,13 +114,24 @@ export class PgvectorDriver implements VectorStore {
96
114
  if (documents.length === 0) return
97
115
  // pgvector accepts the vector as a stringified array literal —
98
116
  // `[0.12,0.34,...]` — cast with `::vector` at the boundary.
117
+ //
118
+ // Tenant scoping: the `tenant_id` column on `rag_vector` is
119
+ // NOT NULL with no default, so apps wrapping the call in
120
+ // `tenants.withTenant(...)` need a value supplied. We read
121
+ // `current_setting('app.tenant_id')` inside the SQL itself —
122
+ // the same session var the RLS policy reads — so the INSERT
123
+ // works under tenant scope without the driver knowing the PK
124
+ // type ahead of time. The `true` second arg makes the
125
+ // setting return NULL (not throw) outside `withTenant`; the
126
+ // INSERT then fails the NOT NULL constraint with a clear
127
+ // error message that nudges the app toward the right wrap.
99
128
  for (const doc of documents) {
100
129
  const id = doc.id ?? crypto.randomUUID()
101
130
  const embeddingLiteral = `[${doc.embedding.join(',')}]`
102
- await this.db.execute(
131
+ await this.exec().execute(
103
132
  `INSERT INTO "${this.table}"
104
- ("id", "collection", "source_id", "content", "metadata", "embedding", "created_at")
105
- VALUES ($1, $2, $3, $4, $5::jsonb, $6::vector, NOW())
133
+ ("id", "tenant_id", "collection", "source_id", "content", "metadata", "embedding", "created_at")
134
+ VALUES ($1, current_setting('app.tenant_id', true), $2, $3, $4, $5::jsonb, $6::vector, NOW())
106
135
  ON CONFLICT ("id") DO UPDATE SET
107
136
  "collection" = EXCLUDED."collection",
108
137
  "source_id" = EXCLUDED."source_id",
@@ -124,21 +153,21 @@ export class PgvectorDriver implements VectorStore {
124
153
  async delete(collection: string, ids: readonly string[]): Promise<void> {
125
154
  if (ids.length === 0) return
126
155
  const placeholders = ids.map((_, i) => `$${i + 2}`).join(', ')
127
- await this.db.execute(
156
+ await this.exec().execute(
128
157
  `DELETE FROM "${this.table}" WHERE "collection" = $1 AND "id" IN (${placeholders})`,
129
158
  [collection, ...ids],
130
159
  )
131
160
  }
132
161
 
133
162
  async deleteBySource(collection: string, sourceId: string): Promise<void> {
134
- await this.db.execute(
163
+ await this.exec().execute(
135
164
  `DELETE FROM "${this.table}" WHERE "collection" = $1 AND "source_id" = $2`,
136
165
  [collection, sourceId],
137
166
  )
138
167
  }
139
168
 
140
169
  async flush(collection: string): Promise<void> {
141
- await this.db.execute(
170
+ await this.exec().execute(
142
171
  `DELETE FROM "${this.table}" WHERE "collection" = $1`,
143
172
  [collection],
144
173
  )
@@ -190,7 +219,7 @@ export class PgvectorDriver implements VectorStore {
190
219
  score: number | string
191
220
  }>
192
221
  try {
193
- rows = await this.db.query(sql, params)
222
+ rows = await this.exec().query(sql, params)
194
223
  } catch (cause) {
195
224
  throw new VectorQueryError(
196
225
  `pgvector query failed for collection "${collection}".`,
package/src/index.ts CHANGED
@@ -32,8 +32,14 @@ export {
32
32
  type RagManagerOptions,
33
33
  type StoreFactory,
34
34
  } from './rag_manager.ts'
35
+ export {
36
+ RagConsoleProvider,
37
+ RagFlush,
38
+ RagList,
39
+ } from './console/index.ts'
35
40
  export { RagProvider } from './rag_provider.ts'
36
41
  export { ragVectorSchema } from './rag_vector_schema.ts'
42
+ export { retrievable } from './retrievable.ts'
37
43
  export type {
38
44
  Chunk,
39
45
  Chunker,
@@ -28,7 +28,7 @@ import { PostgresDatabase } from '@strav/database'
28
28
  // biome-ignore lint/style/useImportType: BrainManager value import for @inject() param-type metadata.
29
29
  import { BrainManager } from '@strav/brain'
30
30
  // biome-ignore lint/style/useImportType: Application value import for the container handle.
31
- import { Application, inject } from '@strav/kernel'
31
+ import { Application, inject, ulid } from '@strav/kernel'
32
32
  import { createChunker } from './chunking/chunker.ts'
33
33
  import { MemoryDriver } from './drivers/memory_driver.ts'
34
34
  import { PgvectorDriver } from './drivers/pgvector_driver.ts'
@@ -204,9 +204,8 @@ export class RagManager {
204
204
  )
205
205
  }
206
206
 
207
- const baseId = crypto.randomUUID()
208
207
  const documents: VectorDocument[] = chunks.map((chunk, i) => ({
209
- id: `${baseId}_${i}`,
208
+ id: ulid(),
210
209
  ...(options.sourceId !== undefined ? { sourceId: options.sourceId } : {}),
211
210
  content: chunk.content,
212
211
  embedding: embeddings[i]!,
@@ -0,0 +1,270 @@
1
+ /**
2
+ * `retrievable(Repository)` — class mixin that bolts vector-index
3
+ * methods onto a Repository so apps can re-index a row and search
4
+ * its collection without juggling `RagManager` calls by hand.
5
+ *
6
+ * ```ts
7
+ * @inject()
8
+ * export class ArticleRepository extends retrievable(Repository<Article>) {
9
+ * static override readonly schema = articleSchema
10
+ * static override readonly model = Article
11
+ *
12
+ * constructor(db: PostgresDatabase, events: EventBus, rag: RagManager) {
13
+ * super(db, events)
14
+ * this.rag = rag
15
+ * }
16
+ *
17
+ * // Override the extension points as needed:
18
+ * protected override toContent(a: Article): string {
19
+ * return `${a.title}\n\n${a.body}`
20
+ * }
21
+ *
22
+ * protected override toMetadata(a: Article): Record<string, unknown> {
23
+ * return { authorId: a.author_id, tags: a.tags }
24
+ * }
25
+ * }
26
+ * ```
27
+ *
28
+ * Usage:
29
+ *
30
+ * ```ts
31
+ * const article = await articles.create(...)
32
+ * await articles.vectorize(article) // index it
33
+ *
34
+ * const { matches } = await articles.retrieve('query') // search
35
+ *
36
+ * await articles.delete(article)
37
+ * await articles.vectorRemove(article) // drop from index
38
+ * ```
39
+ *
40
+ * Why not auto-vectorize on `create` / `update`?
41
+ *
42
+ * V1 ships the explicit pattern. An auto-hook tied to repository
43
+ * events would couple persistence to the embedding provider's
44
+ * availability — a transient rate-limit on the embedder would
45
+ * fail the create call. Apps that want auto-vectorize wire it
46
+ * themselves via `events.on('article.created', m =>
47
+ * articles.vectorize(m))` so they control the failure mode
48
+ * (fire-and-forget vs awaited vs queued via `@strav/queue`).
49
+ *
50
+ * Extension points (all optional overrides):
51
+ *
52
+ * - `collectionName()` — defaults to the table name from the
53
+ * schema. Override when the collection should differ from the
54
+ * table, or to compose a per-tenant suffix dynamically.
55
+ *
56
+ * - `toContent(model)` — defaults to concatenating every string
57
+ * field on the model with `\n`. The default works for simple
58
+ * row shapes; apps with structured content override.
59
+ *
60
+ * - `toMetadata(model)` — defaults to `{}`. Apps return fields
61
+ * they want to filter on (e.g. `author_id`, `lang`, `kind`).
62
+ *
63
+ * - `shouldRetrieve(model)` — gates indexing. Return `false` for
64
+ * draft / soft-deleted / private rows. The default is `true`.
65
+ */
66
+
67
+ import type { Repository } from '@strav/database'
68
+ import type { RagManager } from './rag_manager.ts'
69
+ import type {
70
+ RetrieveOptions,
71
+ RetrieveResult,
72
+ VectorMatch,
73
+ } from './types.ts'
74
+
75
+ /** Minimal constructor type we can mix into. Wider than `typeof Repository` so subclasses with extra ctor args still type-check. */
76
+ // biome-ignore lint/suspicious/noExplicitAny: mixin constructor signatures intentionally accept any[]; the user-side subclass narrows.
77
+ type RepositoryConstructor<TModel extends object> = new (...args: any[]) => Repository<TModel>
78
+
79
+ /**
80
+ * Returns a subclass that extends `Base` with `vectorize` /
81
+ * `vectorRemove` / `retrieve` plus override-points
82
+ * (`collectionName`, `toContent`, `toMetadata`,
83
+ * `shouldRetrieve`). The user-side class declares an explicit
84
+ * constructor that calls `super(...)` and assigns `this.rag`.
85
+ */
86
+ export function retrievable<TModel extends object, TBase extends RepositoryConstructor<TModel>>(
87
+ Base: TBase,
88
+ ) {
89
+ abstract class RetrievableRepository extends Base {
90
+ /**
91
+ * The framework's `RagManager`. Assigned by the user-side
92
+ * subclass constructor. Public on purpose — apps that want to
93
+ * drop down to raw `rag.store()` / `rag.ingest(...)` access
94
+ * have a hook.
95
+ */
96
+ rag!: RagManager
97
+
98
+ /**
99
+ * Collection name for vector storage. Defaults to the table
100
+ * name from `static schema`. Override to point at a different
101
+ * collection (or to compose per-tenant / per-env suffixes).
102
+ */
103
+ protected collectionName(): string {
104
+ const ctor = this.constructor as unknown as { schema: { name: string } }
105
+ return ctor.schema.name
106
+ }
107
+
108
+ /**
109
+ * Build the indexable text from a model row. Default
110
+ * concatenates every non-underscore string field with `\n`.
111
+ * Apps with structured content override this — typically
112
+ * something like `` `${a.title}\n\n${a.body}` ``.
113
+ */
114
+ protected toContent(model: TModel): string {
115
+ const parts: string[] = []
116
+ for (const [key, value] of Object.entries(model as Record<string, unknown>)) {
117
+ if (key.startsWith('_')) continue
118
+ if (typeof value === 'string' && value.length > 0) parts.push(value)
119
+ }
120
+ return parts.join('\n')
121
+ }
122
+
123
+ /**
124
+ * Build the metadata bag attached to every chunk. Apps return
125
+ * fields they want to filter retrievals on. The framework
126
+ * automatically adds `chunkIndex`, `startOffset`, `endOffset`
127
+ * — overrides shouldn't try to re-add those.
128
+ */
129
+ protected toMetadata(_model: TModel): Record<string, unknown> {
130
+ return {}
131
+ }
132
+
133
+ /**
134
+ * Whether the model should currently be indexed. Override to
135
+ * skip drafts, soft-deleted rows, private records, etc. The
136
+ * default `true` indexes every model — fine for the common
137
+ * case.
138
+ */
139
+ protected shouldRetrieve(_model: TModel): boolean {
140
+ return true
141
+ }
142
+
143
+ /**
144
+ * (Re-)index a single model. Drops any existing chunks for
145
+ * the model's id, then ingests fresh chunks of the current
146
+ * content. When `shouldRetrieve(model)` returns `false`, the
147
+ * chunks are dropped without re-ingest — apps don't need a
148
+ * separate "this just became private" path.
149
+ *
150
+ * Returns the vector ids written. Empty array when content
151
+ * was empty or `shouldRetrieve` returned `false`.
152
+ */
153
+ async vectorize(model: TModel): Promise<string[]> {
154
+ const collection = this.collectionName()
155
+ const id = modelId(model)
156
+
157
+ // Drop existing chunks for this source first so updates
158
+ // replace cleanly. (RagManager.ingest writes fresh ids per
159
+ // call; without this step every re-vectorize would
160
+ // duplicate.)
161
+ await this.rag
162
+ .store()
163
+ .deleteBySource(this.rag.collectionName(collection), id)
164
+
165
+ if (!this.shouldRetrieve(model)) return []
166
+
167
+ const content = this.toContent(model)
168
+ if (!content) return []
169
+
170
+ return this.rag.ingest(collection, content, {
171
+ sourceId: id,
172
+ metadata: this.toMetadata(model),
173
+ })
174
+ }
175
+
176
+ /**
177
+ * Drop every chunk for one model. Apps call this after
178
+ * `delete(model)` in their domain code. The mixin doesn't
179
+ * auto-hook the delete lifecycle for the same reason it
180
+ * doesn't auto-hook create/update — keeps embedding-provider
181
+ * availability out of the persistence path.
182
+ */
183
+ async vectorRemove(model: TModel): Promise<void> {
184
+ const collection = this.collectionName()
185
+ const id = modelId(model)
186
+ await this.rag
187
+ .store()
188
+ .deleteBySource(this.rag.collectionName(collection), id)
189
+ }
190
+
191
+ /**
192
+ * Semantic search over this repository's collection. Default
193
+ * `collection` is the mixin's `collectionName()` — apps that
194
+ * want to retrieve from another collection pass it explicitly.
195
+ */
196
+ async retrieve(
197
+ query: string,
198
+ options: Omit<RetrieveOptions, 'collection'> & { collection?: string } = {},
199
+ ): Promise<RetrieveResult> {
200
+ return this.rag.retrieve(query, {
201
+ ...options,
202
+ collection: options.collection ?? this.collectionName(),
203
+ })
204
+ }
205
+
206
+ /**
207
+ * Re-index every row in the repository. Walks rows in batches
208
+ * of `batchSize` and vectorizes each. Useful for backfilling
209
+ * a new collection or recovering after a schema change.
210
+ *
211
+ * The CLI's `rag:reindex <repository>` doesn't ship in V1 —
212
+ * apps that want one wire it as their own console command
213
+ * pointing at this method.
214
+ *
215
+ * Returns the total count of rows processed (NOT the chunk
216
+ * count — chunks per row vary with content size).
217
+ */
218
+ async reindexAll(batchSize: number = 100): Promise<number> {
219
+ let processed = 0
220
+ let offset = 0
221
+ while (true) {
222
+ const rows = await this.query().orderBy('id', 'asc').limit(batchSize).offset(offset).get()
223
+ if (rows.length === 0) break
224
+ for (const row of rows) await this.vectorize(row)
225
+ processed += rows.length
226
+ offset += rows.length
227
+ if (rows.length < batchSize) break
228
+ }
229
+ return processed
230
+ }
231
+
232
+ /**
233
+ * Match-to-models helper. Takes the `matches` array from
234
+ * `retrieve(...)` and hydrates the source rows by id, in
235
+ * match order. Matches whose `sourceId` doesn't resolve to a
236
+ * row (deleted between index time + retrieval) are dropped.
237
+ */
238
+ async resolveMatches(matches: readonly VectorMatch[]): Promise<TModel[]> {
239
+ const ids = [...new Set(matches.map((m) => m.sourceId).filter((s): s is string => !!s))]
240
+ if (ids.length === 0) return []
241
+ const found = await this.findMany(ids as unknown as readonly string[])
242
+ const byId = new Map<string, TModel>(
243
+ found.map((m) => [modelId(m), m]),
244
+ )
245
+ const out: TModel[] = []
246
+ for (const match of matches) {
247
+ if (!match.sourceId) continue
248
+ const row = byId.get(match.sourceId)
249
+ if (row) out.push(row)
250
+ }
251
+ return out
252
+ }
253
+ }
254
+ return RetrievableRepository
255
+ }
256
+
257
+ /**
258
+ * Coerce a model's `id` to a string. Repositories use ULID or UUID
259
+ * ids by default, both of which round-trip through `String(...)`
260
+ * cleanly; integer PKs (bigSerial) coerce the same way.
261
+ */
262
+ function modelId(model: object): string {
263
+ const id = (model as { id?: unknown }).id
264
+ if (id === undefined || id === null) {
265
+ throw new Error(
266
+ `retrievable: model has no \`id\` to use as a vector sourceId. The mixin only works on models with a single-column id.`,
267
+ )
268
+ }
269
+ return String(id)
270
+ }