Skip to content

feat: implement high-order tier vector search, embedding backfill, an… #784

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
Tanmay-008 wants to merge 1 commit into
base: main
Choose a base branch
from
Open

feat: implement high-order tier vector search, embedding backfill, an… #784

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion AGENTS.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -117,7 +117,7 @@ Hook scripts in `src/hooks/` are standalone Node.js scripts (no iii-sdk import).
## Current Stats (v0.9.16)

- 53 MCP tools (8 visible by default, `AGENTMEMORY_TOOLS=all` for all)
- 125 REST endpoints
- 126 REST endpoints
- 6 MCP resources, 3 MCP prompts
- 12 hooks, 4 skills
- 50+ iii functions
Expand Down
2 changes: 1 addition & 1 deletion README.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1429,7 +1429,7 @@ Create `~/.agentmemory/.env`:

<h2 id="api"><picture><source media="(prefers-color-scheme: dark)" srcset="assets/tags/light/section-api.svg"><img src="assets/tags/section-api.svg" alt="API" height="32" /></picture></h2>

125 endpoints on port `3111`. The REST API binds to `127.0.0.1` by default. Protected endpoints require `Authorization: Bearer <secret>` when `AGENTMEMORY_SECRET` is set, and mesh sync endpoints require `AGENTMEMORY_SECRET` on both peers.
126 endpoints on port `3111`. The REST API binds to `127.0.0.1` by default. Protected endpoints require `Authorization: Bearer <secret>` when `AGENTMEMORY_SECRET` is set, and mesh sync endpoints require `AGENTMEMORY_SECRET` on both peers.

<details>
<summary>Key endpoints</summary>
Expand Down
14 changes: 14 additions & 0 deletions src/config.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -375,6 +375,20 @@ export function getConsolidationDecayDays(): number {
return safeParseInt(getMergedEnv()["CONSOLIDATION_DECAY_DAYS"], 30);
}

export function isHighOrderSearchEnabled(): boolean {
const env = getMergedEnv();
const explicit = env["AGENTMEMORY_HIGH_ORDER_SEARCH"];
if (explicit === "false" || explicit === "0") return false;
if (explicit === "true" || explicit === "1") return true;
return isConsolidationEnabled();
}

export function getHighOrderConfidenceFloor(): number {
const env = getMergedEnv();
const raw = parseFloat(env["AGENTMEMORY_SMART_SEARCH_CONFIDENCE_FLOOR"] || "0.3");
return Number.isFinite(raw) && raw >= 0 && raw <= 1 ? raw : 0.3;
}

export function isStandaloneMcp(): boolean {
return getMergedEnv()["STANDALONE_MCP"] === "true";
}
Expand Down
41 changes: 40 additions & 1 deletion src/functions/consolidation-pipeline.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 +17,8 @@ import {
import { recordAudit } from "./audit.js";
import { getConsolidationDecayDays, isConsolidationEnabled } from "../config.js";
import { logger } from "../logger.js";

import { getEmbeddingProvider } from "./search.js";
import { float32ToBase64 } from "../state/vector-index.js";
function applyDecay(
items: Array<{
strength: number;
Expand Down Expand Up @@ -93,6 +94,8 @@ export function registerConsolidationPipelineFunction(
const confidence = Number.isNaN(parsedConf) ? 0.5 : parsedConf;
const fact = match[2].trim();

const ep = getEmbeddingProvider();

const existing = existingSemantic.find(
(s) => s.fact.toLowerCase() === fact.toLowerCase(),
);
Expand All @@ -101,6 +104,14 @@ export function registerConsolidationPipelineFunction(
existing.lastAccessedAt = now;
existing.updatedAt = now;
existing.confidence = Math.max(existing.confidence, confidence);
if (!existing.embedding && ep) {
try {
existing.embedding = float32ToBase64(await ep.embed(existing.fact));
existing.embeddingModel = ep.name;
} catch (e) {
logger.warn("Failed to embed existing semantic fact", { id: existing.id, error: String(e) });
}
}
await kv.set(KV.semantic, existing.id, existing);
} else {
const sem: SemanticMemory = {
Expand All @@ -115,6 +126,14 @@ export function registerConsolidationPipelineFunction(
createdAt: now,
updatedAt: now,
};
if (ep) {
try {
sem.embedding = float32ToBase64(await ep.embed(sem.fact));
sem.embeddingModel = ep.name;
} catch (e) {
logger.warn("Failed to embed new semantic fact", { error: String(e) });
}
}
await kv.set(KV.semantic, sem.id, sem);
newFacts++;
}
Expand Down Expand Up @@ -187,13 +206,24 @@ export function registerConsolidationPipelineFunction(
steps.push(stepMatch[1].trim());
}

const ep = getEmbeddingProvider();

const existing = existingProcs.find(
(p) => p.name.toLowerCase() === name.toLowerCase(),
);
if (existing) {
existing.frequency++;
existing.updatedAt = now;
existing.strength = Math.min(1, existing.strength + 0.1);
if (!existing.embedding && ep) {
try {
const text = `${existing.name} ${existing.triggerCondition} ${existing.steps.join(" ")}`;
existing.embedding = float32ToBase64(await ep.embed(text));
existing.embeddingModel = ep.name;
} catch (e) {
logger.warn("Failed to embed existing procedural skill", { id: existing.id, error: String(e) });
}
}
await kv.set(KV.procedural, existing.id, existing);
} else {
const proc: ProceduralMemory = {
Expand All @@ -207,6 +237,15 @@ export function registerConsolidationPipelineFunction(
createdAt: now,
updatedAt: now,
};
if (ep) {
try {
const text = `${proc.name} ${proc.triggerCondition} ${proc.steps.join(" ")}`;
proc.embedding = float32ToBase64(await ep.embed(text));
proc.embeddingModel = ep.name;
} catch (e) {
logger.warn("Failed to embed new procedural skill", { error: String(e) });
}
}
await kv.set(KV.procedural, proc.id, proc);
newProcs++;
}
Expand Down
15 changes: 15 additions & 0 deletions src/functions/crystallize.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,10 @@ import type { StateKV } from "../state/kv.js";
import { KV, generateId } from "../state/schema.js";
import type { Action, ActionEdge, Crystal, MemoryProvider } from "../types.js";

import { getEmbeddingProvider } from "./search.js";
import { float32ToBase64 } from "../state/vector-index.js";
import { logger } from "../logger.js";

interface CrystalDigest {
narrative: string;
keyOutcomes: string[];
Expand Down Expand Up @@ -52,6 +56,7 @@ export function registerCrystallizeFunction(
);

const prompt = buildChainText(actions, relevantEdges);
const ep = getEmbeddingProvider();

try {
const response = await provider.summarize(CRYSTALLIZE_SYSTEM, prompt);
Expand All @@ -69,6 +74,16 @@ export function registerCrystallizeFunction(
createdAt: new Date().toISOString(),
};

if (ep) {
try {
const text = `${crystal.narrative} ${crystal.lessons.join(" ")}`;
crystal.embedding = float32ToBase64(await ep.embed(text));
crystal.embeddingModel = ep.name;
} catch (e) {
logger.warn("Failed to embed new crystal", { error: String(e) });
}
}

await kv.set(KV.crystals, crystal.id, crystal);

await Promise.all(
Expand Down
122 changes: 122 additions & 0 deletions src/functions/high-order-backfill.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,122 @@
import type { ISdk } from "iii-sdk";
import type { StateKV } from "../state/kv.js";
import { KV } from "../state/schema.js";
import type { SemanticMemory, ProceduralMemory, Crystal, Insight } from "../types.js";
import { getEmbeddingProvider } from "./search.js";
import { float32ToBase64 } from "../state/vector-index.js";
import { logger } from "../logger.js";

const BACKFILL_BATCH_SIZE = 20;

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 };
}
});
Comment on lines +11 to +121
Copy link
Copy Markdown
Contributor

@coderabbitai coderabbitai Bot Jun 2, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟠 Major | ⚡ Quick win

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.

}
Loading