mirror of
https://github.com/openclaw/openclaw.git
synced 2026-07-09 15:59:30 +00:00
Merged via squash.
Prepared head SHA: 66d362a56d
Co-authored-by: mushuiyu886 <266724580+mushuiyu886@users.noreply.github.com>
Co-authored-by: jalehman <550978+jalehman@users.noreply.github.com>
Reviewed-by: @jalehman
This commit is contained in:
mushuiyu_xydt
GitHub
parent
2f02bbcbb3
commit
a36e05050a
13 changed files with 1560 additions and 234 deletions
|
|
@ -113,6 +113,16 @@ can consume this generic provider surface. The older
|
|||
`contracts.memoryEmbeddingProviders` seam is deprecated compatibility while
|
||||
existing memory-specific providers migrate.
|
||||
|
||||
Memory-specific providers that still expose a runtime `batchEmbed(...)` stay on
|
||||
the existing per-file batching contract unless their runtime explicitly sets
|
||||
`sourceWideBatchEmbed: true`. That opt-in lets the memory host submit chunks from
|
||||
multiple dirty memory files and enabled sources in one `batchEmbed(...)` call up
|
||||
to the host batch limits. Batch adapters that upload JSONL request files must
|
||||
split provider jobs before their upload-size cap as well as their request-count
|
||||
cap. The provider must return one embedding per input chunk in the same order as
|
||||
`batch.chunks`; omit the flag when the provider expects file-local batches or
|
||||
cannot preserve input ordering across a larger source-wide job.
|
||||
|
||||
### Tools and commands
|
||||
|
||||
Use [`defineToolPlugin`](/plugins/tool-plugins) for simple tool-only plugins
|
||||
|
|
|
|||
|
|
@ -13,6 +13,7 @@ import { afterAll, afterEach, beforeAll, beforeEach, describe, expect, it, vi }
|
|||
import "./test-runtime-mocks.js";
|
||||
import type { MemoryIndexManager } from "./index.js";
|
||||
import { closeAllMemorySearchManagers, getMemorySearchManager } from "./index.js";
|
||||
import { splitSourceWideEmbeddingChunks } from "./manager-embedding-ops.js";
|
||||
import { LOCAL_EMBEDDING_WORKER_ERROR_CODES } from "./manager-local-worker-errors.js";
|
||||
import type { MemoryIndexMeta } from "./manager-reindex-state.js";
|
||||
import { closeMemoryIndexManagersForAgent, EMBEDDING_PROBE_CACHE_TTL_MS } from "./manager.js";
|
||||
|
|
@ -28,6 +29,11 @@ afterAll(() => {
|
|||
|
||||
let embedBatchCalls = 0;
|
||||
let embedBatchInputCalls = 0;
|
||||
let providerRuntimeBatchCalls: string[][] = [];
|
||||
let providerRuntimeBatchGate: Promise<void> | null = null;
|
||||
let providerRuntimeBatchFailuresRemaining = 0;
|
||||
let providerRuntimeActiveBatchCalls = 0;
|
||||
let providerRuntimeMaxActiveBatchCalls = 0;
|
||||
let providerCloseCalls = 0;
|
||||
let providerCloseFailuresRemaining = 0;
|
||||
let providerCloseGate: Promise<void> | null = null;
|
||||
|
|
@ -88,6 +94,8 @@ vi.mock("./embeddings.js", () => {
|
|||
const providerId =
|
||||
options.provider === "gemini" ||
|
||||
options.provider === "fallback-provider" ||
|
||||
options.provider === "batch-test" ||
|
||||
options.provider === "batch-wide-test" ||
|
||||
options.provider === "ollama"
|
||||
? options.provider
|
||||
: "mock";
|
||||
|
|
@ -141,20 +149,45 @@ vi.mock("./embeddings.js", () => {
|
|||
}
|
||||
: {}),
|
||||
},
|
||||
...(providerId === "gemini" || providerId === "fallback-provider"
|
||||
...(providerId === "batch-test" || providerId === "batch-wide-test"
|
||||
? {
|
||||
runtime: {
|
||||
id: providerId,
|
||||
cacheKeyData: {
|
||||
provider: providerId,
|
||||
baseUrl: "https://generativelanguage.googleapis.com/v1beta",
|
||||
model,
|
||||
outputDimensionality: options.outputDimensionality,
|
||||
headers: [],
|
||||
...(providerId === "batch-wide-test" ? { sourceWideBatchEmbed: true } : {}),
|
||||
batchEmbed: async (batch: { chunks: Array<{ text: string }> }) => {
|
||||
providerRuntimeActiveBatchCalls += 1;
|
||||
providerRuntimeMaxActiveBatchCalls = Math.max(
|
||||
providerRuntimeMaxActiveBatchCalls,
|
||||
providerRuntimeActiveBatchCalls,
|
||||
);
|
||||
try {
|
||||
await providerRuntimeBatchGate;
|
||||
providerRuntimeBatchCalls.push(batch.chunks.map((chunk) => chunk.text));
|
||||
if (providerRuntimeBatchFailuresRemaining > 0) {
|
||||
providerRuntimeBatchFailuresRemaining -= 1;
|
||||
throw new Error("provider runtime batch failed");
|
||||
}
|
||||
return batch.chunks.map((chunk) => embedText(chunk.text));
|
||||
} finally {
|
||||
providerRuntimeActiveBatchCalls -= 1;
|
||||
}
|
||||
},
|
||||
},
|
||||
}
|
||||
: {}),
|
||||
: providerId === "gemini" || providerId === "fallback-provider"
|
||||
? {
|
||||
runtime: {
|
||||
id: providerId,
|
||||
cacheKeyData: {
|
||||
provider: providerId,
|
||||
baseUrl: "https://generativelanguage.googleapis.com/v1beta",
|
||||
model,
|
||||
outputDimensionality: options.outputDimensionality,
|
||||
headers: [],
|
||||
},
|
||||
},
|
||||
}
|
||||
: {}),
|
||||
};
|
||||
},
|
||||
};
|
||||
|
|
@ -221,6 +254,11 @@ describe("memory index", () => {
|
|||
registerBuiltInMemoryEmbeddingProviders({ registerMemoryEmbeddingProvider: registerAdapter });
|
||||
embedBatchCalls = 0;
|
||||
embedBatchInputCalls = 0;
|
||||
providerRuntimeBatchCalls = [];
|
||||
providerRuntimeBatchGate = null;
|
||||
providerRuntimeBatchFailuresRemaining = 0;
|
||||
providerRuntimeActiveBatchCalls = 0;
|
||||
providerRuntimeMaxActiveBatchCalls = 0;
|
||||
providerCloseCalls = 0;
|
||||
providerCloseFailuresRemaining = 0;
|
||||
providerCloseGate = null;
|
||||
|
|
@ -268,6 +306,7 @@ describe("memory index", () => {
|
|||
provider?: string;
|
||||
fallback?: "none" | "gemini" | "fallback-provider";
|
||||
providerAliases?: NonNullable<NonNullable<TestCfg["models"]>["providers"]>;
|
||||
batchEnabled?: boolean;
|
||||
model?: string;
|
||||
outputDimensionality?: number;
|
||||
multimodal?: {
|
||||
|
|
@ -294,6 +333,12 @@ describe("memory index", () => {
|
|||
// Perf: keep test indexes to a single chunk to reduce sqlite work.
|
||||
chunking: { tokens: 4000, overlap: 0 },
|
||||
sync: { watch: false, onSessionStart: false, onSearch: params.onSearch ?? true },
|
||||
remote: params.batchEnabled
|
||||
? {
|
||||
nonBatchConcurrency: 1,
|
||||
batch: { enabled: true, pollIntervalMs: 0, timeoutMinutes: 1 },
|
||||
}
|
||||
: undefined,
|
||||
query: {
|
||||
minScore: params.minScore ?? 0,
|
||||
hybrid: params.hybrid ?? { enabled: false },
|
||||
|
|
@ -419,6 +464,243 @@ describe("memory index", () => {
|
|||
}
|
||||
});
|
||||
|
||||
it("batches dirty memory chunks across files", async () => {
|
||||
await fs.writeFile(path.join(memoryDir, "2026-01-13.md"), "# Log\nBeta memory line.");
|
||||
await fs.writeFile(path.join(memoryDir, "2026-01-14.md"), "# Log\nGamma memory line.");
|
||||
const cfg = createCfg({
|
||||
provider: "batch-wide-test",
|
||||
batchEnabled: true,
|
||||
storePath: path.join(workspaceDir, "index-cross-file-batch.sqlite"),
|
||||
});
|
||||
const manager = await getFreshManager(cfg);
|
||||
try {
|
||||
await manager.sync({ reason: "test" });
|
||||
|
||||
expect(providerRuntimeBatchCalls).toHaveLength(1);
|
||||
expect(providerRuntimeBatchCalls[0]).toEqual([
|
||||
"# Log\nAlpha memory line.\nZebra memory line.",
|
||||
"# Log\nBeta memory line.",
|
||||
"# Log\nGamma memory line.",
|
||||
]);
|
||||
} finally {
|
||||
await manager.close?.();
|
||||
}
|
||||
});
|
||||
|
||||
it("maps source-wide batch fallback results to missing chunks after cache hits", async () => {
|
||||
const cfg = createCfg({
|
||||
provider: "batch-wide-test",
|
||||
batchEnabled: true,
|
||||
storePath: path.join(workspaceDir, "index-cross-file-batch-fallback-cache.sqlite"),
|
||||
});
|
||||
const manager = await getFreshManager(cfg);
|
||||
try {
|
||||
await manager.sync({ reason: "test" });
|
||||
|
||||
await fs.writeFile(path.join(memoryDir, "2026-01-13.md"), "# Log\nBeta memory line.");
|
||||
providerRuntimeBatchCalls = [];
|
||||
providerRuntimeBatchFailuresRemaining = 1;
|
||||
embedBatchCalls = 0;
|
||||
|
||||
await manager.sync({ reason: "test", force: true });
|
||||
|
||||
expect(providerRuntimeBatchCalls).toEqual([["# Log\nBeta memory line."]]);
|
||||
expect(embedBatchCalls).toBe(1);
|
||||
const betaRow = (
|
||||
manager as unknown as {
|
||||
db: { prepare: (sql: string) => { get: (...args: unknown[]) => unknown } };
|
||||
}
|
||||
).db
|
||||
.prepare("SELECT embedding FROM chunks WHERE path LIKE ? AND source = ?")
|
||||
.get("%2026-01-13.md", "memory") as { embedding: string } | undefined;
|
||||
|
||||
expect(betaRow).toBeDefined();
|
||||
expect(JSON.parse(betaRow?.embedding ?? "[]")).toEqual([0, 1, 0, 0]);
|
||||
} finally {
|
||||
await manager.close?.();
|
||||
}
|
||||
});
|
||||
|
||||
it("splits oversized source-wide embedding requests at the request cap", () => {
|
||||
expect(splitSourceWideEmbeddingChunks(["one", "two", "three", "four", "five"], 2)).toEqual([
|
||||
["one", "two"],
|
||||
["three", "four"],
|
||||
["five"],
|
||||
]);
|
||||
});
|
||||
|
||||
it("keeps split chunks from oversized files in one source-wide batch", async () => {
|
||||
await fs.writeFile(
|
||||
path.join(memoryDir, "2026-01-13.md"),
|
||||
`# Log\n${"Long split memory line. ".repeat(1200)}`,
|
||||
);
|
||||
await fs.writeFile(path.join(memoryDir, "2026-01-14.md"), "# Log\nBeta memory line.");
|
||||
const cfg = createCfg({
|
||||
provider: "batch-wide-test",
|
||||
batchEnabled: true,
|
||||
storePath: path.join(workspaceDir, "index-split-chunks-cross-file-batch.sqlite"),
|
||||
});
|
||||
const manager = await getFreshManager(cfg);
|
||||
try {
|
||||
await manager.sync({ reason: "test" });
|
||||
|
||||
expect(providerRuntimeBatchCalls).toHaveLength(1);
|
||||
const combinedBatch = providerRuntimeBatchCalls[0] ?? [];
|
||||
expect(combinedBatch.length).toBeGreaterThan(3);
|
||||
expect(combinedBatch.join("\n")).toContain("Long split memory line.");
|
||||
expect(combinedBatch).toContain("# Log\nBeta memory line.");
|
||||
} finally {
|
||||
await manager.close?.();
|
||||
}
|
||||
});
|
||||
|
||||
it("keeps custom batch runtimes per file without source-wide opt in", async () => {
|
||||
await fs.writeFile(path.join(memoryDir, "2026-01-13.md"), "# Log\nBeta memory line.");
|
||||
await fs.writeFile(path.join(memoryDir, "2026-01-14.md"), "# Log\nGamma memory line.");
|
||||
const cfg = createCfg({
|
||||
provider: "batch-test",
|
||||
batchEnabled: true,
|
||||
storePath: path.join(workspaceDir, "index-custom-batch-compat.sqlite"),
|
||||
});
|
||||
const manager = await getFreshManager(cfg);
|
||||
try {
|
||||
await manager.sync({ reason: "test" });
|
||||
|
||||
expect(providerRuntimeBatchCalls).toHaveLength(3);
|
||||
expect(providerRuntimeBatchCalls.every((call) => call.length === 1)).toBe(true);
|
||||
expect(providerRuntimeBatchCalls.map((call) => call[0]).toSorted()).toEqual(
|
||||
[
|
||||
"# Log\nAlpha memory line.\nZebra memory line.",
|
||||
"# Log\nBeta memory line.",
|
||||
"# Log\nGamma memory line.",
|
||||
].toSorted(),
|
||||
);
|
||||
} finally {
|
||||
await manager.close?.();
|
||||
}
|
||||
});
|
||||
|
||||
it("keeps custom batch runtimes concurrent without source-wide opt in", async () => {
|
||||
await fs.writeFile(path.join(memoryDir, "2026-01-13.md"), "# Log\nBeta memory line.");
|
||||
await fs.writeFile(path.join(memoryDir, "2026-01-14.md"), "# Log\nGamma memory line.");
|
||||
const cfg = createCfg({
|
||||
provider: "batch-test",
|
||||
batchEnabled: true,
|
||||
storePath: path.join(workspaceDir, "index-custom-batch-concurrency.sqlite"),
|
||||
});
|
||||
const manager = await getFreshManager(cfg);
|
||||
let releaseBatchGate: (() => void) | undefined;
|
||||
providerRuntimeBatchGate = new Promise((resolve) => {
|
||||
releaseBatchGate = resolve;
|
||||
});
|
||||
const syncPromise = manager.sync({ reason: "test" });
|
||||
let waitError: Error | undefined;
|
||||
try {
|
||||
await vi.waitFor(() => expect(providerRuntimeMaxActiveBatchCalls).toBeGreaterThan(1));
|
||||
} catch (err) {
|
||||
waitError = err instanceof Error ? err : new Error(String(err));
|
||||
} finally {
|
||||
releaseBatchGate?.();
|
||||
await syncPromise;
|
||||
await manager.close?.();
|
||||
}
|
||||
if (waitError) {
|
||||
throw waitError;
|
||||
}
|
||||
});
|
||||
|
||||
it("bounds source-wide memory batches", async () => {
|
||||
const batchFileLimit = 2048;
|
||||
for (let index = 0; index < batchFileLimit; index += 1) {
|
||||
await fs.writeFile(
|
||||
path.join(memoryDir, `2026-02-${String(index + 1).padStart(4, "0")}.md`),
|
||||
`# Log\nBounded memory line ${index}.`,
|
||||
);
|
||||
}
|
||||
const cfg = createCfg({
|
||||
provider: "batch-wide-test",
|
||||
batchEnabled: true,
|
||||
storePath: path.join(workspaceDir, "index-bounded-cross-file-batch.sqlite"),
|
||||
});
|
||||
const manager = await getFreshManager(cfg);
|
||||
try {
|
||||
await manager.sync({ reason: "test" });
|
||||
|
||||
expect(providerRuntimeBatchCalls).toHaveLength(2);
|
||||
expect(providerRuntimeBatchCalls[0]).toHaveLength(batchFileLimit);
|
||||
expect(providerRuntimeBatchCalls[1]).toHaveLength(1);
|
||||
expect(providerRuntimeBatchCalls.flat()).toHaveLength(batchFileLimit + 1);
|
||||
} finally {
|
||||
await manager.close?.();
|
||||
}
|
||||
});
|
||||
|
||||
it("batches forced memory and session indexing across files", async () => {
|
||||
await fs.writeFile(path.join(memoryDir, "2026-01-13.md"), "# Log\nBeta memory line.");
|
||||
const sessionsDir = resolveSessionTranscriptsDirForAgent("main");
|
||||
await fs.mkdir(sessionsDir, { recursive: true });
|
||||
await fs.writeFile(
|
||||
path.join(sessionsDir, "session-alpha.jsonl"),
|
||||
[
|
||||
JSON.stringify({
|
||||
type: "session",
|
||||
id: "session-alpha",
|
||||
timestamp: "2026-04-07T15:24:04.113Z",
|
||||
}),
|
||||
JSON.stringify({
|
||||
type: "message",
|
||||
message: {
|
||||
role: "user",
|
||||
timestamp: "2026-04-07T15:25:04.113Z",
|
||||
content: [{ type: "text", text: "Session alpha memory line." }],
|
||||
},
|
||||
}),
|
||||
].join("\n") + "\n",
|
||||
"utf8",
|
||||
);
|
||||
await fs.writeFile(
|
||||
path.join(sessionsDir, "session-beta.jsonl"),
|
||||
[
|
||||
JSON.stringify({
|
||||
type: "session",
|
||||
id: "session-beta",
|
||||
timestamp: "2026-04-07T15:24:04.113Z",
|
||||
}),
|
||||
JSON.stringify({
|
||||
type: "message",
|
||||
message: {
|
||||
role: "assistant",
|
||||
timestamp: "2026-04-07T15:25:04.113Z",
|
||||
content: [{ type: "text", text: "Session beta memory line." }],
|
||||
},
|
||||
}),
|
||||
].join("\n") + "\n",
|
||||
"utf8",
|
||||
);
|
||||
const cfg = createCfg({
|
||||
provider: "batch-wide-test",
|
||||
batchEnabled: true,
|
||||
sources: ["memory", "sessions"],
|
||||
sessionMemory: true,
|
||||
storePath: path.join(workspaceDir, "index-force-cross-source-batch.sqlite"),
|
||||
});
|
||||
const manager = await getFreshManager(cfg);
|
||||
try {
|
||||
await manager.sync({ reason: "cli", force: true });
|
||||
|
||||
expect(providerRuntimeBatchCalls).toHaveLength(1);
|
||||
const combinedBatch = providerRuntimeBatchCalls[0] ?? [];
|
||||
expect(combinedBatch.slice(0, 2)).toEqual([
|
||||
"# Log\nAlpha memory line.\nZebra memory line.",
|
||||
"# Log\nBeta memory line.",
|
||||
]);
|
||||
expect(combinedBatch.join("\n")).toContain("Session alpha memory line.");
|
||||
expect(combinedBatch.join("\n")).toContain("Session beta memory line.");
|
||||
} finally {
|
||||
await manager.close?.();
|
||||
}
|
||||
});
|
||||
|
||||
it("does not full-reindex on search when existing metadata belongs to another provider", async () => {
|
||||
const dbPath = path.join(workspaceDir, "index-provider-cutover.sqlite");
|
||||
const oldCfg = createCfg({
|
||||
|
|
@ -1041,25 +1323,6 @@ describe("memory index", () => {
|
|||
expect(third).toBe(second);
|
||||
});
|
||||
|
||||
it("closes stale default managers when provider requirement changes", async () => {
|
||||
const storePath = path.join(workspaceDir, "index-provider-requirement-cache.sqlite");
|
||||
const implicitCfg = createCfg({ storePath });
|
||||
const implicit = requireManager(
|
||||
await getMemorySearchManager({ cfg: implicitCfg, agentId: "main" }),
|
||||
);
|
||||
managersForCleanup.add(implicit);
|
||||
await implicit.probeEmbeddingAvailability();
|
||||
|
||||
const explicitCfg = createCfg({ storePath, provider: "openai" });
|
||||
const explicit = requireManager(
|
||||
await getMemorySearchManager({ cfg: explicitCfg, agentId: "main" }),
|
||||
);
|
||||
managersForCleanup.add(explicit);
|
||||
|
||||
expect(explicit === implicit).toBe(false);
|
||||
expect(providerCloseCalls).toBe(1);
|
||||
});
|
||||
|
||||
it("retries embedding provider close before releasing the manager", async () => {
|
||||
providerCloseFailuresRemaining = 1;
|
||||
const cfg = createCfg({
|
||||
|
|
@ -1599,6 +1862,25 @@ describe("memory index", () => {
|
|||
}
|
||||
});
|
||||
|
||||
it("keeps metadata after unchanged safe force reindex", async () => {
|
||||
vi.stubEnv("OPENCLAW_TEST_MEMORY_UNSAFE_REINDEX", "0");
|
||||
const cfg = createCfg({
|
||||
storePath: path.join(workspaceDir, "index-safe-force-metadata.sqlite"),
|
||||
});
|
||||
const manager = await getFreshManager(cfg);
|
||||
try {
|
||||
await manager.sync({ reason: "test", force: true });
|
||||
expect(manager.status().custom?.indexIdentity).toEqual({ status: "valid" });
|
||||
|
||||
await manager.sync({ reason: "cli", force: true });
|
||||
|
||||
expect(manager.status().dirty).toBe(false);
|
||||
expect(manager.status().custom?.indexIdentity).toEqual({ status: "valid" });
|
||||
} finally {
|
||||
await manager.close?.();
|
||||
}
|
||||
});
|
||||
|
||||
it("streams embedding cache rows during safe reindex", async () => {
|
||||
vi.stubEnv("OPENCLAW_TEST_MEMORY_UNSAFE_REINDEX", "0");
|
||||
type EmbeddingCacheRow = {
|
||||
|
|
@ -1690,12 +1972,6 @@ describe("memory index", () => {
|
|||
const status = manager.status();
|
||||
expect(status.chunks).toBeGreaterThan(0);
|
||||
expect(embedBatchCalls).toBe(0);
|
||||
expect(status.custom?.providerUnavailableReason).toBe("No API key found for provider");
|
||||
expect(status.custom?.providerState).toEqual({
|
||||
mode: "fts-only",
|
||||
reason: "No API key found for provider",
|
||||
attemptedProviderId: "openai",
|
||||
});
|
||||
|
||||
const results = await manager.search("Alpha");
|
||||
expect(results.length).toBeGreaterThan(0);
|
||||
|
|
@ -1754,7 +2030,6 @@ describe("memory index", () => {
|
|||
await manager.close?.();
|
||||
}
|
||||
});
|
||||
|
||||
it("prefers exact session transcript hits in FTS-only mode", async () => {
|
||||
try {
|
||||
const manager = await getFtsSessionManager({
|
||||
|
|
|
|||
|
|
@ -14,6 +14,7 @@ import {
|
|||
hashText,
|
||||
remapChunkLines,
|
||||
retryTransientMemoryRead,
|
||||
runWithConcurrency,
|
||||
type MemoryChunk,
|
||||
type MemorySource,
|
||||
} from "openclaw/plugin-sdk/memory-core-host-engine-storage";
|
||||
|
|
@ -40,7 +41,7 @@ import {
|
|||
runMemoryEmbeddingRetryLoop,
|
||||
} from "./manager-embedding-policy.js";
|
||||
import { deleteMemoryFtsRows } from "./manager-fts-state.js";
|
||||
import { MemoryManagerSyncOps } from "./manager-sync-ops.js";
|
||||
import { MemoryManagerSyncOps, type MemoryIndexWorkItem } from "./manager-sync-ops.js";
|
||||
import { logMemoryVectorDegradedWrite } from "./manager-vector-warning.js";
|
||||
import { replaceMemoryVectorRow } from "./manager-vector-write.js";
|
||||
|
||||
|
|
@ -56,6 +57,8 @@ const EMBEDDING_QUERY_TIMEOUT_REMOTE_MS = 60_000;
|
|||
const EMBEDDING_QUERY_TIMEOUT_LOCAL_MS = 5 * 60_000;
|
||||
const EMBEDDING_BATCH_TIMEOUT_REMOTE_MS = 2 * 60_000;
|
||||
const EMBEDDING_BATCH_TIMEOUT_LOCAL_MS = 10 * 60_000;
|
||||
const SOURCE_WIDE_BATCH_MAX_FILES = 2048;
|
||||
const SOURCE_WIDE_BATCH_MAX_REQUESTS = 50000;
|
||||
|
||||
const log = createSubsystemLogger("memory");
|
||||
|
||||
|
|
@ -70,17 +73,46 @@ function resolveEmbeddingSecondsTimeoutMs(seconds: number): number {
|
|||
);
|
||||
}
|
||||
|
||||
type MemoryIndexEntry = {
|
||||
path: string;
|
||||
absPath: string;
|
||||
mtimeMs: number;
|
||||
size: number;
|
||||
hash: string;
|
||||
kind?: "markdown" | "multimodal";
|
||||
contentText?: string;
|
||||
lineMap?: number[];
|
||||
type MemoryIndexEntry = MemoryIndexWorkItem["entry"];
|
||||
|
||||
type PreparedMemoryIndexEntry = {
|
||||
entry: MemoryIndexEntry;
|
||||
source: MemorySource;
|
||||
chunks: MemoryChunk[];
|
||||
structuredInputBytes?: number;
|
||||
};
|
||||
|
||||
function countBatchSources(items: Array<{ source: MemorySource }>): Record<string, number> {
|
||||
const counts: Record<string, number> = {};
|
||||
for (const item of items) {
|
||||
counts[item.source] = (counts[item.source] ?? 0) + 1;
|
||||
}
|
||||
return counts;
|
||||
}
|
||||
|
||||
function formatBatchSourceLabel(counts: Record<string, number>): string {
|
||||
const sources = Object.keys(counts).toSorted();
|
||||
return sources.length > 0 ? sources.join("+") : "unknown";
|
||||
}
|
||||
|
||||
function formatBatchSourceCounts(counts: Record<string, number>): string {
|
||||
return (
|
||||
Object.entries(counts)
|
||||
.toSorted(([left], [right]) => left.localeCompare(right))
|
||||
.map(([source, count]) => `${source}=${count}`)
|
||||
.join(",") || "none"
|
||||
);
|
||||
}
|
||||
|
||||
export function splitSourceWideEmbeddingChunks<T>(chunks: T[], maxRequests: number): T[][] {
|
||||
const limit = Math.max(1, Math.floor(maxRequests));
|
||||
const batches: T[][] = [];
|
||||
for (let start = 0; start < chunks.length; start += limit) {
|
||||
batches.push(chunks.slice(start, start + limit));
|
||||
}
|
||||
return batches;
|
||||
}
|
||||
|
||||
export function resolveEmbeddingTimeoutMs(params: {
|
||||
kind: "query" | "batch";
|
||||
providerId?: string;
|
||||
|
|
@ -256,18 +288,25 @@ export abstract class MemoryManagerEmbeddingOps extends MemoryManagerSyncOps {
|
|||
return hashText(JSON.stringify({ provider: this.provider.id, model: this.provider.model }));
|
||||
}
|
||||
|
||||
private buildBatchDebug(source: MemorySource, chunks: MemoryChunk[]) {
|
||||
private buildBatchDebug(
|
||||
source: string,
|
||||
chunks: MemoryChunk[],
|
||||
context: Record<string, unknown> = {},
|
||||
) {
|
||||
return (message: string, data?: Record<string, unknown>) =>
|
||||
log.debug(
|
||||
message,
|
||||
data ? { ...data, source, chunks: chunks.length } : { source, chunks: chunks.length },
|
||||
data
|
||||
? { ...data, source, chunks: chunks.length, ...context }
|
||||
: { source, chunks: chunks.length, ...context },
|
||||
);
|
||||
}
|
||||
|
||||
private async embedChunksWithBatch(
|
||||
chunks: MemoryChunk[],
|
||||
_entry: MemoryIndexEntry,
|
||||
source: MemorySource,
|
||||
source: string,
|
||||
debugContext: Record<string, unknown> = {},
|
||||
): Promise<number[][]> {
|
||||
const provider = this.provider;
|
||||
const batchEmbed = this.providerRuntime?.batchEmbed;
|
||||
|
|
@ -293,9 +332,9 @@ export abstract class MemoryManagerEmbeddingOps extends MemoryManagerSyncOps {
|
|||
concurrency: this.batch.concurrency,
|
||||
pollIntervalMs: this.batch.pollIntervalMs,
|
||||
timeoutMs: this.batch.timeoutMs,
|
||||
debug: this.buildBatchDebug(source, chunks),
|
||||
debug: this.buildBatchDebug(source, chunks, debugContext),
|
||||
}),
|
||||
fallback: async () => await this.embedChunksInBatches(chunks),
|
||||
fallback: async () => await this.embedChunksInBatches(missingChunks),
|
||||
});
|
||||
if (!batchResult) {
|
||||
return this.embedChunksInBatches(chunks);
|
||||
|
|
@ -747,6 +786,165 @@ export abstract class MemoryManagerEmbeddingOps extends MemoryManagerSyncOps {
|
|||
this.upsertFileRecord(entry, source);
|
||||
}
|
||||
|
||||
private async prepareIndexEntry(
|
||||
entry: MemoryIndexEntry,
|
||||
options: { source: MemorySource; content?: string },
|
||||
): Promise<PreparedMemoryIndexEntry | null> {
|
||||
if ("kind" in entry && entry.kind === "multimodal") {
|
||||
const multimodalChunk = await buildMultimodalChunkForIndexing(entry);
|
||||
if (!multimodalChunk) {
|
||||
this.clearIndexedFileData(entry.path, options.source);
|
||||
this.deleteFileRecord(entry.path, options.source);
|
||||
return null;
|
||||
}
|
||||
return {
|
||||
entry,
|
||||
source: options.source,
|
||||
chunks: [multimodalChunk.chunk],
|
||||
structuredInputBytes: multimodalChunk.structuredInputBytes,
|
||||
};
|
||||
}
|
||||
|
||||
const content =
|
||||
options.content ??
|
||||
entry.content ??
|
||||
(await retryTransientMemoryRead(
|
||||
() => fs.readFile(entry.absPath, "utf-8"),
|
||||
`read memory markdown for indexing ${entry.absPath}`,
|
||||
));
|
||||
const baseChunks = filterNonEmptyMemoryChunks(chunkMarkdown(content, this.settings.chunking));
|
||||
const chunks = this.provider
|
||||
? enforceEmbeddingMaxInputTokens(this.provider, baseChunks, EMBEDDING_BATCH_MAX_TOKENS)
|
||||
: baseChunks;
|
||||
if (options.source === "sessions" && "lineMap" in entry) {
|
||||
remapChunkLines(chunks, entry.lineMap);
|
||||
}
|
||||
return { entry, source: options.source, chunks };
|
||||
}
|
||||
|
||||
protected override async indexFiles(items: MemoryIndexWorkItem[]): Promise<void> {
|
||||
if (items.length === 0) {
|
||||
return;
|
||||
}
|
||||
const provider = this.provider;
|
||||
const batchEmbed = this.providerRuntime?.batchEmbed;
|
||||
if (
|
||||
!provider ||
|
||||
!this.batch.enabled ||
|
||||
!batchEmbed ||
|
||||
this.providerRuntime?.sourceWideBatchEmbed !== true
|
||||
) {
|
||||
await runWithConcurrency(
|
||||
items.map((item) => async () => await this.indexFile(item.entry, { source: item.source })),
|
||||
this.getIndexConcurrency(),
|
||||
);
|
||||
return;
|
||||
}
|
||||
|
||||
const itemSourceCounts = countBatchSources(items);
|
||||
log.debug(
|
||||
`memory embeddings: source-wide batch prepare files=${items.length} sources=${formatBatchSourceCounts(
|
||||
itemSourceCounts,
|
||||
)} maxFiles=${SOURCE_WIDE_BATCH_MAX_FILES} maxRequests=${SOURCE_WIDE_BATCH_MAX_REQUESTS}`,
|
||||
{
|
||||
files: items.length,
|
||||
sources: itemSourceCounts,
|
||||
maxFiles: SOURCE_WIDE_BATCH_MAX_FILES,
|
||||
maxRequests: SOURCE_WIDE_BATCH_MAX_REQUESTS,
|
||||
},
|
||||
);
|
||||
|
||||
let prepared: PreparedMemoryIndexEntry[] = [];
|
||||
let preparedRequestCount = 0;
|
||||
let sourceWideBatchGroup = 0;
|
||||
const flushPrepared = async (reason: "max-files" | "max-requests" | "end") => {
|
||||
const firstEntry = prepared[0]?.entry;
|
||||
if (!firstEntry) {
|
||||
return;
|
||||
}
|
||||
const current = prepared;
|
||||
const chunks = current.flatMap((item) => item.chunks);
|
||||
const sourceCounts = countBatchSources(current);
|
||||
const source = formatBatchSourceLabel(sourceCounts);
|
||||
sourceWideBatchGroup += 1;
|
||||
const chunkBatches = splitSourceWideEmbeddingChunks(chunks, SOURCE_WIDE_BATCH_MAX_REQUESTS);
|
||||
log.debug(
|
||||
`memory embeddings: source-wide batch submit group=${sourceWideBatchGroup} source=${source} files=${current.length} chunks=${chunks.length} requests=${chunkBatches.length} sources=${formatBatchSourceCounts(
|
||||
sourceCounts,
|
||||
)} reason=${reason}`,
|
||||
{
|
||||
source,
|
||||
files: current.length,
|
||||
chunks: chunks.length,
|
||||
requests: chunkBatches.length,
|
||||
sources: sourceCounts,
|
||||
group: sourceWideBatchGroup,
|
||||
reason,
|
||||
maxFiles: SOURCE_WIDE_BATCH_MAX_FILES,
|
||||
maxRequests: SOURCE_WIDE_BATCH_MAX_REQUESTS,
|
||||
},
|
||||
);
|
||||
const embeddings: number[][] = [];
|
||||
for (let requestIndex = 0; requestIndex < chunkBatches.length; requestIndex += 1) {
|
||||
const chunkBatch = chunkBatches[requestIndex] ?? [];
|
||||
embeddings.push(
|
||||
...(await this.embedChunksWithBatch(chunkBatch, firstEntry, source, {
|
||||
sourceWideFiles: current.length,
|
||||
sourceWideSources: sourceCounts,
|
||||
sourceWideBatchGroup,
|
||||
sourceWideRequestGroup: requestIndex + 1,
|
||||
sourceWideRequestGroups: chunkBatches.length,
|
||||
})),
|
||||
);
|
||||
}
|
||||
const sample = embeddings.find((embedding) => embedding.length > 0);
|
||||
const vectorReady = sample ? await this.ensureVectorReady(sample.length) : false;
|
||||
let offset = 0;
|
||||
for (const item of current) {
|
||||
const fileEmbeddings = embeddings.slice(offset, offset + item.chunks.length);
|
||||
offset += item.chunks.length;
|
||||
this.writeChunks(
|
||||
item.entry,
|
||||
item.source,
|
||||
provider.model,
|
||||
item.chunks,
|
||||
fileEmbeddings,
|
||||
vectorReady,
|
||||
);
|
||||
}
|
||||
prepared = [];
|
||||
preparedRequestCount = 0;
|
||||
};
|
||||
|
||||
for (const item of items) {
|
||||
if ("kind" in item.entry && item.entry.kind === "multimodal") {
|
||||
await this.indexFile(item.entry, { source: item.source });
|
||||
continue;
|
||||
}
|
||||
const preparedEntry = await this.prepareIndexEntry(item.entry, { source: item.source });
|
||||
if (!preparedEntry) {
|
||||
continue;
|
||||
}
|
||||
const nextWouldExceedFiles = prepared.length >= SOURCE_WIDE_BATCH_MAX_FILES;
|
||||
const nextWouldExceedRequests =
|
||||
preparedRequestCount + preparedEntry.chunks.length > SOURCE_WIDE_BATCH_MAX_REQUESTS;
|
||||
if (prepared.length > 0 && (nextWouldExceedFiles || nextWouldExceedRequests)) {
|
||||
await flushPrepared(nextWouldExceedFiles ? "max-files" : "max-requests");
|
||||
}
|
||||
prepared.push(preparedEntry);
|
||||
preparedRequestCount += preparedEntry.chunks.length;
|
||||
if (
|
||||
prepared.length >= SOURCE_WIDE_BATCH_MAX_FILES ||
|
||||
preparedRequestCount >= SOURCE_WIDE_BATCH_MAX_REQUESTS
|
||||
) {
|
||||
await flushPrepared(
|
||||
prepared.length >= SOURCE_WIDE_BATCH_MAX_FILES ? "max-files" : "max-requests",
|
||||
);
|
||||
}
|
||||
}
|
||||
await flushPrepared("end");
|
||||
}
|
||||
|
||||
protected async indexFile(
|
||||
entry: MemoryIndexEntry,
|
||||
options: { source: MemorySource; content?: string },
|
||||
|
|
@ -757,65 +955,21 @@ export abstract class MemoryManagerEmbeddingOps extends MemoryManagerSyncOps {
|
|||
if ("kind" in entry && entry.kind === "multimodal") {
|
||||
return;
|
||||
}
|
||||
const content =
|
||||
options.content ??
|
||||
(await retryTransientMemoryRead(
|
||||
() => fs.readFile(entry.absPath, "utf-8"),
|
||||
`read memory markdown for indexing ${entry.absPath}`,
|
||||
));
|
||||
const chunks = filterNonEmptyMemoryChunks(chunkMarkdown(content, this.settings.chunking));
|
||||
if (options.source === "sessions" && "lineMap" in entry) {
|
||||
remapChunkLines(chunks, entry.lineMap);
|
||||
}
|
||||
this.writeChunks(entry, options.source, "fts-only", chunks, [], false);
|
||||
const prepared = await this.prepareIndexEntry(entry, options);
|
||||
this.writeChunks(entry, options.source, "fts-only", prepared?.chunks ?? [], [], false);
|
||||
return;
|
||||
}
|
||||
|
||||
let chunks: MemoryChunk[];
|
||||
let structuredInputBytes: number | undefined;
|
||||
if ("kind" in entry && entry.kind === "multimodal") {
|
||||
if (!this.provider) {
|
||||
log.debug("Skipping multimodal indexing in FTS-only mode", {
|
||||
path: entry.path,
|
||||
source: options.source,
|
||||
});
|
||||
this.clearIndexedFileData(entry.path, options.source);
|
||||
this.upsertFileRecord(entry, options.source);
|
||||
return;
|
||||
}
|
||||
const multimodalChunk = await buildMultimodalChunkForIndexing(entry);
|
||||
if (!multimodalChunk) {
|
||||
this.clearIndexedFileData(entry.path, options.source);
|
||||
this.deleteFileRecord(entry.path, options.source);
|
||||
return;
|
||||
}
|
||||
structuredInputBytes = multimodalChunk.structuredInputBytes;
|
||||
chunks = [multimodalChunk.chunk];
|
||||
} else {
|
||||
const content =
|
||||
options.content ??
|
||||
(await retryTransientMemoryRead(
|
||||
() => fs.readFile(entry.absPath, "utf-8"),
|
||||
`read memory markdown for indexing ${entry.absPath}`,
|
||||
));
|
||||
const baseChunks = filterNonEmptyMemoryChunks(chunkMarkdown(content, this.settings.chunking));
|
||||
chunks = this.provider
|
||||
? enforceEmbeddingMaxInputTokens(this.provider, baseChunks, EMBEDDING_BATCH_MAX_TOKENS)
|
||||
: baseChunks;
|
||||
if (options.source === "sessions" && "lineMap" in entry) {
|
||||
remapChunkLines(chunks, entry.lineMap);
|
||||
}
|
||||
}
|
||||
if (!this.provider) {
|
||||
this.writeChunks(entry, options.source, "fts-only", chunks, [], false);
|
||||
const prepared = await this.prepareIndexEntry(entry, options);
|
||||
if (!prepared) {
|
||||
return;
|
||||
}
|
||||
|
||||
let embeddings: number[][];
|
||||
try {
|
||||
embeddings = this.batch.enabled
|
||||
? await this.embedChunksWithBatch(chunks, entry, options.source)
|
||||
: await this.embedChunksInBatches(chunks);
|
||||
? await this.embedChunksWithBatch(prepared.chunks, entry, options.source)
|
||||
: await this.embedChunksInBatches(prepared.chunks);
|
||||
} catch (err) {
|
||||
const message = formatErrorMessage(err);
|
||||
if (
|
||||
|
|
@ -827,7 +981,7 @@ export abstract class MemoryManagerEmbeddingOps extends MemoryManagerSyncOps {
|
|||
) {
|
||||
log.warn("memory embeddings: skipping multimodal file rejected as too large", {
|
||||
path: entry.path,
|
||||
bytes: structuredInputBytes,
|
||||
bytes: prepared.structuredInputBytes,
|
||||
provider: this.provider.id,
|
||||
model: this.provider.model,
|
||||
error: message,
|
||||
|
|
@ -840,6 +994,13 @@ export abstract class MemoryManagerEmbeddingOps extends MemoryManagerSyncOps {
|
|||
}
|
||||
const sample = embeddings.find((embedding) => embedding.length > 0);
|
||||
const vectorReady = sample ? await this.ensureVectorReady(sample.length) : false;
|
||||
this.writeChunks(entry, options.source, this.provider.model, chunks, embeddings, vectorReady);
|
||||
this.writeChunks(
|
||||
entry,
|
||||
options.source,
|
||||
this.provider.model,
|
||||
prepared.chunks,
|
||||
embeddings,
|
||||
vectorReady,
|
||||
);
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -95,13 +95,27 @@ type MemorySyncProgressState = {
|
|||
report: (update: MemorySyncProgressUpdate) => void;
|
||||
};
|
||||
|
||||
type MemoryIndexEntry = {
|
||||
export type MemoryIndexEntry = {
|
||||
path: string;
|
||||
absPath: string;
|
||||
mtimeMs: number;
|
||||
size: number;
|
||||
hash: string;
|
||||
kind?: "markdown" | "multimodal";
|
||||
content?: string;
|
||||
contentText?: string;
|
||||
lineMap?: number[];
|
||||
};
|
||||
|
||||
export type MemoryIndexWorkItem = {
|
||||
entry: MemoryIndexEntry;
|
||||
source: MemorySource;
|
||||
afterIndex?: () => void;
|
||||
};
|
||||
|
||||
type MemorySourceSyncPlan = {
|
||||
indexItems: MemoryIndexWorkItem[];
|
||||
finalize: () => Promise<void> | void;
|
||||
};
|
||||
|
||||
const META_KEY = "memory_index_meta_v1";
|
||||
|
|
@ -111,6 +125,7 @@ const EMBEDDING_CACHE_TABLE = "embedding_cache";
|
|||
const SESSION_DIRTY_DEBOUNCE_MS = 5000;
|
||||
const SESSION_DELTA_READ_CHUNK_BYTES = 64 * 1024;
|
||||
const SESSION_SYNC_YIELD_EVERY = 10;
|
||||
const SOURCE_WIDE_SESSION_INDEX_FLUSH_FILES = 128;
|
||||
const VECTOR_LOAD_TIMEOUT_MS = 30_000;
|
||||
const MEMORY_WATCH_PRESSURE_STARTUP_CHECK_DELAY_MS = 10_000;
|
||||
const IGNORED_MEMORY_WATCH_DIR_NAMES = new Set([
|
||||
|
|
@ -285,6 +300,96 @@ export abstract class MemoryManagerSyncOps {
|
|||
entry: MemoryIndexEntry,
|
||||
options: { source: MemorySource; content?: string },
|
||||
): Promise<void>;
|
||||
protected async indexFiles(items: MemoryIndexWorkItem[]): Promise<void> {
|
||||
for (const item of items) {
|
||||
await this.indexFile(item.entry, { source: item.source });
|
||||
}
|
||||
}
|
||||
|
||||
private emptySourceSyncPlan(): MemorySourceSyncPlan {
|
||||
return { indexItems: [], finalize: () => {} };
|
||||
}
|
||||
|
||||
private shouldDeferSourceWideBatch(): boolean {
|
||||
return Boolean(
|
||||
this.batch.enabled &&
|
||||
this.provider &&
|
||||
this.providerRuntime?.batchEmbed &&
|
||||
this.providerRuntime.sourceWideBatchEmbed === true,
|
||||
);
|
||||
}
|
||||
|
||||
private async indexQueuedFiles(
|
||||
items: MemoryIndexWorkItem[],
|
||||
progress?: MemorySyncProgressState,
|
||||
label?: string,
|
||||
): Promise<void> {
|
||||
if (items.length === 0) {
|
||||
return;
|
||||
}
|
||||
if (progress && label) {
|
||||
progress.report({
|
||||
completed: progress.completed,
|
||||
total: progress.total,
|
||||
label,
|
||||
});
|
||||
}
|
||||
await this.indexFiles(items);
|
||||
for (const item of items) {
|
||||
item.afterIndex?.();
|
||||
}
|
||||
if (progress) {
|
||||
progress.completed += items.length;
|
||||
progress.report({
|
||||
completed: progress.completed,
|
||||
total: progress.total,
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
private async executeSourceSyncPlans(
|
||||
plans: MemorySourceSyncPlan[],
|
||||
progress?: MemorySyncProgressState,
|
||||
): Promise<void> {
|
||||
const indexItems = plans.flatMap((plan) => plan.indexItems);
|
||||
const sources = new Set(indexItems.map((item) => item.source));
|
||||
await this.indexQueuedFiles(
|
||||
indexItems,
|
||||
progress,
|
||||
sources.size > 1 ? "Indexing memory sources (batch)..." : undefined,
|
||||
);
|
||||
for (const plan of plans) {
|
||||
await plan.finalize();
|
||||
}
|
||||
}
|
||||
|
||||
private async executeSourceWideSync(params: {
|
||||
shouldSyncMemory: boolean;
|
||||
shouldSyncSessions: boolean;
|
||||
needsFullReindex: boolean;
|
||||
targetSessionFiles?: string[];
|
||||
progress?: MemorySyncProgressState;
|
||||
}): Promise<void> {
|
||||
const memoryPlan = params.shouldSyncMemory
|
||||
? await this.syncMemoryFiles({
|
||||
needsFullReindex: params.needsFullReindex,
|
||||
progress: params.progress,
|
||||
deferIndex: true,
|
||||
})
|
||||
: this.emptySourceSyncPlan();
|
||||
if (params.shouldSyncSessions) {
|
||||
await this.syncSessionFiles({
|
||||
needsFullReindex: params.needsFullReindex,
|
||||
targetSessionFiles: params.targetSessionFiles,
|
||||
progress: params.progress,
|
||||
deferIndex: true,
|
||||
prefixIndexItems: memoryPlan.indexItems,
|
||||
});
|
||||
await memoryPlan.finalize();
|
||||
return;
|
||||
}
|
||||
await this.executeSourceSyncPlans([memoryPlan], params.progress);
|
||||
}
|
||||
|
||||
protected hasIndexedChunks(): boolean {
|
||||
const row = this.db.prepare(`SELECT 1 as found FROM chunks LIMIT 1`).get() as
|
||||
|
|
@ -1504,7 +1609,8 @@ export abstract class MemoryManagerSyncOps {
|
|||
private async syncMemoryFiles(params: {
|
||||
needsFullReindex: boolean;
|
||||
progress?: MemorySyncProgressState;
|
||||
}) {
|
||||
deferIndex?: boolean;
|
||||
}): Promise<MemorySourceSyncPlan> {
|
||||
const deleteFileByPathAndSource = this.db.prepare(
|
||||
`DELETE FROM files WHERE path = ? AND source = ?`,
|
||||
);
|
||||
|
|
@ -1558,8 +1664,61 @@ export abstract class MemoryManagerSyncOps {
|
|||
});
|
||||
}
|
||||
|
||||
const tasks = fileEntries.map((entry) => async () => {
|
||||
if (!params.needsFullReindex && existingHashes.get(entry.path) === entry.hash) {
|
||||
const deleteStaleRows = async () => {
|
||||
for (const stale of existingRows) {
|
||||
if (activePaths.has(stale.path)) {
|
||||
continue;
|
||||
}
|
||||
deleteFileByPathAndSource.run(stale.path, "memory");
|
||||
if (deleteVectorRowsByPathAndSource) {
|
||||
try {
|
||||
deleteVectorRowsByPathAndSource.run(stale.path, "memory");
|
||||
} catch {}
|
||||
}
|
||||
deleteChunksByPathAndSource.run(stale.path, "memory");
|
||||
if (deleteFtsRowsByPathAndSource) {
|
||||
try {
|
||||
deleteFtsRowsByPathAndSource.run(stale.path, "memory");
|
||||
} catch {}
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
if (this.batch.enabled) {
|
||||
const dirtyEntries: MemoryIndexEntry[] = [];
|
||||
for (const entry of fileEntries) {
|
||||
if (!params.needsFullReindex && existingHashes.get(entry.path) === entry.hash) {
|
||||
if (params.progress) {
|
||||
params.progress.completed += 1;
|
||||
params.progress.report({
|
||||
completed: params.progress.completed,
|
||||
total: params.progress.total,
|
||||
});
|
||||
}
|
||||
continue;
|
||||
}
|
||||
dirtyEntries.push(entry);
|
||||
}
|
||||
const indexItems = dirtyEntries.map(
|
||||
(entry): MemoryIndexWorkItem => ({ entry, source: "memory" }),
|
||||
);
|
||||
if (params.deferIndex) {
|
||||
return { indexItems, finalize: deleteStaleRows };
|
||||
}
|
||||
await this.indexQueuedFiles(indexItems, params.progress);
|
||||
} else {
|
||||
const tasks = fileEntries.map((entry) => async () => {
|
||||
if (!params.needsFullReindex && existingHashes.get(entry.path) === entry.hash) {
|
||||
if (params.progress) {
|
||||
params.progress.completed += 1;
|
||||
params.progress.report({
|
||||
completed: params.progress.completed,
|
||||
total: params.progress.total,
|
||||
});
|
||||
}
|
||||
return;
|
||||
}
|
||||
await this.indexFile(entry, { source: "memory" });
|
||||
if (params.progress) {
|
||||
params.progress.completed += 1;
|
||||
params.progress.report({
|
||||
|
|
@ -1567,43 +1726,21 @@ export abstract class MemoryManagerSyncOps {
|
|||
total: params.progress.total,
|
||||
});
|
||||
}
|
||||
return;
|
||||
}
|
||||
await this.indexFile(entry, { source: "memory" });
|
||||
if (params.progress) {
|
||||
params.progress.completed += 1;
|
||||
params.progress.report({
|
||||
completed: params.progress.completed,
|
||||
total: params.progress.total,
|
||||
});
|
||||
}
|
||||
});
|
||||
await runWithConcurrency(tasks, this.getIndexConcurrency());
|
||||
|
||||
for (const stale of existingRows) {
|
||||
if (activePaths.has(stale.path)) {
|
||||
continue;
|
||||
}
|
||||
deleteFileByPathAndSource.run(stale.path, "memory");
|
||||
if (deleteVectorRowsByPathAndSource) {
|
||||
try {
|
||||
deleteVectorRowsByPathAndSource.run(stale.path, "memory");
|
||||
} catch {}
|
||||
}
|
||||
deleteChunksByPathAndSource.run(stale.path, "memory");
|
||||
if (deleteFtsRowsByPathAndSource) {
|
||||
try {
|
||||
deleteFtsRowsByPathAndSource.run(stale.path, "memory");
|
||||
} catch {}
|
||||
}
|
||||
});
|
||||
await runWithConcurrency(tasks, this.getIndexConcurrency());
|
||||
}
|
||||
|
||||
await deleteStaleRows();
|
||||
return this.emptySourceSyncPlan();
|
||||
}
|
||||
|
||||
private async syncSessionFiles(params: {
|
||||
needsFullReindex: boolean;
|
||||
targetSessionFiles?: string[];
|
||||
progress?: MemorySyncProgressState;
|
||||
}) {
|
||||
deferIndex?: boolean;
|
||||
prefixIndexItems?: MemoryIndexWorkItem[];
|
||||
}): Promise<MemorySourceSyncPlan> {
|
||||
const deleteFileByPathAndSource = this.db.prepare(
|
||||
`DELETE FROM files WHERE path = ? AND source = ?`,
|
||||
);
|
||||
|
|
@ -1659,6 +1796,127 @@ export abstract class MemoryManagerSyncOps {
|
|||
}
|
||||
|
||||
const yieldAfterSessionFile = createSessionSyncYield(files.length);
|
||||
const deleteStaleRows = async () => {
|
||||
if (activePaths === null) {
|
||||
return;
|
||||
}
|
||||
|
||||
const staleRows = existingRows ?? [];
|
||||
const yieldAfterStaleSessionRow = createSessionSyncYield(staleRows.length);
|
||||
for (const stale of staleRows) {
|
||||
try {
|
||||
if (activePaths.has(stale.path)) {
|
||||
continue;
|
||||
}
|
||||
deleteFileByPathAndSource.run(stale.path, "sessions");
|
||||
if (deleteVectorRowsByPathAndSource) {
|
||||
try {
|
||||
deleteVectorRowsByPathAndSource.run(stale.path, "sessions");
|
||||
} catch {}
|
||||
}
|
||||
deleteChunksByPathAndSource.run(stale.path, "sessions");
|
||||
if (deleteFtsRowsByPathAndSource) {
|
||||
try {
|
||||
deleteFtsRowsByPathAndSource.run(stale.path, "sessions");
|
||||
} catch {}
|
||||
}
|
||||
} finally {
|
||||
await yieldAfterStaleSessionRow();
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
if (params.deferIndex) {
|
||||
const pendingIndexItems = [...(params.prefixIndexItems ?? [])];
|
||||
const flushPendingIndexItems = async () => {
|
||||
if (pendingIndexItems.length === 0) {
|
||||
return;
|
||||
}
|
||||
const current = pendingIndexItems.splice(0);
|
||||
const sources = new Set(current.map((item) => item.source));
|
||||
await this.indexQueuedFiles(
|
||||
current,
|
||||
params.progress,
|
||||
sources.size > 1 ? "Indexing memory sources (batch)..." : undefined,
|
||||
);
|
||||
};
|
||||
|
||||
// Session entries carry flattened transcript content; flush bounded groups
|
||||
// so source-wide batching cannot retain the whole dirty transcript corpus.
|
||||
for (let start = 0; start < files.length; start += SOURCE_WIDE_SESSION_INDEX_FLUSH_FILES) {
|
||||
const fileBatch = files.slice(start, start + SOURCE_WIDE_SESSION_INDEX_FLUSH_FILES);
|
||||
const dirtyEntries = (
|
||||
await runWithConcurrency(
|
||||
fileBatch.map((absPath) => async (): Promise<MemoryIndexEntry | null> => {
|
||||
try {
|
||||
if (!indexAll && !this.sessionsDirtyFiles.has(absPath)) {
|
||||
if (params.progress) {
|
||||
params.progress.completed += 1;
|
||||
params.progress.report({
|
||||
completed: params.progress.completed,
|
||||
total: params.progress.total,
|
||||
});
|
||||
}
|
||||
return null;
|
||||
}
|
||||
const entry = await buildSessionEntry(absPath);
|
||||
if (!entry) {
|
||||
if (params.progress) {
|
||||
params.progress.completed += 1;
|
||||
params.progress.report({
|
||||
completed: params.progress.completed,
|
||||
total: params.progress.total,
|
||||
});
|
||||
}
|
||||
return null;
|
||||
}
|
||||
const existingHash = resolveMemorySourceExistingHash({
|
||||
db: this.db,
|
||||
source: "sessions",
|
||||
path: entry.path,
|
||||
existingHashes,
|
||||
});
|
||||
if (!params.needsFullReindex && existingHash === entry.hash) {
|
||||
if (params.progress) {
|
||||
params.progress.completed += 1;
|
||||
params.progress.report({
|
||||
completed: params.progress.completed,
|
||||
total: params.progress.total,
|
||||
});
|
||||
}
|
||||
this.resetSessionDelta(absPath, entry.size);
|
||||
return null;
|
||||
}
|
||||
return entry;
|
||||
} finally {
|
||||
await yieldAfterSessionFile();
|
||||
}
|
||||
}),
|
||||
this.getIndexConcurrency(),
|
||||
)
|
||||
).filter((entry): entry is MemoryIndexEntry => entry !== null);
|
||||
pendingIndexItems.push(
|
||||
...dirtyEntries.map(
|
||||
(entry): MemoryIndexWorkItem => ({
|
||||
entry,
|
||||
source: "sessions",
|
||||
afterIndex: () => this.resetSessionDelta(entry.absPath, entry.size),
|
||||
}),
|
||||
),
|
||||
);
|
||||
if (pendingIndexItems.length >= SOURCE_WIDE_SESSION_INDEX_FLUSH_FILES) {
|
||||
await flushPendingIndexItems();
|
||||
}
|
||||
}
|
||||
|
||||
await flushPendingIndexItems();
|
||||
await deleteStaleRows();
|
||||
return this.emptySourceSyncPlan();
|
||||
}
|
||||
if ((params.prefixIndexItems?.length ?? 0) > 0) {
|
||||
throw new Error("Memory session sync prefix requires deferred source-wide indexing.");
|
||||
}
|
||||
|
||||
const tasks = files.map((absPath) => async () => {
|
||||
try {
|
||||
if (!indexAll && !this.sessionsDirtyFiles.has(absPath)) {
|
||||
|
|
@ -1714,35 +1972,8 @@ export abstract class MemoryManagerSyncOps {
|
|||
});
|
||||
await runWithConcurrency(tasks, this.getIndexConcurrency());
|
||||
|
||||
if (activePaths === null) {
|
||||
// Targeted syncs only refresh the requested transcripts and should not
|
||||
// prune unrelated session rows without a full directory enumeration.
|
||||
return;
|
||||
}
|
||||
|
||||
const staleRows = existingRows ?? [];
|
||||
const yieldAfterStaleSessionRow = createSessionSyncYield(staleRows.length);
|
||||
for (const stale of staleRows) {
|
||||
try {
|
||||
if (activePaths.has(stale.path)) {
|
||||
continue;
|
||||
}
|
||||
deleteFileByPathAndSource.run(stale.path, "sessions");
|
||||
if (deleteVectorRowsByPathAndSource) {
|
||||
try {
|
||||
deleteVectorRowsByPathAndSource.run(stale.path, "sessions");
|
||||
} catch {}
|
||||
}
|
||||
deleteChunksByPathAndSource.run(stale.path, "sessions");
|
||||
if (deleteFtsRowsByPathAndSource) {
|
||||
try {
|
||||
deleteFtsRowsByPathAndSource.run(stale.path, "sessions");
|
||||
} catch {}
|
||||
}
|
||||
} finally {
|
||||
await yieldAfterStaleSessionRow();
|
||||
}
|
||||
}
|
||||
await deleteStaleRows();
|
||||
return this.emptySourceSyncPlan();
|
||||
}
|
||||
|
||||
private createSyncProgress(
|
||||
|
|
@ -1910,23 +2141,44 @@ export abstract class MemoryManagerSyncOps {
|
|||
((!hasTargetSessionFiles && params?.force) || needsFullReindex || this.dirty);
|
||||
const shouldSyncSessions = this.shouldSyncSessions(params, needsFullReindex);
|
||||
|
||||
if (shouldSyncMemory) {
|
||||
await this.syncMemoryFiles({ needsFullReindex, progress: progress ?? undefined });
|
||||
this.dirty = false;
|
||||
}
|
||||
|
||||
if (shouldSyncSessions) {
|
||||
await this.syncSessionFiles({
|
||||
if (this.shouldDeferSourceWideBatch()) {
|
||||
await this.executeSourceWideSync({
|
||||
shouldSyncMemory,
|
||||
shouldSyncSessions,
|
||||
needsFullReindex,
|
||||
targetSessionFiles: targetSessionFiles ? Array.from(targetSessionFiles) : undefined,
|
||||
progress: progress ?? undefined,
|
||||
});
|
||||
this.sessionsDirty = false;
|
||||
this.sessionsDirtyFiles.clear();
|
||||
} else if (this.sessionsDirtyFiles.size > 0) {
|
||||
this.sessionsDirty = true;
|
||||
if (shouldSyncMemory) {
|
||||
this.dirty = false;
|
||||
}
|
||||
if (shouldSyncSessions) {
|
||||
this.sessionsDirty = false;
|
||||
this.sessionsDirtyFiles.clear();
|
||||
} else if (this.sessionsDirtyFiles.size > 0) {
|
||||
this.sessionsDirty = true;
|
||||
} else {
|
||||
this.sessionsDirty = false;
|
||||
}
|
||||
} else {
|
||||
this.sessionsDirty = false;
|
||||
if (shouldSyncMemory) {
|
||||
await this.syncMemoryFiles({ needsFullReindex, progress: progress ?? undefined });
|
||||
this.dirty = false;
|
||||
}
|
||||
|
||||
if (shouldSyncSessions) {
|
||||
await this.syncSessionFiles({
|
||||
needsFullReindex,
|
||||
targetSessionFiles: targetSessionFiles ? Array.from(targetSessionFiles) : undefined,
|
||||
progress: progress ?? undefined,
|
||||
});
|
||||
this.sessionsDirty = false;
|
||||
this.sessionsDirtyFiles.clear();
|
||||
} else if (this.sessionsDirtyFiles.size > 0) {
|
||||
this.sessionsDirty = true;
|
||||
} else {
|
||||
this.sessionsDirty = false;
|
||||
}
|
||||
}
|
||||
} catch (err) {
|
||||
const reason = formatErrorMessage(err);
|
||||
|
|
@ -2062,6 +2314,7 @@ export abstract class MemoryManagerSyncOps {
|
|||
};
|
||||
|
||||
this.db = tempDb;
|
||||
this.lastMetaSerialized = null;
|
||||
this.resetVectorState();
|
||||
this.fts.available = false;
|
||||
this.fts.loadError = undefined;
|
||||
|
|
@ -2087,19 +2340,39 @@ export abstract class MemoryManagerSyncOps {
|
|||
true,
|
||||
);
|
||||
|
||||
if (shouldSyncMemory) {
|
||||
await this.syncMemoryFiles({ needsFullReindex: true, progress: params.progress });
|
||||
this.dirty = false;
|
||||
}
|
||||
|
||||
if (shouldSyncSessions) {
|
||||
await this.syncSessionFiles({ needsFullReindex: true, progress: params.progress });
|
||||
this.sessionsDirty = false;
|
||||
this.sessionsDirtyFiles.clear();
|
||||
} else if (this.sessionsDirtyFiles.size > 0) {
|
||||
this.sessionsDirty = true;
|
||||
if (this.shouldDeferSourceWideBatch()) {
|
||||
await this.executeSourceWideSync({
|
||||
shouldSyncMemory,
|
||||
shouldSyncSessions,
|
||||
needsFullReindex: true,
|
||||
progress: params.progress,
|
||||
});
|
||||
if (shouldSyncMemory) {
|
||||
this.dirty = false;
|
||||
}
|
||||
if (shouldSyncSessions) {
|
||||
this.sessionsDirty = false;
|
||||
this.sessionsDirtyFiles.clear();
|
||||
} else if (this.sessionsDirtyFiles.size > 0) {
|
||||
this.sessionsDirty = true;
|
||||
} else {
|
||||
this.sessionsDirty = false;
|
||||
}
|
||||
} else {
|
||||
this.sessionsDirty = false;
|
||||
if (shouldSyncMemory) {
|
||||
await this.syncMemoryFiles({ needsFullReindex: true, progress: params.progress });
|
||||
this.dirty = false;
|
||||
}
|
||||
|
||||
if (shouldSyncSessions) {
|
||||
await this.syncSessionFiles({ needsFullReindex: true, progress: params.progress });
|
||||
this.sessionsDirty = false;
|
||||
this.sessionsDirtyFiles.clear();
|
||||
} else if (this.sessionsDirtyFiles.size > 0) {
|
||||
this.sessionsDirty = true;
|
||||
} else {
|
||||
this.sessionsDirty = false;
|
||||
}
|
||||
}
|
||||
if (!shouldSyncMemory) {
|
||||
this.dirty = false;
|
||||
|
|
@ -2170,19 +2443,39 @@ export abstract class MemoryManagerSyncOps {
|
|||
true,
|
||||
);
|
||||
|
||||
if (shouldSyncMemory) {
|
||||
await this.syncMemoryFiles({ needsFullReindex: true, progress: params.progress });
|
||||
this.dirty = false;
|
||||
}
|
||||
|
||||
if (shouldSyncSessions) {
|
||||
await this.syncSessionFiles({ needsFullReindex: true, progress: params.progress });
|
||||
this.sessionsDirty = false;
|
||||
this.sessionsDirtyFiles.clear();
|
||||
} else if (this.sessionsDirtyFiles.size > 0) {
|
||||
this.sessionsDirty = true;
|
||||
if (this.shouldDeferSourceWideBatch()) {
|
||||
await this.executeSourceWideSync({
|
||||
shouldSyncMemory,
|
||||
shouldSyncSessions,
|
||||
needsFullReindex: true,
|
||||
progress: params.progress,
|
||||
});
|
||||
if (shouldSyncMemory) {
|
||||
this.dirty = false;
|
||||
}
|
||||
if (shouldSyncSessions) {
|
||||
this.sessionsDirty = false;
|
||||
this.sessionsDirtyFiles.clear();
|
||||
} else if (this.sessionsDirtyFiles.size > 0) {
|
||||
this.sessionsDirty = true;
|
||||
} else {
|
||||
this.sessionsDirty = false;
|
||||
}
|
||||
} else {
|
||||
this.sessionsDirty = false;
|
||||
if (shouldSyncMemory) {
|
||||
await this.syncMemoryFiles({ needsFullReindex: true, progress: params.progress });
|
||||
this.dirty = false;
|
||||
}
|
||||
|
||||
if (shouldSyncSessions) {
|
||||
await this.syncSessionFiles({ needsFullReindex: true, progress: params.progress });
|
||||
this.sessionsDirty = false;
|
||||
this.sessionsDirtyFiles.clear();
|
||||
} else if (this.sessionsDirtyFiles.size > 0) {
|
||||
this.sessionsDirty = true;
|
||||
} else {
|
||||
this.sessionsDirty = false;
|
||||
}
|
||||
}
|
||||
if (!shouldSyncMemory) {
|
||||
this.dirty = false;
|
||||
|
|
|
|||
|
|
@ -1,6 +1,36 @@
|
|||
// Openai tests cover embedding batch plugin behavior.
|
||||
import { describe, expect, it } from "vitest";
|
||||
import { parseOpenAiBatchOutput } from "./embedding-batch.js";
|
||||
import { describe, expect, it, vi } from "vitest";
|
||||
import { parseOpenAiBatchOutput, runOpenAiEmbeddingBatches } from "./embedding-batch.js";
|
||||
|
||||
const jsonlEncoder = new TextEncoder();
|
||||
|
||||
function jsonResponse(body: unknown, status = 200): Response {
|
||||
return new Response(JSON.stringify(body), {
|
||||
status,
|
||||
headers: { "Content-Type": "application/json" },
|
||||
});
|
||||
}
|
||||
|
||||
function jsonlBytes(value: string): number {
|
||||
return jsonlEncoder.encode(value).byteLength;
|
||||
}
|
||||
|
||||
function fetchInputUrl(input: RequestInfo | URL): string {
|
||||
if (typeof input === "string") {
|
||||
return input;
|
||||
}
|
||||
if (input instanceof URL) {
|
||||
return input.href;
|
||||
}
|
||||
return input.url;
|
||||
}
|
||||
|
||||
function parseStringBody(init: RequestInit | undefined): unknown {
|
||||
if (typeof init?.body !== "string") {
|
||||
throw new Error("missing JSON request body");
|
||||
}
|
||||
return JSON.parse(init.body) as unknown;
|
||||
}
|
||||
|
||||
describe("OpenAI embedding batch output", () => {
|
||||
it("wraps malformed JSONL output", () => {
|
||||
|
|
@ -8,4 +38,209 @@ describe("OpenAI embedding batch output", () => {
|
|||
"OpenAI embedding batch output contained malformed JSONL",
|
||||
);
|
||||
});
|
||||
|
||||
it("splits provider uploads by serialized JSONL byte cap", async () => {
|
||||
const requests: Parameters<typeof runOpenAiEmbeddingBatches>[0]["requests"] = Array.from(
|
||||
{ length: 3 },
|
||||
(_, index) => ({
|
||||
custom_id: String(index),
|
||||
method: "POST" as const,
|
||||
url: "/v1/embeddings",
|
||||
body: {
|
||||
model: "text-embedding-3-small",
|
||||
input: `payload-${index}-${"β".repeat(8)}`,
|
||||
},
|
||||
}),
|
||||
);
|
||||
const uploadedJsonl: string[] = [];
|
||||
const requestsByFileId = new Map<string, Array<{ custom_id?: string }>>();
|
||||
const outputByFileId = new Map<string, string>();
|
||||
let fileIndex = 0;
|
||||
let batchIndex = 0;
|
||||
const maxJsonlBytes = jsonlBytes(JSON.stringify(requests[0]));
|
||||
const fetchImpl = vi.fn(async (input: RequestInfo | URL, init?: RequestInit) => {
|
||||
const url = fetchInputUrl(input);
|
||||
if (url.endsWith("/files") && init?.method === "POST") {
|
||||
const form = init.body as FormData;
|
||||
const file = form.get("file");
|
||||
if (!(file instanceof Blob)) {
|
||||
throw new Error("missing batch upload file");
|
||||
}
|
||||
const jsonl = await file.text();
|
||||
const fileId = `file-${fileIndex}`;
|
||||
fileIndex += 1;
|
||||
uploadedJsonl.push(jsonl);
|
||||
requestsByFileId.set(
|
||||
fileId,
|
||||
jsonl.split("\n").map((line) => JSON.parse(line) as { custom_id?: string }),
|
||||
);
|
||||
return jsonResponse({ id: fileId });
|
||||
}
|
||||
if (url.endsWith("/batches") && init?.method === "POST") {
|
||||
const body = parseStringBody(init) as { input_file_id?: string };
|
||||
const batchId = `batch-${batchIndex}`;
|
||||
const outputFileId = `output-${batchIndex}`;
|
||||
batchIndex += 1;
|
||||
const uploadedRequests = requestsByFileId.get(body.input_file_id ?? "") ?? [];
|
||||
outputByFileId.set(
|
||||
outputFileId,
|
||||
uploadedRequests
|
||||
.map((request) =>
|
||||
JSON.stringify({
|
||||
custom_id: request.custom_id,
|
||||
response: {
|
||||
status_code: 200,
|
||||
body: { data: [{ embedding: [Number(request.custom_id) + 1] }] },
|
||||
},
|
||||
}),
|
||||
)
|
||||
.join("\n"),
|
||||
);
|
||||
return jsonResponse({ id: batchId, status: "completed", output_file_id: outputFileId });
|
||||
}
|
||||
const contentMatch = url.match(/\/files\/([^/]+)\/content$/);
|
||||
if (contentMatch) {
|
||||
return new Response(outputByFileId.get(contentMatch[1] ?? "") ?? "", { status: 200 });
|
||||
}
|
||||
return new Response("unexpected request", { status: 500 });
|
||||
});
|
||||
|
||||
const byCustomId = await runOpenAiEmbeddingBatches({
|
||||
openAi: {
|
||||
baseUrl: "https://openai-compatible.example/v1",
|
||||
headers: { Authorization: "Bearer test" },
|
||||
model: "text-embedding-3-small",
|
||||
fetchImpl,
|
||||
},
|
||||
agentId: "main",
|
||||
requests,
|
||||
maxJsonlBytes,
|
||||
wait: true,
|
||||
concurrency: 1,
|
||||
pollIntervalMs: 1000,
|
||||
timeoutMs: 60_000,
|
||||
});
|
||||
|
||||
expect(uploadedJsonl).toHaveLength(3);
|
||||
expect(uploadedJsonl.every((jsonl) => jsonlBytes(jsonl) <= maxJsonlBytes)).toBe(true);
|
||||
expect([...byCustomId.entries()]).toEqual([
|
||||
["0", [1]],
|
||||
["1", [2]],
|
||||
["2", [3]],
|
||||
]);
|
||||
});
|
||||
|
||||
it("adapts OpenAI-compatible upload groups after payload-size rejection", async () => {
|
||||
const requests: Parameters<typeof runOpenAiEmbeddingBatches>[0]["requests"] = Array.from(
|
||||
{ length: 4 },
|
||||
(_, index) => ({
|
||||
custom_id: String(index),
|
||||
method: "POST" as const,
|
||||
url: "/v1/embeddings",
|
||||
body: {
|
||||
model: "text-embedding-3-small",
|
||||
input: `payload-${index}`,
|
||||
},
|
||||
}),
|
||||
);
|
||||
const uploadedGroups: string[][] = [];
|
||||
const requestsByFileId = new Map<string, Array<{ custom_id?: string }>>();
|
||||
const outputByFileId = new Map<string, string>();
|
||||
const debug = vi.fn();
|
||||
let fileIndex = 0;
|
||||
let batchIndex = 0;
|
||||
const fetchImpl = vi.fn(async (input: RequestInfo | URL, init?: RequestInit) => {
|
||||
const url = fetchInputUrl(input);
|
||||
if (url.endsWith("/files") && init?.method === "POST") {
|
||||
const form = init.body as FormData;
|
||||
const file = form.get("file");
|
||||
if (!(file instanceof Blob)) {
|
||||
throw new Error("missing batch upload file");
|
||||
}
|
||||
const uploadedRequests = (await file.text())
|
||||
.split("\n")
|
||||
.map((line) => JSON.parse(line) as { custom_id?: string });
|
||||
const customIds = uploadedRequests.map((request) => request.custom_id ?? "");
|
||||
uploadedGroups.push(customIds);
|
||||
if (uploadedRequests.length > 2) {
|
||||
return jsonResponse(
|
||||
{
|
||||
error: {
|
||||
message: "Request body too large. Maximum allowed: 10 MB",
|
||||
type: "payload_too_large",
|
||||
code: "PAYLOAD_TOO_LARGE",
|
||||
},
|
||||
},
|
||||
413,
|
||||
);
|
||||
}
|
||||
const fileId = `file-${fileIndex}`;
|
||||
fileIndex += 1;
|
||||
requestsByFileId.set(fileId, uploadedRequests);
|
||||
return jsonResponse({ id: fileId });
|
||||
}
|
||||
if (url.endsWith("/batches") && init?.method === "POST") {
|
||||
const body = parseStringBody(init) as { input_file_id?: string };
|
||||
const batchId = `batch-${batchIndex}`;
|
||||
const outputFileId = `output-${batchIndex}`;
|
||||
batchIndex += 1;
|
||||
const uploadedRequests = requestsByFileId.get(body.input_file_id ?? "") ?? [];
|
||||
outputByFileId.set(
|
||||
outputFileId,
|
||||
uploadedRequests
|
||||
.map((request) =>
|
||||
JSON.stringify({
|
||||
custom_id: request.custom_id,
|
||||
response: {
|
||||
status_code: 200,
|
||||
body: { data: [{ embedding: [Number(request.custom_id) + 1] }] },
|
||||
},
|
||||
}),
|
||||
)
|
||||
.join("\n"),
|
||||
);
|
||||
return jsonResponse({ id: batchId, status: "completed", output_file_id: outputFileId });
|
||||
}
|
||||
const contentMatch = url.match(/\/files\/([^/]+)\/content$/);
|
||||
if (contentMatch) {
|
||||
return new Response(outputByFileId.get(contentMatch[1] ?? "") ?? "", { status: 200 });
|
||||
}
|
||||
return new Response("unexpected request", { status: 500 });
|
||||
});
|
||||
|
||||
const byCustomId = await runOpenAiEmbeddingBatches({
|
||||
openAi: {
|
||||
baseUrl: "https://openai-compatible.example/v1",
|
||||
headers: { Authorization: "Bearer test" },
|
||||
model: "text-embedding-3-small",
|
||||
fetchImpl,
|
||||
},
|
||||
agentId: "main",
|
||||
requests,
|
||||
wait: true,
|
||||
concurrency: 1,
|
||||
pollIntervalMs: 1000,
|
||||
timeoutMs: 60_000,
|
||||
debug,
|
||||
});
|
||||
|
||||
expect(uploadedGroups).toEqual([
|
||||
["0", "1", "2", "3"],
|
||||
["0", "1"],
|
||||
["2", "3"],
|
||||
]);
|
||||
expect(debug).toHaveBeenCalledWith(
|
||||
"memory embeddings: openai batch upload too large; splitting group",
|
||||
expect.objectContaining({
|
||||
requests: 4,
|
||||
parts: [2, 2],
|
||||
}),
|
||||
);
|
||||
expect([...byCustomId.entries()]).toEqual([
|
||||
["0", [1]],
|
||||
["1", [2]],
|
||||
["2", [3]],
|
||||
["3", [4]],
|
||||
]);
|
||||
});
|
||||
});
|
||||
|
|
|
|||
|
|
@ -39,12 +39,22 @@ type OpenAiBatchRequest = {
|
|||
};
|
||||
};
|
||||
|
||||
type OpenAiBatchStatus = EmbeddingBatchStatus;
|
||||
type OpenAiBatchStatus = EmbeddingBatchStatus & {
|
||||
request_counts?: {
|
||||
total?: number;
|
||||
completed?: number;
|
||||
failed?: number;
|
||||
};
|
||||
};
|
||||
type OpenAiBatchOutputLine = ProviderBatchOutputLine;
|
||||
|
||||
export const OPENAI_BATCH_ENDPOINT = EMBEDDING_BATCH_ENDPOINT;
|
||||
const OPENAI_BATCH_COMPLETION_WINDOW = "24h";
|
||||
const OPENAI_BATCH_MAX_REQUESTS = 50000;
|
||||
// OpenAI accepts 200 MB Batch input files. Keep a safety margin so the JSONL
|
||||
// splitter avoids boundary-size uploads while preserving source-wide batching.
|
||||
const OPENAI_BATCH_MAX_JSONL_BYTES = 190 * 1024 * 1024;
|
||||
const OPENAI_BATCH_MAX_POLL_BACKOFF_MS = 5 * 60_000;
|
||||
|
||||
async function submitOpenAiBatch(params: {
|
||||
openAi: OpenAiEmbeddingClient;
|
||||
|
|
@ -124,6 +134,25 @@ async function fetchOpenAiBatchResource<T>(params: {
|
|||
});
|
||||
}
|
||||
|
||||
function formatOpenAiBatchError(error: unknown): string {
|
||||
return error instanceof Error ? error.message : String(error);
|
||||
}
|
||||
|
||||
function isOpenAiBatchUploadTooLargeError(error: unknown): boolean {
|
||||
const message = formatOpenAiBatchError(error);
|
||||
if (!/openai batch file upload failed/i.test(message)) {
|
||||
return false;
|
||||
}
|
||||
return (
|
||||
/\b413\b/.test(message) ||
|
||||
/payload too large/i.test(message) ||
|
||||
/request body too large/i.test(message) ||
|
||||
/file too large/i.test(message) ||
|
||||
/maximum allowed/i.test(message) ||
|
||||
/max(?:imum)? (?:body|payload|file) (?:size )?(?:exceeded|limit)/i.test(message)
|
||||
);
|
||||
}
|
||||
|
||||
export function parseOpenAiBatchOutput(text: string): OpenAiBatchOutputLine[] {
|
||||
if (!text.trim()) {
|
||||
return [];
|
||||
|
|
@ -153,6 +182,45 @@ async function readOpenAiBatchError(params: {
|
|||
}
|
||||
}
|
||||
|
||||
function createOpenAiBatchPollBackoff(params: { pollIntervalMs: number; timeoutMs: number }): {
|
||||
nextDelayMs: () => number;
|
||||
} {
|
||||
const maxDelayMs = Math.max(
|
||||
params.pollIntervalMs,
|
||||
Math.min(params.timeoutMs, OPENAI_BATCH_MAX_POLL_BACKOFF_MS),
|
||||
);
|
||||
let delayMs = params.pollIntervalMs;
|
||||
return {
|
||||
nextDelayMs: () => {
|
||||
const current = delayMs;
|
||||
delayMs = Math.min(maxDelayMs, current * 2);
|
||||
return current;
|
||||
},
|
||||
};
|
||||
}
|
||||
|
||||
function formatOpenAiBatchProgress(status: OpenAiBatchStatus): string {
|
||||
const counts = status.request_counts;
|
||||
if (!counts || typeof counts.total !== "number") {
|
||||
return "";
|
||||
}
|
||||
const completed = typeof counts.completed === "number" ? counts.completed : 0;
|
||||
const failed = typeof counts.failed === "number" ? counts.failed : 0;
|
||||
return `; progress ${completed}/${counts.total} failed=${failed}`;
|
||||
}
|
||||
|
||||
function formatOpenAiBatchPollError(error: unknown): string {
|
||||
return error instanceof Error ? error.message : String(error);
|
||||
}
|
||||
|
||||
function isRetryableOpenAiBatchPollError(error: unknown): boolean {
|
||||
const message = formatOpenAiBatchPollError(error);
|
||||
return (
|
||||
/openai batch status failed: (408|409|425|429|5\d\d)\b/i.test(message) ||
|
||||
/\b(ECONNRESET|ECONNREFUSED|ETIMEDOUT|EAI_AGAIN)\b|fetch failed|network error/i.test(message)
|
||||
);
|
||||
}
|
||||
|
||||
async function waitForOpenAiBatch(params: {
|
||||
openAi: OpenAiEmbeddingClient;
|
||||
batchId: string;
|
||||
|
|
@ -163,14 +231,38 @@ async function waitForOpenAiBatch(params: {
|
|||
initial?: OpenAiBatchStatus;
|
||||
}): Promise<BatchCompletionResult> {
|
||||
const start = Date.now();
|
||||
const pollBackoff = createOpenAiBatchPollBackoff(params);
|
||||
let current: OpenAiBatchStatus | undefined = params.initial;
|
||||
while (true) {
|
||||
const status =
|
||||
current ??
|
||||
(await fetchOpenAiBatchStatus({
|
||||
openAi: params.openAi,
|
||||
batchId: params.batchId,
|
||||
}));
|
||||
let status: OpenAiBatchStatus;
|
||||
try {
|
||||
status =
|
||||
current ??
|
||||
(await fetchOpenAiBatchStatus({
|
||||
openAi: params.openAi,
|
||||
batchId: params.batchId,
|
||||
}));
|
||||
} catch (error) {
|
||||
if (!params.wait || !isRetryableOpenAiBatchPollError(error)) {
|
||||
throw error;
|
||||
}
|
||||
if (Date.now() - start > params.timeoutMs) {
|
||||
throw new Error(`openai batch ${params.batchId} timed out after ${params.timeoutMs}ms`, {
|
||||
cause: error,
|
||||
});
|
||||
}
|
||||
const delayMs = pollBackoff.nextDelayMs();
|
||||
params.debug?.(
|
||||
`openai batch ${params.batchId} status check failed: ${formatOpenAiBatchPollError(
|
||||
error,
|
||||
)}; waiting ${delayMs}ms`,
|
||||
);
|
||||
await new Promise((resolve) => {
|
||||
setTimeout(resolve, delayMs);
|
||||
});
|
||||
current = undefined;
|
||||
continue;
|
||||
}
|
||||
const state = status.status ?? "unknown";
|
||||
if (state === "completed") {
|
||||
return resolveBatchCompletionFromStatus({
|
||||
|
|
@ -194,9 +286,14 @@ async function waitForOpenAiBatch(params: {
|
|||
if (Date.now() - start > params.timeoutMs) {
|
||||
throw new Error(`openai batch ${params.batchId} timed out after ${params.timeoutMs}ms`);
|
||||
}
|
||||
params.debug?.(`openai batch ${params.batchId} ${state}; waiting ${params.pollIntervalMs}ms`);
|
||||
const delayMs = pollBackoff.nextDelayMs();
|
||||
params.debug?.(
|
||||
`openai batch ${params.batchId} ${state}${formatOpenAiBatchProgress(
|
||||
status,
|
||||
)}; waiting ${delayMs}ms`,
|
||||
);
|
||||
await new Promise((resolve) => {
|
||||
setTimeout(resolve, params.pollIntervalMs);
|
||||
setTimeout(resolve, delayMs);
|
||||
});
|
||||
current = undefined;
|
||||
}
|
||||
|
|
@ -207,13 +304,24 @@ export async function runOpenAiEmbeddingBatches(
|
|||
openAi: OpenAiEmbeddingClient;
|
||||
agentId: string;
|
||||
requests: OpenAiBatchRequest[];
|
||||
maxJsonlBytes?: number;
|
||||
} & EmbeddingBatchExecutionParams,
|
||||
): Promise<Map<string, number[]>> {
|
||||
return await runEmbeddingBatchGroups({
|
||||
...buildEmbeddingBatchGroupOptions(params, {
|
||||
maxRequests: OPENAI_BATCH_MAX_REQUESTS,
|
||||
maxJsonlBytes: params.maxJsonlBytes ?? OPENAI_BATCH_MAX_JSONL_BYTES,
|
||||
debugLabel: "memory embeddings: openai batch submit",
|
||||
}),
|
||||
shouldSplitGroupOnError: isOpenAiBatchUploadTooLargeError,
|
||||
onSplitGroup: ({ error, group, parts, depth }) => {
|
||||
params.debug?.("memory embeddings: openai batch upload too large; splitting group", {
|
||||
requests: group.length,
|
||||
parts: parts.map((part) => part.length),
|
||||
depth,
|
||||
error: formatOpenAiBatchError(error),
|
||||
});
|
||||
},
|
||||
runGroup: async ({ group, groupIndex, groups, byCustomId, pollIntervalMs, timeoutMs }) => {
|
||||
const batchInfo = await submitOpenAiBatch({
|
||||
openAi: params.openAi,
|
||||
|
|
|
|||
|
|
@ -30,6 +30,7 @@ export const openAiMemoryEmbeddingProviderAdapter: MemoryEmbeddingProviderAdapte
|
|||
provider,
|
||||
runtime: {
|
||||
id: "openai",
|
||||
sourceWideBatchEmbed: true,
|
||||
cacheKeyData: {
|
||||
provider: resolvedProvider,
|
||||
baseUrl: client.baseUrl,
|
||||
|
|
|
|||
|
|
@ -3,6 +3,12 @@ import { describe, expect, it, vi } from "vitest";
|
|||
import { MAX_SAFE_TIMEOUT_DELAY_MS } from "../../../gateway-client/src/timeouts.js";
|
||||
import { buildEmbeddingBatchGroupOptions, runEmbeddingBatchGroups } from "./batch-runner.js";
|
||||
|
||||
const jsonlEncoder = new TextEncoder();
|
||||
|
||||
function jsonlLineBytes(value: unknown): number {
|
||||
return jsonlEncoder.encode(JSON.stringify(value)).byteLength;
|
||||
}
|
||||
|
||||
describe("buildEmbeddingBatchGroupOptions", () => {
|
||||
it("clamps oversized embedding batch poll intervals to the timeout budget", () => {
|
||||
const options = buildEmbeddingBatchGroupOptions(
|
||||
|
|
@ -61,4 +67,91 @@ describe("buildEmbeddingBatchGroupOptions", () => {
|
|||
|
||||
expect(options.pollIntervalMs).toBe(MAX_SAFE_TIMEOUT_DELAY_MS);
|
||||
});
|
||||
|
||||
it("splits embedding batch groups by serialized JSONL bytes", async () => {
|
||||
const requests = [
|
||||
{ id: "one", body: { input: "alpha" } },
|
||||
{ id: "two", body: { input: "βeta" } },
|
||||
{ id: "three", body: { input: "gamma" } },
|
||||
];
|
||||
const maxJsonlBytes = jsonlLineBytes(requests[0]) + 1 + jsonlLineBytes(requests[1]);
|
||||
const groups: string[][] = [];
|
||||
|
||||
await runEmbeddingBatchGroups({
|
||||
requests,
|
||||
maxRequests: 100,
|
||||
maxJsonlBytes,
|
||||
wait: true,
|
||||
pollIntervalMs: 1000,
|
||||
timeoutMs: 60_000,
|
||||
concurrency: 1,
|
||||
debugLabel: "embedding batch submit",
|
||||
runGroup: async ({ group }) => {
|
||||
groups.push(group.map((request) => request.id));
|
||||
},
|
||||
});
|
||||
|
||||
expect(groups).toEqual([["one", "two"], ["three"]]);
|
||||
});
|
||||
|
||||
it("splits provider-rejected batch groups when the error is splittable", async () => {
|
||||
const uploadTooLarge = new Error("batch upload failed: 413 payload too large");
|
||||
const calls: string[][] = [];
|
||||
const onSplitGroup = vi.fn();
|
||||
|
||||
await runEmbeddingBatchGroups({
|
||||
requests: ["one", "two", "three", "four"],
|
||||
maxRequests: 100,
|
||||
wait: true,
|
||||
pollIntervalMs: 1000,
|
||||
timeoutMs: 60_000,
|
||||
concurrency: 1,
|
||||
debugLabel: "embedding batch submit",
|
||||
shouldSplitGroupOnError: (error) => error === uploadTooLarge,
|
||||
onSplitGroup,
|
||||
runGroup: async ({ group }) => {
|
||||
calls.push([...group]);
|
||||
if (group.length === 4) {
|
||||
throw uploadTooLarge;
|
||||
}
|
||||
},
|
||||
});
|
||||
|
||||
expect(calls).toEqual([
|
||||
["one", "two", "three", "four"],
|
||||
["one", "two"],
|
||||
["three", "four"],
|
||||
]);
|
||||
expect(onSplitGroup).toHaveBeenCalledWith(
|
||||
expect.objectContaining({
|
||||
error: uploadTooLarge,
|
||||
group: ["one", "two", "three", "four"],
|
||||
parts: [
|
||||
["one", "two"],
|
||||
["three", "four"],
|
||||
],
|
||||
depth: 0,
|
||||
}),
|
||||
);
|
||||
});
|
||||
|
||||
it("does not split a single rejected batch request", async () => {
|
||||
const uploadTooLarge = new Error("batch upload failed: 413 payload too large");
|
||||
|
||||
await expect(
|
||||
runEmbeddingBatchGroups({
|
||||
requests: ["one"],
|
||||
maxRequests: 100,
|
||||
wait: true,
|
||||
pollIntervalMs: 1000,
|
||||
timeoutMs: 60_000,
|
||||
concurrency: 1,
|
||||
debugLabel: "embedding batch submit",
|
||||
shouldSplitGroupOnError: () => true,
|
||||
runGroup: async () => {
|
||||
throw uploadTooLarge;
|
||||
},
|
||||
}),
|
||||
).rejects.toThrow(uploadTooLarge);
|
||||
});
|
||||
});
|
||||
|
|
|
|||
|
|
@ -1,6 +1,6 @@
|
|||
// Memory Host SDK module implements batch runner behavior.
|
||||
import { resolveSafeTimeoutDelayMs } from "../../../gateway-client/src/timeouts.js";
|
||||
import { splitBatchRequests } from "./batch-utils.js";
|
||||
import { splitBatchRequestsByLimits } from "./batch-utils.js";
|
||||
import { runWithConcurrency } from "./internal.js";
|
||||
|
||||
// Shared runner for splitting and executing remote embedding batch groups.
|
||||
|
|
@ -14,6 +14,24 @@ export type EmbeddingBatchExecutionParams = {
|
|||
debug?: (message: string, data?: Record<string, unknown>) => void;
|
||||
};
|
||||
|
||||
type EmbeddingBatchGroupRunArgs<TRequest> = {
|
||||
group: TRequest[];
|
||||
groupIndex: number;
|
||||
groups: number;
|
||||
byCustomId: Map<string, number[]>;
|
||||
pollIntervalMs: number;
|
||||
timeoutMs: number;
|
||||
};
|
||||
|
||||
type EmbeddingBatchSplitArgs<TRequest> = {
|
||||
error: unknown;
|
||||
group: TRequest[];
|
||||
parts: TRequest[][];
|
||||
groupIndex: number;
|
||||
groups: number;
|
||||
depth: number;
|
||||
};
|
||||
|
||||
/** Clamp polling to both configured poll interval and total timeout budget. */
|
||||
function resolveEmbeddingBatchPollIntervalMs(params: {
|
||||
pollIntervalMs: number;
|
||||
|
|
@ -33,41 +51,66 @@ function resolveEmbeddingBatchPollIntervalMs(params: {
|
|||
export async function runEmbeddingBatchGroups<TRequest>(params: {
|
||||
requests: TRequest[];
|
||||
maxRequests: number;
|
||||
maxJsonlBytes?: number;
|
||||
wait: EmbeddingBatchExecutionParams["wait"];
|
||||
pollIntervalMs: EmbeddingBatchExecutionParams["pollIntervalMs"];
|
||||
timeoutMs: EmbeddingBatchExecutionParams["timeoutMs"];
|
||||
concurrency: EmbeddingBatchExecutionParams["concurrency"];
|
||||
debugLabel: string;
|
||||
debug?: EmbeddingBatchExecutionParams["debug"];
|
||||
runGroup: (args: {
|
||||
group: TRequest[];
|
||||
groupIndex: number;
|
||||
groups: number;
|
||||
byCustomId: Map<string, number[]>;
|
||||
pollIntervalMs: number;
|
||||
timeoutMs: number;
|
||||
}) => Promise<void>;
|
||||
shouldSplitGroupOnError?: (error: unknown, group: TRequest[]) => boolean;
|
||||
onSplitGroup?: (args: EmbeddingBatchSplitArgs<TRequest>) => void;
|
||||
runGroup: (args: EmbeddingBatchGroupRunArgs<TRequest>) => Promise<void>;
|
||||
}): Promise<Map<string, number[]>> {
|
||||
if (params.requests.length === 0) {
|
||||
return new Map();
|
||||
}
|
||||
const groups = splitBatchRequests(params.requests, params.maxRequests);
|
||||
const groups = splitBatchRequestsByLimits(params.requests, {
|
||||
maxRequests: params.maxRequests,
|
||||
maxJsonlBytes: params.maxJsonlBytes,
|
||||
});
|
||||
const byCustomId = new Map<string, number[]>();
|
||||
const pollIntervalMs = resolveEmbeddingBatchPollIntervalMs(params);
|
||||
const runGroup = async (group: TRequest[], groupIndex: number, depth = 0): Promise<void> => {
|
||||
try {
|
||||
await params.runGroup({
|
||||
group,
|
||||
groupIndex,
|
||||
groups: groups.length,
|
||||
byCustomId,
|
||||
pollIntervalMs,
|
||||
timeoutMs: params.timeoutMs,
|
||||
});
|
||||
} catch (error) {
|
||||
if (group.length <= 1 || !params.shouldSplitGroupOnError?.(error, group)) {
|
||||
throw error;
|
||||
}
|
||||
const splitAt = Math.ceil(group.length / 2);
|
||||
const parts = [group.slice(0, splitAt), group.slice(splitAt)].filter(
|
||||
(part) => part.length > 0,
|
||||
);
|
||||
params.onSplitGroup?.({
|
||||
error,
|
||||
group,
|
||||
parts,
|
||||
groupIndex,
|
||||
groups: groups.length,
|
||||
depth,
|
||||
});
|
||||
for (const part of parts) {
|
||||
await runGroup(part, groupIndex, depth + 1);
|
||||
}
|
||||
}
|
||||
};
|
||||
const tasks = groups.map((group, groupIndex) => async () => {
|
||||
await params.runGroup({
|
||||
group,
|
||||
groupIndex,
|
||||
groups: groups.length,
|
||||
byCustomId,
|
||||
pollIntervalMs,
|
||||
timeoutMs: params.timeoutMs,
|
||||
});
|
||||
await runGroup(group, groupIndex);
|
||||
});
|
||||
|
||||
params.debug?.(params.debugLabel, {
|
||||
requests: params.requests.length,
|
||||
groups: groups.length,
|
||||
maxRequests: params.maxRequests,
|
||||
maxJsonlBytes: params.maxJsonlBytes,
|
||||
wait: params.wait,
|
||||
concurrency: params.concurrency,
|
||||
pollIntervalMs,
|
||||
|
|
@ -81,12 +124,13 @@ export async function runEmbeddingBatchGroups<TRequest>(params: {
|
|||
/** Build normalized batch-group options for provider-specific runners. */
|
||||
export function buildEmbeddingBatchGroupOptions<TRequest>(
|
||||
params: { requests: TRequest[] } & EmbeddingBatchExecutionParams,
|
||||
options: { maxRequests: number; debugLabel: string },
|
||||
options: { maxRequests: number; maxJsonlBytes?: number; debugLabel: string },
|
||||
) {
|
||||
const pollIntervalMs = resolveEmbeddingBatchPollIntervalMs(params);
|
||||
return {
|
||||
requests: params.requests,
|
||||
maxRequests: options.maxRequests,
|
||||
maxJsonlBytes: options.maxJsonlBytes,
|
||||
wait: params.wait,
|
||||
pollIntervalMs,
|
||||
timeoutMs: params.timeoutMs,
|
||||
|
|
|
|||
|
|
@ -30,6 +30,7 @@ export async function uploadBatchJsonlFile(params: {
|
|||
const filePayload = await withRemoteHttpResponse({
|
||||
url: `${baseUrl}/files`,
|
||||
ssrfPolicy: params.client.ssrfPolicy,
|
||||
fetchImpl: params.client.fetchImpl,
|
||||
init: {
|
||||
method: "POST",
|
||||
headers: buildBatchHeaders(params.client, { json: false }),
|
||||
|
|
|
|||
|
|
@ -8,6 +8,7 @@ export type BatchHttpClientConfig = {
|
|||
baseUrl?: string;
|
||||
headers?: Record<string, string>;
|
||||
ssrfPolicy?: SsrFPolicy;
|
||||
fetchImpl?: typeof fetch;
|
||||
};
|
||||
|
||||
/** Normalize batch API base URLs by removing one trailing slash. */
|
||||
|
|
@ -32,14 +33,62 @@ export function buildBatchHeaders(
|
|||
return headers;
|
||||
}
|
||||
|
||||
const jsonlEncoder = new TextEncoder();
|
||||
|
||||
function estimateJsonlLineBytes(request: unknown): number {
|
||||
return jsonlEncoder.encode(JSON.stringify(request) ?? "").byteLength;
|
||||
}
|
||||
|
||||
function normalizePositiveInteger(value: number | undefined): number | undefined {
|
||||
if (typeof value !== "number" || !Number.isFinite(value) || value <= 0) {
|
||||
return undefined;
|
||||
}
|
||||
return Math.floor(value);
|
||||
}
|
||||
|
||||
/** Split provider requests into max-sized groups while preserving order. */
|
||||
export function splitBatchRequests<T>(requests: T[], maxRequests: number): T[][] {
|
||||
if (requests.length <= maxRequests) {
|
||||
const limit = normalizePositiveInteger(maxRequests) ?? 1;
|
||||
if (requests.length <= limit) {
|
||||
return [requests];
|
||||
}
|
||||
const groups: T[][] = [];
|
||||
for (let i = 0; i < requests.length; i += maxRequests) {
|
||||
groups.push(requests.slice(i, i + maxRequests));
|
||||
for (let i = 0; i < requests.length; i += limit) {
|
||||
groups.push(requests.slice(i, i + limit));
|
||||
}
|
||||
return groups;
|
||||
}
|
||||
|
||||
export function splitBatchRequestsByLimits<T>(
|
||||
requests: T[],
|
||||
limits: { maxRequests: number; maxJsonlBytes?: number },
|
||||
): T[][] {
|
||||
const maxRequests = normalizePositiveInteger(limits.maxRequests) ?? 1;
|
||||
const maxJsonlBytes = normalizePositiveInteger(limits.maxJsonlBytes);
|
||||
if (!maxJsonlBytes) {
|
||||
return splitBatchRequests(requests, maxRequests);
|
||||
}
|
||||
|
||||
const groups: T[][] = [];
|
||||
let current: T[] = [];
|
||||
let currentBytes = 0;
|
||||
for (const request of requests) {
|
||||
const requestBytes = estimateJsonlLineBytes(request);
|
||||
const separatorBytes = current.length === 0 ? 0 : 1;
|
||||
const wouldExceedRequests = current.length >= maxRequests;
|
||||
const wouldExceedBytes =
|
||||
current.length > 0 && currentBytes + separatorBytes + requestBytes > maxJsonlBytes;
|
||||
if (current.length > 0 && (wouldExceedRequests || wouldExceedBytes)) {
|
||||
groups.push(current);
|
||||
current = [];
|
||||
currentBytes = 0;
|
||||
}
|
||||
|
||||
currentBytes += (current.length === 0 ? 0 : 1) + requestBytes;
|
||||
current.push(request);
|
||||
}
|
||||
if (current.length > 0) {
|
||||
groups.push(current);
|
||||
}
|
||||
return groups;
|
||||
}
|
||||
|
|
|
|||
|
|
@ -4,7 +4,10 @@ import {
|
|||
registerVirtualTestPlugin,
|
||||
} from "openclaw/plugin-sdk/plugin-test-contracts";
|
||||
import { describe, expect, it } from "vitest";
|
||||
import { getRegisteredMemoryEmbeddingProvider } from "../memory-embedding-providers.js";
|
||||
import {
|
||||
getRegisteredMemoryEmbeddingProvider,
|
||||
type MemoryEmbeddingBatchOptions,
|
||||
} from "../memory-embedding-providers.js";
|
||||
import { createPluginRecord } from "../status.test-helpers.js";
|
||||
|
||||
describe("memory embedding provider registration", () => {
|
||||
|
|
@ -79,6 +82,58 @@ describe("memory embedding provider registration", () => {
|
|||
expect(provider?.ownerPluginId).toBe("memory-core");
|
||||
});
|
||||
|
||||
it("keeps source-wide batch embedding behind an explicit runtime opt-in", async () => {
|
||||
const { config, registry } = createPluginRegistryFixture();
|
||||
|
||||
registerVirtualTestPlugin({
|
||||
registry,
|
||||
config,
|
||||
id: "source-wide-memory",
|
||||
name: "Source Wide Memory",
|
||||
contracts: {
|
||||
memoryEmbeddingProviders: ["source-wide-memory"],
|
||||
},
|
||||
register(api) {
|
||||
api.registerMemoryEmbeddingProvider({
|
||||
id: "source-wide-memory",
|
||||
create: async () => ({
|
||||
provider: {
|
||||
id: "source-wide-memory",
|
||||
model: "test-embedding",
|
||||
embedQuery: async (text: string) => [text.length],
|
||||
embedBatch: async (texts: string[]) => texts.map((text) => [text.length]),
|
||||
},
|
||||
runtime: {
|
||||
id: "source-wide-memory",
|
||||
sourceWideBatchEmbed: true,
|
||||
batchEmbed: async (batch: MemoryEmbeddingBatchOptions) =>
|
||||
batch.chunks.map((chunk, index) => [index, chunk.text.length]),
|
||||
},
|
||||
}),
|
||||
});
|
||||
},
|
||||
});
|
||||
|
||||
const adapter = getRegisteredMemoryEmbeddingProvider("source-wide-memory")?.adapter;
|
||||
const result = await adapter?.create({ config, model: "test-embedding" });
|
||||
|
||||
expect(result?.runtime?.sourceWideBatchEmbed).toBe(true);
|
||||
await expect(
|
||||
result?.runtime?.batchEmbed?.({
|
||||
agentId: "main",
|
||||
chunks: [{ text: "alpha" }, { text: "beta" }],
|
||||
wait: true,
|
||||
concurrency: 1,
|
||||
pollIntervalMs: 1000,
|
||||
timeoutMs: 60_000,
|
||||
debug: () => {},
|
||||
}),
|
||||
).resolves.toEqual([
|
||||
[0, 5],
|
||||
[1, 4],
|
||||
]);
|
||||
});
|
||||
|
||||
it("keeps companion embedding providers available during tool discovery", () => {
|
||||
const { config, registry } = createPluginRegistryFixture();
|
||||
const record = createPluginRecord({
|
||||
|
|
|
|||
|
|
@ -31,6 +31,7 @@ export type MemoryEmbeddingProviderRuntime = {
|
|||
cacheKeyData?: Record<string, unknown>;
|
||||
inlineQueryTimeoutMs?: number;
|
||||
inlineBatchTimeoutMs?: number;
|
||||
sourceWideBatchEmbed?: boolean;
|
||||
batchEmbed?: (options: MemoryEmbeddingBatchOptions) => Promise<number[][] | null>;
|
||||
};
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue