@strav/rag 0.1.0

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 ADDED
@@ -0,0 +1,33 @@
1
+ {
2
+ "name": "@strav/rag",
3
+ "version": "0.1.0",
4
+ "type": "module",
5
+ "description": "Vector retrieval framework for RAG in the Strav framework",
6
+ "license": "MIT",
7
+ "exports": {
8
+ ".": "./src/index.ts",
9
+ "./*": "./src/*.ts"
10
+ },
11
+ "strav": {
12
+ "commands": "src/commands"
13
+ },
14
+ "files": [
15
+ "src/",
16
+ "stubs/",
17
+ "package.json",
18
+ "tsconfig.json"
19
+ ],
20
+ "peerDependencies": {
21
+ "@strav/kernel": "0.1.0",
22
+ "@strav/brain": "0.1.0",
23
+ "@strav/database": "0.1.0",
24
+ "@strav/cli": "0.1.0"
25
+ },
26
+ "scripts": {
27
+ "test": "bun test tests/",
28
+ "typecheck": "tsc --noEmit"
29
+ },
30
+ "devDependencies": {
31
+ "commander": "^14.0.3"
32
+ }
33
+ }
@@ -0,0 +1,14 @@
1
+ import type { Chunker, ChunkingConfig } from '../types.ts'
2
+ import { FixedSizeChunker } from './fixed_size_chunker.ts'
3
+ import { RecursiveChunker } from './recursive_chunker.ts'
4
+
5
+ export function createChunker(config: ChunkingConfig): Chunker {
6
+ switch (config.strategy) {
7
+ case 'fixed':
8
+ return new FixedSizeChunker(config.chunkSize, config.overlap)
9
+ case 'recursive':
10
+ return new RecursiveChunker(config.chunkSize, config.overlap, config.separators)
11
+ default:
12
+ return new RecursiveChunker(config.chunkSize, config.overlap)
13
+ }
14
+ }
@@ -0,0 +1,33 @@
1
+ import type { Chunk, Chunker } from '../types.ts'
2
+
3
+ export class FixedSizeChunker implements Chunker {
4
+ constructor(
5
+ private readonly chunkSize: number = 512,
6
+ private readonly overlap: number = 64
7
+ ) {}
8
+
9
+ chunk(content: string): Chunk[] {
10
+ if (!content) return []
11
+
12
+ const chunks: Chunk[] = []
13
+ const step = Math.max(1, this.chunkSize - this.overlap)
14
+
15
+ let start = 0
16
+ let index = 0
17
+
18
+ while (start < content.length) {
19
+ const end = Math.min(start + this.chunkSize, content.length)
20
+ chunks.push({
21
+ content: content.slice(start, end),
22
+ index,
23
+ startOffset: start,
24
+ endOffset: end,
25
+ })
26
+ index++
27
+ start += step
28
+ if (end === content.length) break
29
+ }
30
+
31
+ return chunks
32
+ }
33
+ }
@@ -0,0 +1,83 @@
1
+ import type { Chunk, Chunker } from '../types.ts'
2
+
3
+ const DEFAULT_SEPARATORS = ['\n\n', '\n', '. ', ' ']
4
+
5
+ export class RecursiveChunker implements Chunker {
6
+ private readonly separators: string[]
7
+
8
+ constructor(
9
+ private readonly chunkSize: number = 512,
10
+ private readonly overlap: number = 64,
11
+ separators?: string[]
12
+ ) {
13
+ this.separators = separators ?? DEFAULT_SEPARATORS
14
+ }
15
+
16
+ chunk(content: string): Chunk[] {
17
+ if (!content) return []
18
+ const pieces = this.splitRecursive(content, 0)
19
+ return this.buildChunks(content, pieces)
20
+ }
21
+
22
+ private splitRecursive(text: string, separatorIndex: number): string[] {
23
+ if (text.length <= this.chunkSize) return [text]
24
+
25
+ const separator = this.separators[separatorIndex]
26
+ if (!separator) {
27
+ const result: string[] = []
28
+ for (let i = 0; i < text.length; i += this.chunkSize) {
29
+ result.push(text.slice(i, i + this.chunkSize))
30
+ }
31
+ return result
32
+ }
33
+
34
+ const parts = text.split(separator)
35
+ const merged: string[] = []
36
+ let current = ''
37
+
38
+ for (const part of parts) {
39
+ const candidate = current ? current + separator + part : part
40
+
41
+ if (candidate.length <= this.chunkSize) {
42
+ current = candidate
43
+ } else {
44
+ if (current) merged.push(current)
45
+ if (part.length > this.chunkSize) {
46
+ merged.push(...this.splitRecursive(part, separatorIndex + 1))
47
+ current = ''
48
+ } else {
49
+ current = part
50
+ }
51
+ }
52
+ }
53
+ if (current) merged.push(current)
54
+
55
+ return merged
56
+ }
57
+
58
+ private buildChunks(original: string, pieces: string[]): Chunk[] {
59
+ const chunks: Chunk[] = []
60
+ let searchFrom = 0
61
+
62
+ for (let i = 0; i < pieces.length; i++) {
63
+ const piece = pieces[i]!
64
+ const foundAt = original.indexOf(piece, searchFrom)
65
+ const startOffset = foundAt >= 0 ? foundAt : searchFrom
66
+ const pieceEnd = startOffset + piece.length
67
+
68
+ const overlapEnd = Math.min(pieceEnd + this.overlap, original.length)
69
+ const chunkContent = original.slice(startOffset, overlapEnd)
70
+
71
+ chunks.push({
72
+ content: chunkContent,
73
+ index: i,
74
+ startOffset,
75
+ endOffset: overlapEnd,
76
+ })
77
+
78
+ searchFrom = pieceEnd
79
+ }
80
+
81
+ return chunks
82
+ }
83
+ }
@@ -0,0 +1,41 @@
1
+ import type { Command } from 'commander'
2
+ import chalk from 'chalk'
3
+ import { bootstrap, shutdown } from '@stravigor/cli'
4
+ import { BaseModel } from '@stravigor/database'
5
+ import RagManager from '../rag_manager.ts'
6
+
7
+ export function register(program: Command): void {
8
+ program
9
+ .command('rag:flush <model>')
10
+ .description("Flush all vectors from a model's vector collection")
11
+ .action(async (modelPath: string) => {
12
+ let db
13
+ try {
14
+ const { db: database, config } = await bootstrap()
15
+ db = database
16
+
17
+ new BaseModel(db)
18
+ new RagManager(config)
19
+
20
+ const resolved = require.resolve(`${process.cwd()}/${modelPath}`)
21
+ const module = await import(resolved)
22
+ const ModelClass = module.default ?? (Object.values(module)[0] as any)
23
+
24
+ if (typeof ModelClass?.flushVectors !== 'function') {
25
+ console.error(chalk.red(`Model "${modelPath}" does not use the retrievable() mixin.`))
26
+ process.exit(1)
27
+ }
28
+
29
+ const collectionName = ModelClass.retrievableAs()
30
+ console.log(chalk.dim(`Flushing "${collectionName}"...`))
31
+
32
+ await ModelClass.flushVectors()
33
+ console.log(chalk.green(`Flushed all vectors from "${collectionName}".`))
34
+ } catch (err) {
35
+ console.error(chalk.red(`Error: ${err instanceof Error ? err.message : err}`))
36
+ process.exit(1)
37
+ } finally {
38
+ if (db) await shutdown(db)
39
+ }
40
+ })
41
+ }
@@ -0,0 +1,45 @@
1
+ import type { Command } from 'commander'
2
+ import chalk from 'chalk'
3
+ import { bootstrap, shutdown } from '@stravigor/cli'
4
+ import { BaseModel } from '@stravigor/database'
5
+ import { BrainManager } from '@stravigor/brain'
6
+ import RagManager from '../rag_manager.ts'
7
+
8
+ export function register(program: Command): void {
9
+ program
10
+ .command('rag:ingest <model>')
11
+ .description('Vectorize all records for a model into the vector store')
12
+ .option('--chunk <size>', 'Records per batch', '100')
13
+ .action(async (modelPath: string, options: { chunk: string }) => {
14
+ let db
15
+ try {
16
+ const { db: database, config } = await bootstrap()
17
+ db = database
18
+
19
+ new BaseModel(db)
20
+ new RagManager(config)
21
+ new BrainManager(config)
22
+
23
+ const resolved = require.resolve(`${process.cwd()}/${modelPath}`)
24
+ const module = await import(resolved)
25
+ const ModelClass = module.default ?? (Object.values(module)[0] as any)
26
+
27
+ if (typeof ModelClass?.importAll !== 'function') {
28
+ console.error(chalk.red(`Model "${modelPath}" does not use the retrievable() mixin.`))
29
+ process.exit(1)
30
+ }
31
+
32
+ const chunkSize = parseInt(options.chunk, 10)
33
+ const collectionName = ModelClass.retrievableAs()
34
+ console.log(chalk.dim(`Vectorizing ${ModelClass.name} into "${collectionName}"...`))
35
+
36
+ const count = await ModelClass.importAll(chunkSize)
37
+ console.log(chalk.green(`Vectorized ${count} record(s) into "${collectionName}".`))
38
+ } catch (err) {
39
+ console.error(chalk.red(`Error: ${err instanceof Error ? err.message : err}`))
40
+ process.exit(1)
41
+ } finally {
42
+ if (db) await shutdown(db)
43
+ }
44
+ })
45
+ }
@@ -0,0 +1,135 @@
1
+ import type { VectorStore } from '../vector_store.ts'
2
+ import type { VectorDocument, QueryOptions, QueryResult, VectorMatch } from '../types.ts'
3
+
4
+ export class MemoryDriver implements VectorStore {
5
+ readonly name = 'memory'
6
+ private collections = new Map<string, VectorDocument[]>()
7
+
8
+ async createCollection(collection: string, _dimension: number): Promise<void> {
9
+ if (!this.collections.has(collection)) {
10
+ this.collections.set(collection, [])
11
+ }
12
+ }
13
+
14
+ async deleteCollection(collection: string): Promise<void> {
15
+ this.collections.delete(collection)
16
+ }
17
+
18
+ async upsert(collection: string, documents: VectorDocument[]): Promise<void> {
19
+ let docs = this.collections.get(collection)
20
+ if (!docs) {
21
+ docs = []
22
+ this.collections.set(collection, docs)
23
+ }
24
+
25
+ for (const doc of documents) {
26
+ if (doc.id != null) {
27
+ const existingIndex = docs.findIndex(d => d.id === doc.id)
28
+ if (existingIndex >= 0) {
29
+ docs[existingIndex] = doc
30
+ } else {
31
+ docs.push(doc)
32
+ }
33
+ } else {
34
+ docs.push(doc)
35
+ }
36
+ }
37
+ }
38
+
39
+ async delete(collection: string, ids: (string | number)[]): Promise<void> {
40
+ const docs = this.collections.get(collection)
41
+ if (!docs) return
42
+
43
+ const idSet = new Set(ids.map(String))
44
+ this.collections.set(
45
+ collection,
46
+ docs.filter(d => !idSet.has(String(d.id)))
47
+ )
48
+ }
49
+
50
+ async deleteBySource(collection: string, sourceId: string | number): Promise<void> {
51
+ const docs = this.collections.get(collection)
52
+ if (!docs) return
53
+
54
+ const sourceStr = String(sourceId)
55
+ this.collections.set(
56
+ collection,
57
+ docs.filter(d => String(d.sourceId) !== sourceStr)
58
+ )
59
+ }
60
+
61
+ async flush(collection: string): Promise<void> {
62
+ if (this.collections.has(collection)) {
63
+ this.collections.set(collection, [])
64
+ }
65
+ }
66
+
67
+ async query(
68
+ collection: string,
69
+ vector: number[],
70
+ options?: QueryOptions
71
+ ): Promise<QueryResult> {
72
+ const start = performance.now()
73
+ const docs = this.collections.get(collection)
74
+ if (!docs || docs.length === 0) {
75
+ return { matches: [], processingTimeMs: performance.now() - start }
76
+ }
77
+
78
+ const topK = options?.topK ?? 5
79
+ const threshold = options?.threshold ?? 0
80
+
81
+ let scored: VectorMatch[] = docs.map(doc => ({
82
+ id: doc.id ?? 0,
83
+ content: doc.content,
84
+ score: cosineSimilarity(vector, doc.embedding),
85
+ metadata: doc.metadata ?? {},
86
+ }))
87
+
88
+ if (options?.filter) {
89
+ scored = scored.filter(m => matchesFilter(m.metadata, options.filter!))
90
+ }
91
+
92
+ if (threshold > 0) {
93
+ scored = scored.filter(m => m.score >= threshold)
94
+ }
95
+
96
+ scored.sort((a, b) => b.score - a.score)
97
+ const matches = scored.slice(0, topK)
98
+
99
+ return {
100
+ matches,
101
+ processingTimeMs: performance.now() - start,
102
+ }
103
+ }
104
+
105
+ getCollection(collection: string): VectorDocument[] {
106
+ return this.collections.get(collection) ?? []
107
+ }
108
+ }
109
+
110
+ function cosineSimilarity(a: number[], b: number[]): number {
111
+ let dot = 0
112
+ let magA = 0
113
+ let magB = 0
114
+
115
+ for (let i = 0; i < a.length; i++) {
116
+ const ai = a[i]!
117
+ const bi = b[i]!
118
+ dot += ai * bi
119
+ magA += ai * ai
120
+ magB += bi * bi
121
+ }
122
+
123
+ const denom = Math.sqrt(magA) * Math.sqrt(magB)
124
+ return denom === 0 ? 0 : dot / denom
125
+ }
126
+
127
+ function matchesFilter(
128
+ metadata: Record<string, unknown>,
129
+ filter: Record<string, unknown>
130
+ ): boolean {
131
+ for (const [key, value] of Object.entries(filter)) {
132
+ if (metadata[key] !== value) return false
133
+ }
134
+ return true
135
+ }
@@ -0,0 +1,21 @@
1
+ import type { VectorStore } from '../vector_store.ts'
2
+ import type { VectorDocument, QueryOptions, QueryResult } from '../types.ts'
3
+
4
+ export class NullDriver implements VectorStore {
5
+ readonly name = 'null'
6
+
7
+ async createCollection(_collection: string, _dimension: number): Promise<void> {}
8
+ async deleteCollection(_collection: string): Promise<void> {}
9
+ async upsert(_collection: string, _documents: VectorDocument[]): Promise<void> {}
10
+ async delete(_collection: string, _ids: (string | number)[]): Promise<void> {}
11
+ async deleteBySource(_collection: string, _sourceId: string | number): Promise<void> {}
12
+ async flush(_collection: string): Promise<void> {}
13
+
14
+ async query(
15
+ _collection: string,
16
+ _vector: number[],
17
+ _options?: QueryOptions
18
+ ): Promise<QueryResult> {
19
+ return { matches: [] }
20
+ }
21
+ }
@@ -0,0 +1,157 @@
1
+ import { Database } from '@stravigor/database'
2
+ import type { VectorStore } from '../vector_store.ts'
3
+ import type { VectorDocument, QueryOptions, QueryResult, VectorMatch, StoreConfig } from '../types.ts'
4
+ import { VectorQueryError } from '../errors.ts'
5
+
6
+ export class PgvectorDriver implements VectorStore {
7
+ readonly name = 'pgvector'
8
+ private initialized = false
9
+
10
+ constructor(_config: StoreConfig) {}
11
+
12
+ async createCollection(collection: string, dimension: number): Promise<void> {
13
+ await this.ensureTable(dimension)
14
+
15
+ const indexName = `idx_strav_vectors_hnsw_${collection.replace(/[^a-z0-9_]/gi, '_')}`
16
+ try {
17
+ await Database.raw.unsafe(
18
+ `CREATE INDEX IF NOT EXISTS "${indexName}"
19
+ ON _strav_vectors USING hnsw (embedding vector_cosine_ops)
20
+ WHERE collection = '${collection}'`
21
+ )
22
+ } catch {
23
+ // Index may already exist
24
+ }
25
+ }
26
+
27
+ async deleteCollection(collection: string): Promise<void> {
28
+ await Database.raw.unsafe(
29
+ `DELETE FROM _strav_vectors WHERE collection = $1`,
30
+ [collection]
31
+ )
32
+ }
33
+
34
+ async upsert(collection: string, documents: VectorDocument[]): Promise<void> {
35
+ const sql = Database.raw
36
+
37
+ for (const doc of documents) {
38
+ const embeddingStr = `[${doc.embedding.join(',')}]`
39
+ const metadata = JSON.stringify(doc.metadata ?? {})
40
+ const id = doc.id != null ? String(doc.id) : crypto.randomUUID()
41
+ const sourceId = doc.sourceId != null ? String(doc.sourceId) : null
42
+
43
+ await sql.unsafe(
44
+ `INSERT INTO _strav_vectors (collection, source_id, content, metadata, embedding)
45
+ VALUES ($1, $2, $3, $4::jsonb, $5::vector)`,
46
+ [collection, sourceId, doc.content, metadata, embeddingStr]
47
+ )
48
+ }
49
+ }
50
+
51
+ async delete(collection: string, ids: (string | number)[]): Promise<void> {
52
+ if (ids.length === 0) return
53
+ const placeholders = ids.map((_, i) => `$${i + 2}`).join(', ')
54
+ await Database.raw.unsafe(
55
+ `DELETE FROM _strav_vectors WHERE collection = $1 AND id IN (${placeholders})`,
56
+ [collection, ...ids]
57
+ )
58
+ }
59
+
60
+ async deleteBySource(collection: string, sourceId: string | number): Promise<void> {
61
+ await Database.raw.unsafe(
62
+ `DELETE FROM _strav_vectors WHERE collection = $1 AND source_id = $2`,
63
+ [collection, String(sourceId)]
64
+ )
65
+ }
66
+
67
+ async flush(collection: string): Promise<void> {
68
+ await Database.raw.unsafe(
69
+ `DELETE FROM _strav_vectors WHERE collection = $1`,
70
+ [collection]
71
+ )
72
+ }
73
+
74
+ async query(
75
+ collection: string,
76
+ vector: number[],
77
+ options?: QueryOptions
78
+ ): Promise<QueryResult> {
79
+ const start = performance.now()
80
+ const topK = options?.topK ?? 5
81
+ const threshold = options?.threshold ?? 0
82
+ const embeddingStr = `[${vector.join(',')}]`
83
+
84
+ let whereClause = 'collection = $1'
85
+ const params: unknown[] = [collection]
86
+ let paramIndex = 2
87
+
88
+ if (options?.filter) {
89
+ for (const [key, value] of Object.entries(options.filter)) {
90
+ whereClause += ` AND metadata->>'${key}' = $${paramIndex}`
91
+ params.push(String(value))
92
+ paramIndex++
93
+ }
94
+ }
95
+
96
+ if (threshold > 0) {
97
+ whereClause += ` AND (embedding <=> $${paramIndex}::vector) <= $${paramIndex + 1}`
98
+ params.push(embeddingStr, 1 - threshold)
99
+ paramIndex += 2
100
+ }
101
+
102
+ try {
103
+ const rows = (await Database.raw.unsafe(
104
+ `SELECT id, source_id, content, metadata,
105
+ 1 - (embedding <=> $${paramIndex}::vector) AS score
106
+ FROM _strav_vectors
107
+ WHERE ${whereClause}
108
+ ORDER BY embedding <=> $${paramIndex}::vector
109
+ LIMIT $${paramIndex + 1}`,
110
+ [...params, embeddingStr, topK]
111
+ )) as any[]
112
+
113
+ const matches: VectorMatch[] = rows.map((row: any) => ({
114
+ id: row.source_id ?? row.id,
115
+ content: row.content,
116
+ score: parseFloat(row.score),
117
+ metadata: typeof row.metadata === 'string' ? JSON.parse(row.metadata) : row.metadata ?? {},
118
+ }))
119
+
120
+ return {
121
+ matches,
122
+ processingTimeMs: performance.now() - start,
123
+ }
124
+ } catch (err) {
125
+ throw new VectorQueryError(collection, err instanceof Error ? err.message : String(err))
126
+ }
127
+ }
128
+
129
+ private async ensureTable(dimension: number): Promise<void> {
130
+ if (this.initialized) return
131
+
132
+ const sql = Database.raw
133
+
134
+ await sql.unsafe(`CREATE EXTENSION IF NOT EXISTS vector`)
135
+
136
+ await sql.unsafe(`
137
+ CREATE TABLE IF NOT EXISTS _strav_vectors (
138
+ id BIGSERIAL PRIMARY KEY,
139
+ collection VARCHAR(255) NOT NULL,
140
+ source_id VARCHAR(255),
141
+ content TEXT NOT NULL,
142
+ metadata JSONB DEFAULT '{}',
143
+ embedding vector(${dimension}),
144
+ created_at TIMESTAMPTZ DEFAULT NOW()
145
+ )
146
+ `)
147
+
148
+ await sql.unsafe(
149
+ `CREATE INDEX IF NOT EXISTS idx_strav_vectors_collection ON _strav_vectors(collection)`
150
+ )
151
+ await sql.unsafe(
152
+ `CREATE INDEX IF NOT EXISTS idx_strav_vectors_source ON _strav_vectors(collection, source_id)`
153
+ )
154
+
155
+ this.initialized = true
156
+ }
157
+ }
package/src/errors.ts ADDED
@@ -0,0 +1,21 @@
1
+ import { StravError } from '@stravigor/kernel'
2
+
3
+ export class RagError extends StravError {}
4
+
5
+ export class CollectionNotFoundError extends RagError {
6
+ constructor(collection: string) {
7
+ super(`Vector collection "${collection}" not found.`)
8
+ }
9
+ }
10
+
11
+ export class VectorQueryError extends RagError {
12
+ constructor(collection: string, cause?: string) {
13
+ super(`Vector query on "${collection}" failed${cause ? `: ${cause}` : ''}.`)
14
+ }
15
+ }
16
+
17
+ export class EmbeddingError extends RagError {
18
+ constructor(cause?: string) {
19
+ super(`Embedding generation failed${cause ? `: ${cause}` : ''}.`)
20
+ }
21
+ }
package/src/helpers.ts ADDED
@@ -0,0 +1,160 @@
1
+ import { brain } from '@stravigor/brain'
2
+ import RagManager from './rag_manager.ts'
3
+ import type { VectorStore } from './vector_store.ts'
4
+ import type {
5
+ RetrieveOptions,
6
+ RetrieveResult,
7
+ RetrievedDocument,
8
+ VectorDocument,
9
+ StoreConfig,
10
+ } from './types.ts'
11
+ import { createChunker } from './chunking/chunker.ts'
12
+ import { EmbeddingError } from './errors.ts'
13
+
14
+ export interface IngestOptions {
15
+ metadata?: Record<string, unknown>
16
+ sourceId?: string | number
17
+ chunkSize?: number
18
+ overlap?: number
19
+ strategy?: string
20
+ }
21
+
22
+ export const rag = {
23
+ store(name?: string): VectorStore {
24
+ return RagManager.store(name)
25
+ },
26
+
27
+ extend(name: string, factory: (config: StoreConfig) => VectorStore): void {
28
+ RagManager.extend(name, factory)
29
+ },
30
+
31
+ async ingest(
32
+ collection: string,
33
+ content: string,
34
+ options: IngestOptions = {}
35
+ ): Promise<string[]> {
36
+ const config = RagManager.config
37
+ const fullCollection = RagManager.collectionName(collection)
38
+
39
+ const chunkerConfig = {
40
+ strategy: options.strategy ?? config.chunking.strategy,
41
+ chunkSize: options.chunkSize ?? config.chunking.chunkSize,
42
+ overlap: options.overlap ?? config.chunking.overlap,
43
+ separators: config.chunking.separators,
44
+ }
45
+ const chunker = createChunker(chunkerConfig)
46
+ const chunks = chunker.chunk(content)
47
+
48
+ if (chunks.length === 0) return []
49
+
50
+ const chunkTexts = chunks.map(c => c.content)
51
+ let embeddings: number[][]
52
+ try {
53
+ embeddings = await brain.embed(chunkTexts, {
54
+ provider: config.embedding.provider,
55
+ model: config.embedding.model,
56
+ })
57
+ } catch (err) {
58
+ throw new EmbeddingError(err instanceof Error ? err.message : String(err))
59
+ }
60
+
61
+ const baseId = crypto.randomUUID()
62
+ const documents: VectorDocument[] = chunks.map((chunk, i) => ({
63
+ id: `${baseId}_${i}`,
64
+ sourceId: options.sourceId,
65
+ content: chunk.content,
66
+ embedding: embeddings[i]!,
67
+ metadata: {
68
+ ...options.metadata,
69
+ chunkIndex: chunk.index,
70
+ startOffset: chunk.startOffset,
71
+ endOffset: chunk.endOffset,
72
+ },
73
+ }))
74
+
75
+ await RagManager.store().upsert(fullCollection, documents)
76
+ return documents.map(d => String(d.id))
77
+ },
78
+
79
+ async retrieve(query: string, options: RetrieveOptions = {}): Promise<RetrieveResult> {
80
+ const start = performance.now()
81
+ const config = RagManager.config
82
+ const collection = RagManager.collectionName(options.collection ?? 'default')
83
+
84
+ let queryVector: number[]
85
+ try {
86
+ const vectors = await brain.embed(query, {
87
+ provider: config.embedding.provider,
88
+ model: config.embedding.model,
89
+ })
90
+ queryVector = vectors[0]!
91
+ } catch (err) {
92
+ throw new EmbeddingError(err instanceof Error ? err.message : String(err))
93
+ }
94
+
95
+ const queryResult = await RagManager.store().query(collection, queryVector, {
96
+ topK: options.topK,
97
+ threshold: options.threshold,
98
+ filter: options.filter,
99
+ })
100
+
101
+ let matches: RetrievedDocument[] = queryResult.matches.map(m => ({
102
+ id: m.id,
103
+ content: m.content,
104
+ score: m.score,
105
+ similarity: m.score,
106
+ metadata: m.metadata,
107
+ }))
108
+
109
+ if (options.rerank) {
110
+ const {
111
+ similarityWeight = 0.6,
112
+ authorityWeight = 0.2,
113
+ recencyWeight = 0.2,
114
+ } = options.rerank
115
+
116
+ matches = matches.map(m => {
117
+ const authority =
118
+ typeof m.metadata.authority === 'number' ? m.metadata.authority : 0
119
+ const createdAt = m.metadata.createdAt
120
+ const recencyScore = createdAt
121
+ ? 1 / (1 + daysSince(new Date(createdAt as string)) / 30)
122
+ : 0.5
123
+
124
+ const finalScore =
125
+ m.similarity * similarityWeight +
126
+ authority * authorityWeight +
127
+ recencyScore * recencyWeight
128
+
129
+ return { ...m, score: finalScore }
130
+ })
131
+
132
+ matches.sort((a, b) => b.score - a.score)
133
+ }
134
+
135
+ return {
136
+ matches,
137
+ query,
138
+ processingTimeMs: performance.now() - start,
139
+ }
140
+ },
141
+
142
+ async delete(collection: string, ids: (string | number)[]): Promise<void> {
143
+ const fullCollection = RagManager.collectionName(collection)
144
+ await RagManager.store().delete(fullCollection, ids)
145
+ },
146
+
147
+ async deleteBySource(collection: string, sourceId: string | number): Promise<void> {
148
+ const fullCollection = RagManager.collectionName(collection)
149
+ await RagManager.store().deleteBySource(fullCollection, sourceId)
150
+ },
151
+
152
+ async flush(collection: string): Promise<void> {
153
+ const fullCollection = RagManager.collectionName(collection)
154
+ await RagManager.store().flush(fullCollection)
155
+ },
156
+ }
157
+
158
+ function daysSince(date: Date): number {
159
+ return (Date.now() - date.getTime()) / (1000 * 60 * 60 * 24)
160
+ }
package/src/index.ts ADDED
@@ -0,0 +1,48 @@
1
+ // Manager
2
+ export { default, default as RagManager } from './rag_manager.ts'
3
+
4
+ // Provider
5
+ export { default as RagProvider } from './rag_provider.ts'
6
+
7
+ // Store interface
8
+ export type { VectorStore } from './vector_store.ts'
9
+
10
+ // Drivers
11
+ export { NullDriver } from './drivers/null_driver.ts'
12
+ export { MemoryDriver } from './drivers/memory_driver.ts'
13
+ export { PgvectorDriver } from './drivers/pgvector_driver.ts'
14
+
15
+ // Mixin
16
+ export { retrievable } from './retrievable.ts'
17
+ export type { RetrievableInstance, RetrievableModel } from './retrievable.ts'
18
+
19
+ // Helper
20
+ export { rag } from './helpers.ts'
21
+
22
+ // Chunking
23
+ export { createChunker } from './chunking/chunker.ts'
24
+ export { FixedSizeChunker } from './chunking/fixed_size_chunker.ts'
25
+ export { RecursiveChunker } from './chunking/recursive_chunker.ts'
26
+
27
+ // Errors
28
+ export { RagError, CollectionNotFoundError, VectorQueryError, EmbeddingError } from './errors.ts'
29
+
30
+ // Types
31
+ export type {
32
+ RagConfig,
33
+ StoreConfig,
34
+ EmbeddingConfig,
35
+ ChunkingConfig,
36
+ VectorDocument,
37
+ QueryOptions,
38
+ QueryResult,
39
+ VectorMatch,
40
+ RetrieveOptions,
41
+ RerankOptions,
42
+ RetrieveResult,
43
+ RetrievedDocument,
44
+ Chunk,
45
+ Chunker,
46
+ } from './types.ts'
47
+
48
+ export type { IngestOptions } from './helpers.ts'
@@ -0,0 +1,98 @@
1
+ import { inject, Configuration, ConfigurationError } from '@stravigor/kernel'
2
+ import type { VectorStore } from './vector_store.ts'
3
+ import type { RagConfig, StoreConfig, EmbeddingConfig, ChunkingConfig } from './types.ts'
4
+ import { NullDriver } from './drivers/null_driver.ts'
5
+ import { MemoryDriver } from './drivers/memory_driver.ts'
6
+ import { PgvectorDriver } from './drivers/pgvector_driver.ts'
7
+
8
+ @inject
9
+ export default class RagManager {
10
+ private static _config: RagConfig
11
+ private static _stores = new Map<string, VectorStore>()
12
+ private static _extensions = new Map<string, (config: StoreConfig) => VectorStore>()
13
+
14
+ constructor(config: Configuration) {
15
+ RagManager._config = {
16
+ default: config.get('rag.default', 'null') as string,
17
+ prefix: config.get('rag.prefix', '') as string,
18
+ embedding: config.get('rag.embedding', {
19
+ provider: 'openai',
20
+ model: 'text-embedding-3-small',
21
+ dimension: 1536,
22
+ }) as EmbeddingConfig,
23
+ chunking: config.get('rag.chunking', {
24
+ strategy: 'recursive',
25
+ chunkSize: 512,
26
+ overlap: 64,
27
+ }) as ChunkingConfig,
28
+ stores: config.get('rag.stores', {}) as Record<string, StoreConfig>,
29
+ }
30
+ }
31
+
32
+ static get config(): RagConfig {
33
+ if (!RagManager._config) {
34
+ throw new ConfigurationError(
35
+ 'RagManager not configured. Resolve it through the container first.'
36
+ )
37
+ }
38
+ return RagManager._config
39
+ }
40
+
41
+ static store(name?: string): VectorStore {
42
+ const key = name ?? RagManager.config.default
43
+
44
+ let store = RagManager._stores.get(key)
45
+ if (store) return store
46
+
47
+ const storeConfig = RagManager.config.stores[key]
48
+ if (!storeConfig) {
49
+ throw new ConfigurationError(`RAG store "${key}" is not configured.`)
50
+ }
51
+
52
+ store = RagManager.createStore(key, storeConfig)
53
+ RagManager._stores.set(key, store)
54
+ return store
55
+ }
56
+
57
+ static get prefix(): string {
58
+ return RagManager._config?.prefix ?? ''
59
+ }
60
+
61
+ static collectionName(name: string): string {
62
+ return RagManager.prefix ? `${RagManager.prefix}${name}` : name
63
+ }
64
+
65
+ static extend(name: string, factory: (config: StoreConfig) => VectorStore): void {
66
+ RagManager._extensions.set(name, factory)
67
+ }
68
+
69
+ static useStore(store: VectorStore): void {
70
+ RagManager._stores.set(store.name, store)
71
+ }
72
+
73
+ static reset(): void {
74
+ RagManager._stores.clear()
75
+ RagManager._extensions.clear()
76
+ RagManager._config = undefined as any
77
+ }
78
+
79
+ private static createStore(name: string, config: StoreConfig): VectorStore {
80
+ const driverName = config.driver ?? name
81
+
82
+ const extension = RagManager._extensions.get(driverName)
83
+ if (extension) return extension(config)
84
+
85
+ switch (driverName) {
86
+ case 'pgvector':
87
+ return new PgvectorDriver(config)
88
+ case 'memory':
89
+ return new MemoryDriver()
90
+ case 'null':
91
+ return new NullDriver()
92
+ default:
93
+ throw new ConfigurationError(
94
+ `Unknown RAG driver "${driverName}". Register it with RagManager.extend().`
95
+ )
96
+ }
97
+ }
98
+ }
@@ -0,0 +1,16 @@
1
+ import { ServiceProvider } from '@stravigor/kernel'
2
+ import type { Application } from '@stravigor/kernel'
3
+ import RagManager from './rag_manager.ts'
4
+
5
+ export default class RagProvider extends ServiceProvider {
6
+ readonly name = 'rag'
7
+ override readonly dependencies = ['config']
8
+
9
+ override register(app: Application): void {
10
+ app.singleton(RagManager)
11
+ }
12
+
13
+ override boot(app: Application): void {
14
+ app.resolve(RagManager)
15
+ }
16
+ }
@@ -0,0 +1,179 @@
1
+ import type { BaseModel } from '@stravigor/database'
2
+ import type { NormalizeConstructor } from '@stravigor/kernel'
3
+ import { Emitter } from '@stravigor/kernel'
4
+ import { brain } from '@stravigor/brain'
5
+ import RagManager from './rag_manager.ts'
6
+ import { createChunker } from './chunking/chunker.ts'
7
+ import type { VectorDocument, RetrieveOptions, RetrieveResult } from './types.ts'
8
+
9
+ export function retrievable<T extends NormalizeConstructor<typeof BaseModel>>(Base: T) {
10
+ return class Retrievable extends Base {
11
+ private static _retrievalBooted = false
12
+
13
+ static retrievableAs(): string {
14
+ return (this as unknown as typeof BaseModel).tableName
15
+ }
16
+
17
+ toRetrievableContent(): string {
18
+ const parts: string[] = []
19
+ for (const key of Object.keys(this)) {
20
+ if (key.startsWith('_')) continue
21
+ const val = (this as any)[key]
22
+ if (typeof val === 'string' && val.length > 0) parts.push(val)
23
+ }
24
+ return parts.join('\n')
25
+ }
26
+
27
+ toRetrievableMetadata(): Record<string, unknown> {
28
+ return {}
29
+ }
30
+
31
+ shouldBeRetrievable(): boolean {
32
+ return true
33
+ }
34
+
35
+ // ── Instance methods ──────────────────────────────────────────────
36
+
37
+ async vectorize(): Promise<void> {
38
+ if (!this.shouldBeRetrievable()) return
39
+
40
+ const ctor = this.constructor as typeof Retrievable
41
+ const collection = RagManager.collectionName(ctor.retrievableAs())
42
+ const config = RagManager.config
43
+ const pkProp = (ctor as unknown as typeof BaseModel).primaryKeyProperty
44
+ const id = (this as any)[pkProp]
45
+
46
+ const content = this.toRetrievableContent()
47
+ if (!content) return
48
+
49
+ // Remove existing chunks for this model instance
50
+ await RagManager.store().deleteBySource(collection, id)
51
+
52
+ const chunker = createChunker(config.chunking)
53
+ const chunks = chunker.chunk(content)
54
+ if (chunks.length === 0) return
55
+
56
+ const texts = chunks.map(c => c.content)
57
+ const embeddings = await brain.embed(texts, {
58
+ provider: config.embedding.provider,
59
+ model: config.embedding.model,
60
+ })
61
+
62
+ const metadata = this.toRetrievableMetadata()
63
+ const documents: VectorDocument[] = chunks.map((chunk, i) => ({
64
+ id: `${id}_${i}`,
65
+ sourceId: id,
66
+ content: chunk.content,
67
+ embedding: embeddings[i]!,
68
+ metadata: {
69
+ ...metadata,
70
+ modelId: id,
71
+ chunkIndex: chunk.index,
72
+ },
73
+ }))
74
+
75
+ await RagManager.store().upsert(collection, documents)
76
+ }
77
+
78
+ async vectorRemove(): Promise<void> {
79
+ const ctor = this.constructor as typeof Retrievable
80
+ const collection = RagManager.collectionName(ctor.retrievableAs())
81
+ const pkProp = (ctor as unknown as typeof BaseModel).primaryKeyProperty
82
+ const id = (this as any)[pkProp]
83
+ await RagManager.store().deleteBySource(collection, id)
84
+ }
85
+
86
+ // ── Static methods ────────────────────────────────────────────────
87
+
88
+ static async retrieve(query: string, options?: RetrieveOptions): Promise<RetrieveResult> {
89
+ const { rag } = await import('./helpers.ts')
90
+ return rag.retrieve(query, {
91
+ ...options,
92
+ collection: options?.collection ?? this.retrievableAs(),
93
+ })
94
+ }
95
+
96
+ static async importAll(chunkSize: number = 100): Promise<number> {
97
+ const ModelCtor = this as unknown as typeof BaseModel & typeof Retrievable
98
+ const collection = RagManager.collectionName(this.retrievableAs())
99
+ const config = RagManager.config
100
+ const db = ModelCtor.db
101
+ const table = ModelCtor.tableName
102
+ const pkCol = ModelCtor.primaryKeyColumn
103
+
104
+ await RagManager.store().createCollection(collection, config.embedding.dimension)
105
+
106
+ let imported = 0
107
+ let offset = 0
108
+
109
+ while (true) {
110
+ const rows = (await db.sql.unsafe(
111
+ `SELECT * FROM "${table}" ORDER BY "${pkCol}" LIMIT $1 OFFSET $2`,
112
+ [chunkSize, offset]
113
+ )) as Record<string, unknown>[]
114
+
115
+ if (rows.length === 0) break
116
+
117
+ for (const row of rows) {
118
+ const instance = ModelCtor.hydrate(row) as InstanceType<typeof Retrievable>
119
+ if (instance.shouldBeRetrievable()) {
120
+ try {
121
+ await instance.vectorize()
122
+ imported++
123
+ } catch {
124
+ // Vectorization is secondary — continue on failure
125
+ }
126
+ }
127
+ }
128
+
129
+ offset += chunkSize
130
+ if (rows.length < chunkSize) break
131
+ }
132
+
133
+ return imported
134
+ }
135
+
136
+ static async flushVectors(): Promise<void> {
137
+ const collection = RagManager.collectionName(this.retrievableAs())
138
+ await RagManager.store().flush(collection)
139
+ }
140
+
141
+ static async createVectorCollection(): Promise<void> {
142
+ const collection = RagManager.collectionName(this.retrievableAs())
143
+ await RagManager.store().createCollection(collection, RagManager.config.embedding.dimension)
144
+ }
145
+
146
+ static bootRetrieval(eventPrefix: string): void {
147
+ if (this._retrievalBooted) return
148
+ this._retrievalBooted = true
149
+
150
+ const vectorizeFn = async (model: unknown) => {
151
+ if (model && typeof (model as any).vectorize === 'function') {
152
+ try {
153
+ await (model as any).vectorize()
154
+ } catch {
155
+ // Vectorization is secondary — failures should not break the event pipeline
156
+ }
157
+ }
158
+ }
159
+
160
+ const removeFn = async (model: unknown) => {
161
+ if (model && typeof (model as any).vectorRemove === 'function') {
162
+ try {
163
+ await (model as any).vectorRemove()
164
+ } catch {
165
+ // Vector removal is secondary
166
+ }
167
+ }
168
+ }
169
+
170
+ Emitter.on(`${eventPrefix}.created`, vectorizeFn)
171
+ Emitter.on(`${eventPrefix}.updated`, vectorizeFn)
172
+ Emitter.on(`${eventPrefix}.synced`, vectorizeFn)
173
+ Emitter.on(`${eventPrefix}.deleted`, removeFn)
174
+ }
175
+ }
176
+ }
177
+
178
+ export type RetrievableInstance = InstanceType<ReturnType<typeof retrievable>>
179
+ export type RetrievableModel = ReturnType<typeof retrievable>
package/src/types.ts ADDED
@@ -0,0 +1,100 @@
1
+ // ── Vector Documents ─────────────────────────────────────────────────────
2
+
3
+ export interface VectorDocument {
4
+ id?: string | number
5
+ sourceId?: string | number
6
+ content: string
7
+ embedding: number[]
8
+ metadata?: Record<string, unknown>
9
+ }
10
+
11
+ // ── Query Options & Results ──────────────────────────────────────────────
12
+
13
+ export interface QueryOptions {
14
+ topK?: number
15
+ threshold?: number
16
+ filter?: Record<string, unknown>
17
+ }
18
+
19
+ export interface QueryResult {
20
+ matches: VectorMatch[]
21
+ processingTimeMs?: number
22
+ }
23
+
24
+ export interface VectorMatch {
25
+ id: string | number
26
+ content: string
27
+ score: number
28
+ metadata: Record<string, unknown>
29
+ }
30
+
31
+ // ── Retrieval (high-level pipeline) ──────────────────────────────────────
32
+
33
+ export interface RetrieveOptions {
34
+ collection?: string
35
+ topK?: number
36
+ threshold?: number
37
+ filter?: Record<string, unknown>
38
+ rerank?: RerankOptions
39
+ }
40
+
41
+ export interface RerankOptions {
42
+ authorityWeight?: number
43
+ recencyWeight?: number
44
+ similarityWeight?: number
45
+ }
46
+
47
+ export interface RetrieveResult {
48
+ matches: RetrievedDocument[]
49
+ query: string
50
+ processingTimeMs: number
51
+ }
52
+
53
+ export interface RetrievedDocument {
54
+ id: string | number
55
+ content: string
56
+ score: number
57
+ similarity: number
58
+ metadata: Record<string, unknown>
59
+ }
60
+
61
+ // ── Chunking ─────────────────────────────────────────────────────────────
62
+
63
+ export interface Chunk {
64
+ content: string
65
+ index: number
66
+ startOffset: number
67
+ endOffset: number
68
+ }
69
+
70
+ export interface Chunker {
71
+ chunk(content: string): Chunk[]
72
+ }
73
+
74
+ // ── Configuration ────────────────────────────────────────────────────────
75
+
76
+ export interface RagConfig {
77
+ default: string
78
+ prefix: string
79
+ embedding: EmbeddingConfig
80
+ chunking: ChunkingConfig
81
+ stores: Record<string, StoreConfig>
82
+ }
83
+
84
+ export interface EmbeddingConfig {
85
+ provider: string
86
+ model: string
87
+ dimension: number
88
+ }
89
+
90
+ export interface ChunkingConfig {
91
+ strategy: string
92
+ chunkSize: number
93
+ overlap: number
94
+ separators?: string[]
95
+ }
96
+
97
+ export interface StoreConfig {
98
+ driver: string
99
+ [key: string]: unknown
100
+ }
@@ -0,0 +1,15 @@
1
+ import type { VectorDocument, QueryOptions, QueryResult } from './types.ts'
2
+
3
+ export interface VectorStore {
4
+ readonly name: string
5
+
6
+ createCollection(collection: string, dimension: number): Promise<void>
7
+ deleteCollection(collection: string): Promise<void>
8
+
9
+ upsert(collection: string, documents: VectorDocument[]): Promise<void>
10
+ delete(collection: string, ids: (string | number)[]): Promise<void>
11
+ deleteBySource(collection: string, sourceId: string | number): Promise<void>
12
+ flush(collection: string): Promise<void>
13
+
14
+ query(collection: string, vector: number[], options?: QueryOptions): Promise<QueryResult>
15
+ }
@@ -0,0 +1,33 @@
1
+ import { env } from '@stravigor/kernel'
2
+
3
+ export default {
4
+ default: env('RAG_DRIVER', 'pgvector'),
5
+
6
+ prefix: env('RAG_PREFIX', ''),
7
+
8
+ embedding: {
9
+ provider: env('RAG_EMBEDDING_PROVIDER', 'openai'),
10
+ model: env('RAG_EMBEDDING_MODEL', 'text-embedding-3-small'),
11
+ dimension: 1536,
12
+ },
13
+
14
+ chunking: {
15
+ strategy: 'recursive',
16
+ chunkSize: 512,
17
+ overlap: 64,
18
+ },
19
+
20
+ stores: {
21
+ pgvector: {
22
+ driver: 'pgvector',
23
+ },
24
+
25
+ memory: {
26
+ driver: 'memory',
27
+ },
28
+
29
+ null: {
30
+ driver: 'null',
31
+ },
32
+ },
33
+ }
package/tsconfig.json ADDED
@@ -0,0 +1,5 @@
1
+ {
2
+ "extends": "../../tsconfig.json",
3
+ "include": ["src/**/*.ts"],
4
+ "exclude": ["node_modules", "tests"]
5
+ }