feat: expand secure assistant workflows
Harden runtime execution and add local knowledge, Smart Heartbeat, usage visibility, responsive product surfaces, and cross-platform packaging support. Co-authored-by: factory-droid[bot] <138933559+factory-droid[bot]@users.noreply.github.com>
This commit is contained in:
co-authored by
factory-droid[bot] <138933559+factory-droid[bot]@users.noreply.github.com>
parent
6ef1795b81
commit
b3fdf96962
@@ -1,4 +1,5 @@
|
||||
import { mkdtemp, rm } from 'node:fs/promises'
|
||||
import { createHash } from 'node:crypto'
|
||||
import { tmpdir } from 'node:os'
|
||||
import { join } from 'node:path'
|
||||
import { DatabaseSync } from 'node:sqlite'
|
||||
@@ -85,12 +86,12 @@ describe('KnowledgeDatabase', () => {
|
||||
const inspection = new DatabaseSync(path)
|
||||
expect(
|
||||
inspection.prepare('PRAGMA user_version').get()
|
||||
).toEqual({ user_version: 1 })
|
||||
).toEqual({ user_version: 2 })
|
||||
expect(
|
||||
inspection
|
||||
.prepare('SELECT version FROM schema_migrations ORDER BY version')
|
||||
.all()
|
||||
).toEqual([{ version: 1 }])
|
||||
).toEqual([{ version: 1 }, { version: 2 }])
|
||||
inspection.close()
|
||||
|
||||
const reopened = new KnowledgeDatabase(path)
|
||||
@@ -107,6 +108,53 @@ describe('KnowledgeDatabase', () => {
|
||||
.toHaveLength(1)
|
||||
})
|
||||
|
||||
it('upgrades an existing v1 database to vector schema v2', async () => {
|
||||
const { database, path } = await createDatabase()
|
||||
const knowledgeBase = database.createKnowledgeBase({
|
||||
name: 'Version one data',
|
||||
storageMode: 'reference'
|
||||
})
|
||||
seedDocument(database, knowledgeBase.id, 'version-one')
|
||||
database.close()
|
||||
|
||||
const downgrade = new DatabaseSync(path)
|
||||
downgrade.exec(`
|
||||
DROP TABLE embedding_index_state;
|
||||
DROP TABLE chunk_embeddings;
|
||||
DELETE FROM schema_migrations WHERE version = 2;
|
||||
PRAGMA user_version = 1;
|
||||
`)
|
||||
downgrade.close()
|
||||
|
||||
const upgraded = new KnowledgeDatabase(path)
|
||||
openDatabases.push(upgraded)
|
||||
upgraded.initialize()
|
||||
const inspection = new DatabaseSync(path)
|
||||
expect(inspection.prepare('PRAGMA user_version').get()).toEqual({
|
||||
user_version: 2
|
||||
})
|
||||
expect(
|
||||
inspection
|
||||
.prepare(
|
||||
`SELECT name FROM sqlite_master
|
||||
WHERE type = 'table' AND name IN
|
||||
('chunk_embeddings', 'embedding_index_state')
|
||||
ORDER BY name`
|
||||
)
|
||||
.all()
|
||||
).toEqual([
|
||||
{ name: 'chunk_embeddings' },
|
||||
{ name: 'embedding_index_state' }
|
||||
])
|
||||
inspection.close()
|
||||
expect(
|
||||
upgraded.search({
|
||||
knowledgeBaseId: knowledgeBase.id,
|
||||
query: 'lighthouse'
|
||||
})
|
||||
).toHaveLength(1)
|
||||
})
|
||||
|
||||
it('isolates FTS results by knowledge base and replaces indexed chunks', async () => {
|
||||
const { database } = await createDatabase()
|
||||
const first = database.createKnowledgeBase({
|
||||
@@ -298,6 +346,266 @@ describe('KnowledgeDatabase', () => {
|
||||
expect(database.deleteEntity(other.id)).toBe(true)
|
||||
})
|
||||
|
||||
it('persists isolated Float32 embeddings and clears stale index state transactionally', async () => {
|
||||
const created = await createDatabase()
|
||||
let database = created.database
|
||||
const knowledgeBase = database.createKnowledgeBase({
|
||||
name: 'Vectors',
|
||||
storageMode: 'reference'
|
||||
})
|
||||
const alpha = seedDocument(database, knowledgeBase.id, 'alpha-vector')
|
||||
const beta = seedDocument(database, knowledgeBase.id, 'beta-vector')
|
||||
const checksum = (value: string): string =>
|
||||
createHash('sha256').update(value).digest('hex')
|
||||
const alphaContent =
|
||||
'alpha-vector contains the searchable lighthouse phrase'
|
||||
const betaContent =
|
||||
'beta-vector contains the searchable lighthouse phrase'
|
||||
|
||||
expect(
|
||||
database.replaceDocumentEmbeddings(
|
||||
alpha.documentId,
|
||||
'ollama',
|
||||
'test-model',
|
||||
[
|
||||
{
|
||||
chunkId: alpha.chunkId,
|
||||
contentChecksum: checksum(alphaContent),
|
||||
vector: [1, 0]
|
||||
}
|
||||
]
|
||||
)
|
||||
).toMatchObject({ status: 'ready', dimensions: 2 })
|
||||
database.replaceDocumentEmbeddings(
|
||||
beta.documentId,
|
||||
'ollama',
|
||||
'test-model',
|
||||
[
|
||||
{
|
||||
chunkId: beta.chunkId,
|
||||
contentChecksum: checksum(betaContent),
|
||||
vector: [0, 1]
|
||||
}
|
||||
]
|
||||
)
|
||||
database.close()
|
||||
database = new KnowledgeDatabase(created.path)
|
||||
openDatabases.push(database)
|
||||
database.initialize()
|
||||
|
||||
expect(
|
||||
database.vectorSearch({
|
||||
knowledgeBaseId: knowledgeBase.id,
|
||||
provider: 'ollama',
|
||||
model: 'test-model',
|
||||
vector: [0.9, 0.1]
|
||||
}).map((result) => result.chunk.id)
|
||||
).toEqual([alpha.chunkId, beta.chunkId])
|
||||
expect(
|
||||
database.vectorSearch({
|
||||
knowledgeBaseId: knowledgeBase.id,
|
||||
provider: 'ollama',
|
||||
model: 'other-model',
|
||||
vector: [0.9, 0.1]
|
||||
})
|
||||
).toEqual([])
|
||||
expect(() =>
|
||||
database.replaceDocumentEmbeddings(
|
||||
alpha.documentId,
|
||||
'ollama',
|
||||
'test-model',
|
||||
[
|
||||
{
|
||||
chunkId: alpha.chunkId,
|
||||
contentChecksum: '0'.repeat(64),
|
||||
vector: [1, 0]
|
||||
}
|
||||
]
|
||||
)
|
||||
).toThrow('checksum')
|
||||
expect(
|
||||
database.vectorSearch({
|
||||
knowledgeBaseId: knowledgeBase.id,
|
||||
provider: 'ollama',
|
||||
model: 'test-model',
|
||||
vector: [1, 0],
|
||||
limit: 1
|
||||
})[0]?.chunk.id
|
||||
).toBe(alpha.chunkId)
|
||||
|
||||
database.upsertDocument(
|
||||
{
|
||||
id: alpha.documentId,
|
||||
knowledgeBaseId: knowledgeBase.id,
|
||||
sourceId: alpha.sourceId,
|
||||
externalId: 'alpha-vector',
|
||||
title: 'alpha-vector'
|
||||
},
|
||||
[{ id: alpha.chunkId, ordinal: 0, content: 'fresh lexical fallback' }]
|
||||
)
|
||||
expect(
|
||||
database.getEmbeddingIndexState(
|
||||
alpha.documentId,
|
||||
'ollama',
|
||||
'test-model'
|
||||
)
|
||||
).toBeUndefined()
|
||||
expect(
|
||||
database.vectorSearch({
|
||||
knowledgeBaseId: knowledgeBase.id,
|
||||
provider: 'ollama',
|
||||
model: 'test-model',
|
||||
vector: [1, 0]
|
||||
}).map((result) => result.chunk.id)
|
||||
).not.toContain(alpha.chunkId)
|
||||
expect(
|
||||
database.search({
|
||||
knowledgeBaseId: knowledgeBase.id,
|
||||
query: 'fallback'
|
||||
})
|
||||
).toHaveLength(1)
|
||||
})
|
||||
|
||||
it('fuses FTS and vector ranks while isolating providers and libraries', async () => {
|
||||
const { database } = await createDatabase()
|
||||
const first = database.createKnowledgeBase({
|
||||
name: 'Hybrid one',
|
||||
storageMode: 'reference',
|
||||
graphEnabled: false
|
||||
})
|
||||
const second = database.createKnowledgeBase({
|
||||
name: 'Hybrid two',
|
||||
storageMode: 'reference',
|
||||
graphEnabled: false
|
||||
})
|
||||
const firstSeed = seedDocument(database, first.id, 'hybrid-first')
|
||||
const secondSeed = seedDocument(database, second.id, 'hybrid-second')
|
||||
for (const [databaseId, seed, marker] of [
|
||||
[first.id, firstSeed, 'hybrid-first'],
|
||||
[second.id, secondSeed, 'hybrid-second']
|
||||
] as const) {
|
||||
const content = `${marker} contains the searchable lighthouse phrase`
|
||||
database.replaceDocumentEmbeddings(
|
||||
seed.documentId,
|
||||
'ollama',
|
||||
'hybrid-model',
|
||||
[
|
||||
{
|
||||
chunkId: seed.chunkId,
|
||||
contentChecksum: createHash('sha256').update(content).digest('hex'),
|
||||
vector: [1, 0, 0]
|
||||
}
|
||||
]
|
||||
)
|
||||
expect(databaseId).toBeTruthy()
|
||||
}
|
||||
|
||||
const results = database.hybridSearch({
|
||||
knowledgeBaseId: first.id,
|
||||
query: 'lighthouse',
|
||||
provider: 'ollama',
|
||||
model: 'hybrid-model',
|
||||
vector: [1, 0, 0],
|
||||
graphEnabled: false
|
||||
})
|
||||
expect(results).toHaveLength(1)
|
||||
expect(results[0]?.chunk.id).toBe(firstSeed.chunkId)
|
||||
expect(results[0]?.retrieval.channels).toEqual(['fts', 'vector'])
|
||||
expect(results[0]?.retrieval.similarity).toBeCloseTo(1)
|
||||
expect(results.map((result) => result.chunk.id)).not.toContain(
|
||||
secondSeed.chunkId
|
||||
)
|
||||
})
|
||||
|
||||
it('expands persisted graph seeds only through evidence-backed same-library paths', async () => {
|
||||
const { database } = await createDatabase()
|
||||
const first = database.createKnowledgeBase({
|
||||
name: 'GraphRAG one',
|
||||
storageMode: 'reference'
|
||||
})
|
||||
const second = database.createKnowledgeBase({
|
||||
name: 'GraphRAG two',
|
||||
storageMode: 'reference'
|
||||
})
|
||||
const goodBuddy = seedDocument(database, first.id, 'GoodBuddy')
|
||||
const electron = seedDocument(database, first.id, 'Electron')
|
||||
const foreign = seedDocument(database, second.id, 'GoodBuddy-foreign')
|
||||
const source = database.createEntity({
|
||||
knowledgeBaseId: first.id,
|
||||
name: 'GoodBuddy',
|
||||
type: 'product'
|
||||
})
|
||||
const target = database.createEntity({
|
||||
knowledgeBaseId: first.id,
|
||||
name: 'Electron',
|
||||
type: 'framework'
|
||||
})
|
||||
const relation = database.createRelation({
|
||||
knowledgeBaseId: first.id,
|
||||
sourceEntityId: source.id,
|
||||
targetEntityId: target.id,
|
||||
type: 'uses'
|
||||
})
|
||||
expect(() =>
|
||||
database.createEvidence({
|
||||
knowledgeBaseId: first.id,
|
||||
entityId: source.id,
|
||||
documentId: goodBuddy.documentId,
|
||||
chunkId: electron.chunkId
|
||||
})
|
||||
).toThrow('must belong')
|
||||
database.createEvidence({
|
||||
knowledgeBaseId: first.id,
|
||||
entityId: source.id,
|
||||
documentId: goodBuddy.documentId,
|
||||
chunkId: goodBuddy.chunkId
|
||||
})
|
||||
database.createEvidence({
|
||||
knowledgeBaseId: first.id,
|
||||
entityId: target.id,
|
||||
documentId: electron.documentId,
|
||||
chunkId: electron.chunkId
|
||||
})
|
||||
database.createEvidence({
|
||||
knowledgeBaseId: first.id,
|
||||
relationId: relation.id,
|
||||
documentId: goodBuddy.documentId,
|
||||
chunkId: goodBuddy.chunkId
|
||||
})
|
||||
const unbacked = database.createEntity({
|
||||
knowledgeBaseId: first.id,
|
||||
name: 'Unbacked',
|
||||
type: 'concept'
|
||||
})
|
||||
const foreignEntity = database.createEntity({
|
||||
knowledgeBaseId: second.id,
|
||||
name: 'GoodBuddy',
|
||||
type: 'foreign'
|
||||
})
|
||||
database.createEvidence({
|
||||
knowledgeBaseId: second.id,
|
||||
entityId: foreignEntity.id,
|
||||
documentId: foreign.documentId,
|
||||
chunkId: foreign.chunkId
|
||||
})
|
||||
|
||||
const graphResults = database.graphSearch(first.id, 'GoodBuddy', 10, 1)
|
||||
expect(graphResults.map((result) => result.chunk.id)).toEqual(
|
||||
expect.arrayContaining([goodBuddy.chunkId, electron.chunkId])
|
||||
)
|
||||
expect(
|
||||
graphResults.every(
|
||||
(result) =>
|
||||
result.retrieval.channels[0] === 'graph' &&
|
||||
result.retrieval.evidenceIds.length > 0
|
||||
)
|
||||
).toBe(true)
|
||||
expect(graphResults.map((result) => result.chunk.id)).not.toContain(
|
||||
foreign.chunkId
|
||||
)
|
||||
expect(database.graphSearch(first.id, unbacked.name)).toEqual([])
|
||||
})
|
||||
|
||||
it('bounds inputs and rejects API keys in extensible metadata', async () => {
|
||||
const { database } = await createDatabase()
|
||||
expect(() =>
|
||||
|
||||
@@ -1,16 +1,20 @@
|
||||
import { randomUUID } from 'node:crypto'
|
||||
import { createHash, randomUUID } from 'node:crypto'
|
||||
import { DatabaseSync, type StatementSync } from 'node:sqlite'
|
||||
import type {
|
||||
Chunk,
|
||||
ChunkEmbeddingInput,
|
||||
CreateEvidenceInput,
|
||||
CreateGraphEntityInput,
|
||||
CreateGraphRelationInput,
|
||||
CreateKnowledgeBaseInput,
|
||||
Document,
|
||||
Evidence,
|
||||
EmbeddingIndexState,
|
||||
GraphEntity,
|
||||
GraphRelation,
|
||||
GraphStrategy,
|
||||
HybridSearchOptions,
|
||||
HybridSearchResult,
|
||||
JsonObject,
|
||||
KnowledgeBase,
|
||||
KnowledgeSource,
|
||||
@@ -25,10 +29,11 @@ import type {
|
||||
UpdateGraphRelationInput,
|
||||
UpdateKnowledgeBaseInput,
|
||||
UpsertDocumentInput,
|
||||
UpsertKnowledgeSourceInput
|
||||
UpsertKnowledgeSourceInput,
|
||||
VectorSearchOptions
|
||||
} from './types'
|
||||
|
||||
const DATABASE_VERSION = 1
|
||||
const DATABASE_VERSION = 2
|
||||
const MAX_ID_LENGTH = 128
|
||||
const MAX_NAME_LENGTH = 512
|
||||
const MAX_LOCATION_LENGTH = 8192
|
||||
@@ -42,6 +47,19 @@ const MAX_JSON_ARRAY_ITEMS = 1_000
|
||||
const MAX_JSON_DEPTH = 20
|
||||
const MAX_JSON_NODES = 10_000
|
||||
const MAX_JSON_STRING_LENGTH = 32_768
|
||||
const MAX_EMBEDDING_DIMENSIONS = 8_192
|
||||
const MAX_EMBEDDING_PROVIDER_LENGTH = 128
|
||||
const MAX_EMBEDDING_MODEL_LENGTH = 512
|
||||
const MAX_EMBEDDING_ERROR_LENGTH = 2_000
|
||||
const MAX_GRAPH_DEPTH = 3
|
||||
const MAX_VECTOR_CANDIDATES = 5_000
|
||||
const RRF_CONSTANT = 60
|
||||
|
||||
type ScoredSearchResult = {
|
||||
result: SearchResult
|
||||
similarity?: number
|
||||
evidenceIds?: string[]
|
||||
}
|
||||
|
||||
type Row = Record<string, null | number | bigint | string | Uint8Array>
|
||||
|
||||
@@ -221,6 +239,102 @@ function asNumber(row: Row, key: string): number {
|
||||
return row[key] as number
|
||||
}
|
||||
|
||||
function asBytes(row: Row, key: string): Uint8Array {
|
||||
return row[key] as Uint8Array
|
||||
}
|
||||
|
||||
function contentChecksum(content: string): string {
|
||||
return createHash('sha256').update(content).digest('hex')
|
||||
}
|
||||
|
||||
function normalizedChecksum(value: string, field: string): string {
|
||||
const checksum = requiredString(value, field, 64).toLowerCase()
|
||||
if (!/^[a-f0-9]{64}$/u.test(checksum)) {
|
||||
throw new RangeError(`${field} must be a SHA-256 checksum`)
|
||||
}
|
||||
return checksum
|
||||
}
|
||||
|
||||
function normalizeVector(
|
||||
value: readonly number[],
|
||||
field: string
|
||||
): {
|
||||
bytes: Buffer
|
||||
dimensions: number
|
||||
magnitude: number
|
||||
values: number[]
|
||||
} {
|
||||
if (
|
||||
!Array.isArray(value) ||
|
||||
value.length < 1 ||
|
||||
value.length > MAX_EMBEDDING_DIMENSIONS
|
||||
) {
|
||||
throw new RangeError(
|
||||
`${field} must contain between 1 and ${MAX_EMBEDDING_DIMENSIONS} dimensions`
|
||||
)
|
||||
}
|
||||
const bytes = Buffer.allocUnsafe(value.length * Float32Array.BYTES_PER_ELEMENT)
|
||||
const values: number[] = []
|
||||
let magnitudeSquared = 0
|
||||
for (let index = 0; index < value.length; index += 1) {
|
||||
const component = value[index]
|
||||
if (typeof component !== 'number' || !Number.isFinite(component)) {
|
||||
throw new TypeError(`${field} must contain only finite numbers`)
|
||||
}
|
||||
const storedComponent = Math.fround(component)
|
||||
if (!Number.isFinite(storedComponent)) {
|
||||
throw new RangeError(`${field} components must fit in Float32`)
|
||||
}
|
||||
bytes.writeFloatLE(
|
||||
storedComponent,
|
||||
index * Float32Array.BYTES_PER_ELEMENT
|
||||
)
|
||||
values.push(storedComponent)
|
||||
magnitudeSquared += storedComponent * storedComponent
|
||||
}
|
||||
const magnitude = Math.sqrt(magnitudeSquared)
|
||||
if (!Number.isFinite(magnitude) || magnitude <= 0) {
|
||||
throw new RangeError(`${field} must have a finite non-zero norm`)
|
||||
}
|
||||
return { bytes, dimensions: value.length, magnitude, values }
|
||||
}
|
||||
|
||||
function cosineSimilarity(
|
||||
left: readonly number[],
|
||||
leftMagnitude: number,
|
||||
rightBytes: Uint8Array,
|
||||
dimensions: number,
|
||||
rightMagnitude: number
|
||||
): number | undefined {
|
||||
if (
|
||||
left.length !== dimensions ||
|
||||
rightBytes.byteLength !== dimensions * Float32Array.BYTES_PER_ELEMENT ||
|
||||
!Number.isFinite(rightMagnitude) ||
|
||||
rightMagnitude <= 0
|
||||
) {
|
||||
return undefined
|
||||
}
|
||||
const buffer = Buffer.from(
|
||||
rightBytes.buffer,
|
||||
rightBytes.byteOffset,
|
||||
rightBytes.byteLength
|
||||
)
|
||||
let dot = 0
|
||||
for (let index = 0; index < dimensions; index += 1) {
|
||||
const component = buffer.readFloatLE(
|
||||
index * Float32Array.BYTES_PER_ELEMENT
|
||||
)
|
||||
if (!Number.isFinite(component)) {
|
||||
return undefined
|
||||
}
|
||||
dot += (left[index] ?? 0) * component
|
||||
}
|
||||
const similarity = dot / (leftMagnitude * rightMagnitude)
|
||||
return Number.isFinite(similarity)
|
||||
? Math.max(-1, Math.min(1, similarity))
|
||||
: undefined
|
||||
}
|
||||
|
||||
function mapKnowledgeBase(row: Row): KnowledgeBase {
|
||||
return {
|
||||
id: asString(row, 'id'),
|
||||
@@ -325,6 +439,21 @@ function mapEvidence(row: Row): Evidence {
|
||||
}
|
||||
}
|
||||
|
||||
function mapEmbeddingIndexState(row: Row): EmbeddingIndexState {
|
||||
return {
|
||||
documentId: asString(row, 'document_id'),
|
||||
knowledgeBaseId: asString(row, 'knowledge_base_id'),
|
||||
provider: asString(row, 'provider'),
|
||||
model: asString(row, 'model'),
|
||||
dimensions:
|
||||
row.dimensions === null ? undefined : asNumber(row, 'dimensions'),
|
||||
contentChecksum: asString(row, 'content_checksum'),
|
||||
status: asString(row, 'status') as EmbeddingIndexState['status'],
|
||||
lastError: asOptionalString(row, 'last_error'),
|
||||
updatedAt: asString(row, 'updated_at')
|
||||
}
|
||||
}
|
||||
|
||||
export class KnowledgeDatabase {
|
||||
private database?: DatabaseSync
|
||||
|
||||
@@ -688,6 +817,9 @@ export class KnowledgeDatabase {
|
||||
now,
|
||||
now
|
||||
)
|
||||
database
|
||||
.prepare('DELETE FROM embedding_index_state WHERE document_id = ?')
|
||||
.run(id)
|
||||
database.prepare('DELETE FROM chunks WHERE document_id = ?').run(id)
|
||||
const insertChunk = database.prepare(
|
||||
`INSERT INTO chunks
|
||||
@@ -768,7 +900,7 @@ export class KnowledgeDatabase {
|
||||
'documentId',
|
||||
MAX_ID_LENGTH
|
||||
)
|
||||
boundedInteger(limit, 'limit', 1, MAX_LIST_LIMIT)
|
||||
boundedInteger(limit, 'limit', 1, MAX_CHUNKS)
|
||||
return this.requireDatabase()
|
||||
.prepare(
|
||||
`SELECT * FROM chunks WHERE document_id = ?
|
||||
@@ -778,6 +910,365 @@ export class KnowledgeDatabase {
|
||||
.map(mapChunk)
|
||||
}
|
||||
|
||||
replaceDocumentEmbeddings(
|
||||
documentId: string,
|
||||
provider: string,
|
||||
model: string,
|
||||
embeddings: readonly ChunkEmbeddingInput[]
|
||||
): EmbeddingIndexState {
|
||||
const normalizedDocumentId = requiredString(
|
||||
documentId,
|
||||
'documentId',
|
||||
MAX_ID_LENGTH
|
||||
)
|
||||
const normalizedProvider = requiredString(
|
||||
provider,
|
||||
'provider',
|
||||
MAX_EMBEDDING_PROVIDER_LENGTH
|
||||
)
|
||||
const normalizedModel = requiredString(
|
||||
model,
|
||||
'model',
|
||||
MAX_EMBEDDING_MODEL_LENGTH
|
||||
)
|
||||
if (!Array.isArray(embeddings) || embeddings.length > MAX_CHUNKS) {
|
||||
throw new RangeError(`embeddings must contain at most ${MAX_CHUNKS} items`)
|
||||
}
|
||||
const database = this.requireDatabase()
|
||||
const document = database
|
||||
.prepare('SELECT id, knowledge_base_id FROM documents WHERE id = ?')
|
||||
.get(normalizedDocumentId)
|
||||
if (!document) {
|
||||
throw new Error(`Document not found: ${normalizedDocumentId}`)
|
||||
}
|
||||
const chunks = database
|
||||
.prepare(
|
||||
`SELECT id, content FROM chunks
|
||||
WHERE document_id = ? ORDER BY ordinal ASC, id ASC`
|
||||
)
|
||||
.all(normalizedDocumentId)
|
||||
if (embeddings.length !== chunks.length) {
|
||||
throw new Error('Embeddings must cover every current document chunk')
|
||||
}
|
||||
const chunksById = new Map(
|
||||
chunks.map((row) => [asString(row, 'id'), asString(row, 'content')])
|
||||
)
|
||||
const seen = new Set<string>()
|
||||
let dimensions: number | undefined
|
||||
const normalized = embeddings.map((embedding, index) => {
|
||||
const chunkId = requiredString(
|
||||
embedding.chunkId,
|
||||
`embeddings[${index}].chunkId`,
|
||||
MAX_ID_LENGTH
|
||||
)
|
||||
const content = chunksById.get(chunkId)
|
||||
if (content === undefined || seen.has(chunkId)) {
|
||||
throw new Error('Embeddings must reference unique chunks in the document')
|
||||
}
|
||||
seen.add(chunkId)
|
||||
const checksum = normalizedChecksum(
|
||||
embedding.contentChecksum,
|
||||
`embeddings[${index}].contentChecksum`
|
||||
)
|
||||
if (checksum !== contentChecksum(content)) {
|
||||
throw new Error('Embedding content checksum does not match the chunk')
|
||||
}
|
||||
const vector = normalizeVector(
|
||||
embedding.vector,
|
||||
`embeddings[${index}].vector`
|
||||
)
|
||||
if (dimensions === undefined) {
|
||||
dimensions = vector.dimensions
|
||||
} else if (dimensions !== vector.dimensions) {
|
||||
throw new Error('Document embeddings must have consistent dimensions')
|
||||
}
|
||||
return { chunkId, checksum, ...vector }
|
||||
})
|
||||
const indexChecksum = createHash('sha256')
|
||||
.update(
|
||||
normalized
|
||||
.map((item) => `${item.chunkId}\0${item.checksum}`)
|
||||
.sort()
|
||||
.join('\n')
|
||||
)
|
||||
.digest('hex')
|
||||
const now = new Date().toISOString()
|
||||
this.transaction(database, () => {
|
||||
database
|
||||
.prepare(
|
||||
`DELETE FROM chunk_embeddings
|
||||
WHERE provider = ? AND model = ? AND chunk_id IN
|
||||
(SELECT id FROM chunks WHERE document_id = ?)`
|
||||
)
|
||||
.run(normalizedProvider, normalizedModel, normalizedDocumentId)
|
||||
const insert = database.prepare(
|
||||
`INSERT INTO chunk_embeddings
|
||||
(chunk_id, knowledge_base_id, provider, model, dimensions,
|
||||
content_checksum, vector, magnitude, created_at, updated_at)
|
||||
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)`
|
||||
)
|
||||
for (const item of normalized) {
|
||||
insert.run(
|
||||
item.chunkId,
|
||||
asString(document, 'knowledge_base_id'),
|
||||
normalizedProvider,
|
||||
normalizedModel,
|
||||
item.dimensions,
|
||||
item.checksum,
|
||||
item.bytes,
|
||||
item.magnitude,
|
||||
now,
|
||||
now
|
||||
)
|
||||
}
|
||||
database
|
||||
.prepare(
|
||||
`INSERT INTO embedding_index_state
|
||||
(document_id, knowledge_base_id, provider, model, dimensions,
|
||||
content_checksum, status, last_error, updated_at)
|
||||
VALUES (?, ?, ?, ?, ?, ?, 'ready', NULL, ?)
|
||||
ON CONFLICT(document_id, provider, model) DO UPDATE SET
|
||||
knowledge_base_id = excluded.knowledge_base_id,
|
||||
dimensions = excluded.dimensions,
|
||||
content_checksum = excluded.content_checksum,
|
||||
status = 'ready',
|
||||
last_error = NULL,
|
||||
updated_at = excluded.updated_at`
|
||||
)
|
||||
.run(
|
||||
normalizedDocumentId,
|
||||
asString(document, 'knowledge_base_id'),
|
||||
normalizedProvider,
|
||||
normalizedModel,
|
||||
dimensions ?? null,
|
||||
indexChecksum,
|
||||
now
|
||||
)
|
||||
})
|
||||
return this.requiredEmbeddingIndexState(
|
||||
normalizedDocumentId,
|
||||
normalizedProvider,
|
||||
normalizedModel
|
||||
)
|
||||
}
|
||||
|
||||
recordEmbeddingIndexError(
|
||||
documentId: string,
|
||||
provider: string,
|
||||
model: string,
|
||||
error: string
|
||||
): EmbeddingIndexState {
|
||||
const normalizedDocumentId = requiredString(
|
||||
documentId,
|
||||
'documentId',
|
||||
MAX_ID_LENGTH
|
||||
)
|
||||
const normalizedProvider = requiredString(
|
||||
provider,
|
||||
'provider',
|
||||
MAX_EMBEDDING_PROVIDER_LENGTH
|
||||
)
|
||||
const normalizedModel = requiredString(
|
||||
model,
|
||||
'model',
|
||||
MAX_EMBEDDING_MODEL_LENGTH
|
||||
)
|
||||
const normalizedError = requiredString(
|
||||
error,
|
||||
'error',
|
||||
MAX_EMBEDDING_ERROR_LENGTH,
|
||||
false
|
||||
)
|
||||
const database = this.requireDatabase()
|
||||
const document = database
|
||||
.prepare('SELECT knowledge_base_id FROM documents WHERE id = ?')
|
||||
.get(normalizedDocumentId)
|
||||
if (!document) {
|
||||
throw new Error(`Document not found: ${normalizedDocumentId}`)
|
||||
}
|
||||
database
|
||||
.prepare(
|
||||
`INSERT INTO embedding_index_state
|
||||
(document_id, knowledge_base_id, provider, model, dimensions,
|
||||
content_checksum, status, last_error, updated_at)
|
||||
VALUES (?, ?, ?, ?, NULL, '', 'error', ?, ?)
|
||||
ON CONFLICT(document_id, provider, model) DO UPDATE SET
|
||||
status = 'error',
|
||||
last_error = excluded.last_error,
|
||||
updated_at = excluded.updated_at`
|
||||
)
|
||||
.run(
|
||||
normalizedDocumentId,
|
||||
asString(document, 'knowledge_base_id'),
|
||||
normalizedProvider,
|
||||
normalizedModel,
|
||||
normalizedError,
|
||||
new Date().toISOString()
|
||||
)
|
||||
return this.requiredEmbeddingIndexState(
|
||||
normalizedDocumentId,
|
||||
normalizedProvider,
|
||||
normalizedModel
|
||||
)
|
||||
}
|
||||
|
||||
getEmbeddingIndexState(
|
||||
documentId: string,
|
||||
provider: string,
|
||||
model: string
|
||||
): EmbeddingIndexState | undefined {
|
||||
const row = this.requireDatabase()
|
||||
.prepare(
|
||||
`SELECT * FROM embedding_index_state
|
||||
WHERE document_id = ? AND provider = ? AND model = ?`
|
||||
)
|
||||
.get(
|
||||
requiredString(documentId, 'documentId', MAX_ID_LENGTH),
|
||||
requiredString(
|
||||
provider,
|
||||
'provider',
|
||||
MAX_EMBEDDING_PROVIDER_LENGTH
|
||||
),
|
||||
requiredString(model, 'model', MAX_EMBEDDING_MODEL_LENGTH)
|
||||
)
|
||||
return row ? mapEmbeddingIndexState(row) : undefined
|
||||
}
|
||||
|
||||
vectorSearch(options: VectorSearchOptions): SearchResult[] {
|
||||
return this.vectorSearchScored(options).map((item) => item.result)
|
||||
}
|
||||
|
||||
graphSearch(
|
||||
knowledgeBaseId: string,
|
||||
query: string,
|
||||
limit = 20,
|
||||
maximumDepth = 1
|
||||
): HybridSearchResult[] {
|
||||
boundedInteger(limit, 'limit', 1, 100)
|
||||
boundedInteger(maximumDepth, 'maximumDepth', 0, MAX_GRAPH_DEPTH)
|
||||
return this.graphSearchScored(
|
||||
requiredString(knowledgeBaseId, 'knowledgeBaseId', MAX_ID_LENGTH),
|
||||
requiredString(query, 'query', 512),
|
||||
limit,
|
||||
maximumDepth
|
||||
).map((item, index) => ({
|
||||
...item.result,
|
||||
retrieval: {
|
||||
score: 1 / (RRF_CONSTANT + index + 1),
|
||||
channels: ['graph'],
|
||||
graphRank: index + 1,
|
||||
evidenceIds: item.evidenceIds ?? []
|
||||
}
|
||||
}))
|
||||
}
|
||||
|
||||
hybridSearch(options: HybridSearchOptions): HybridSearchResult[] {
|
||||
const knowledgeBaseId = requiredString(
|
||||
options.knowledgeBaseId,
|
||||
'knowledgeBaseId',
|
||||
MAX_ID_LENGTH
|
||||
)
|
||||
const query = requiredString(options.query, 'query', 512)
|
||||
const limit = options.limit ?? 20
|
||||
boundedInteger(limit, 'limit', 1, 100)
|
||||
const lexical = this.search({
|
||||
knowledgeBaseId,
|
||||
query,
|
||||
limit: Math.min(100, Math.max(limit * 4, limit))
|
||||
})
|
||||
const vector =
|
||||
options.vector && options.provider && options.model
|
||||
? this.vectorSearchScored({
|
||||
knowledgeBaseId,
|
||||
provider: options.provider,
|
||||
model: options.model,
|
||||
vector: options.vector,
|
||||
limit: options.vectorLimit ?? Math.min(100, limit * 4)
|
||||
})
|
||||
: []
|
||||
const graph = options.graphEnabled === false
|
||||
? []
|
||||
: this.graphSearchScored(
|
||||
knowledgeBaseId,
|
||||
query,
|
||||
Math.min(100, Math.max(limit * 4, limit)),
|
||||
boundedInteger(
|
||||
options.graphDepth ?? 1,
|
||||
'graphDepth',
|
||||
0,
|
||||
MAX_GRAPH_DEPTH
|
||||
)
|
||||
)
|
||||
const fused = new Map<
|
||||
string,
|
||||
{
|
||||
result: SearchResult
|
||||
score: number
|
||||
channels: Set<'fts' | 'vector' | 'graph'>
|
||||
lexicalRank?: number
|
||||
vectorRank?: number
|
||||
graphRank?: number
|
||||
similarity?: number
|
||||
evidenceIds: Set<string>
|
||||
}
|
||||
>()
|
||||
const add = (
|
||||
channel: 'fts' | 'vector' | 'graph',
|
||||
candidates: readonly ScoredSearchResult[],
|
||||
weight: number
|
||||
): void => {
|
||||
candidates.forEach((candidate, index) => {
|
||||
const current = fused.get(candidate.result.chunk.id) ?? {
|
||||
result: candidate.result,
|
||||
score: 0,
|
||||
channels: new Set<'fts' | 'vector' | 'graph'>(),
|
||||
evidenceIds: new Set<string>()
|
||||
}
|
||||
current.score += weight / (RRF_CONSTANT + index + 1)
|
||||
current.channels.add(channel)
|
||||
if (channel === 'fts') {
|
||||
current.lexicalRank = index + 1
|
||||
} else if (channel === 'vector') {
|
||||
current.vectorRank = index + 1
|
||||
current.similarity = candidate.similarity
|
||||
} else {
|
||||
current.graphRank = index + 1
|
||||
for (const evidenceId of candidate.evidenceIds ?? []) {
|
||||
current.evidenceIds.add(evidenceId)
|
||||
}
|
||||
}
|
||||
fused.set(candidate.result.chunk.id, current)
|
||||
})
|
||||
}
|
||||
add(
|
||||
'fts',
|
||||
lexical.map((result) => ({ result })),
|
||||
1
|
||||
)
|
||||
add('vector', vector, 1)
|
||||
add('graph', graph, 0.8)
|
||||
return [...fused.values()]
|
||||
.sort(
|
||||
(left, right) =>
|
||||
right.score - left.score ||
|
||||
left.result.chunk.id.localeCompare(right.result.chunk.id)
|
||||
)
|
||||
.slice(0, limit)
|
||||
.map((item) => ({
|
||||
...item.result,
|
||||
rank: -item.score,
|
||||
retrieval: {
|
||||
score: item.score,
|
||||
channels: [...item.channels],
|
||||
lexicalRank: item.lexicalRank,
|
||||
vectorRank: item.vectorRank,
|
||||
graphRank: item.graphRank,
|
||||
similarity: item.similarity,
|
||||
evidenceIds: [...item.evidenceIds]
|
||||
}
|
||||
}))
|
||||
}
|
||||
|
||||
search(options: SearchOptions): SearchResult[] {
|
||||
const knowledgeBaseId = requiredString(
|
||||
options.knowledgeBaseId,
|
||||
@@ -1318,6 +1809,301 @@ export class KnowledgeDatabase {
|
||||
return this.requiredEntity(target.id)
|
||||
}
|
||||
|
||||
private vectorSearchScored(
|
||||
options: VectorSearchOptions
|
||||
): ScoredSearchResult[] {
|
||||
const knowledgeBaseId = requiredString(
|
||||
options.knowledgeBaseId,
|
||||
'knowledgeBaseId',
|
||||
MAX_ID_LENGTH
|
||||
)
|
||||
const provider = requiredString(
|
||||
options.provider,
|
||||
'provider',
|
||||
MAX_EMBEDDING_PROVIDER_LENGTH
|
||||
)
|
||||
const model = requiredString(
|
||||
options.model,
|
||||
'model',
|
||||
MAX_EMBEDDING_MODEL_LENGTH
|
||||
)
|
||||
const queryVector = normalizeVector(options.vector, 'vector')
|
||||
const limit = options.limit ?? 20
|
||||
boundedInteger(limit, 'limit', 1, 100)
|
||||
const minimumSimilarity = options.minimumSimilarity ?? -1
|
||||
if (
|
||||
typeof minimumSimilarity !== 'number' ||
|
||||
!Number.isFinite(minimumSimilarity) ||
|
||||
minimumSimilarity < -1 ||
|
||||
minimumSimilarity > 1
|
||||
) {
|
||||
throw new RangeError('minimumSimilarity must be between -1 and 1')
|
||||
}
|
||||
const rows = this.requireDatabase()
|
||||
.prepare(
|
||||
`SELECT
|
||||
ce.chunk_id, ce.vector AS embedding_vector,
|
||||
ce.dimensions AS embedding_dimensions,
|
||||
ce.magnitude AS embedding_magnitude
|
||||
FROM chunk_embeddings ce
|
||||
JOIN embedding_index_state eis
|
||||
ON eis.knowledge_base_id = ce.knowledge_base_id
|
||||
AND eis.provider = ce.provider
|
||||
AND eis.model = ce.model
|
||||
AND eis.dimensions = ce.dimensions
|
||||
AND eis.status = 'ready'
|
||||
JOIN chunks c
|
||||
ON c.id = ce.chunk_id
|
||||
AND c.document_id = eis.document_id
|
||||
AND c.knowledge_base_id = ce.knowledge_base_id
|
||||
WHERE ce.knowledge_base_id = ?
|
||||
AND ce.provider = ? AND ce.model = ? AND ce.dimensions = ?
|
||||
AND length(ce.vector) = ce.dimensions * 4
|
||||
AND ce.content_checksum <> ''
|
||||
ORDER BY ce.chunk_id ASC LIMIT ?`
|
||||
)
|
||||
.all(
|
||||
knowledgeBaseId,
|
||||
provider,
|
||||
model,
|
||||
queryVector.dimensions,
|
||||
MAX_VECTOR_CANDIDATES + 1
|
||||
)
|
||||
if (rows.length > MAX_VECTOR_CANDIDATES) {
|
||||
return []
|
||||
}
|
||||
const winners = rows
|
||||
.map((row): { chunkId: string; similarity: number } | undefined => {
|
||||
const similarity = cosineSimilarity(
|
||||
queryVector.values,
|
||||
queryVector.magnitude,
|
||||
asBytes(row, 'embedding_vector'),
|
||||
asNumber(row, 'embedding_dimensions'),
|
||||
asNumber(row, 'embedding_magnitude')
|
||||
)
|
||||
if (similarity === undefined || similarity < minimumSimilarity) {
|
||||
return undefined
|
||||
}
|
||||
return { chunkId: asString(row, 'chunk_id'), similarity }
|
||||
})
|
||||
.filter(
|
||||
(item): item is { chunkId: string; similarity: number } =>
|
||||
item !== undefined
|
||||
)
|
||||
.sort(
|
||||
(left, right) =>
|
||||
right.similarity - left.similarity ||
|
||||
left.chunkId.localeCompare(right.chunkId)
|
||||
)
|
||||
.slice(0, limit)
|
||||
if (winners.length === 0) {
|
||||
return []
|
||||
}
|
||||
const placeholders = winners.map(() => '?').join(', ')
|
||||
const hydratedRows = this.requireDatabase()
|
||||
.prepare(
|
||||
`SELECT
|
||||
c.*, substr(c.content, 1, 600) AS snippet,
|
||||
d.id AS d_id, d.knowledge_base_id AS d_knowledge_base_id,
|
||||
d.source_id AS d_source_id, d.external_id AS d_external_id,
|
||||
d.title AS d_title, d.mime_type AS d_mime_type,
|
||||
d.source_location AS d_source_location, d.checksum AS d_checksum,
|
||||
d.metadata AS d_metadata, d.created_at AS d_created_at,
|
||||
d.updated_at AS d_updated_at,
|
||||
s.id AS s_id, s.knowledge_base_id AS s_knowledge_base_id,
|
||||
s.type AS s_type, s.location AS s_location,
|
||||
s.display_name AS s_display_name, s.status AS s_status,
|
||||
s.last_error AS s_last_error, s.metadata AS s_metadata,
|
||||
s.created_at AS s_created_at, s.updated_at AS s_updated_at
|
||||
FROM chunks c
|
||||
JOIN documents d ON d.id = c.document_id
|
||||
JOIN knowledge_sources s ON s.id = d.source_id
|
||||
WHERE c.id IN (${placeholders})
|
||||
AND c.knowledge_base_id = ?
|
||||
AND d.knowledge_base_id = ?
|
||||
AND s.knowledge_base_id = ?`
|
||||
)
|
||||
.all(
|
||||
...winners.map((winner) => winner.chunkId),
|
||||
knowledgeBaseId,
|
||||
knowledgeBaseId,
|
||||
knowledgeBaseId
|
||||
) as Row[]
|
||||
const rowsByChunkId = new Map(
|
||||
hydratedRows.map((row) => [asString(row, 'id'), row])
|
||||
)
|
||||
return winners.flatMap((winner) => {
|
||||
const row = rowsByChunkId.get(winner.chunkId)
|
||||
return row
|
||||
? [
|
||||
{
|
||||
similarity: winner.similarity,
|
||||
result: {
|
||||
chunk: mapChunk(row),
|
||||
document: mapDocument(this.prefixedRow(row, 'd_')),
|
||||
source: mapSource(this.prefixedRow(row, 's_')),
|
||||
snippet: asString(row, 'snippet'),
|
||||
rank: -winner.similarity
|
||||
}
|
||||
}
|
||||
]
|
||||
: []
|
||||
})
|
||||
}
|
||||
|
||||
private graphSearchScored(
|
||||
knowledgeBaseId: string,
|
||||
query: string,
|
||||
limit: number,
|
||||
maximumDepth: number
|
||||
): ScoredSearchResult[] {
|
||||
const terms = [
|
||||
...new Set(
|
||||
[query, ...query.split(/[^\p{L}\p{N}_.$/@-]+/u)]
|
||||
.map((term) => term.normalize('NFKC').trim().toLowerCase())
|
||||
.filter((term) => term.length > 1)
|
||||
)
|
||||
]
|
||||
.sort((left, right) => right.length - left.length)
|
||||
.slice(0, 8)
|
||||
if (terms.length === 0) {
|
||||
return []
|
||||
}
|
||||
const conditions = terms
|
||||
.map(
|
||||
() =>
|
||||
`(lower(ge.name) LIKE ? ESCAPE '\\' OR lower(ge.aliases) LIKE ? ESCAPE '\\' OR lower(ge.type) LIKE ? ESCAPE '\\')`
|
||||
)
|
||||
.join(' OR ')
|
||||
const patterns = terms.flatMap((term) => {
|
||||
const escaped = term.replaceAll('\\', '\\\\').replaceAll('%', '\\%')
|
||||
.replaceAll('_', '\\_')
|
||||
return [`%${escaped}%`, `%${escaped}%`, `%${escaped}%`]
|
||||
})
|
||||
const rows = this.requireDatabase()
|
||||
.prepare(
|
||||
`WITH RECURSIVE
|
||||
seed(id, depth) AS (
|
||||
SELECT ge.id, 0
|
||||
FROM graph_entities ge
|
||||
WHERE ge.knowledge_base_id = ? AND (${conditions})
|
||||
AND EXISTS (
|
||||
SELECT 1 FROM graph_evidence ev
|
||||
JOIN chunks ec ON ec.id = ev.chunk_id
|
||||
WHERE ev.entity_id = ge.id
|
||||
AND ev.knowledge_base_id = ge.knowledge_base_id
|
||||
AND ec.knowledge_base_id = ge.knowledge_base_id
|
||||
AND ec.document_id = ev.document_id
|
||||
)
|
||||
ORDER BY ge.name COLLATE NOCASE ASC, ge.id ASC
|
||||
LIMIT 24
|
||||
),
|
||||
reachable(id, depth) AS (
|
||||
SELECT id, depth FROM seed
|
||||
UNION
|
||||
SELECT
|
||||
CASE
|
||||
WHEN gr.source_entity_id = reachable.id
|
||||
THEN gr.target_entity_id
|
||||
ELSE gr.source_entity_id
|
||||
END,
|
||||
reachable.depth + 1
|
||||
FROM reachable
|
||||
JOIN graph_relations gr
|
||||
ON gr.knowledge_base_id = ?
|
||||
AND (gr.source_entity_id = reachable.id
|
||||
OR gr.target_entity_id = reachable.id)
|
||||
WHERE reachable.depth < ?
|
||||
AND EXISTS (
|
||||
SELECT 1 FROM graph_evidence rev
|
||||
JOIN chunks rc ON rc.id = rev.chunk_id
|
||||
WHERE rev.relation_id = gr.id
|
||||
AND rev.knowledge_base_id = gr.knowledge_base_id
|
||||
AND rc.knowledge_base_id = gr.knowledge_base_id
|
||||
AND rc.document_id = rev.document_id
|
||||
)
|
||||
AND EXISTS (
|
||||
SELECT 1 FROM graph_evidence nev
|
||||
JOIN chunks nc ON nc.id = nev.chunk_id
|
||||
WHERE nev.entity_id = CASE
|
||||
WHEN gr.source_entity_id = reachable.id
|
||||
THEN gr.target_entity_id
|
||||
ELSE gr.source_entity_id
|
||||
END
|
||||
AND nev.knowledge_base_id = gr.knowledge_base_id
|
||||
AND nc.knowledge_base_id = gr.knowledge_base_id
|
||||
AND nc.document_id = nev.document_id
|
||||
)
|
||||
),
|
||||
reached(id, depth) AS (
|
||||
SELECT id, MIN(depth) FROM reachable GROUP BY id
|
||||
),
|
||||
backed_evidence AS (
|
||||
SELECT ev.*, reached.depth AS graph_depth
|
||||
FROM graph_evidence ev
|
||||
JOIN reached ON reached.id = ev.entity_id
|
||||
WHERE ev.knowledge_base_id = ? AND ev.chunk_id IS NOT NULL
|
||||
UNION ALL
|
||||
SELECT ev.*, MAX(source.depth, target.depth) AS graph_depth
|
||||
FROM graph_evidence ev
|
||||
JOIN graph_relations gr ON gr.id = ev.relation_id
|
||||
JOIN reached source ON source.id = gr.source_entity_id
|
||||
JOIN reached target ON target.id = gr.target_entity_id
|
||||
WHERE ev.knowledge_base_id = ? AND gr.knowledge_base_id = ?
|
||||
AND ev.chunk_id IS NOT NULL
|
||||
)
|
||||
SELECT
|
||||
c.*, substr(c.content, 1, 600) AS snippet,
|
||||
200 + MIN(backed_evidence.graph_depth) AS rank,
|
||||
group_concat(DISTINCT backed_evidence.id) AS evidence_ids,
|
||||
d.id AS d_id, d.knowledge_base_id AS d_knowledge_base_id,
|
||||
d.source_id AS d_source_id, d.external_id AS d_external_id,
|
||||
d.title AS d_title, d.mime_type AS d_mime_type,
|
||||
d.source_location AS d_source_location, d.checksum AS d_checksum,
|
||||
d.metadata AS d_metadata, d.created_at AS d_created_at,
|
||||
d.updated_at AS d_updated_at,
|
||||
s.id AS s_id, s.knowledge_base_id AS s_knowledge_base_id,
|
||||
s.type AS s_type, s.location AS s_location,
|
||||
s.display_name AS s_display_name, s.status AS s_status,
|
||||
s.last_error AS s_last_error, s.metadata AS s_metadata,
|
||||
s.created_at AS s_created_at, s.updated_at AS s_updated_at
|
||||
FROM backed_evidence
|
||||
JOIN chunks c
|
||||
ON c.id = backed_evidence.chunk_id
|
||||
AND c.document_id = backed_evidence.document_id
|
||||
JOIN documents d ON d.id = c.document_id
|
||||
JOIN knowledge_sources s ON s.id = d.source_id
|
||||
WHERE c.knowledge_base_id = ? AND d.knowledge_base_id = ?
|
||||
AND s.knowledge_base_id = ?
|
||||
GROUP BY c.id
|
||||
ORDER BY MIN(backed_evidence.graph_depth) ASC, c.id ASC
|
||||
LIMIT ?`
|
||||
)
|
||||
.all(
|
||||
knowledgeBaseId,
|
||||
...patterns,
|
||||
knowledgeBaseId,
|
||||
maximumDepth,
|
||||
knowledgeBaseId,
|
||||
knowledgeBaseId,
|
||||
knowledgeBaseId,
|
||||
knowledgeBaseId,
|
||||
knowledgeBaseId,
|
||||
knowledgeBaseId,
|
||||
limit
|
||||
)
|
||||
return rows.map((row) => ({
|
||||
evidenceIds: asString(row, 'evidence_ids').split(','),
|
||||
result: {
|
||||
chunk: mapChunk(row),
|
||||
document: mapDocument(this.prefixedRow(row, 'd_')),
|
||||
source: mapSource(this.prefixedRow(row, 's_')),
|
||||
snippet: asString(row, 'snippet'),
|
||||
rank: asNumber(row, 'rank')
|
||||
}
|
||||
}))
|
||||
}
|
||||
|
||||
private assertFts5(database: DatabaseSync): void {
|
||||
try {
|
||||
database.exec(`
|
||||
@@ -1358,6 +2144,14 @@ export class KnowledgeDatabase {
|
||||
)
|
||||
.run(1, new Date().toISOString())
|
||||
}
|
||||
if (currentVersion < 2) {
|
||||
this.migrateToVersion2(database)
|
||||
database
|
||||
.prepare(
|
||||
'INSERT INTO schema_migrations (version, applied_at) VALUES (?, ?)'
|
||||
)
|
||||
.run(2, new Date().toISOString())
|
||||
}
|
||||
database.exec(`PRAGMA user_version = ${DATABASE_VERSION}`)
|
||||
database.exec('COMMIT')
|
||||
} catch (error) {
|
||||
@@ -1492,6 +2286,52 @@ export class KnowledgeDatabase {
|
||||
`)
|
||||
}
|
||||
|
||||
private migrateToVersion2(database: DatabaseSync): void {
|
||||
database.exec(`
|
||||
CREATE TABLE chunk_embeddings (
|
||||
chunk_id TEXT NOT NULL REFERENCES chunks(id) ON DELETE CASCADE,
|
||||
knowledge_base_id TEXT NOT NULL
|
||||
REFERENCES knowledge_bases(id) ON DELETE CASCADE,
|
||||
provider TEXT NOT NULL,
|
||||
model TEXT NOT NULL,
|
||||
dimensions INTEGER NOT NULL
|
||||
CHECK (dimensions >= 1 AND dimensions <= 8192),
|
||||
content_checksum TEXT NOT NULL
|
||||
CHECK (length(content_checksum) = 64),
|
||||
vector BLOB NOT NULL
|
||||
CHECK (length(vector) = dimensions * 4),
|
||||
magnitude REAL NOT NULL CHECK (magnitude > 0),
|
||||
created_at TEXT NOT NULL,
|
||||
updated_at TEXT NOT NULL,
|
||||
PRIMARY KEY (chunk_id, provider, model)
|
||||
);
|
||||
CREATE INDEX chunk_embeddings_lookup_idx
|
||||
ON chunk_embeddings(
|
||||
knowledge_base_id, provider, model, dimensions, chunk_id
|
||||
);
|
||||
|
||||
CREATE TABLE embedding_index_state (
|
||||
document_id TEXT NOT NULL REFERENCES documents(id) ON DELETE CASCADE,
|
||||
knowledge_base_id TEXT NOT NULL
|
||||
REFERENCES knowledge_bases(id) ON DELETE CASCADE,
|
||||
provider TEXT NOT NULL,
|
||||
model TEXT NOT NULL,
|
||||
dimensions INTEGER
|
||||
CHECK (dimensions IS NULL OR
|
||||
(dimensions >= 1 AND dimensions <= 8192)),
|
||||
content_checksum TEXT NOT NULL,
|
||||
status TEXT NOT NULL CHECK (status IN ('ready', 'error')),
|
||||
last_error TEXT,
|
||||
updated_at TEXT NOT NULL,
|
||||
PRIMARY KEY (document_id, provider, model)
|
||||
);
|
||||
CREATE INDEX embedding_index_state_lookup_idx
|
||||
ON embedding_index_state(
|
||||
knowledge_base_id, provider, model, status, document_id
|
||||
);
|
||||
`)
|
||||
}
|
||||
|
||||
private normalizeChunks(chunks: ReplaceChunkInput[]): Array<{
|
||||
id: string
|
||||
ordinal: number
|
||||
@@ -1603,6 +2443,17 @@ export class KnowledgeDatabase {
|
||||
statement.get(id, value.knowledgeBaseId) as Row,
|
||||
'count'
|
||||
) === 1
|
||||
const chunkMatches =
|
||||
value.chunkId === undefined ||
|
||||
asNumber(
|
||||
database
|
||||
.prepare(
|
||||
`SELECT COUNT(*) AS count FROM chunks
|
||||
WHERE id = ? AND knowledge_base_id = ? AND document_id = ?`
|
||||
)
|
||||
.get(value.chunkId, value.knowledgeBaseId, value.documentId) as Row,
|
||||
'count'
|
||||
) === 1
|
||||
if (
|
||||
!matches(
|
||||
database.prepare(
|
||||
@@ -1622,12 +2473,7 @@ export class KnowledgeDatabase {
|
||||
),
|
||||
value.documentId
|
||||
) ||
|
||||
!matches(
|
||||
database.prepare(
|
||||
'SELECT COUNT(*) AS count FROM chunks WHERE id = ? AND knowledge_base_id = ?'
|
||||
),
|
||||
value.chunkId
|
||||
)
|
||||
!chunkMatches
|
||||
) {
|
||||
throw new Error('Evidence targets must belong to the evidence knowledge base')
|
||||
}
|
||||
@@ -1714,4 +2560,16 @@ export class KnowledgeDatabase {
|
||||
}
|
||||
return mapEvidence(row)
|
||||
}
|
||||
|
||||
private requiredEmbeddingIndexState(
|
||||
documentId: string,
|
||||
provider: string,
|
||||
model: string
|
||||
): EmbeddingIndexState {
|
||||
const value = this.getEmbeddingIndexState(documentId, provider, model)
|
||||
if (!value) {
|
||||
throw new Error(`Embedding index state not found: ${documentId}`)
|
||||
}
|
||||
return value
|
||||
}
|
||||
}
|
||||
|
||||
@@ -7,22 +7,25 @@ import {
|
||||
} from 'node:fs/promises'
|
||||
import { tmpdir } from 'node:os'
|
||||
import { join } from 'node:path'
|
||||
import { afterEach, describe, expect, it } from 'vitest'
|
||||
import { afterEach, describe, expect, it, vi } from 'vitest'
|
||||
import { KnowledgeService } from './knowledge-service'
|
||||
import type { EmbeddingProvider } from './types'
|
||||
import { UrlImporter } from './url-importer'
|
||||
|
||||
const temporaryDirectories: string[] = []
|
||||
const services: KnowledgeService[] = []
|
||||
|
||||
async function createService(
|
||||
urlImporter?: UrlImporter
|
||||
urlImporter?: UrlImporter,
|
||||
embeddingProvider?: EmbeddingProvider
|
||||
): Promise<{ directory: string; service: KnowledgeService }> {
|
||||
const directory = await mkdtemp(join(tmpdir(), 'goodbuddy-knowledge-service-'))
|
||||
temporaryDirectories.push(directory)
|
||||
const service = new KnowledgeService({
|
||||
databasePath: join(directory, 'knowledge.sqlite'),
|
||||
managedRoot: join(directory, 'managed'),
|
||||
urlImporter
|
||||
urlImporter,
|
||||
embeddingProvider
|
||||
})
|
||||
await service.initialize()
|
||||
services.push(service)
|
||||
@@ -144,4 +147,143 @@ describe('KnowledgeService', () => {
|
||||
expect(snapshot.evidence.length).toBeGreaterThan(0)
|
||||
await service.dispose()
|
||||
})
|
||||
|
||||
it('indexes optional embeddings and performs vector-backed hybrid search', async () => {
|
||||
const provider: EmbeddingProvider = {
|
||||
provider: 'test-provider',
|
||||
model: 'test-model',
|
||||
embed: async (input) =>
|
||||
input.map((text) =>
|
||||
text.includes('orbital') || text === 'related meaning'
|
||||
? [1, 0]
|
||||
: [0, 1]
|
||||
)
|
||||
}
|
||||
const { directory, service } = await createService(undefined, provider)
|
||||
const sourcePath = join(directory, 'vectors.txt')
|
||||
await writeFile(sourcePath, 'orbital telescope notes', 'utf8')
|
||||
const library = service.createLibrary({
|
||||
name: 'Vector knowledge',
|
||||
storageMode: 'reference',
|
||||
graphEnabled: false
|
||||
})
|
||||
|
||||
await service.importPaths(library.id, [sourcePath])
|
||||
const document = service.snapshot(library.id).documents[0]
|
||||
if (!document) {
|
||||
throw new Error('Indexed document missing')
|
||||
}
|
||||
expect(
|
||||
service.database.getEmbeddingIndexState(
|
||||
document.id,
|
||||
provider.provider,
|
||||
provider.model
|
||||
)
|
||||
).toMatchObject({ status: 'ready', dimensions: 2 })
|
||||
|
||||
const results = await service.searchHybrid(
|
||||
library.id,
|
||||
'related meaning'
|
||||
)
|
||||
expect(results[0]?.document.id).toBe(document.id)
|
||||
expect(results[0]?.retrieval.channels).toContain('vector')
|
||||
})
|
||||
|
||||
it('keeps FTS available and records diagnostics when embeddings fail', async () => {
|
||||
const provider: EmbeddingProvider = {
|
||||
provider: 'failing-provider',
|
||||
model: 'failing-model',
|
||||
embed: async () => {
|
||||
throw new Error('synthetic provider outage')
|
||||
}
|
||||
}
|
||||
const { directory, service } = await createService(undefined, provider)
|
||||
const sourcePath = join(directory, 'fallback.txt')
|
||||
await writeFile(sourcePath, 'lexical fallback remains searchable', 'utf8')
|
||||
const library = service.createLibrary({
|
||||
name: 'Fallback knowledge',
|
||||
storageMode: 'reference',
|
||||
graphEnabled: false
|
||||
})
|
||||
|
||||
await service.importPaths(library.id, [sourcePath])
|
||||
const document = service.snapshot(library.id).documents[0]
|
||||
if (!document) {
|
||||
throw new Error('Indexed document missing')
|
||||
}
|
||||
expect(service.search(library.id, 'fallback')).toHaveLength(1)
|
||||
expect(
|
||||
service.database.getEmbeddingIndexState(
|
||||
document.id,
|
||||
provider.provider,
|
||||
provider.model
|
||||
)
|
||||
).toMatchObject({
|
||||
status: 'error',
|
||||
lastError: 'synthetic provider outage'
|
||||
})
|
||||
const results = await service.searchHybrid(library.id, 'fallback')
|
||||
expect(results[0]?.retrieval.channels).toContain('fts')
|
||||
})
|
||||
|
||||
it('reindexes existing documents when an embedding provider is enabled', async () => {
|
||||
const { directory, service } = await createService()
|
||||
const sourcePath = join(directory, 'existing.txt')
|
||||
await writeFile(sourcePath, 'existing semantic content', 'utf8')
|
||||
const library = service.createLibrary({
|
||||
name: 'Existing knowledge',
|
||||
storageMode: 'reference',
|
||||
graphEnabled: false
|
||||
})
|
||||
await service.importPaths(library.id, [sourcePath])
|
||||
const document = service.snapshot(library.id).documents[0]!
|
||||
const provider: EmbeddingProvider = {
|
||||
provider: 'late-provider',
|
||||
model: 'late-model',
|
||||
embed: async (input) => input.map(() => [0.5, 0.5])
|
||||
}
|
||||
|
||||
await service.setEmbeddingProvider(provider)
|
||||
|
||||
expect(
|
||||
service.database.getEmbeddingIndexState(
|
||||
document.id,
|
||||
provider.provider,
|
||||
provider.model
|
||||
)
|
||||
).toMatchObject({ status: 'ready', dimensions: 2 })
|
||||
})
|
||||
|
||||
it('embeds a hybrid query once across multiple libraries', async () => {
|
||||
const embed = vi.fn<EmbeddingProvider['embed']>(
|
||||
async (input) => input.map(() => [1, 0])
|
||||
)
|
||||
const provider: EmbeddingProvider = {
|
||||
provider: 'shared-query-provider',
|
||||
model: 'shared-query-model',
|
||||
embed
|
||||
}
|
||||
const { directory, service } = await createService(undefined, provider)
|
||||
const libraryIds: string[] = []
|
||||
for (const index of [1, 2]) {
|
||||
const sourcePath = join(directory, `library-${index}.txt`)
|
||||
await writeFile(sourcePath, `shared topic ${index}`, 'utf8')
|
||||
const library = service.createLibrary({
|
||||
name: `Library ${index}`,
|
||||
storageMode: 'reference',
|
||||
graphEnabled: false
|
||||
})
|
||||
libraryIds.push(library.id)
|
||||
await service.importPaths(library.id, [sourcePath])
|
||||
}
|
||||
embed.mockClear()
|
||||
|
||||
const results = await service.searchHybridMany(
|
||||
libraryIds,
|
||||
'shared topic'
|
||||
)
|
||||
|
||||
expect(embed).toHaveBeenCalledOnce()
|
||||
expect(results).toHaveLength(2)
|
||||
})
|
||||
})
|
||||
|
||||
@@ -31,6 +31,8 @@ import type {
|
||||
GraphStrategy,
|
||||
GraphEntity,
|
||||
GraphRelation,
|
||||
EmbeddingProvider,
|
||||
HybridSearchResult,
|
||||
KnowledgeBase,
|
||||
KnowledgeSource,
|
||||
SearchResult
|
||||
@@ -76,12 +78,15 @@ export type KnowledgeServiceOptions = {
|
||||
managedRoot: string
|
||||
extractStructured?: ExtractStructured
|
||||
urlImporter?: UrlImporter
|
||||
embeddingProvider?: EmbeddingProvider
|
||||
embeddingBatchSize?: number
|
||||
}
|
||||
|
||||
const supportedExtensions = new Set<string>(supportedDocumentExtensions)
|
||||
const maximumFileBytes = 20 * 1024 * 1024
|
||||
const maximumSourceBytes = 500 * 1024 * 1024
|
||||
const maximumFilesPerSource = 2_000
|
||||
const maximumEmbeddingChunksPerBatch = 32
|
||||
|
||||
function isInside(root: string, candidate: string): boolean {
|
||||
const path = relative(resolve(root), resolve(candidate))
|
||||
@@ -114,15 +119,30 @@ export class KnowledgeService {
|
||||
private readonly managedRoot: string
|
||||
private readonly extractStructured?: ExtractStructured
|
||||
private readonly urlImporter: UrlImporter
|
||||
private embeddingProvider?: EmbeddingProvider
|
||||
private readonly embeddingBatchSize: number
|
||||
private readonly watchers = new Map<string, FSWatcher>()
|
||||
private readonly syncTimers = new Map<string, ReturnType<typeof setTimeout>>()
|
||||
private readonly activeSyncs = new Map<string, Promise<void>>()
|
||||
private readonly lifecycleController = new AbortController()
|
||||
|
||||
constructor(options: KnowledgeServiceOptions) {
|
||||
this.database = new KnowledgeDatabase(options.databasePath)
|
||||
this.managedRoot = resolve(options.managedRoot)
|
||||
this.extractStructured = options.extractStructured
|
||||
this.urlImporter = options.urlImporter ?? new UrlImporter()
|
||||
this.embeddingProvider = options.embeddingProvider
|
||||
const embeddingBatchSize = options.embeddingBatchSize ?? 16
|
||||
if (
|
||||
!Number.isSafeInteger(embeddingBatchSize) ||
|
||||
embeddingBatchSize < 1 ||
|
||||
embeddingBatchSize > maximumEmbeddingChunksPerBatch
|
||||
) {
|
||||
throw new RangeError(
|
||||
`embeddingBatchSize must be between 1 and ${maximumEmbeddingChunksPerBatch}`
|
||||
)
|
||||
}
|
||||
this.embeddingBatchSize = embeddingBatchSize
|
||||
}
|
||||
|
||||
async initialize(): Promise<void> {
|
||||
@@ -142,6 +162,9 @@ export class KnowledgeService {
|
||||
}
|
||||
|
||||
async dispose(): Promise<void> {
|
||||
this.lifecycleController.abort(
|
||||
new Error('Knowledge service is shutting down')
|
||||
)
|
||||
for (const timer of this.syncTimers.values()) {
|
||||
clearTimeout(timer)
|
||||
}
|
||||
@@ -154,6 +177,57 @@ export class KnowledgeService {
|
||||
this.database.close()
|
||||
}
|
||||
|
||||
setEmbeddingProvider(provider?: EmbeddingProvider): Promise<void> {
|
||||
if (
|
||||
this.embeddingProvider === provider ||
|
||||
(this.embeddingProvider?.fingerprint !== undefined &&
|
||||
this.embeddingProvider.fingerprint === provider?.fingerprint)
|
||||
) {
|
||||
this.embeddingProvider = provider
|
||||
return Promise.resolve()
|
||||
}
|
||||
this.embeddingProvider = provider
|
||||
if (!provider) {
|
||||
return Promise.resolve()
|
||||
}
|
||||
const reindex = this.reindexEmbeddings(provider)
|
||||
this.activeSyncs.set('embedding-reindex', reindex)
|
||||
void reindex.then(
|
||||
() => {
|
||||
if (this.activeSyncs.get('embedding-reindex') === reindex) {
|
||||
this.activeSyncs.delete('embedding-reindex')
|
||||
}
|
||||
},
|
||||
() => {
|
||||
if (this.activeSyncs.get('embedding-reindex') === reindex) {
|
||||
this.activeSyncs.delete('embedding-reindex')
|
||||
}
|
||||
}
|
||||
)
|
||||
return reindex
|
||||
}
|
||||
|
||||
private async reindexEmbeddings(
|
||||
provider: EmbeddingProvider
|
||||
): Promise<void> {
|
||||
for (const library of this.database.listKnowledgeBases(100)) {
|
||||
if (this.embeddingProvider !== provider) {
|
||||
return
|
||||
}
|
||||
for (const document of this.database.listDocuments(
|
||||
library.id,
|
||||
500
|
||||
)) {
|
||||
if (this.embeddingProvider !== provider) {
|
||||
return
|
||||
}
|
||||
if (document.metadata.status === 'ready') {
|
||||
await this.indexDocumentEmbeddings(document, provider)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
createLibrary(input: CreateKnowledgeBaseInput): KnowledgeBase {
|
||||
return this.database.createKnowledgeBase(input)
|
||||
}
|
||||
@@ -256,6 +330,76 @@ export class KnowledgeService {
|
||||
})
|
||||
}
|
||||
|
||||
async searchHybrid(
|
||||
knowledgeBaseId: string,
|
||||
query: string,
|
||||
limit = 6,
|
||||
signal?: AbortSignal
|
||||
): Promise<HybridSearchResult[]> {
|
||||
const library = this.requireLibrary(knowledgeBaseId)
|
||||
const vector = await this.embedQuery(query, signal)
|
||||
return this.database.hybridSearch({
|
||||
knowledgeBaseId,
|
||||
query,
|
||||
limit,
|
||||
provider: vector ? this.embeddingProvider?.provider : undefined,
|
||||
model: vector ? this.embeddingProvider?.model : undefined,
|
||||
vector,
|
||||
graphEnabled: library.graphEnabled
|
||||
})
|
||||
}
|
||||
|
||||
async searchHybridMany(
|
||||
knowledgeBaseIds: readonly string[],
|
||||
query: string,
|
||||
limitPerLibrary = 6,
|
||||
signal?: AbortSignal
|
||||
): Promise<
|
||||
Array<{ knowledgeBaseId: string; result: HybridSearchResult }>
|
||||
> {
|
||||
const vector = await this.embedQuery(query, signal)
|
||||
return knowledgeBaseIds.flatMap((knowledgeBaseId) => {
|
||||
const library = this.requireLibrary(knowledgeBaseId)
|
||||
return this.database
|
||||
.hybridSearch({
|
||||
knowledgeBaseId,
|
||||
query,
|
||||
limit: limitPerLibrary,
|
||||
provider: vector
|
||||
? this.embeddingProvider?.provider
|
||||
: undefined,
|
||||
model: vector ? this.embeddingProvider?.model : undefined,
|
||||
vector,
|
||||
graphEnabled: library.graphEnabled
|
||||
})
|
||||
.map((result) => ({ knowledgeBaseId, result }))
|
||||
})
|
||||
}
|
||||
|
||||
private async embedQuery(
|
||||
query: string,
|
||||
signal?: AbortSignal
|
||||
): Promise<readonly number[] | undefined> {
|
||||
if (!this.embeddingProvider) {
|
||||
return undefined
|
||||
}
|
||||
const effectiveSignal = signal
|
||||
? AbortSignal.any([signal, this.lifecycleController.signal])
|
||||
: this.lifecycleController.signal
|
||||
try {
|
||||
const result = await this.embeddingProvider.embed(
|
||||
[query],
|
||||
effectiveSignal
|
||||
)
|
||||
return result.length === 1 ? result[0] : undefined
|
||||
} catch {
|
||||
if (effectiveSignal.aborted) {
|
||||
throw effectiveSignal.reason
|
||||
}
|
||||
return undefined
|
||||
}
|
||||
}
|
||||
|
||||
async importPaths(
|
||||
knowledgeBaseId: string,
|
||||
selectedPaths: string[],
|
||||
@@ -343,7 +487,12 @@ export class KnowledgeService {
|
||||
const effectiveLibrary = graphStrategy
|
||||
? { ...library, graphStrategy }
|
||||
: library
|
||||
const result = await this.urlImporter.import(input, signal)
|
||||
const effectiveSignal = AbortSignal.any([
|
||||
signal,
|
||||
this.lifecycleController.signal,
|
||||
AbortSignal.timeout(60_000)
|
||||
])
|
||||
const result = await this.urlImporter.import(input, effectiveSignal)
|
||||
let source = this.database.upsertSource({
|
||||
id: sourceId,
|
||||
knowledgeBaseId,
|
||||
@@ -381,6 +530,7 @@ export class KnowledgeService {
|
||||
location: chunk.locator
|
||||
}))
|
||||
)
|
||||
await this.indexDocumentEmbeddings(document)
|
||||
await this.extractGraph(effectiveLibrary, document)
|
||||
source = this.database.upsertSource({
|
||||
...source,
|
||||
@@ -452,7 +602,7 @@ export class KnowledgeService {
|
||||
await this.importUrl(
|
||||
library.id,
|
||||
source.location,
|
||||
new AbortController().signal,
|
||||
this.lifecycleController.signal,
|
||||
source.id
|
||||
)
|
||||
return
|
||||
@@ -543,6 +693,7 @@ export class KnowledgeService {
|
||||
}))
|
||||
)
|
||||
this.database.removeEvidenceForDocument(document.id)
|
||||
await this.indexDocumentEmbeddings(document)
|
||||
await this.extractGraph(library, document)
|
||||
} catch (error) {
|
||||
failures.push(
|
||||
@@ -567,6 +718,83 @@ export class KnowledgeService {
|
||||
}
|
||||
}
|
||||
|
||||
private async indexDocumentEmbeddings(
|
||||
document: Document,
|
||||
requestedProvider?: EmbeddingProvider
|
||||
): Promise<void> {
|
||||
const provider = requestedProvider ?? this.embeddingProvider
|
||||
if (!provider) {
|
||||
return
|
||||
}
|
||||
try {
|
||||
const chunks = this.database.listChunks(document.id, 10_000)
|
||||
const embeddings: Array<{
|
||||
chunkId: string
|
||||
contentChecksum: string
|
||||
vector: readonly number[]
|
||||
}> = []
|
||||
let expectedDimensions: number | undefined
|
||||
for (
|
||||
let offset = 0;
|
||||
offset < chunks.length;
|
||||
offset += this.embeddingBatchSize
|
||||
) {
|
||||
const batch = chunks.slice(offset, offset + this.embeddingBatchSize)
|
||||
const vectors = await provider.embed(
|
||||
batch.map((chunk) => chunk.content),
|
||||
this.lifecycleController.signal
|
||||
)
|
||||
if (vectors.length !== batch.length) {
|
||||
throw new Error('Embedding provider returned an invalid result count')
|
||||
}
|
||||
for (let index = 0; index < batch.length; index += 1) {
|
||||
const chunk = batch[index]
|
||||
const vector = vectors[index]
|
||||
if (!chunk || !vector) {
|
||||
throw new Error('Embedding provider returned an incomplete batch')
|
||||
}
|
||||
if (expectedDimensions === undefined) {
|
||||
expectedDimensions = vector.length
|
||||
} else if (vector.length !== expectedDimensions) {
|
||||
throw new Error('Embedding provider returned inconsistent dimensions')
|
||||
}
|
||||
embeddings.push({
|
||||
chunkId: chunk.id,
|
||||
contentChecksum: createHash('sha256')
|
||||
.update(chunk.content)
|
||||
.digest('hex'),
|
||||
vector
|
||||
})
|
||||
}
|
||||
}
|
||||
if (this.embeddingProvider !== provider) {
|
||||
return
|
||||
}
|
||||
this.database.replaceDocumentEmbeddings(
|
||||
document.id,
|
||||
provider.provider,
|
||||
provider.model,
|
||||
embeddings
|
||||
)
|
||||
} catch (error) {
|
||||
if (this.lifecycleController.signal.aborted) {
|
||||
return
|
||||
}
|
||||
const message =
|
||||
error instanceof Error ? error.message : 'Embedding indexing failed'
|
||||
try {
|
||||
this.database.recordEmbeddingIndexError(
|
||||
document.id,
|
||||
provider.provider,
|
||||
provider.model,
|
||||
message.slice(0, 2_000)
|
||||
)
|
||||
} catch {
|
||||
// FTS indexing is authoritative; embedding diagnostics are best effort.
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
private async extractGraph(
|
||||
library: KnowledgeBase,
|
||||
document: Document
|
||||
|
||||
@@ -0,0 +1,146 @@
|
||||
import { describe, expect, it, vi } from 'vitest'
|
||||
import { OllamaEmbeddingClient } from './ollama-embedding-client'
|
||||
|
||||
describe('OllamaEmbeddingClient', () => {
|
||||
it('batches bounded embed requests and validates consistent vectors', async () => {
|
||||
const transport = vi.fn<typeof fetch>(async (_input, init) => {
|
||||
const body = JSON.parse(String(init?.body)) as {
|
||||
input: string[]
|
||||
model: string
|
||||
}
|
||||
return new Response(
|
||||
JSON.stringify({
|
||||
embeddings: body.input.map((_, index) => [index + 1, 2, 3])
|
||||
}),
|
||||
{
|
||||
status: 200,
|
||||
headers: { 'content-type': 'application/json' }
|
||||
}
|
||||
)
|
||||
})
|
||||
const client = new OllamaEmbeddingClient({
|
||||
url: 'http://embedding.test:11434',
|
||||
model: 'synthetic-model',
|
||||
batchSize: 2,
|
||||
fetch: transport
|
||||
})
|
||||
|
||||
const result = await client.embed(['alpha', 'beta', 'gamma'])
|
||||
|
||||
expect(result).toEqual([
|
||||
[1, 2, 3],
|
||||
[2, 2, 3],
|
||||
[1, 2, 3]
|
||||
])
|
||||
expect(transport).toHaveBeenCalledTimes(2)
|
||||
expect(transport.mock.calls[0]?.[0]).toBe(
|
||||
'http://embedding.test:11434/api/embed'
|
||||
)
|
||||
expect(JSON.parse(String(transport.mock.calls[0]?.[1]?.body))).toEqual({
|
||||
model: 'synthetic-model',
|
||||
input: ['alpha', 'beta'],
|
||||
truncate: true
|
||||
})
|
||||
})
|
||||
|
||||
it('rejects invalid inputs and malformed or oversized responses', async () => {
|
||||
expect(
|
||||
() =>
|
||||
new OllamaEmbeddingClient({
|
||||
url: 'file:///tmp/ollama.sock',
|
||||
model: 'model'
|
||||
})
|
||||
).toThrow('HTTP or HTTPS')
|
||||
|
||||
const malformed = new OllamaEmbeddingClient({
|
||||
url: 'https://embedding.test',
|
||||
model: 'model',
|
||||
fetch: async () =>
|
||||
new Response(JSON.stringify({ embeddings: [[1, Number.NaN]] }))
|
||||
})
|
||||
await expect(malformed.embed(['safe synthetic input'])).rejects.toThrow(
|
||||
'finite numbers'
|
||||
)
|
||||
|
||||
const oversized = new OllamaEmbeddingClient({
|
||||
url: 'https://embedding.test',
|
||||
model: 'model',
|
||||
fetch: async () =>
|
||||
new Response('ignored', {
|
||||
headers: { 'content-length': String(16 * 1024 * 1024 + 1) }
|
||||
})
|
||||
})
|
||||
await expect(oversized.embed(['safe synthetic input'])).rejects.toThrow(
|
||||
'too large'
|
||||
)
|
||||
await expect(
|
||||
malformed.embed(['x'.repeat(16_001)])
|
||||
).rejects.toThrow('at most 16000')
|
||||
})
|
||||
|
||||
it('honors caller cancellation without exposing request input', async () => {
|
||||
const controller = new AbortController()
|
||||
controller.abort()
|
||||
const transport = vi.fn<typeof fetch>()
|
||||
const client = new OllamaEmbeddingClient({
|
||||
url: 'https://embedding.test',
|
||||
model: 'model',
|
||||
fetch: transport
|
||||
})
|
||||
|
||||
await expect(
|
||||
client.embed(['synthetic cancellation text'], controller.signal)
|
||||
).rejects.toBeDefined()
|
||||
expect(transport).not.toHaveBeenCalled()
|
||||
})
|
||||
|
||||
it.runIf(
|
||||
['1', 'true'].includes(
|
||||
process.env.GOODBUDDY_OLLAMA_INTEGRATION?.toLowerCase() ?? ''
|
||||
)
|
||||
)(
|
||||
'embeds synthetic text against an explicitly configured Ollama instance',
|
||||
async () => {
|
||||
const url = process.env.GOODBUDDY_OLLAMA_URL
|
||||
const model = process.env.GOODBUDDY_OLLAMA_MODEL
|
||||
if (!url || !model) {
|
||||
throw new Error(
|
||||
'GOODBUDDY_OLLAMA_URL and GOODBUDDY_OLLAMA_MODEL are required'
|
||||
)
|
||||
}
|
||||
const client = new OllamaEmbeddingClient({
|
||||
url,
|
||||
model,
|
||||
timeoutMs: 30_000
|
||||
})
|
||||
const vectors = await client.embed([
|
||||
'A cat is sleeping peacefully on a sunny windowsill.',
|
||||
'A database transaction uses indexes and rollback logs.',
|
||||
'Where is the sleeping cat resting?'
|
||||
])
|
||||
const cosine = (left: number[], right: number[]): number => {
|
||||
const dot = left.reduce(
|
||||
(total, value, index) =>
|
||||
total + value * (right[index] ?? 0),
|
||||
0
|
||||
)
|
||||
const magnitude = (vector: number[]): number =>
|
||||
Math.sqrt(
|
||||
vector.reduce(
|
||||
(total, value) => total + value * value,
|
||||
0
|
||||
)
|
||||
)
|
||||
return dot / (magnitude(left) * magnitude(right))
|
||||
}
|
||||
expect(vectors).toHaveLength(3)
|
||||
expect(vectors[0]?.length).toBeGreaterThan(0)
|
||||
expect(vectors[1]?.length).toBe(vectors[0]?.length)
|
||||
expect(vectors[2]?.length).toBe(vectors[0]?.length)
|
||||
expect(cosine(vectors[2]!, vectors[0]!)).toBeGreaterThan(
|
||||
cosine(vectors[2]!, vectors[1]!)
|
||||
)
|
||||
},
|
||||
40_000
|
||||
)
|
||||
})
|
||||
@@ -0,0 +1,262 @@
|
||||
import type { EmbeddingProvider } from './types'
|
||||
|
||||
const MAX_INPUTS = 256
|
||||
const MAX_BATCH_SIZE = 32
|
||||
const MAX_INPUT_LENGTH = 16_000
|
||||
const MAX_BATCH_CHARACTERS = 128_000
|
||||
const MAX_MODEL_LENGTH = 256
|
||||
const MAX_URL_LENGTH = 2_048
|
||||
const MAX_DIMENSIONS = 8_192
|
||||
const MAX_RESPONSE_BYTES = 16 * 1024 * 1024
|
||||
const MIN_TIMEOUT_MS = 100
|
||||
const MAX_TIMEOUT_MS = 120_000
|
||||
|
||||
export interface OllamaEmbeddingClientOptions {
|
||||
url: string
|
||||
model: string
|
||||
batchSize?: number
|
||||
timeoutMs?: number
|
||||
fetch?: typeof fetch
|
||||
}
|
||||
|
||||
function boundedInteger(
|
||||
value: number,
|
||||
field: string,
|
||||
minimum: number,
|
||||
maximum: number
|
||||
): number {
|
||||
if (!Number.isSafeInteger(value) || value < minimum || value > maximum) {
|
||||
throw new RangeError(
|
||||
`${field} must be an integer between ${minimum} and ${maximum}`
|
||||
)
|
||||
}
|
||||
return value
|
||||
}
|
||||
|
||||
function requiredString(value: string, field: string, maximum: number): string {
|
||||
if (typeof value !== 'string' || value.trim().length === 0) {
|
||||
throw new TypeError(`${field} must be a non-empty string`)
|
||||
}
|
||||
const normalized = value.trim()
|
||||
if (normalized.length > maximum) {
|
||||
throw new RangeError(`${field} must be at most ${maximum} characters`)
|
||||
}
|
||||
return normalized
|
||||
}
|
||||
|
||||
function endpointFor(input: string): string {
|
||||
const value = requiredString(input, 'url', MAX_URL_LENGTH)
|
||||
const url = new URL(value)
|
||||
if (!['http:', 'https:'].includes(url.protocol)) {
|
||||
throw new RangeError('url must use HTTP or HTTPS')
|
||||
}
|
||||
if (url.username || url.password) {
|
||||
throw new RangeError('url must not contain credentials')
|
||||
}
|
||||
url.search = ''
|
||||
url.hash = ''
|
||||
url.pathname = `${url.pathname.replace(/\/+$/u, '')}/api/embed`
|
||||
return url.toString()
|
||||
}
|
||||
|
||||
async function readBoundedJson(response: Response): Promise<unknown> {
|
||||
const declaredLength = response.headers.get('content-length')
|
||||
if (
|
||||
declaredLength !== null &&
|
||||
Number(declaredLength) > MAX_RESPONSE_BYTES
|
||||
) {
|
||||
throw new RangeError('Ollama embedding response is too large')
|
||||
}
|
||||
if (!response.body) {
|
||||
throw new Error('Ollama embedding response has no body')
|
||||
}
|
||||
const reader = response.body.getReader()
|
||||
const chunks: Uint8Array[] = []
|
||||
let length = 0
|
||||
while (true) {
|
||||
const result = await reader.read()
|
||||
if (result.done) {
|
||||
break
|
||||
}
|
||||
length += result.value.byteLength
|
||||
if (length > MAX_RESPONSE_BYTES) {
|
||||
await reader.cancel()
|
||||
throw new RangeError('Ollama embedding response is too large')
|
||||
}
|
||||
chunks.push(result.value)
|
||||
}
|
||||
const bytes = new Uint8Array(length)
|
||||
let offset = 0
|
||||
for (const chunk of chunks) {
|
||||
bytes.set(chunk, offset)
|
||||
offset += chunk.byteLength
|
||||
}
|
||||
try {
|
||||
return JSON.parse(new TextDecoder().decode(bytes)) as unknown
|
||||
} catch {
|
||||
throw new Error('Ollama embedding response is not valid JSON')
|
||||
}
|
||||
}
|
||||
|
||||
function validateEmbeddings(value: unknown, expected: number): number[][] {
|
||||
if (
|
||||
typeof value !== 'object' ||
|
||||
value === null ||
|
||||
!('embeddings' in value) ||
|
||||
!Array.isArray(value.embeddings) ||
|
||||
value.embeddings.length !== expected
|
||||
) {
|
||||
throw new Error('Ollama embedding response has an invalid result count')
|
||||
}
|
||||
let dimensions: number | undefined
|
||||
return value.embeddings.map((candidate, embeddingIndex) => {
|
||||
if (
|
||||
!Array.isArray(candidate) ||
|
||||
candidate.length < 1 ||
|
||||
candidate.length > MAX_DIMENSIONS
|
||||
) {
|
||||
throw new RangeError(
|
||||
`Ollama embedding ${embeddingIndex} has invalid dimensions`
|
||||
)
|
||||
}
|
||||
if (dimensions === undefined) {
|
||||
dimensions = candidate.length
|
||||
} else if (candidate.length !== dimensions) {
|
||||
throw new Error('Ollama embeddings have inconsistent dimensions')
|
||||
}
|
||||
let magnitudeSquared = 0
|
||||
const vector = candidate.map((component) => {
|
||||
if (typeof component !== 'number' || !Number.isFinite(component)) {
|
||||
throw new TypeError('Ollama embeddings must contain finite numbers')
|
||||
}
|
||||
magnitudeSquared += component * component
|
||||
return component
|
||||
})
|
||||
if (!Number.isFinite(magnitudeSquared) || magnitudeSquared <= 0) {
|
||||
throw new RangeError('Ollama embeddings must have a finite non-zero norm')
|
||||
}
|
||||
return vector
|
||||
})
|
||||
}
|
||||
|
||||
export class OllamaEmbeddingClient implements EmbeddingProvider {
|
||||
readonly provider = 'ollama'
|
||||
readonly model: string
|
||||
readonly fingerprint: string
|
||||
private readonly endpoint: string
|
||||
private readonly batchSize: number
|
||||
private readonly timeoutMs: number
|
||||
private readonly transport: typeof fetch
|
||||
|
||||
constructor(options: OllamaEmbeddingClientOptions) {
|
||||
this.endpoint = endpointFor(options.url)
|
||||
this.model = requiredString(options.model, 'model', MAX_MODEL_LENGTH)
|
||||
this.fingerprint = `${this.provider}:${this.endpoint}:${this.model}`
|
||||
this.batchSize = boundedInteger(
|
||||
options.batchSize ?? 16,
|
||||
'batchSize',
|
||||
1,
|
||||
MAX_BATCH_SIZE
|
||||
)
|
||||
this.timeoutMs = boundedInteger(
|
||||
options.timeoutMs ?? 15_000,
|
||||
'timeoutMs',
|
||||
MIN_TIMEOUT_MS,
|
||||
MAX_TIMEOUT_MS
|
||||
)
|
||||
this.transport = options.fetch ?? globalThis.fetch
|
||||
if (typeof this.transport !== 'function') {
|
||||
throw new Error('A Fetch API implementation is required')
|
||||
}
|
||||
}
|
||||
|
||||
async embed(
|
||||
input: readonly string[],
|
||||
signal?: AbortSignal
|
||||
): Promise<number[][]> {
|
||||
if (!Array.isArray(input) || input.length < 1 || input.length > MAX_INPUTS) {
|
||||
throw new RangeError(`input must contain between 1 and ${MAX_INPUTS} items`)
|
||||
}
|
||||
const normalized = input.map((item, index) => {
|
||||
if (typeof item !== 'string' || item.length < 1) {
|
||||
throw new TypeError(`input[${index}] must be a non-empty string`)
|
||||
}
|
||||
if (item.length > MAX_INPUT_LENGTH) {
|
||||
throw new RangeError(
|
||||
`input[${index}] must be at most ${MAX_INPUT_LENGTH} characters`
|
||||
)
|
||||
}
|
||||
return item
|
||||
})
|
||||
|
||||
const embeddings: number[][] = []
|
||||
let offset = 0
|
||||
let expectedDimensions: number | undefined
|
||||
while (offset < normalized.length) {
|
||||
let end = offset
|
||||
let characters = 0
|
||||
while (end < normalized.length && end - offset < this.batchSize) {
|
||||
const next = normalized[end]
|
||||
if (next === undefined) {
|
||||
break
|
||||
}
|
||||
if (end > offset && characters + next.length > MAX_BATCH_CHARACTERS) {
|
||||
break
|
||||
}
|
||||
characters += next.length
|
||||
end += 1
|
||||
}
|
||||
const batch = normalized.slice(offset, end)
|
||||
const vectors = await this.embedBatch(batch, signal)
|
||||
for (const vector of vectors) {
|
||||
if (expectedDimensions === undefined) {
|
||||
expectedDimensions = vector.length
|
||||
} else if (vector.length !== expectedDimensions) {
|
||||
throw new Error('Ollama embedding batches have inconsistent dimensions')
|
||||
}
|
||||
embeddings.push(vector)
|
||||
}
|
||||
offset = end
|
||||
}
|
||||
return embeddings
|
||||
}
|
||||
|
||||
private async embedBatch(
|
||||
input: readonly string[],
|
||||
signal?: AbortSignal
|
||||
): Promise<number[][]> {
|
||||
if (signal?.aborted) {
|
||||
throw signal.reason
|
||||
}
|
||||
const timeout = AbortSignal.timeout(this.timeoutMs)
|
||||
const requestSignal = signal ? AbortSignal.any([signal, timeout]) : timeout
|
||||
let response: Response
|
||||
try {
|
||||
response = await this.transport(this.endpoint, {
|
||||
method: 'POST',
|
||||
headers: {
|
||||
accept: 'application/json',
|
||||
'content-type': 'application/json'
|
||||
},
|
||||
body: JSON.stringify({
|
||||
model: this.model,
|
||||
input,
|
||||
truncate: true
|
||||
}),
|
||||
redirect: 'error',
|
||||
signal: requestSignal
|
||||
})
|
||||
} catch (error) {
|
||||
if (requestSignal.aborted) {
|
||||
const abortError = new Error('Ollama embedding request was cancelled')
|
||||
abortError.name = 'AbortError'
|
||||
throw abortError
|
||||
}
|
||||
throw new Error('Ollama embedding request failed', { cause: error })
|
||||
}
|
||||
if (!response.ok) {
|
||||
throw new Error(`Ollama embedding request failed with HTTP ${response.status}`)
|
||||
}
|
||||
return validateEmbeddings(await readBoundedJson(response), input.length)
|
||||
}
|
||||
}
|
||||
@@ -132,6 +132,63 @@ export interface SearchResult {
|
||||
rank: number
|
||||
}
|
||||
|
||||
export interface EmbeddingProvider {
|
||||
readonly provider: string
|
||||
readonly model: string
|
||||
readonly fingerprint?: string
|
||||
embed(input: readonly string[], signal?: AbortSignal): Promise<number[][]>
|
||||
}
|
||||
|
||||
export interface ChunkEmbeddingInput {
|
||||
chunkId: string
|
||||
contentChecksum: string
|
||||
vector: readonly number[]
|
||||
}
|
||||
|
||||
export interface EmbeddingIndexState {
|
||||
documentId: string
|
||||
knowledgeBaseId: string
|
||||
provider: string
|
||||
model: string
|
||||
dimensions?: number
|
||||
contentChecksum: string
|
||||
status: 'ready' | 'error'
|
||||
lastError?: string
|
||||
updatedAt: string
|
||||
}
|
||||
|
||||
export interface VectorSearchOptions {
|
||||
knowledgeBaseId: string
|
||||
provider: string
|
||||
model: string
|
||||
vector: readonly number[]
|
||||
limit?: number
|
||||
minimumSimilarity?: number
|
||||
}
|
||||
|
||||
export interface HybridSearchOptions extends SearchOptions {
|
||||
provider?: string
|
||||
model?: string
|
||||
vector?: readonly number[]
|
||||
graphEnabled?: boolean
|
||||
vectorLimit?: number
|
||||
graphDepth?: number
|
||||
}
|
||||
|
||||
export interface RetrievalMetadata {
|
||||
score: number
|
||||
channels: Array<'fts' | 'vector' | 'graph'>
|
||||
lexicalRank?: number
|
||||
vectorRank?: number
|
||||
graphRank?: number
|
||||
similarity?: number
|
||||
evidenceIds: string[]
|
||||
}
|
||||
|
||||
export interface HybridSearchResult extends SearchResult {
|
||||
retrieval: RetrievalMetadata
|
||||
}
|
||||
|
||||
export interface GraphEntity {
|
||||
id: string
|
||||
knowledgeBaseId: string
|
||||
|
||||
Reference in New Issue
Block a user