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:
lofyer
2026-08-02 10:04:59 +08:00
co-authored by factory-droid[bot] <138933559+factory-droid[bot]@users.noreply.github.com>
parent 6ef1795b81
commit b3fdf96962
82 changed files with 17608 additions and 825 deletions
+310 -2
View File
@@ -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(() =>
+868 -10
View File
@@ -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
}
}
+145 -3
View File
@@ -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)
})
})
+230 -2
View File
@@ -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)
}
}
+57
View File
@@ -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