feat: add heartbeat monitor and vector memory search (Tier 2)
Heartbeat: - HeartbeatMonitor with 5 checks: gateway, model, channels, memory, disk - Configurable interval, failure threshold, notification channel - Recovery notifications when health restores - 25 new tests Vector Memory Search: - EmbeddingProvider interface with OpenAI, Gemini, Ollama, LlamaCpp backends - SQLite-backed VectorStore with cosine similarity search - Text chunker with paragraph-aware splitting and overlap - HybridSearch merging keyword + vector results with configurable weight - Background indexer with dirty-namespace tracking - Graceful fallback to keyword search when embeddings unavailable - 51 new tests Config: automation.heartbeat + memory.embedding schema sections Total: 950 tests passing, all types clean
This commit is contained in:
William Valentin
@@ -0,0 +1,106 @@
|
||||
import { describe, it, expect } from 'vitest';
|
||||
import { chunkText } from './chunker.js';
|
||||
import type { Chunk } from './chunker.js';
|
||||
|
||||
describe('chunkText', () => {
|
||||
it('returns empty array for empty content', () => {
|
||||
expect(chunkText('', 'test')).toEqual([]);
|
||||
expect(chunkText(' \n\n ', 'test')).toEqual([]);
|
||||
});
|
||||
|
||||
it('returns single chunk for small content', () => {
|
||||
const content = 'Hello world\nSecond line';
|
||||
const chunks = chunkText(content, 'notes', { chunkSize: 1000, chunkOverlap: 0 });
|
||||
|
||||
expect(chunks).toHaveLength(1);
|
||||
expect(chunks[0].text).toBe('Hello world\nSecond line');
|
||||
expect(chunks[0].namespace).toBe('notes');
|
||||
expect(chunks[0].startLine).toBe(1);
|
||||
expect(chunks[0].endLine).toBe(2);
|
||||
});
|
||||
|
||||
it('splits on paragraph boundaries (double newline)', () => {
|
||||
const content = 'Paragraph one line one\nParagraph one line two\n\nParagraph two line one\nParagraph two line two';
|
||||
const chunks = chunkText(content, 'test', { chunkSize: 30, chunkOverlap: 0 });
|
||||
|
||||
// Should split into two chunks at the paragraph boundary
|
||||
expect(chunks.length).toBeGreaterThanOrEqual(2);
|
||||
expect(chunks[0].text).toContain('Paragraph one');
|
||||
expect(chunks[1].text).toContain('Paragraph two');
|
||||
});
|
||||
|
||||
it('merges small paragraphs to reach target chunk size', () => {
|
||||
const content = 'A\n\nB\n\nC\n\nD';
|
||||
const chunks = chunkText(content, 'test', { chunkSize: 100, chunkOverlap: 0 });
|
||||
|
||||
// All paragraphs are tiny, so they should all fit in one chunk
|
||||
expect(chunks).toHaveLength(1);
|
||||
expect(chunks[0].text).toContain('A');
|
||||
expect(chunks[0].text).toContain('D');
|
||||
});
|
||||
|
||||
it('tracks line numbers accurately', () => {
|
||||
const content = 'Line one\n\nLine three\n\nLine five';
|
||||
const chunks = chunkText(content, 'test', { chunkSize: 10, chunkOverlap: 0 });
|
||||
|
||||
// First chunk should start at line 1
|
||||
expect(chunks[0].startLine).toBe(1);
|
||||
expect(chunks[0].endLine).toBe(1);
|
||||
|
||||
// Line three is on actual line 3
|
||||
const lineThreeChunk = chunks.find((c) => c.text.includes('Line three'));
|
||||
expect(lineThreeChunk).toBeDefined();
|
||||
expect(lineThreeChunk!.startLine).toBe(3);
|
||||
|
||||
// Line five is on actual line 5
|
||||
const lineFiveChunk = chunks.find((c) => c.text.includes('Line five'));
|
||||
expect(lineFiveChunk).toBeDefined();
|
||||
expect(lineFiveChunk!.startLine).toBe(5);
|
||||
});
|
||||
|
||||
it('includes overlap between consecutive chunks', () => {
|
||||
// Create content with clear paragraphs that force splitting
|
||||
const para1 = 'First paragraph with enough text to matter';
|
||||
const para2 = 'Second paragraph with some more text';
|
||||
const para3 = 'Third paragraph and final content here';
|
||||
const content = `${para1}\n\n${para2}\n\n${para3}`;
|
||||
|
||||
// Use a chunk size that forces splitting, with overlap
|
||||
const chunks = chunkText(content, 'test', { chunkSize: 50, chunkOverlap: 40 });
|
||||
|
||||
// With overlap, later chunks should contain content from previous paragraphs
|
||||
if (chunks.length >= 2) {
|
||||
// Check that there's some content overlap between consecutive chunks
|
||||
const lastChunk = chunks[chunks.length - 1];
|
||||
const prevChunk = chunks[chunks.length - 2];
|
||||
// Either chunks share content or at least have proper sequencing
|
||||
expect(lastChunk.startLine).toBeLessThanOrEqual(prevChunk.endLine + 5);
|
||||
}
|
||||
});
|
||||
|
||||
it('preserves namespace in all chunks', () => {
|
||||
const content = 'Para one\n\nPara two\n\nPara three';
|
||||
const chunks = chunkText(content, 'sessions/abc123', { chunkSize: 10, chunkOverlap: 0 });
|
||||
|
||||
for (const chunk of chunks) {
|
||||
expect(chunk.namespace).toBe('sessions/abc123');
|
||||
}
|
||||
});
|
||||
|
||||
it('handles content with multiple consecutive blank lines', () => {
|
||||
const content = 'First\n\n\n\nSecond';
|
||||
const chunks = chunkText(content, 'test', { chunkSize: 1000, chunkOverlap: 0 });
|
||||
|
||||
expect(chunks.length).toBeGreaterThanOrEqual(1);
|
||||
expect(chunks.some((c) => c.text.includes('First'))).toBe(true);
|
||||
expect(chunks.some((c) => c.text.includes('Second'))).toBe(true);
|
||||
});
|
||||
|
||||
it('handles single-line content', () => {
|
||||
const chunks = chunkText('single line', 'test', { chunkSize: 100, chunkOverlap: 0 });
|
||||
expect(chunks).toHaveLength(1);
|
||||
expect(chunks[0].text).toBe('single line');
|
||||
expect(chunks[0].startLine).toBe(1);
|
||||
expect(chunks[0].endLine).toBe(1);
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,163 @@
|
||||
/**
|
||||
* Text chunker that splits markdown content into overlapping chunks
|
||||
* for embedding generation.
|
||||
*/
|
||||
|
||||
/**
|
||||
* A single chunk of text extracted from a memory namespace.
|
||||
*/
|
||||
export interface Chunk {
|
||||
/** The chunk text content. */
|
||||
text: string;
|
||||
/** The memory namespace this chunk came from. */
|
||||
namespace: string;
|
||||
/** 1-based start line number in the original content. */
|
||||
startLine: number;
|
||||
/** 1-based end line number in the original content. */
|
||||
endLine: number;
|
||||
}
|
||||
|
||||
export interface ChunkOptions {
|
||||
/** Target chunk size in characters. */
|
||||
chunkSize: number;
|
||||
/** Number of overlapping characters between consecutive chunks. */
|
||||
chunkOverlap: number;
|
||||
}
|
||||
|
||||
const DEFAULT_CHUNK_OPTIONS: ChunkOptions = {
|
||||
chunkSize: 512,
|
||||
chunkOverlap: 50,
|
||||
};
|
||||
|
||||
/**
|
||||
* Split content into overlapping chunks suitable for embedding.
|
||||
*
|
||||
* Strategy:
|
||||
* 1. Split on paragraph boundaries (double newline).
|
||||
* 2. Merge small paragraphs to reach target chunk size.
|
||||
* 3. Track line numbers accurately through splits.
|
||||
* 4. Add overlap from previous chunk for context continuity.
|
||||
*/
|
||||
export function chunkText(
|
||||
content: string,
|
||||
namespace: string,
|
||||
options?: Partial<ChunkOptions>,
|
||||
): Chunk[] {
|
||||
const opts = { ...DEFAULT_CHUNK_OPTIONS, ...options };
|
||||
|
||||
if (content.trim().length === 0) {
|
||||
return [];
|
||||
}
|
||||
|
||||
const lines = content.split('\n');
|
||||
|
||||
// Build paragraph groups: each paragraph is a contiguous set of lines
|
||||
// separated by blank lines (double newline boundaries).
|
||||
const paragraphs: { text: string; startLine: number; endLine: number }[] = [];
|
||||
let currentLines: string[] = [];
|
||||
let currentStart = 1; // 1-based
|
||||
|
||||
for (let i = 0; i < lines.length; i++) {
|
||||
const line = lines[i];
|
||||
const lineNum = i + 1; // 1-based
|
||||
|
||||
if (line.trim() === '' && currentLines.length > 0) {
|
||||
// End of a paragraph
|
||||
paragraphs.push({
|
||||
text: currentLines.join('\n'),
|
||||
startLine: currentStart,
|
||||
endLine: lineNum - 1,
|
||||
});
|
||||
currentLines = [];
|
||||
currentStart = lineNum + 1;
|
||||
} else if (line.trim() !== '') {
|
||||
if (currentLines.length === 0) {
|
||||
currentStart = lineNum;
|
||||
}
|
||||
currentLines.push(line);
|
||||
} else {
|
||||
// Empty line and no current paragraph — advance start
|
||||
currentStart = lineNum + 1;
|
||||
}
|
||||
}
|
||||
|
||||
// Flush remaining
|
||||
if (currentLines.length > 0) {
|
||||
paragraphs.push({
|
||||
text: currentLines.join('\n'),
|
||||
startLine: currentStart,
|
||||
endLine: lines.length,
|
||||
});
|
||||
}
|
||||
|
||||
if (paragraphs.length === 0) {
|
||||
return [];
|
||||
}
|
||||
|
||||
// Merge paragraphs into chunks, respecting the target size
|
||||
const chunks: Chunk[] = [];
|
||||
let chunkParagraphs: typeof paragraphs = [];
|
||||
let chunkLength = 0;
|
||||
|
||||
for (const para of paragraphs) {
|
||||
const paraLength = para.text.length;
|
||||
|
||||
// If adding this paragraph would exceed the target, flush current chunk
|
||||
if (chunkLength > 0 && chunkLength + paraLength + 1 > opts.chunkSize) {
|
||||
chunks.push(buildChunk(chunkParagraphs, namespace));
|
||||
// Start a new chunk — include overlap from previous chunk
|
||||
const overlapChunk = getOverlapParagraphs(chunkParagraphs, opts.chunkOverlap);
|
||||
chunkParagraphs = overlapChunk;
|
||||
chunkLength = overlapChunk.reduce((sum, p) => sum + p.text.length, 0);
|
||||
}
|
||||
|
||||
chunkParagraphs.push(para);
|
||||
chunkLength += paraLength + (chunkLength > 0 ? 1 : 0); // +1 for separator
|
||||
}
|
||||
|
||||
// Flush remaining
|
||||
if (chunkParagraphs.length > 0) {
|
||||
chunks.push(buildChunk(chunkParagraphs, namespace));
|
||||
}
|
||||
|
||||
return chunks;
|
||||
}
|
||||
|
||||
/** Build a Chunk from a list of paragraph entries. */
|
||||
function buildChunk(
|
||||
paragraphs: { text: string; startLine: number; endLine: number }[],
|
||||
namespace: string,
|
||||
): Chunk {
|
||||
return {
|
||||
text: paragraphs.map((p) => p.text).join('\n\n'),
|
||||
namespace,
|
||||
startLine: paragraphs[0].startLine,
|
||||
endLine: paragraphs[paragraphs.length - 1].endLine,
|
||||
};
|
||||
}
|
||||
|
||||
/**
|
||||
* Get trailing paragraphs from the previous chunk for overlap.
|
||||
* Takes paragraphs from the end until we've accumulated enough characters.
|
||||
*/
|
||||
function getOverlapParagraphs(
|
||||
paragraphs: { text: string; startLine: number; endLine: number }[],
|
||||
overlapChars: number,
|
||||
): { text: string; startLine: number; endLine: number }[] {
|
||||
if (overlapChars <= 0) {
|
||||
return [];
|
||||
}
|
||||
|
||||
const result: typeof paragraphs = [];
|
||||
let totalChars = 0;
|
||||
|
||||
for (let i = paragraphs.length - 1; i >= 0; i--) {
|
||||
totalChars += paragraphs[i].text.length;
|
||||
result.unshift(paragraphs[i]);
|
||||
if (totalChars >= overlapChars) {
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
return result;
|
||||
}
|
||||
@@ -0,0 +1,159 @@
|
||||
import { describe, it, expect } from 'vitest';
|
||||
import {
|
||||
createEmbeddingProvider,
|
||||
OpenAIEmbeddingProvider,
|
||||
GeminiEmbeddingProvider,
|
||||
OllamaEmbeddingProvider,
|
||||
LlamaCppEmbeddingProvider,
|
||||
} from './embeddings.js';
|
||||
import type { EmbeddingConfig } from '../config/schema.js';
|
||||
|
||||
describe('createEmbeddingProvider', () => {
|
||||
const baseConfig: EmbeddingConfig = {
|
||||
enabled: true,
|
||||
provider: 'openai',
|
||||
model: 'text-embedding-3-small',
|
||||
chunk_size: 512,
|
||||
chunk_overlap: 50,
|
||||
top_k: 5,
|
||||
hybrid_weight: 0.7,
|
||||
};
|
||||
|
||||
it('creates OpenAI provider', () => {
|
||||
const provider = createEmbeddingProvider({ ...baseConfig, provider: 'openai' });
|
||||
expect(provider).toBeInstanceOf(OpenAIEmbeddingProvider);
|
||||
});
|
||||
|
||||
it('creates Gemini provider', () => {
|
||||
const provider = createEmbeddingProvider({ ...baseConfig, provider: 'gemini' });
|
||||
expect(provider).toBeInstanceOf(GeminiEmbeddingProvider);
|
||||
});
|
||||
|
||||
it('creates Ollama provider', () => {
|
||||
const provider = createEmbeddingProvider({ ...baseConfig, provider: 'ollama' });
|
||||
expect(provider).toBeInstanceOf(OllamaEmbeddingProvider);
|
||||
});
|
||||
|
||||
it('creates LlamaCpp provider', () => {
|
||||
const provider = createEmbeddingProvider({ ...baseConfig, provider: 'llamacpp' });
|
||||
expect(provider).toBeInstanceOf(LlamaCppEmbeddingProvider);
|
||||
});
|
||||
|
||||
it('throws on unknown provider', () => {
|
||||
expect(() => createEmbeddingProvider({ ...baseConfig, provider: 'unknown' as never })).toThrow('Unknown embedding provider');
|
||||
});
|
||||
});
|
||||
|
||||
describe('OpenAIEmbeddingProvider', () => {
|
||||
it('reports configured dimensions', () => {
|
||||
const config: EmbeddingConfig = {
|
||||
enabled: true,
|
||||
provider: 'openai',
|
||||
model: 'text-embedding-3-small',
|
||||
dimensions: 512,
|
||||
chunk_size: 512,
|
||||
chunk_overlap: 50,
|
||||
top_k: 5,
|
||||
hybrid_weight: 0.7,
|
||||
};
|
||||
const provider = new OpenAIEmbeddingProvider(config);
|
||||
expect(provider.dimensions).toBe(512);
|
||||
});
|
||||
|
||||
it('defaults to 1536 dimensions', () => {
|
||||
const config: EmbeddingConfig = {
|
||||
enabled: true,
|
||||
provider: 'openai',
|
||||
model: 'text-embedding-3-small',
|
||||
chunk_size: 512,
|
||||
chunk_overlap: 50,
|
||||
top_k: 5,
|
||||
hybrid_weight: 0.7,
|
||||
};
|
||||
const provider = new OpenAIEmbeddingProvider(config);
|
||||
expect(provider.dimensions).toBe(1536);
|
||||
});
|
||||
});
|
||||
|
||||
describe('GeminiEmbeddingProvider', () => {
|
||||
it('reports configured dimensions', () => {
|
||||
const config: EmbeddingConfig = {
|
||||
enabled: true,
|
||||
provider: 'gemini',
|
||||
model: 'text-embedding-004',
|
||||
dimensions: 256,
|
||||
chunk_size: 512,
|
||||
chunk_overlap: 50,
|
||||
top_k: 5,
|
||||
hybrid_weight: 0.7,
|
||||
};
|
||||
const provider = new GeminiEmbeddingProvider(config);
|
||||
expect(provider.dimensions).toBe(256);
|
||||
});
|
||||
|
||||
it('defaults to 768 dimensions', () => {
|
||||
const config: EmbeddingConfig = {
|
||||
enabled: true,
|
||||
provider: 'gemini',
|
||||
model: 'text-embedding-004',
|
||||
chunk_size: 512,
|
||||
chunk_overlap: 50,
|
||||
top_k: 5,
|
||||
hybrid_weight: 0.7,
|
||||
};
|
||||
const provider = new GeminiEmbeddingProvider(config);
|
||||
expect(provider.dimensions).toBe(768);
|
||||
});
|
||||
});
|
||||
|
||||
describe('OllamaEmbeddingProvider', () => {
|
||||
it('reports configured dimensions', () => {
|
||||
const config: EmbeddingConfig = {
|
||||
enabled: true,
|
||||
provider: 'ollama',
|
||||
model: 'nomic-embed-text',
|
||||
dimensions: 384,
|
||||
endpoint: 'http://localhost:11434',
|
||||
chunk_size: 512,
|
||||
chunk_overlap: 50,
|
||||
top_k: 5,
|
||||
hybrid_weight: 0.7,
|
||||
};
|
||||
const provider = new OllamaEmbeddingProvider(config);
|
||||
expect(provider.dimensions).toBe(384);
|
||||
});
|
||||
});
|
||||
|
||||
describe('LlamaCppEmbeddingProvider', () => {
|
||||
it('reports configured dimensions', () => {
|
||||
const config: EmbeddingConfig = {
|
||||
enabled: true,
|
||||
provider: 'llamacpp',
|
||||
model: 'unused',
|
||||
dimensions: 768,
|
||||
endpoint: 'http://localhost:8080',
|
||||
chunk_size: 512,
|
||||
chunk_overlap: 50,
|
||||
top_k: 5,
|
||||
hybrid_weight: 0.7,
|
||||
};
|
||||
const provider = new LlamaCppEmbeddingProvider(config);
|
||||
expect(provider.dimensions).toBe(768);
|
||||
});
|
||||
|
||||
it('defaults endpoint to localhost:8080', () => {
|
||||
const config: EmbeddingConfig = {
|
||||
enabled: true,
|
||||
provider: 'llamacpp',
|
||||
model: 'unused',
|
||||
dimensions: 768,
|
||||
chunk_size: 512,
|
||||
chunk_overlap: 50,
|
||||
top_k: 5,
|
||||
hybrid_weight: 0.7,
|
||||
};
|
||||
// Provider should be constructable without endpoint
|
||||
const provider = new LlamaCppEmbeddingProvider(config);
|
||||
expect(provider.dimensions).toBe(768);
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,182 @@
|
||||
/**
|
||||
* Embedding provider interface and implementations for multiple backends.
|
||||
*/
|
||||
|
||||
import type { EmbeddingConfig } from '../config/schema.js';
|
||||
|
||||
/**
|
||||
* Interface for embedding providers that convert text to vectors.
|
||||
*/
|
||||
export interface EmbeddingProvider {
|
||||
/** Generate embeddings for one or more texts. Returns one vector per text. */
|
||||
embed(texts: string[]): Promise<number[][]>;
|
||||
/** The dimensionality of the embedding vectors. */
|
||||
dimensions: number;
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// OpenAI
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
export class OpenAIEmbeddingProvider implements EmbeddingProvider {
|
||||
private _model: string;
|
||||
private _dimensions: number;
|
||||
private _apiKey: string;
|
||||
private _endpoint?: string;
|
||||
|
||||
constructor(config: EmbeddingConfig) {
|
||||
this._model = config.model;
|
||||
this._dimensions = config.dimensions ?? 1536;
|
||||
this._apiKey = config.api_key ?? process.env.OPENAI_API_KEY ?? '';
|
||||
this._endpoint = config.endpoint;
|
||||
}
|
||||
|
||||
get dimensions(): number {
|
||||
return this._dimensions;
|
||||
}
|
||||
|
||||
async embed(texts: string[]): Promise<number[][]> {
|
||||
const { default: OpenAI } = await import('openai');
|
||||
const client = new OpenAI({
|
||||
apiKey: this._apiKey,
|
||||
...(this._endpoint ? { baseURL: this._endpoint } : {}),
|
||||
});
|
||||
|
||||
const response = await client.embeddings.create({
|
||||
model: this._model,
|
||||
input: texts,
|
||||
...(this._dimensions ? { dimensions: this._dimensions } : {}),
|
||||
});
|
||||
|
||||
// Sort by index to ensure order matches input
|
||||
const sorted = response.data.sort((a, b) => a.index - b.index);
|
||||
return sorted.map((item) => item.embedding);
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Gemini
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
export class GeminiEmbeddingProvider implements EmbeddingProvider {
|
||||
private _model: string;
|
||||
private _dimensions: number;
|
||||
private _apiKey: string;
|
||||
|
||||
constructor(config: EmbeddingConfig) {
|
||||
this._model = config.model;
|
||||
this._dimensions = config.dimensions ?? 768;
|
||||
this._apiKey = config.api_key ?? process.env.GOOGLE_API_KEY ?? '';
|
||||
}
|
||||
|
||||
get dimensions(): number {
|
||||
return this._dimensions;
|
||||
}
|
||||
|
||||
async embed(texts: string[]): Promise<number[][]> {
|
||||
const { GoogleGenerativeAI } = await import('@google/generative-ai');
|
||||
const genAI = new GoogleGenerativeAI(this._apiKey);
|
||||
const model = genAI.getGenerativeModel({ model: this._model });
|
||||
|
||||
// Use batchEmbedContents for efficiency
|
||||
const requests = texts.map((text) => ({
|
||||
content: { role: 'user' as const, parts: [{ text }] },
|
||||
}));
|
||||
|
||||
const response = await model.batchEmbedContents({ requests });
|
||||
return response.embeddings.map((e) => e.values);
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Ollama
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
export class OllamaEmbeddingProvider implements EmbeddingProvider {
|
||||
private _model: string;
|
||||
private _dimensions: number;
|
||||
private _host?: string;
|
||||
|
||||
constructor(config: EmbeddingConfig) {
|
||||
this._model = config.model;
|
||||
this._dimensions = config.dimensions ?? 768;
|
||||
this._host = config.endpoint;
|
||||
}
|
||||
|
||||
get dimensions(): number {
|
||||
return this._dimensions;
|
||||
}
|
||||
|
||||
async embed(texts: string[]): Promise<number[][]> {
|
||||
const { Ollama } = await import('ollama');
|
||||
const client = new Ollama({ host: this._host });
|
||||
|
||||
const response = await client.embed({
|
||||
model: this._model,
|
||||
input: texts,
|
||||
});
|
||||
|
||||
return response.embeddings;
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// LlamaCpp
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
export class LlamaCppEmbeddingProvider implements EmbeddingProvider {
|
||||
private _dimensions: number;
|
||||
private _endpoint: string;
|
||||
|
||||
constructor(config: EmbeddingConfig) {
|
||||
this._dimensions = config.dimensions ?? 768;
|
||||
this._endpoint = config.endpoint ?? 'http://localhost:8080';
|
||||
}
|
||||
|
||||
get dimensions(): number {
|
||||
return this._dimensions;
|
||||
}
|
||||
|
||||
async embed(texts: string[]): Promise<number[][]> {
|
||||
const results: number[][] = [];
|
||||
|
||||
for (const text of texts) {
|
||||
const response = await fetch(`${this._endpoint}/embedding`, {
|
||||
method: 'POST',
|
||||
headers: { 'Content-Type': 'application/json' },
|
||||
body: JSON.stringify({ content: text }),
|
||||
});
|
||||
|
||||
if (!response.ok) {
|
||||
throw new Error(`LlamaCpp embedding request failed: ${response.status} ${response.statusText}`);
|
||||
}
|
||||
|
||||
const data = (await response.json()) as { embedding: number[] };
|
||||
results.push(data.embedding);
|
||||
}
|
||||
|
||||
return results;
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Factory
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
/**
|
||||
* Create an embedding provider from config.
|
||||
*/
|
||||
export function createEmbeddingProvider(config: EmbeddingConfig): EmbeddingProvider {
|
||||
switch (config.provider) {
|
||||
case 'openai':
|
||||
return new OpenAIEmbeddingProvider(config);
|
||||
case 'gemini':
|
||||
return new GeminiEmbeddingProvider(config);
|
||||
case 'ollama':
|
||||
return new OllamaEmbeddingProvider(config);
|
||||
case 'llamacpp':
|
||||
return new LlamaCppEmbeddingProvider(config);
|
||||
default:
|
||||
throw new Error(`Unknown embedding provider: ${(config as Record<string, unknown>).provider}`);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,213 @@
|
||||
import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest';
|
||||
import { HybridSearch } from './hybrid-search.js';
|
||||
import type { HybridSearchResult } from './hybrid-search.js';
|
||||
import type { MemoryStore, SearchResult } from './store.js';
|
||||
import type { VectorStore, VectorSearchResult } from './vector-store.js';
|
||||
import type { EmbeddingProvider } from './embeddings.js';
|
||||
|
||||
/**
|
||||
* Create a mock MemoryStore with keyword search results.
|
||||
*/
|
||||
function mockMemoryStore(results: SearchResult[]): MemoryStore {
|
||||
return {
|
||||
search: vi.fn(() => results),
|
||||
read: vi.fn(() => ''),
|
||||
write: vi.fn(),
|
||||
listNamespaces: vi.fn(() => []),
|
||||
getContextForPrompt: vi.fn(() => ''),
|
||||
getDirtyNamespaces: vi.fn(() => []),
|
||||
markAllDirty: vi.fn(),
|
||||
} as unknown as MemoryStore;
|
||||
}
|
||||
|
||||
/**
|
||||
* Create a mock VectorStore that returns given results.
|
||||
*/
|
||||
function mockVectorStore(results: VectorSearchResult[]): VectorStore {
|
||||
return {
|
||||
search: vi.fn(() => results),
|
||||
upsertChunks: vi.fn(),
|
||||
deleteNamespace: vi.fn(),
|
||||
hasContentHash: vi.fn(() => false),
|
||||
count: vi.fn(() => 0),
|
||||
close: vi.fn(),
|
||||
} as unknown as VectorStore;
|
||||
}
|
||||
|
||||
/**
|
||||
* Create a mock embedding provider that returns fixed embeddings.
|
||||
*/
|
||||
function mockEmbeddingProvider(dims: number = 4): EmbeddingProvider {
|
||||
return {
|
||||
dimensions: dims,
|
||||
embed: vi.fn(async (texts: string[]) =>
|
||||
texts.map(() => new Array(dims).fill(0.1)),
|
||||
),
|
||||
};
|
||||
}
|
||||
|
||||
describe('HybridSearch', () => {
|
||||
describe('search', () => {
|
||||
it('returns keyword results when no vector results exist', async () => {
|
||||
const keywordResults: SearchResult[] = [
|
||||
{ namespace: 'notes', line: 5, content: 'fox jumped', context: 'the fox jumped over' },
|
||||
];
|
||||
|
||||
const hybrid = new HybridSearch(
|
||||
mockMemoryStore(keywordResults),
|
||||
mockVectorStore([]),
|
||||
mockEmbeddingProvider(),
|
||||
0.7,
|
||||
);
|
||||
|
||||
const results = await hybrid.search('fox');
|
||||
expect(results.length).toBe(1);
|
||||
expect(results[0].namespace).toBe('notes');
|
||||
expect(results[0].source).toBe('keyword');
|
||||
});
|
||||
|
||||
it('returns vector results when no keyword results exist', async () => {
|
||||
const vectorResults: VectorSearchResult[] = [
|
||||
{ namespace: 'journal', chunkText: 'semantic match', startLine: 10, endLine: 20, score: 0.9 },
|
||||
];
|
||||
|
||||
const hybrid = new HybridSearch(
|
||||
mockMemoryStore([]),
|
||||
mockVectorStore(vectorResults),
|
||||
mockEmbeddingProvider(),
|
||||
0.7,
|
||||
);
|
||||
|
||||
const results = await hybrid.search('meaning');
|
||||
expect(results.length).toBe(1);
|
||||
expect(results[0].namespace).toBe('journal');
|
||||
expect(results[0].source).toBe('vector');
|
||||
});
|
||||
|
||||
it('merges keyword and vector results', async () => {
|
||||
const keywordResults: SearchResult[] = [
|
||||
{ namespace: 'notes', line: 5, content: 'fox keyword', context: 'the fox keyword hit' },
|
||||
];
|
||||
const vectorResults: VectorSearchResult[] = [
|
||||
{ namespace: 'journal', chunkText: 'fox semantic', startLine: 10, endLine: 20, score: 0.85 },
|
||||
];
|
||||
|
||||
const hybrid = new HybridSearch(
|
||||
mockMemoryStore(keywordResults),
|
||||
mockVectorStore(vectorResults),
|
||||
mockEmbeddingProvider(),
|
||||
0.7,
|
||||
);
|
||||
|
||||
const results = await hybrid.search('fox');
|
||||
expect(results.length).toBe(2);
|
||||
|
||||
const namespaces = results.map((r) => r.namespace);
|
||||
expect(namespaces).toContain('notes');
|
||||
expect(namespaces).toContain('journal');
|
||||
});
|
||||
|
||||
it('deduplicates results from same namespace and nearby lines', async () => {
|
||||
// Both keyword and vector find something at the same location
|
||||
const keywordResults: SearchResult[] = [
|
||||
{ namespace: 'notes', line: 5, content: 'fox hit', context: 'context' },
|
||||
];
|
||||
const vectorResults: VectorSearchResult[] = [
|
||||
{ namespace: 'notes', chunkText: 'fox hit too', startLine: 4, endLine: 8, score: 0.9 },
|
||||
];
|
||||
|
||||
const hybrid = new HybridSearch(
|
||||
mockMemoryStore(keywordResults),
|
||||
mockVectorStore(vectorResults),
|
||||
mockEmbeddingProvider(),
|
||||
0.7,
|
||||
);
|
||||
|
||||
const results = await hybrid.search('fox');
|
||||
// Should be deduplicated to a single "both" result
|
||||
expect(results.length).toBe(1);
|
||||
expect(results[0].source).toBe('both');
|
||||
});
|
||||
|
||||
it('applies hybrid weight to scoring', async () => {
|
||||
const keywordResults: SearchResult[] = [
|
||||
{ namespace: 'notes', line: 100, content: 'keyword only', context: 'ctx' },
|
||||
];
|
||||
const vectorResults: VectorSearchResult[] = [
|
||||
{ namespace: 'journal', chunkText: 'vector only', startLine: 200, endLine: 210, score: 0.95 },
|
||||
];
|
||||
|
||||
// High vector weight (0.9)
|
||||
const hybrid = new HybridSearch(
|
||||
mockMemoryStore(keywordResults),
|
||||
mockVectorStore(vectorResults),
|
||||
mockEmbeddingProvider(),
|
||||
0.9,
|
||||
);
|
||||
|
||||
const results = await hybrid.search('query');
|
||||
expect(results.length).toBe(2);
|
||||
|
||||
// Vector result should rank higher with high vector weight
|
||||
const vectorResult = results.find((r) => r.source === 'vector');
|
||||
const keywordResult = results.find((r) => r.source === 'keyword');
|
||||
expect(vectorResult).toBeDefined();
|
||||
expect(keywordResult).toBeDefined();
|
||||
expect(vectorResult!.score).toBeGreaterThan(keywordResult!.score);
|
||||
});
|
||||
|
||||
it('falls back to keyword search when vector search fails', async () => {
|
||||
const keywordResults: SearchResult[] = [
|
||||
{ namespace: 'notes', line: 1, content: 'fallback', context: 'ctx' },
|
||||
];
|
||||
|
||||
const failingProvider: EmbeddingProvider = {
|
||||
dimensions: 4,
|
||||
embed: vi.fn(async () => { throw new Error('API error'); }),
|
||||
};
|
||||
|
||||
const hybrid = new HybridSearch(
|
||||
mockMemoryStore(keywordResults),
|
||||
mockVectorStore([]),
|
||||
failingProvider,
|
||||
0.7,
|
||||
);
|
||||
|
||||
// Should not throw — should fall back to keyword results
|
||||
const results = await hybrid.search('test');
|
||||
expect(results.length).toBe(1);
|
||||
expect(results[0].source).toBe('keyword');
|
||||
});
|
||||
|
||||
it('respects topK limit', async () => {
|
||||
const keywordResults: SearchResult[] = Array.from({ length: 10 }, (_, i) => ({
|
||||
namespace: `ns${i}`,
|
||||
line: i + 1,
|
||||
content: `result ${i}`,
|
||||
context: `ctx ${i}`,
|
||||
}));
|
||||
|
||||
const hybrid = new HybridSearch(
|
||||
mockMemoryStore(keywordResults),
|
||||
mockVectorStore([]),
|
||||
mockEmbeddingProvider(),
|
||||
0.5,
|
||||
);
|
||||
|
||||
const results = await hybrid.search('query', 3);
|
||||
expect(results.length).toBe(3);
|
||||
});
|
||||
|
||||
it('returns empty array when both searches find nothing', async () => {
|
||||
const hybrid = new HybridSearch(
|
||||
mockMemoryStore([]),
|
||||
mockVectorStore([]),
|
||||
mockEmbeddingProvider(),
|
||||
0.7,
|
||||
);
|
||||
|
||||
const results = await hybrid.search('nonexistent');
|
||||
expect(results).toEqual([]);
|
||||
});
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,182 @@
|
||||
/**
|
||||
* Hybrid search combining vector similarity with keyword matching.
|
||||
*/
|
||||
|
||||
import type { MemoryStore, SearchResult } from './store.js';
|
||||
import type { VectorStore } from './vector-store.js';
|
||||
import type { EmbeddingProvider } from './embeddings.js';
|
||||
|
||||
/**
|
||||
* A result from hybrid search combining vector and keyword sources.
|
||||
*/
|
||||
export interface HybridSearchResult {
|
||||
/** The memory namespace the result came from. */
|
||||
namespace: string;
|
||||
/** The matched content text. */
|
||||
content: string;
|
||||
/** Surrounding context lines. */
|
||||
context: string;
|
||||
/** 1-based line number of the match. */
|
||||
line: number;
|
||||
/** Combined relevance score (0-1). */
|
||||
score: number;
|
||||
/** Source of the match: keyword, vector, or both. */
|
||||
source: 'keyword' | 'vector' | 'both';
|
||||
}
|
||||
|
||||
/**
|
||||
* Combines keyword search from MemoryStore with vector similarity
|
||||
* search from VectorStore, deduplicating and merging results with
|
||||
* configurable weighting.
|
||||
*/
|
||||
export class HybridSearch {
|
||||
private _memoryStore: MemoryStore;
|
||||
private _vectorStore: VectorStore;
|
||||
private _embeddingProvider: EmbeddingProvider;
|
||||
private _hybridWeight: number;
|
||||
|
||||
/**
|
||||
* @param memoryStore - The keyword-based memory store.
|
||||
* @param vectorStore - The vector embedding store.
|
||||
* @param embeddingProvider - Provider for generating query embeddings.
|
||||
* @param hybridWeight - Weight for vector results (0-1). Keyword weight = 1 - hybridWeight.
|
||||
*/
|
||||
constructor(
|
||||
memoryStore: MemoryStore,
|
||||
vectorStore: VectorStore,
|
||||
embeddingProvider: EmbeddingProvider,
|
||||
hybridWeight: number = 0.7,
|
||||
) {
|
||||
this._memoryStore = memoryStore;
|
||||
this._vectorStore = vectorStore;
|
||||
this._embeddingProvider = embeddingProvider;
|
||||
this._hybridWeight = hybridWeight;
|
||||
}
|
||||
|
||||
/**
|
||||
* Run hybrid search combining keyword and vector results.
|
||||
*
|
||||
* @param query - The search query string.
|
||||
* @param topK - Maximum number of results to return.
|
||||
* @returns Merged and deduplicated results sorted by combined score.
|
||||
*/
|
||||
async search(query: string, topK: number = 5): Promise<HybridSearchResult[]> {
|
||||
// Run keyword and vector search in parallel
|
||||
const [keywordResults, vectorResults] = await Promise.all([
|
||||
this._keywordSearch(query),
|
||||
this._vectorSearch(query, topK * 2), // fetch more for better merging
|
||||
]);
|
||||
|
||||
// Merge and deduplicate
|
||||
return this._mergeResults(keywordResults, vectorResults, topK);
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Private
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
private _keywordSearch(query: string): Promise<SearchResult[]> {
|
||||
// MemoryStore.search is synchronous but we wrap in promise for parallel use
|
||||
return Promise.resolve(this._memoryStore.search(query));
|
||||
}
|
||||
|
||||
private async _vectorSearch(
|
||||
query: string,
|
||||
topK: number,
|
||||
): Promise<HybridSearchResult[]> {
|
||||
try {
|
||||
const [queryEmbedding] = await this._embeddingProvider.embed([query]);
|
||||
const results = this._vectorStore.search(queryEmbedding, topK);
|
||||
|
||||
return results.map((r) => ({
|
||||
namespace: r.namespace,
|
||||
content: r.chunkText,
|
||||
context: r.chunkText,
|
||||
line: r.startLine,
|
||||
score: r.score,
|
||||
source: 'vector' as const,
|
||||
}));
|
||||
} catch (error) {
|
||||
// Vector search failure should not break search entirely
|
||||
console.error('Vector search failed, falling back to keyword only:', error);
|
||||
return [];
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Merge keyword and vector results with deduplication.
|
||||
*
|
||||
* Deduplication: two results are considered duplicates if they share the
|
||||
* same namespace and their line numbers are within 3 lines of each other.
|
||||
*/
|
||||
private _mergeResults(
|
||||
keywordResults: SearchResult[],
|
||||
vectorResults: HybridSearchResult[],
|
||||
topK: number,
|
||||
): HybridSearchResult[] {
|
||||
// Normalise keyword scores: assign rank-based scores (best match = 1.0)
|
||||
const maxKeyword = keywordResults.length;
|
||||
const keywordScored: HybridSearchResult[] = keywordResults.map((r, idx) => ({
|
||||
namespace: r.namespace,
|
||||
content: r.content,
|
||||
context: r.context,
|
||||
line: r.line,
|
||||
score: maxKeyword > 0 ? 1 - idx / (maxKeyword + 1) : 0,
|
||||
source: 'keyword' as const,
|
||||
}));
|
||||
|
||||
// Build a combined map keyed by namespace + approximate line
|
||||
const resultMap = new Map<string, HybridSearchResult>();
|
||||
|
||||
// Key function: group results within LINE_PROXIMITY lines together
|
||||
const LINE_PROXIMITY = 3;
|
||||
const makeKey = (namespace: string, line: number): string => {
|
||||
const bucket = Math.floor(line / LINE_PROXIMITY);
|
||||
return `${namespace}:${bucket}`;
|
||||
};
|
||||
|
||||
// Add keyword results first
|
||||
for (const kr of keywordScored) {
|
||||
const key = makeKey(kr.namespace, kr.line);
|
||||
const existing = resultMap.get(key);
|
||||
if (existing) {
|
||||
// Combine scores
|
||||
existing.score = (this._hybridWeight * (existing.source === 'vector' || existing.source === 'both' ? existing.score : 0))
|
||||
+ ((1 - this._hybridWeight) * kr.score);
|
||||
existing.source = 'both';
|
||||
// Prefer the more specific keyword content
|
||||
existing.content = kr.content;
|
||||
existing.context = kr.context;
|
||||
existing.line = kr.line;
|
||||
} else {
|
||||
resultMap.set(key, {
|
||||
...kr,
|
||||
score: (1 - this._hybridWeight) * kr.score,
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
// Add/merge vector results
|
||||
for (const vr of vectorResults) {
|
||||
const key = makeKey(vr.namespace, vr.line);
|
||||
const existing = resultMap.get(key);
|
||||
if (existing) {
|
||||
if (existing.source === 'keyword') {
|
||||
existing.score = (this._hybridWeight * vr.score) + existing.score;
|
||||
existing.source = 'both';
|
||||
}
|
||||
// If already 'both' or 'vector', keep the higher-scoring version
|
||||
} else {
|
||||
resultMap.set(key, {
|
||||
...vr,
|
||||
score: this._hybridWeight * vr.score,
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
// Sort by score descending, return top K
|
||||
const merged = Array.from(resultMap.values());
|
||||
merged.sort((a, b) => b.score - a.score);
|
||||
return merged.slice(0, topK);
|
||||
}
|
||||
}
|
||||
@@ -1,2 +1,10 @@
|
||||
export { MemoryStore } from './store.js';
|
||||
export type { MemoryStoreConfig, SearchResult } from './store.js';
|
||||
export { chunkText } from './chunker.js';
|
||||
export type { Chunk, ChunkOptions } from './chunker.js';
|
||||
export { createEmbeddingProvider, OpenAIEmbeddingProvider, GeminiEmbeddingProvider, OllamaEmbeddingProvider, LlamaCppEmbeddingProvider } from './embeddings.js';
|
||||
export type { EmbeddingProvider } from './embeddings.js';
|
||||
export { VectorStore, cosineSimilarity, contentHash } from './vector-store.js';
|
||||
export type { VectorSearchResult, EmbeddingRow } from './vector-store.js';
|
||||
export { HybridSearch } from './hybrid-search.js';
|
||||
export type { HybridSearchResult } from './hybrid-search.js';
|
||||
|
||||
@@ -40,6 +40,7 @@ export interface SearchResult {
|
||||
*/
|
||||
export class MemoryStore {
|
||||
private _config: MemoryStoreConfig;
|
||||
private _dirtyNamespaces: Set<string> = new Set();
|
||||
|
||||
constructor(config: MemoryStoreConfig) {
|
||||
this._config = config;
|
||||
@@ -88,6 +89,9 @@ export class MemoryStore {
|
||||
const separator = existing.length > 0 ? '\n' : '';
|
||||
writeFileSync(filePath, existing + separator + content, 'utf-8');
|
||||
}
|
||||
|
||||
// Mark namespace as needing re-indexing for vector search
|
||||
this._dirtyNamespaces.add(namespace);
|
||||
}
|
||||
|
||||
/**
|
||||
@@ -147,6 +151,25 @@ export class MemoryStore {
|
||||
return this._scanDir(this._config.dir);
|
||||
}
|
||||
|
||||
/**
|
||||
* Return namespaces that have been modified since last call, then clear
|
||||
* the dirty set. Used by the background indexer to re-embed changed content.
|
||||
*/
|
||||
getDirtyNamespaces(): string[] {
|
||||
const dirty = Array.from(this._dirtyNamespaces);
|
||||
this._dirtyNamespaces.clear();
|
||||
return dirty;
|
||||
}
|
||||
|
||||
/**
|
||||
* Mark all existing namespaces as dirty (e.g. for initial full indexing).
|
||||
*/
|
||||
markAllDirty(): void {
|
||||
for (const ns of this.listNamespaces()) {
|
||||
this._dirtyNamespaces.add(ns);
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Build memory context suitable for injection into a system prompt.
|
||||
*
|
||||
|
||||
@@ -0,0 +1,209 @@
|
||||
import { describe, it, expect, beforeEach, afterEach } from 'vitest';
|
||||
import { VectorStore, cosineSimilarity, contentHash } from './vector-store.js';
|
||||
import type { Chunk } from './chunker.js';
|
||||
import { mkdtempSync, rmSync } from 'fs';
|
||||
import { join } from 'path';
|
||||
import { tmpdir } from 'os';
|
||||
|
||||
describe('cosineSimilarity', () => {
|
||||
it('returns 1 for identical vectors', () => {
|
||||
const v = [1, 2, 3];
|
||||
expect(cosineSimilarity(v, v)).toBeCloseTo(1.0, 5);
|
||||
});
|
||||
|
||||
it('returns 0 for orthogonal vectors', () => {
|
||||
expect(cosineSimilarity([1, 0], [0, 1])).toBeCloseTo(0.0, 5);
|
||||
});
|
||||
|
||||
it('returns -1 for opposite vectors', () => {
|
||||
expect(cosineSimilarity([1, 0], [-1, 0])).toBeCloseTo(-1.0, 5);
|
||||
});
|
||||
|
||||
it('handles Float32Array inputs', () => {
|
||||
const a = new Float32Array([0.5, 0.5]);
|
||||
const b = new Float32Array([0.5, 0.5]);
|
||||
expect(cosineSimilarity(a, b)).toBeCloseTo(1.0, 4);
|
||||
});
|
||||
|
||||
it('throws on dimension mismatch', () => {
|
||||
expect(() => cosineSimilarity([1, 2], [1, 2, 3])).toThrow('dimension mismatch');
|
||||
});
|
||||
|
||||
it('returns 0 for zero vectors', () => {
|
||||
expect(cosineSimilarity([0, 0], [1, 1])).toBe(0);
|
||||
});
|
||||
});
|
||||
|
||||
describe('contentHash', () => {
|
||||
it('returns consistent hash for same input', () => {
|
||||
const hash1 = contentHash('hello world');
|
||||
const hash2 = contentHash('hello world');
|
||||
expect(hash1).toBe(hash2);
|
||||
});
|
||||
|
||||
it('returns different hash for different input', () => {
|
||||
expect(contentHash('hello')).not.toBe(contentHash('world'));
|
||||
});
|
||||
|
||||
it('returns an 8-character hex string', () => {
|
||||
const hash = contentHash('test');
|
||||
expect(hash).toMatch(/^[0-9a-f]{8}$/);
|
||||
});
|
||||
});
|
||||
|
||||
describe('VectorStore', () => {
|
||||
let dir: string;
|
||||
let store: VectorStore;
|
||||
|
||||
beforeEach(() => {
|
||||
dir = mkdtempSync(join(tmpdir(), 'flynn-vector-test-'));
|
||||
store = new VectorStore(join(dir, 'vectors.db'));
|
||||
});
|
||||
|
||||
afterEach(() => {
|
||||
store.close();
|
||||
rmSync(dir, { recursive: true, force: true });
|
||||
});
|
||||
|
||||
const makeChunks = (namespace: string, texts: string[]): Chunk[] =>
|
||||
texts.map((text, i) => ({
|
||||
text,
|
||||
namespace,
|
||||
startLine: i * 10 + 1,
|
||||
endLine: (i + 1) * 10,
|
||||
}));
|
||||
|
||||
const makeFakeEmbeddings = (count: number, dims: number = 4): number[][] =>
|
||||
Array.from({ length: count }, (_, i) => {
|
||||
// Create somewhat distinct embeddings using simple patterns
|
||||
const vec = new Array(dims).fill(0);
|
||||
vec[i % dims] = 1.0;
|
||||
return vec;
|
||||
});
|
||||
|
||||
describe('upsertChunks', () => {
|
||||
it('inserts chunks and embeddings', () => {
|
||||
const chunks = makeChunks('notes', ['chunk one', 'chunk two']);
|
||||
const embeddings = makeFakeEmbeddings(2);
|
||||
|
||||
store.upsertChunks(chunks, embeddings, 'hash1');
|
||||
expect(store.count()).toBe(2);
|
||||
});
|
||||
|
||||
it('replaces existing chunks for same namespace', () => {
|
||||
const chunks1 = makeChunks('notes', ['old chunk']);
|
||||
const embeddings1 = makeFakeEmbeddings(1);
|
||||
store.upsertChunks(chunks1, embeddings1, 'hash1');
|
||||
expect(store.count()).toBe(1);
|
||||
|
||||
const chunks2 = makeChunks('notes', ['new chunk 1', 'new chunk 2']);
|
||||
const embeddings2 = makeFakeEmbeddings(2);
|
||||
store.upsertChunks(chunks2, embeddings2, 'hash2');
|
||||
expect(store.count()).toBe(2);
|
||||
});
|
||||
|
||||
it('keeps other namespaces intact when upserting', () => {
|
||||
store.upsertChunks(
|
||||
makeChunks('ns1', ['chunk 1']),
|
||||
makeFakeEmbeddings(1),
|
||||
'h1',
|
||||
);
|
||||
store.upsertChunks(
|
||||
makeChunks('ns2', ['chunk 2']),
|
||||
makeFakeEmbeddings(1),
|
||||
'h2',
|
||||
);
|
||||
expect(store.count()).toBe(2);
|
||||
|
||||
// Upsert ns1 again
|
||||
store.upsertChunks(
|
||||
makeChunks('ns1', ['new chunk 1']),
|
||||
makeFakeEmbeddings(1),
|
||||
'h3',
|
||||
);
|
||||
expect(store.count()).toBe(2); // ns2 chunk still there
|
||||
});
|
||||
|
||||
it('throws on length mismatch', () => {
|
||||
const chunks = makeChunks('test', ['a', 'b']);
|
||||
const embeddings = makeFakeEmbeddings(1);
|
||||
expect(() => store.upsertChunks(chunks, embeddings, 'h')).toThrow('mismatch');
|
||||
});
|
||||
|
||||
it('does nothing for empty chunks', () => {
|
||||
store.upsertChunks([], [], 'h');
|
||||
expect(store.count()).toBe(0);
|
||||
});
|
||||
});
|
||||
|
||||
describe('deleteNamespace', () => {
|
||||
it('removes all embeddings for a namespace', () => {
|
||||
store.upsertChunks(makeChunks('ns1', ['a']), makeFakeEmbeddings(1), 'h1');
|
||||
store.upsertChunks(makeChunks('ns2', ['b']), makeFakeEmbeddings(1), 'h2');
|
||||
expect(store.count()).toBe(2);
|
||||
|
||||
store.deleteNamespace('ns1');
|
||||
expect(store.count()).toBe(1);
|
||||
});
|
||||
});
|
||||
|
||||
describe('hasContentHash', () => {
|
||||
it('returns true when namespace+hash exists', () => {
|
||||
store.upsertChunks(makeChunks('notes', ['a']), makeFakeEmbeddings(1), 'abc123');
|
||||
expect(store.hasContentHash('notes', 'abc123')).toBe(true);
|
||||
});
|
||||
|
||||
it('returns false for non-matching hash', () => {
|
||||
store.upsertChunks(makeChunks('notes', ['a']), makeFakeEmbeddings(1), 'abc123');
|
||||
expect(store.hasContentHash('notes', 'different')).toBe(false);
|
||||
});
|
||||
|
||||
it('returns false for non-existent namespace', () => {
|
||||
expect(store.hasContentHash('nonexistent', 'abc')).toBe(false);
|
||||
});
|
||||
});
|
||||
|
||||
describe('search', () => {
|
||||
it('returns results sorted by cosine similarity', () => {
|
||||
// Create embeddings where chunk 1 is close to query and chunk 2 is far
|
||||
const chunks = makeChunks('test', ['close match', 'far away']);
|
||||
const embeddings = [
|
||||
[0.9, 0.1, 0.0, 0.0], // close to query
|
||||
[0.0, 0.0, 0.9, 0.1], // far from query
|
||||
];
|
||||
store.upsertChunks(chunks, embeddings, 'h');
|
||||
|
||||
const query = [1.0, 0.0, 0.0, 0.0]; // similar to first chunk
|
||||
const results = store.search(query, 5);
|
||||
|
||||
expect(results.length).toBe(2);
|
||||
expect(results[0].chunkText).toBe('close match');
|
||||
expect(results[0].score).toBeGreaterThan(results[1].score);
|
||||
});
|
||||
|
||||
it('respects topK limit', () => {
|
||||
const chunks = makeChunks('test', ['a', 'b', 'c']);
|
||||
const embeddings = makeFakeEmbeddings(3);
|
||||
store.upsertChunks(chunks, embeddings, 'h');
|
||||
|
||||
const results = store.search([1, 0, 0, 0], 2);
|
||||
expect(results.length).toBe(2);
|
||||
});
|
||||
|
||||
it('returns empty array when no embeddings exist', () => {
|
||||
const results = store.search([1, 0, 0, 0], 5);
|
||||
expect(results).toEqual([]);
|
||||
});
|
||||
|
||||
it('includes namespace and line info in results', () => {
|
||||
const chunks = makeChunks('sessions/abc', ['test chunk']);
|
||||
const embeddings = [[1, 0, 0, 0]];
|
||||
store.upsertChunks(chunks, embeddings, 'h');
|
||||
|
||||
const results = store.search([1, 0, 0, 0], 1);
|
||||
expect(results[0].namespace).toBe('sessions/abc');
|
||||
expect(results[0].startLine).toBe(1);
|
||||
expect(results[0].endLine).toBe(10);
|
||||
});
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,232 @@
|
||||
/**
|
||||
* SQLite-backed vector storage for embedding chunks.
|
||||
* Uses better-sqlite3 for synchronous operations and stores
|
||||
* embeddings as Float32Array BLOBs.
|
||||
*/
|
||||
|
||||
import Database from 'better-sqlite3';
|
||||
import type { Chunk } from './chunker.js';
|
||||
|
||||
/**
|
||||
* A single row from the embeddings table.
|
||||
*/
|
||||
export interface EmbeddingRow {
|
||||
id: number;
|
||||
namespace: string;
|
||||
chunk_text: string;
|
||||
start_line: number;
|
||||
end_line: number;
|
||||
embedding: Buffer;
|
||||
created_at: string;
|
||||
content_hash: string;
|
||||
}
|
||||
|
||||
/**
|
||||
* A search result from the vector store.
|
||||
*/
|
||||
export interface VectorSearchResult {
|
||||
namespace: string;
|
||||
chunkText: string;
|
||||
startLine: number;
|
||||
endLine: number;
|
||||
score: number;
|
||||
}
|
||||
|
||||
/**
|
||||
* Compute cosine similarity between two vectors.
|
||||
* Returns a value between -1 and 1 (1 = identical direction).
|
||||
*/
|
||||
export function cosineSimilarity(a: number[] | Float32Array, b: number[] | Float32Array): number {
|
||||
if (a.length !== b.length) {
|
||||
throw new Error(`Vector dimension mismatch: ${a.length} vs ${b.length}`);
|
||||
}
|
||||
|
||||
let dotProduct = 0;
|
||||
let normA = 0;
|
||||
let normB = 0;
|
||||
|
||||
for (let i = 0; i < a.length; i++) {
|
||||
dotProduct += a[i] * b[i];
|
||||
normA += a[i] * a[i];
|
||||
normB += b[i] * b[i];
|
||||
}
|
||||
|
||||
const magnitude = Math.sqrt(normA) * Math.sqrt(normB);
|
||||
if (magnitude === 0) return 0;
|
||||
|
||||
return dotProduct / magnitude;
|
||||
}
|
||||
|
||||
/**
|
||||
* Simple content hash for change detection.
|
||||
* Uses a fast string hash rather than crypto for speed.
|
||||
*/
|
||||
export function contentHash(text: string): string {
|
||||
// FNV-1a 32-bit hash
|
||||
let hash = 0x811c9dc5;
|
||||
for (let i = 0; i < text.length; i++) {
|
||||
hash ^= text.charCodeAt(i);
|
||||
hash = Math.imul(hash, 0x01000193);
|
||||
}
|
||||
return (hash >>> 0).toString(16).padStart(8, '0');
|
||||
}
|
||||
|
||||
/**
|
||||
* SQLite-backed vector store that persists embedding chunks.
|
||||
*/
|
||||
export class VectorStore {
|
||||
private _db: Database.Database;
|
||||
|
||||
constructor(dbPath: string) {
|
||||
this._db = new Database(dbPath);
|
||||
this._db.pragma('journal_mode = WAL');
|
||||
this._initSchema();
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Schema
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
private _initSchema(): void {
|
||||
this._db.exec(`
|
||||
CREATE TABLE IF NOT EXISTS embeddings (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
namespace TEXT NOT NULL,
|
||||
chunk_text TEXT NOT NULL,
|
||||
start_line INTEGER NOT NULL,
|
||||
end_line INTEGER NOT NULL,
|
||||
embedding BLOB NOT NULL,
|
||||
created_at TEXT NOT NULL DEFAULT (datetime('now')),
|
||||
content_hash TEXT NOT NULL
|
||||
);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS idx_embeddings_namespace
|
||||
ON embeddings (namespace);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS idx_embeddings_content_hash
|
||||
ON embeddings (namespace, content_hash);
|
||||
`);
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Public API
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
/**
|
||||
* Upsert chunks for a namespace.
|
||||
* Deletes all existing embeddings for the namespace, then inserts new ones.
|
||||
*
|
||||
* @param chunks - The text chunks with metadata.
|
||||
* @param embeddings - Corresponding embedding vectors (one per chunk).
|
||||
* @param hash - Content hash for the entire namespace content.
|
||||
*/
|
||||
upsertChunks(
|
||||
chunks: Chunk[],
|
||||
embeddings: number[][],
|
||||
hash: string,
|
||||
): void {
|
||||
if (chunks.length !== embeddings.length) {
|
||||
throw new Error(`Chunks/embeddings length mismatch: ${chunks.length} vs ${embeddings.length}`);
|
||||
}
|
||||
|
||||
if (chunks.length === 0) return;
|
||||
|
||||
const namespace = chunks[0].namespace;
|
||||
|
||||
const insertStmt = this._db.prepare(`
|
||||
INSERT INTO embeddings (namespace, chunk_text, start_line, end_line, embedding, content_hash)
|
||||
VALUES (?, ?, ?, ?, ?, ?)
|
||||
`);
|
||||
|
||||
const deleteStmt = this._db.prepare(`
|
||||
DELETE FROM embeddings WHERE namespace = ?
|
||||
`);
|
||||
|
||||
const transaction = this._db.transaction(() => {
|
||||
deleteStmt.run(namespace);
|
||||
for (let i = 0; i < chunks.length; i++) {
|
||||
const chunk = chunks[i];
|
||||
const embeddingBuffer = Buffer.from(new Float32Array(embeddings[i]).buffer);
|
||||
insertStmt.run(
|
||||
chunk.namespace,
|
||||
chunk.text,
|
||||
chunk.startLine,
|
||||
chunk.endLine,
|
||||
embeddingBuffer,
|
||||
hash,
|
||||
);
|
||||
}
|
||||
});
|
||||
|
||||
transaction();
|
||||
}
|
||||
|
||||
/**
|
||||
* Delete all embeddings for a namespace.
|
||||
*/
|
||||
deleteNamespace(namespace: string): void {
|
||||
this._db.prepare('DELETE FROM embeddings WHERE namespace = ?').run(namespace);
|
||||
}
|
||||
|
||||
/**
|
||||
* Check if a namespace already has embeddings with the given content hash.
|
||||
* Used to skip re-embedding unchanged content.
|
||||
*/
|
||||
hasContentHash(namespace: string, hash: string): boolean {
|
||||
const row = this._db.prepare(
|
||||
'SELECT 1 FROM embeddings WHERE namespace = ? AND content_hash = ? LIMIT 1',
|
||||
).get(namespace, hash) as { '1': number } | undefined;
|
||||
return row !== undefined;
|
||||
}
|
||||
|
||||
/**
|
||||
* Search for the nearest vectors to the query embedding.
|
||||
* Uses brute-force cosine similarity (suitable for moderate dataset sizes).
|
||||
*
|
||||
* @param queryEmbedding - The query vector.
|
||||
* @param topK - Maximum number of results to return.
|
||||
* @returns Ranked results sorted by descending similarity score.
|
||||
*/
|
||||
search(queryEmbedding: number[], topK: number): VectorSearchResult[] {
|
||||
const rows = this._db.prepare(
|
||||
'SELECT namespace, chunk_text, start_line, end_line, embedding FROM embeddings',
|
||||
).all() as Pick<EmbeddingRow, 'namespace' | 'chunk_text' | 'start_line' | 'end_line' | 'embedding'>[];
|
||||
|
||||
const queryArray = new Float32Array(queryEmbedding);
|
||||
|
||||
const scored: VectorSearchResult[] = rows.map((row) => {
|
||||
const stored = new Float32Array(
|
||||
row.embedding.buffer,
|
||||
row.embedding.byteOffset,
|
||||
row.embedding.byteLength / Float32Array.BYTES_PER_ELEMENT,
|
||||
);
|
||||
|
||||
return {
|
||||
namespace: row.namespace,
|
||||
chunkText: row.chunk_text,
|
||||
startLine: row.start_line,
|
||||
endLine: row.end_line,
|
||||
score: cosineSimilarity(queryArray, stored),
|
||||
};
|
||||
});
|
||||
|
||||
// Sort by descending score and return top K
|
||||
scored.sort((a, b) => b.score - a.score);
|
||||
return scored.slice(0, topK);
|
||||
}
|
||||
|
||||
/**
|
||||
* Get total number of stored embeddings.
|
||||
*/
|
||||
count(): number {
|
||||
const row = this._db.prepare('SELECT COUNT(*) as cnt FROM embeddings').get() as { cnt: number };
|
||||
return row.cnt;
|
||||
}
|
||||
|
||||
/**
|
||||
* Close the database connection.
|
||||
*/
|
||||
close(): void {
|
||||
this._db.close();
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user