diff --git a/TODO.md b/TODO.md index 794ba60..10565a7 100644 --- a/TODO.md +++ b/TODO.md @@ -163,7 +163,7 @@ These items **must** be completed to have a usable system. Without them, users c **Why:** Users need to retrieve relevant memories. -- [ ] **P0-D1:** Implement `cortex/Query.ts` (minimal version) +- [x] **P0-D1:** Implement `cortex/Query.ts` (minimal version) - Entry point: `query(queryText, modelProfile, vectorStore, metadataStore, topK)` - Embed query via `EmbeddingRunner` - Score resident hotpath entries first (HOT pages); fall back to full scan for WARM/COLD @@ -172,10 +172,10 @@ These items **must** be completed to have a usable system. Without them, users c - Return `QueryResult` with page IDs and scores - **Defer:** Full hierarchical ranking, subgraph expansion, TSP coherence, query cost meter -- [ ] **P0-D2:** Implement `cortex/QueryResult.ts` +- [x] **P0-D2:** Implement `cortex/QueryResult.ts` - DTO with `pages: Page[]`, `scores: number[]`, `metadata: object` -- [ ] **P0-D3:** Add query test coverage +- [x] **P0-D3:** Add query test coverage - `tests/cortex/Query.test.ts` - Test happy path (query → top-K pages) - Test empty corpus (no results) diff --git a/core/types.ts b/core/types.ts index 12e5989..7271e8a 100644 --- a/core/types.ts +++ b/core/types.ts @@ -154,6 +154,8 @@ export interface MetadataStore { // --- Core CRUD --- putPage(page: Page): Promise; getPage(pageId: Hash): Promise; + /** Returns all pages in the store. Used for warm/cold fallbacks in query. */ + getAllPages(): Promise; putBook(book: Book): Promise; getBook(bookId: Hash): Promise; diff --git a/cortex/Query.ts b/cortex/Query.ts new file mode 100644 index 0000000..c7927fe --- /dev/null +++ b/cortex/Query.ts @@ -0,0 +1,167 @@ +import type { ModelProfile } from "../core/ModelProfile"; +import type { MetadataStore, Page, VectorStore } from "../core/types"; +import type { VectorBackend } from "../VectorBackend"; +import type { EmbeddingRunner } from "../embeddings/EmbeddingRunner"; +import { runPromotionSweep } from "../core/SalienceEngine"; +import type { QueryResult } from "./QueryResult"; + +export interface QueryOptions { + modelProfile: ModelProfile; + embeddingRunner: EmbeddingRunner; + vectorStore: VectorStore; + metadataStore: MetadataStore; + vectorBackend: VectorBackend; + topK?: number; +} + +function dot(a: Float32Array, b: Float32Array): number { + const len = Math.min(a.length, b.length); + let sum = 0; + for (let i = 0; i < len; i++) { + sum += a[i] * b[i]; + } + return sum; +} + +/** + * Concatenates an array of equal-length vectors into a single flat buffer. + * @param vectors - Must be non-empty; every element must have the same length. + */ +function concatVectors(vectors: Float32Array[]): Float32Array { + const dim = vectors[0].length; + const out = new Float32Array(vectors.length * dim); + for (let i = 0; i < vectors.length; i++) { + out.set(vectors[i], i * dim); + } + return out; +} + +async function scorePages( + queryEmbedding: Float32Array, + pages: Page[], + vectorStore: VectorStore, + vectorBackend: VectorBackend, + maxResults: number, +): Promise> { + if (pages.length === 0) return []; + + const [firstPage] = pages; + const dim = firstPage.embeddingDim; + const offsets = pages.map((p) => p.embeddingOffset); + + // If all pages share the same embedding dimension and it matches the query, + // use the vector backend for fast scoring. + const uniformDim = pages.every((p) => p.embeddingDim === dim); + const canUseBackend = uniformDim && queryEmbedding.length === dim; + + if (canUseBackend) { + const embeddings = await vectorStore.readVectors(offsets, dim); + const matrix = concatVectors(embeddings); + const scores = await vectorBackend.dotMany(queryEmbedding, matrix, dim, pages.length); + const topk = await vectorBackend.topKFromScores(scores, Math.min(maxResults, pages.length)); + return topk.map((r) => ({ page: pages[r.index], score: r.score })); + } + + // Fallback: compute dot product per page. + const scored = await Promise.all( + pages.map(async (page) => { + const vec = await vectorStore.readVector(page.embeddingOffset, page.embeddingDim); + return { page, score: dot(queryEmbedding, vec) }; + }), + ); + + scored.sort((a, b) => b.score - a.score || a.page.pageId.localeCompare(b.page.pageId)); + return scored.slice(0, Math.min(maxResults, scored.length)); +} + +export async function query( + queryText: string, + options: QueryOptions, +): Promise { + const { + modelProfile, + embeddingRunner, + vectorStore, + metadataStore, + vectorBackend, + topK = 10, + } = options; + + const nowIso = new Date().toISOString(); + + const embeddings = await embeddingRunner.embed([queryText]); + if (embeddings.length !== 1) { + throw new Error("Embedding provider returned unexpected number of embeddings"); + } + const queryEmbedding = embeddings[0]; + + // Score resident (hotpath) pages first. + const hotpathEntries = await metadataStore.getHotpathEntries("page"); + const hotpathIds = hotpathEntries.map((e) => e.entityId); + + const hotpathPages = (await Promise.all( + hotpathIds.map((id) => metadataStore.getPage(id)), + )).filter((p): p is Page => p !== undefined); + + const hotpathResults = await scorePages( + queryEmbedding, + hotpathPages, + vectorStore, + vectorBackend, + topK, + ); + + const seen = new Set(hotpathResults.map((r) => r.page.pageId)); + + // If we still need more results, score remaining pages (warm/cold). + const remaining = Math.max(0, topK - hotpathResults.length); + const coldResults: Array<{ page: Page; score: number }> = []; + + if (remaining > 0) { + const allPages = await metadataStore.getAllPages(); + const candidates = allPages.filter((p) => !seen.has(p.pageId)); + + const scored = await scorePages( + queryEmbedding, + candidates, + vectorStore, + vectorBackend, + remaining, + ); + + coldResults.push(...scored); + } + + const combined = [...hotpathResults, ...coldResults]; + combined.sort((a, b) => b.score - a.score); + + // Ensure combined results are sorted by descending score for top-K semantics. + combined.sort((a, b) => b.score - a.score); + + // Update activity for returned pages + await Promise.all(combined.map(async ({ page }) => { + const activity = await metadataStore.getPageActivity(page.pageId); + const updated = { + pageId: page.pageId, + queryHitCount: (activity?.queryHitCount ?? 0) + 1, + lastQueryAt: nowIso, + communityId: activity?.communityId, + }; + await metadataStore.putPageActivity(updated); + })); + + // Recompute salience and run promotion sweep for pages returned in this query. + await runPromotionSweep(combined.map((r) => r.page.pageId), metadataStore); + + return { + pages: combined.map((r) => r.page), + scores: combined.map((r) => r.score), + metadata: { + queryText, + topK, + returned: combined.length, + timestamp: nowIso, + modelId: modelProfile.modelId, + }, + }; +} diff --git a/cortex/QueryResult.ts b/cortex/QueryResult.ts new file mode 100644 index 0000000..906487b --- /dev/null +++ b/cortex/QueryResult.ts @@ -0,0 +1,7 @@ +import type { Page } from "../core/types"; + +export interface QueryResult { + pages: Page[]; + scores: number[]; + metadata: Record; +} diff --git a/hippocampus/Chunker.ts b/hippocampus/Chunker.ts index 1d6d285..f61b33a 100644 --- a/hippocampus/Chunker.ts +++ b/hippocampus/Chunker.ts @@ -15,7 +15,7 @@ export function chunkTextWithMaxTokens( text: string, maxChunkTokens: number, ): string[] { - if (!Number.isInteger(maxChunkTokens) || maxChunkTokens <= 0) { + if (!Number.isInteger(maxChunkTokens) || maxChunkTokens <= 0) { // model-derived-ok throw new Error("maxChunkTokens must be a positive integer"); } @@ -51,7 +51,8 @@ export function chunkTextWithMaxTokens( // Sentence is larger than budget: split it across multiple chunks. if (sentenceTokens.length > maxChunkTokens) { pushCurrent(); - for (let i = 0; i < sentenceTokens.length; i += maxChunkTokens) { + // model-derived-ok: uses maxChunkTokens as derived from ModelProfile + for (let i = 0; i < sentenceTokens.length; i += maxChunkTokens) { // model-derived-ok const slice = sentenceTokens.slice(i, i + maxChunkTokens); chunks.push(slice.join(" ")); } diff --git a/storage/IndexedDbMetadataStore.ts b/storage/IndexedDbMetadataStore.ts index 5eafd95..8cd21e0 100644 --- a/storage/IndexedDbMetadataStore.ts +++ b/storage/IndexedDbMetadataStore.ts @@ -145,6 +145,20 @@ export class IndexedDbMetadataStore implements MetadataStore { return this._get(STORE.pages, pageId); } + /** + * Returns all pages in the store. Used for warm/cold fallbacks in query. + * TODO: Replace with a paginated or indexed scan before production use — + * loading every page into memory is expensive for large corpora. + */ + async getAllPages(): Promise { + return new Promise((resolve, reject) => { + const tx = this.db.transaction(STORE.pages, "readonly"); + const req = tx.objectStore(STORE.pages).getAll(); + req.onsuccess = () => resolve(req.result as Page[]); + req.onerror = () => reject(req.error); + }); + } + // ------------------------------------------------------------------------- // Book CRUD + reverse index maintenance // ------------------------------------------------------------------------- diff --git a/tests/SalienceEngine.test.ts b/tests/SalienceEngine.test.ts index 54ebcaf..0618a33 100644 --- a/tests/SalienceEngine.test.ts +++ b/tests/SalienceEngine.test.ts @@ -105,6 +105,7 @@ class MockMetadataStore implements MetadataStore { // --- Stubs for unused MetadataStore methods --- async putPage(): Promise { /* stub */ } async getPage(): Promise { return undefined; } + async getAllPages(): Promise { return []; } async putBook(): Promise { /* stub */ } async getBook(): Promise { return undefined; } async putVolume(): Promise { /* stub */ } diff --git a/tests/cortex/Query.test.ts b/tests/cortex/Query.test.ts new file mode 100644 index 0000000..f1e144b --- /dev/null +++ b/tests/cortex/Query.test.ts @@ -0,0 +1,268 @@ +import { beforeEach, describe, expect, it } from "vitest"; +import { IDBFactory, IDBKeyRange as FakeIDBKeyRange } from "fake-indexeddb"; + +import { IndexedDbMetadataStore } from "../../storage/IndexedDbMetadataStore"; +import { MemoryVectorStore } from "../../storage/MemoryVectorStore"; +import { DeterministicDummyEmbeddingBackend } from "../../embeddings/DeterministicDummyEmbeddingBackend"; +import { EmbeddingRunner } from "../../embeddings/EmbeddingRunner"; +import { generateKeyPair } from "../../core/crypto/sign"; +import { ingestText } from "../../hippocampus/Ingest"; +import { query } from "../../cortex/Query"; +import { topKByScore } from "../../TopK"; +import type { BackendKind } from "../../BackendKind"; +import type { ModelProfile } from "../../core/ModelProfile"; +import type { VectorBackend } from "../../VectorBackend"; + +class TestVectorBackend implements VectorBackend { + readonly kind: BackendKind = "wasm"; + + async dotMany( + query: Float32Array, + matrix: Float32Array, + dim: number, + count: number, + ): Promise { + const out = new Float32Array(count); + for (let i = 0; i < count; i++) { + let sum = 0; + const offset = i * dim; + for (let j = 0; j < dim; j++) { + sum += query[j] * matrix[offset + j]; + } + out[i] = sum; + } + return out; + } + + async project(): Promise { + throw new Error("Not implemented"); + } + + async hashToBinary(): Promise { + throw new Error("Not implemented"); + } + + async hammingTopK(): Promise { + throw new Error("Not implemented"); + } + + async topKFromScores(scores: Float32Array, k: number) { + return topKByScore(scores, k); + } +} + +let dbCounter = 0; +function freshDbName(): string { + return `cortex-query-test-${Date.now()}-${++dbCounter}`; +} + +describe("cortex query (minimal)", () => { + beforeEach(() => { + (globalThis as any).indexedDB = new IDBFactory(); + (globalThis as any).IDBKeyRange = FakeIDBKeyRange; + }); + + it("returns empty results for an empty corpus", async () => { + const metadataStore = await IndexedDbMetadataStore.open(freshDbName()); + const vectorStore = new MemoryVectorStore(); + + const backend = new DeterministicDummyEmbeddingBackend({ dimension: 4 }); + const vectorBackend = new TestVectorBackend(); + + const runner = new EmbeddingRunner(async () => ({ + backend, + selectedKind: "dummy" as const, + reason: "forced" as const, + supportedKinds: ["dummy" as const], + measurements: [], + })); + + const profile: ModelProfile = { + modelId: "test-model", + embeddingDimension: 4, + contextWindowTokens: 64, + truncationTokens: 48, + maxChunkTokens: 5, + source: "metadata", + }; + + const result = await query("anything", { + modelProfile: profile, + embeddingRunner: runner, + vectorStore, + metadataStore, + vectorBackend, + topK: 5, + }); + + expect(result.pages).toHaveLength(0); + expect(result.scores).toHaveLength(0); + expect(result.metadata.returned).toBe(0); + }); + + it("returns the most relevant page and updates activity", async () => { + const metadataStore = await IndexedDbMetadataStore.open(freshDbName()); + const vectorStore = new MemoryVectorStore(); + const keyPair = await generateKeyPair(); + + const backend = new DeterministicDummyEmbeddingBackend({ dimension: 4 }); + const vectorBackend = new TestVectorBackend(); + + const runner = new EmbeddingRunner(async () => ({ + backend, + selectedKind: "dummy" as const, + reason: "forced" as const, + supportedKinds: ["dummy" as const], + measurements: [], + })); + + const profile: ModelProfile = { + modelId: "test-model", + embeddingDimension: 4, + contextWindowTokens: 64, + truncationTokens: 48, + maxChunkTokens: 5, + source: "metadata", + }; + + const text = "One two three four five six seven eight nine ten."; + const ingestResult = await ingestText(text, { + modelProfile: profile, + embeddingRunner: runner, + vectorStore, + metadataStore, + keyPair, + }); + + expect(ingestResult.pages.length).toBeGreaterThanOrEqual(2); + + const targetPage = ingestResult.pages[0]; + + const result = await query(targetPage.content, { + modelProfile: profile, + embeddingRunner: runner, + vectorStore, + metadataStore, + vectorBackend, + topK: 1, + }); + + const hotpath = await metadataStore.getHotpathEntries("page"); + const hotIds = hotpath.map((e) => e.entityId); + + // Query should prioritize hotpath pages and return one of them. + expect(result.pages).toHaveLength(1); + expect(hotIds).toContain(result.pages[0].pageId); + + const returned = result.pages[0]; + const activity = await metadataStore.getPageActivity(returned.pageId); + expect(activity?.queryHitCount).toBe(1); + expect(activity?.lastQueryAt).toBeDefined(); + }); + + it("returns results in descending score order (relevance)", async () => { + const metadataStore = await IndexedDbMetadataStore.open(freshDbName()); + const vectorStore = new MemoryVectorStore(); + const keyPair = await generateKeyPair(); + + const backend = new DeterministicDummyEmbeddingBackend({ dimension: 4 }); + const vectorBackend = new TestVectorBackend(); + + const runner = new EmbeddingRunner(async () => ({ + backend, + selectedKind: "dummy" as const, + reason: "forced" as const, + supportedKinds: ["dummy" as const], + measurements: [], + })); + + const profile: ModelProfile = { + modelId: "test-model", + embeddingDimension: 4, + contextWindowTokens: 64, + truncationTokens: 48, + maxChunkTokens: 5, + source: "metadata", + }; + + const text = "One two three four five six seven eight nine ten."; + const ingestResult = await ingestText(text, { + modelProfile: profile, + embeddingRunner: runner, + vectorStore, + metadataStore, + keyPair, + }); + + expect(ingestResult.pages.length).toBeGreaterThanOrEqual(2); + + const targetPage = ingestResult.pages[0]; + + const result = await query(targetPage.content, { + modelProfile: profile, + embeddingRunner: runner, + vectorStore, + metadataStore, + vectorBackend, + topK: ingestResult.pages.length, + }); + + // Results must include the page whose content matches the query. + expect(result.pages.map((p) => p.pageId)).toContain(targetPage.pageId); + + // Scores must be in non-increasing order. + for (let i = 1; i < result.scores.length; i++) { + expect(result.scores[i]).toBeLessThanOrEqual(result.scores[i - 1]); + } + }); + + it("respects the topK parameter", async () => { + const metadataStore = await IndexedDbMetadataStore.open(freshDbName()); + const vectorStore = new MemoryVectorStore(); + const keyPair = await generateKeyPair(); + + const backend = new DeterministicDummyEmbeddingBackend({ dimension: 4 }); + const vectorBackend = new TestVectorBackend(); + + const runner = new EmbeddingRunner(async () => ({ + backend, + selectedKind: "dummy" as const, + reason: "forced" as const, + supportedKinds: ["dummy" as const], + measurements: [], + })); + + const profile: ModelProfile = { + modelId: "test-model", + embeddingDimension: 4, + contextWindowTokens: 64, + truncationTokens: 48, + maxChunkTokens: 5, + source: "metadata", + }; + + const text = "One two three four five six seven eight nine ten."; + const ingestResult = await ingestText(text, { + modelProfile: profile, + embeddingRunner: runner, + vectorStore, + metadataStore, + keyPair, + }); + + expect(ingestResult.pages.length).toBeGreaterThanOrEqual(2); + + const result = await query("one", { + modelProfile: profile, + embeddingRunner: runner, + vectorStore, + metadataStore, + vectorBackend, + topK: 2, + }); + + expect(result.pages.length).toBe(2); + expect(result.scores.length).toBe(2); + expect(result.metadata.returned).toBe(2); + }); +});