feat: implement high-order tier vector search, embedding backfill, an…#784
Open
Tanmay-008 wants to merge 1 commit into
Open
feat: implement high-order tier vector search, embedding backfill, an…#784Tanmay-008 wants to merge 1 commit into
Tanmay-008 wants to merge 1 commit into
Conversation
|
@Tanmay-008 is attempting to deploy a commit to the rohitg00's projects Team on Vercel. A member of the Team first needs to authorize it. |
Contributor
📝 WalkthroughWalkthroughThis PR implements searchable high-order memory tiers with vector embeddings, backfill capability, and multi-tier ranked retrieval. It adds embedding generation to consolidation, crystallization, and reflection pipelines; introduces a new searchHighOrderTiers function that ranks candidates via BM25-like term matching and RRF fusion; integrates high-order results into smart-search alongside traditional observations; and provides on-demand backfill via a new REST endpoint. ChangesHigh-Order Tier Search Feature
Estimated code review effort🎯 4 (Complex) | ⏱️ ~60 minutes Possibly related issues
Possibly related PRs
Suggested reviewers
Poem
🚥 Pre-merge checks | ✅ 3 | ❌ 2❌ Failed checks (2 warnings)
✅ Passed checks (3 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 2
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (2)
src/state/vector-index.ts (1)
14-21:⚠️ Potential issue | 🟠 Major | 💤 Low valueValidate
base64ToFloat32input to prevent silent embedding truncation
- In
src/state/vector-index.ts(base64ToFloat32, lines 14-21),new Float32Array(..., buf.byteLength / Float32Array.BYTES_PER_ELEMENT)will coerce a non-integer element count and can return a shorter array without throwing whenbuf.byteLengthisn’t a multiple of 4.- Call sites in
src/functions/high-order-search.tsdecodes.embedding/p.embedding/c.embedding/i.embeddingviabase64ToFloat32(...)(around lines 80, 95, 111, 129) without visible guardrails around the decode.- Add explicit length validation (e.g.,
buf.byteLength % 4 === 0) and have the function/callers reject or skip invalid embeddings instead of proceeding with a truncated vector.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/state/vector-index.ts` around lines 14 - 21, The base64ToFloat32 function can silently produce truncated Float32Arrays when the buffer length isn't a multiple of 4; add an explicit validation in base64ToFloat32 (check buf.byteLength % Float32Array.BYTES_PER_ELEMENT === 0) and have it throw or return a clear failure (e.g., null or throw Error) instead of constructing a possibly shortened array, and then update the call sites in high-order-search.ts where s.embedding, p.embedding, c.embedding, and i.embedding are decoded to detect that failure and skip or reject those invalid embeddings (handle the returned null/exception path so invalid base64 embeddings do not proceed into downstream vector logic).src/mcp/server.ts (1)
264-273:⚠️ Potential issue | 🟡 Minor | ⚡ Quick winHandle boolean
falseforincludeHighOrder.Line 272 only disables on the literal string
"false". If an MCP client sends JSON booleanfalse, this still evaluates totrue, so the new opt-out cannot actually opt out for those callers.Suggested fix
const expandIds = parseCsvList(args.expandIds).slice(0, 20); const limit = Math.max(1, Math.min(100, asNumber(args.limit, 10) ?? 10)); + if ( + args.includeHighOrder !== undefined && + args.includeHighOrder !== true && + args.includeHighOrder !== false && + args.includeHighOrder !== "true" && + args.includeHighOrder !== "false" + ) { + return { + status_code: 400, + body: { error: "includeHighOrder must be a boolean or 'true'/'false'" }, + }; + } + const includeHighOrder = + args.includeHighOrder === undefined + ? true + : args.includeHighOrder === true || args.includeHighOrder === "true"; const result = await sdk.trigger({ function_id: "mem::smart-search", payload: { query: args.query, expandIds, limit, - includeHighOrder: args.includeHighOrder !== "false", + includeHighOrder, }, });As per coding guidelines,
src/mcp/server.ts: "MCP tool handler implementation must validate args with typeof checks".🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/mcp/server.ts` around lines 264 - 273, The includeHighOrder flag currently only checks for the string "false" which misses JSON boolean false; update the evaluation in the sdk.trigger payload (near function_id "mem::smart-search" in src/mcp/server.ts) to validate the arg with typeof checks and treat both the boolean false and the string "false" as opt-outs (e.g., set includeHighOrder to false when args.includeHighOrder === false || args.includeHighOrder === "false", otherwise true). Ensure you reference and adjust the use of args.includeHighOrder alongside parseCsvList and asNumber handling so the handler follows the MCP args validation guideline.
🧹 Nitpick comments (2)
src/functions/smart-search.ts (1)
153-154: 💤 Low valueSimplify the defensive Array.isArray check.
searchHighOrderTiersalways returns{ results: CompactHighOrderResult[]; needsBackfill: boolean }, so theArray.isArraycheck is unnecessary. Consider simplifying:-const highOrderResults = Array.isArray(highOrderResponse) ? highOrderResponse : highOrderResponse.results; -const needsBackfill = Array.isArray(highOrderResponse) ? false : highOrderResponse.needsBackfill; +const { results: highOrderResults, needsBackfill } = highOrderResponse;🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/functions/smart-search.ts` around lines 153 - 154, Remove the redundant Array.isArray defensive checks around highOrderResponse since searchHighOrderTiers always returns an object; replace the conditional assignments with direct property access: set highOrderResults = highOrderResponse.results and needsBackfill = highOrderResponse.needsBackfill. Update any related assumptions in the function smart-search.ts (references: highOrderResults, needsBackfill, searchHighOrderTiers, highOrderResponse) and remove the unnecessary Array.isArray branches.test/high-order-search.test.ts (1)
95-353: ⚡ Quick winConsider adding test coverage for the
needsBackfillflag.The PR objectives highlight backfill triggering as a key feature ("smart-search can trigger backfill background tasks when embeddings are missing/outdated"), but the current tests don't verify the
needsBackfillflag behavior. Consider adding tests for:
needsBackfill: truewhen candidates have noembeddingfield.needsBackfill: truewhen candidates haveembeddingbutembeddingModeldiffers from the active provider'sname.needsBackfill: falsewhen all candidates have embeddings matching the active provider model.This would strengthen confidence that the backfill-triggering integration works as designed.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@test/high-order-search.test.ts` around lines 95 - 353, Add unit tests covering the needsBackfill flag for searchHighOrderTiers: create cases using mockKV and helper factories (makeSemantic/makeProcedural/makeCrystal/makeInsight) that assert needsBackfill === true when stored candidates lack an embedding, assert needsBackfill === true when candidates have an embedding but embeddingModel !== active provider name, and assert needsBackfill === false when all candidates include embeddings with embeddingModel equal to the active provider name; reuse existing patterns in the test file (kv := mockKV(), await kv.set(...), call searchHighOrderTiers(...)) and assert the returned object has the expected needsBackfill boolean alongside result checks.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@src/functions/high-order-backfill.ts`:
- Around line 11-121: The backfill function registerHighOrderBackfillFunction
updates multiple KV scopes but never creates an audit trail; update the function
to call recordAudit() after each successful scope backfill (or once after all
scopes complete) including scope name, count backfilled, embeddingModel
(ep.name) and timestamp so changes are recorded; specifically add calls to
recordAudit(...) after the semantic, procedural, crystals and insights loops (or
one consolidated recordAudit at the end using the results object) and ensure
errors still return without emitting a success audit.
In `@src/triggers/api.ts`:
- Around line 1110-1122: Validate and whitelist each expected field from
req.body before calling sdk.trigger: ensure query is a string, expandIds is an
array of strings (or accept a comma-separated string and split/coerce), limit is
a number (parse numeric strings to Number and ignore non-numeric), project and
agentId are strings, and includeLessons/includeHighOrder are booleans; only
include each key in the payload if it passes its type check (otherwise omit or
coerce safely) when constructing the object passed to sdk.trigger so you never
forward raw, unvalidated values from req.body.
---
Outside diff comments:
In `@src/mcp/server.ts`:
- Around line 264-273: The includeHighOrder flag currently only checks for the
string "false" which misses JSON boolean false; update the evaluation in the
sdk.trigger payload (near function_id "mem::smart-search" in src/mcp/server.ts)
to validate the arg with typeof checks and treat both the boolean false and the
string "false" as opt-outs (e.g., set includeHighOrder to false when
args.includeHighOrder === false || args.includeHighOrder === "false", otherwise
true). Ensure you reference and adjust the use of args.includeHighOrder
alongside parseCsvList and asNumber handling so the handler follows the MCP args
validation guideline.
In `@src/state/vector-index.ts`:
- Around line 14-21: The base64ToFloat32 function can silently produce truncated
Float32Arrays when the buffer length isn't a multiple of 4; add an explicit
validation in base64ToFloat32 (check buf.byteLength %
Float32Array.BYTES_PER_ELEMENT === 0) and have it throw or return a clear
failure (e.g., null or throw Error) instead of constructing a possibly shortened
array, and then update the call sites in high-order-search.ts where s.embedding,
p.embedding, c.embedding, and i.embedding are decoded to detect that failure and
skip or reject those invalid embeddings (handle the returned null/exception path
so invalid base64 embeddings do not proceed into downstream vector logic).
---
Nitpick comments:
In `@src/functions/smart-search.ts`:
- Around line 153-154: Remove the redundant Array.isArray defensive checks
around highOrderResponse since searchHighOrderTiers always returns an object;
replace the conditional assignments with direct property access: set
highOrderResults = highOrderResponse.results and needsBackfill =
highOrderResponse.needsBackfill. Update any related assumptions in the function
smart-search.ts (references: highOrderResults, needsBackfill,
searchHighOrderTiers, highOrderResponse) and remove the unnecessary
Array.isArray branches.
In `@test/high-order-search.test.ts`:
- Around line 95-353: Add unit tests covering the needsBackfill flag for
searchHighOrderTiers: create cases using mockKV and helper factories
(makeSemantic/makeProcedural/makeCrystal/makeInsight) that assert needsBackfill
=== true when stored candidates lack an embedding, assert needsBackfill === true
when candidates have an embedding but embeddingModel !== active provider name,
and assert needsBackfill === false when all candidates include embeddings with
embeddingModel equal to the active provider name; reuse existing patterns in the
test file (kv := mockKV(), await kv.set(...), call searchHighOrderTiers(...))
and assert the returned object has the expected needsBackfill boolean alongside
result checks.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: a97f9561-fe90-423c-a829-1f09353c58b7
📒 Files selected for processing (17)
AGENTS.mdREADME.mdsrc/config.tssrc/functions/consolidation-pipeline.tssrc/functions/crystallize.tssrc/functions/high-order-backfill.tssrc/functions/high-order-search.tssrc/functions/reflect.tssrc/functions/smart-search.tssrc/index.tssrc/mcp/server.tssrc/mcp/tools-registry.tssrc/state/vector-index.tssrc/triggers/api.tssrc/types.tstest/high-order-search.test.tstest/smart-search.test.ts
Comment on lines
+11
to
+121
| export function registerHighOrderBackfillFunction(sdk: ISdk, kv: StateKV): void { | ||
| sdk.registerFunction("mem::backfill-embeddings::high-order", async () => { | ||
| const ep = getEmbeddingProvider(); | ||
| if (!ep) { | ||
| return { success: false, error: "No embedding provider available" }; | ||
| } | ||
|
|
||
| const results = { | ||
| semantic: 0, | ||
| procedural: 0, | ||
| crystals: 0, | ||
| insights: 0, | ||
| }; | ||
|
|
||
| try { | ||
| // 1. Semantic Facts | ||
| const semantics = await kv.list<SemanticMemory>(KV.semantic); | ||
| const semToUpdate = semantics.filter( | ||
| (s) => !s.embedding || s.embeddingModel !== ep.name | ||
| ); | ||
| for (let i = 0; i < semToUpdate.length; i += BACKFILL_BATCH_SIZE) { | ||
| const batch = semToUpdate.slice(i, i + BACKFILL_BATCH_SIZE); | ||
| const texts = batch.map((s) => s.fact); | ||
| try { | ||
| const vectors = await ep.embedBatch(texts); | ||
| for (let j = 0; j < batch.length; j++) { | ||
| batch[j].embedding = float32ToBase64(vectors[j]); | ||
| batch[j].embeddingModel = ep.name; | ||
| await kv.set(KV.semantic, batch[j].id, batch[j]); | ||
| } | ||
| results.semantic += batch.length; | ||
| } catch (e) { | ||
| logger.warn("Semantic backfill batch failed", { error: String(e) }); | ||
| } | ||
| } | ||
|
|
||
| // 2. Procedural Skills | ||
| const procedurals = await kv.list<ProceduralMemory>(KV.procedural); | ||
| const procToUpdate = procedurals.filter( | ||
| (p) => !p.embedding || p.embeddingModel !== ep.name | ||
| ); | ||
| for (let i = 0; i < procToUpdate.length; i += BACKFILL_BATCH_SIZE) { | ||
| const batch = procToUpdate.slice(i, i + BACKFILL_BATCH_SIZE); | ||
| const texts = batch.map((p) => `${p.name} ${p.triggerCondition} ${p.steps.join(" ")}`); | ||
| try { | ||
| const vectors = await ep.embedBatch(texts); | ||
| for (let j = 0; j < batch.length; j++) { | ||
| batch[j].embedding = float32ToBase64(vectors[j]); | ||
| batch[j].embeddingModel = ep.name; | ||
| await kv.set(KV.procedural, batch[j].id, batch[j]); | ||
| } | ||
| results.procedural += batch.length; | ||
| } catch (e) { | ||
| logger.warn("Procedural backfill batch failed", { error: String(e) }); | ||
| } | ||
| } | ||
|
|
||
| // 3. Crystals | ||
| const crystals = await kv.list<Crystal>(KV.crystals); | ||
| const crysToUpdate = crystals.filter( | ||
| (c) => !c.embedding || c.embeddingModel !== ep.name | ||
| ); | ||
| for (let i = 0; i < crysToUpdate.length; i += BACKFILL_BATCH_SIZE) { | ||
| const batch = crysToUpdate.slice(i, i + BACKFILL_BATCH_SIZE); | ||
| const texts = batch.map((c) => `${c.narrative} ${c.lessons.join(" ")}`); | ||
| try { | ||
| const vectors = await ep.embedBatch(texts); | ||
| for (let j = 0; j < batch.length; j++) { | ||
| batch[j].embedding = float32ToBase64(vectors[j]); | ||
| batch[j].embeddingModel = ep.name; | ||
| await kv.set(KV.crystals, batch[j].id, batch[j]); | ||
| } | ||
| results.crystals += batch.length; | ||
| } catch (e) { | ||
| logger.warn("Crystal backfill batch failed", { error: String(e) }); | ||
| } | ||
| } | ||
|
|
||
| // 4. Insights | ||
| const insights = await kv.list<Insight>(KV.insights); | ||
| const insToUpdate = insights.filter( | ||
| (ins) => !ins.deleted && (!ins.embedding || ins.embeddingModel !== ep.name) | ||
| ); | ||
| for (let i = 0; i < insToUpdate.length; i += BACKFILL_BATCH_SIZE) { | ||
| const batch = insToUpdate.slice(i, i + BACKFILL_BATCH_SIZE); | ||
| const texts = batch.map((ins) => `${ins.title} ${ins.content}`); | ||
| try { | ||
| const vectors = await ep.embedBatch(texts); | ||
| for (let j = 0; j < batch.length; j++) { | ||
| batch[j].embedding = float32ToBase64(vectors[j]); | ||
| batch[j].embeddingModel = ep.name; | ||
| await kv.set(KV.insights, batch[j].id, batch[j]); | ||
| } | ||
| results.insights += batch.length; | ||
| } catch (e) { | ||
| logger.warn("Insight backfill batch failed", { error: String(e) }); | ||
| } | ||
| } | ||
|
|
||
| const total = results.semantic + results.procedural + results.crystals + results.insights; | ||
| if (total > 0) { | ||
| logger.info("High-order embedding backfill complete", { backfilled: results }); | ||
| } | ||
|
|
||
| return { success: true, backfilled: results }; | ||
| } catch (err) { | ||
| const errorMsg = err instanceof Error ? err.message : String(err); | ||
| logger.error("High-order backfill encountered a fatal error", { error: errorMsg }); | ||
| return { success: false, error: errorMsg }; | ||
| } | ||
| }); |
Contributor
There was a problem hiding this comment.
Add an audit record for high-order backfills.
This function rewrites four KV scopes but never calls recordAudit(), so manual or scheduled backfills leave no audit trail.
Suggested fix
import { float32ToBase64 } from "../state/vector-index.js";
import { logger } from "../logger.js";
+import { recordAudit } from "./audit.js";
@@
const total = results.semantic + results.procedural + results.crystals + results.insights;
if (total > 0) {
logger.info("High-order embedding backfill complete", { backfilled: results });
}
+
+ await recordAudit(
+ kv,
+ "backfill_embeddings",
+ "mem::backfill-embeddings::high-order",
+ [],
+ { backfilled: results, embeddingModel: ep.name },
+ );
return { success: true, backfilled: results };As per coding guidelines, src/functions/**/*.ts: "Use recordAudit() for all state-changing operations".
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@src/functions/high-order-backfill.ts` around lines 11 - 121, The backfill
function registerHighOrderBackfillFunction updates multiple KV scopes but never
creates an audit trail; update the function to call recordAudit() after each
successful scope backfill (or once after all scopes complete) including scope
name, count backfilled, embeddingModel (ep.name) and timestamp so changes are
recorded; specifically add calls to recordAudit(...) after the semantic,
procedural, crystals and insights loops (or one consolidated recordAudit at the
end using the results object) and ensure errors still return without emitting a
success audit.
Comment on lines
+1110
to
+1122
| const body = req.body as Record<string, unknown>; | ||
| const result = await sdk.trigger({ | ||
| function_id: "mem::smart-search", | ||
| payload: { | ||
| ...(body.query !== undefined && { query: body.query }), | ||
| ...(body.expandIds !== undefined && { expandIds: body.expandIds }), | ||
| ...(body.limit !== undefined && { limit: body.limit }), | ||
| ...(body.project !== undefined && { project: body.project }), | ||
| ...(body.includeLessons !== undefined && { includeLessons: body.includeLessons }), | ||
| ...(body.includeHighOrder !== undefined && { includeHighOrder: body.includeHighOrder }), | ||
| ...(body.agentId !== undefined && { agentId: body.agentId }), | ||
| }, | ||
| }); |
Contributor
There was a problem hiding this comment.
Validate smart-search field types before forwarding them.
This whitelist still forwards raw values. A request like { "query": {}, "expandIds": "sem_1", "limit": "10" } now passes the REST boundary unchanged and leaves mem::smart-search to reject or mis-handle it downstream.
As per coding guidelines, src/triggers/**/*.ts: "validate and whitelist fields from request body (never pass raw body to sdk.trigger())", and {src/mcp/server.ts,src/triggers/**/*.ts}: "Input validation must occur at system boundaries (MCP handlers, REST endpoints)".
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@src/triggers/api.ts` around lines 1110 - 1122, Validate and whitelist each
expected field from req.body before calling sdk.trigger: ensure query is a
string, expandIds is an array of strings (or accept a comma-separated string and
split/coerce), limit is a number (parse numeric strings to Number and ignore
non-numeric), project and agentId are strings, and
includeLessons/includeHighOrder are booleans; only include each key in the
payload if it passes its type check (otherwise omit or coerce safely) when
constructing the object passed to sdk.trigger so you never forward raw,
unvalidated values from req.body.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
This PR extends mem::smart-search to query the 4 high-order memory tiers (Semantic facts, Procedural skills, Crystals, and Insights) alongside existing observations and lessons. Previously, these tiers—produced by the consolidation pipeline—sat "dark" in the KV store.
Detailed Features & Changes
src/functions/high-order-search.ts. High-order tiers are loaded in parallel and scored against search query tokens.filterAgentIdis passed, high-order search is skipped entirely. This prevents cross-agent leaks.expandIdsinsrc/functions/smart-search.tsto identify tier prefixes (sem_,proc_/skill_,crys_,ins_) and resolve them from the correct KV scope, unifying the output shape.AGENTMEMORY_HIGH_ORDER_SEARCHandAGENTMEMORY_SMART_SEARCH_CONFIDENCE_FLOOR).embedding(base64 Float32Array) andembeddingModel(version/model stamp) fields toSemanticMemory,ProceduralMemory,Crystal, andInsightinterfaces insrc/types.ts.mem::backfill-embeddings::high-order(src/functions/high-order-backfill.ts) background task and REST endpoint to scan and backfill missing or outdated embeddings. Wiredsmart-searchto detect missing embeddings and trigger this backfill in the background without blocking the user response.high-order-searchto calculate vector cosine similarity using the active embedding model. Fuses BM25 and Vector scores using Reciprocal Rank Fusion (RRF) with a precision of 4 decimal places. Falls back gracefully to BM25-only on model mismatch while scheduling a re-embed task.Fixes #770
Summary by CodeRabbit
New Features
Documentation
Tests