mirror of
https://github.com/danny-avila/LibreChat.git
synced 2026-10-20 17:13:44 +00:00
🗂️ feat: Per-Agent Memory Partitions (#14084)
* feat: per-agent memory partitions (memory_scope)
Adds an optional agentId partition to MemoryEntry so agents can opt into
isolated memory via a new memory_scope field ('user' | 'agent'). Partition
derives from agentId presence ({agentId: null} matches legacy docs, no
migration). Inline set_memory/delete_memory tools, the post-turn memory
agent, the request-scoped memory cache, and context injection are all
partition-aware; context is only injected into agents whose resolved
partition matches. Memory routes accept the partition param, scope
duplicate/token-limit checks per partition, and enrich entries with agent
names. Memories panel gains a partition filter and agent badges; the agent
builder gains an agent-scoped memory toggle.
* fix: address Codex review findings on memory partitions
- strip runtime ____N id suffixes in getMemoryAgentId so added-conversation
runs share the persisted agent's partition
- load each agent's own partition in multi-agent context injection instead
of skipping foreign partitions entirely
- clear memory_scope to 'user' on save when Enable Memory is unchecked
- fall back to 'all' when the selected panel partition no longer exists
- restrict GET /memories agent-name resolution to agents the requester can
VIEW
This commit is contained in:
Danny Avila
GitHub
parent
73c43ded25
commit
3945d293de
27 changed files with 732 additions and 137 deletions
|
|
@ -1,12 +1,14 @@
|
|||
import { Types } from 'mongoose';
|
||||
import { Run, Providers } from '@librechat/agents';
|
||||
import { MemoryScope } from 'librechat-data-provider';
|
||||
import type { IUser } from '@librechat/data-schemas';
|
||||
import type { Response } from 'express';
|
||||
import {
|
||||
processMemory,
|
||||
createMemoryTool,
|
||||
createDeleteMemoryTool,
|
||||
getMemoryAgentId,
|
||||
getRequestMemories,
|
||||
createDeleteMemoryTool,
|
||||
invalidateRequestMemories,
|
||||
agentHasInlineMemoryTools,
|
||||
} from './memory';
|
||||
|
|
@ -706,4 +708,44 @@ describe('getRequestMemories caching', () => {
|
|||
await getRequestMemories({ req, userId: 'user-1', getFormattedMemories });
|
||||
expect(getFormattedMemories).toHaveBeenCalledTimes(2);
|
||||
});
|
||||
|
||||
it('caches and invalidates per partition', async () => {
|
||||
const getFormattedMemories = jest
|
||||
.fn()
|
||||
.mockResolvedValue({ withKeys: '', withoutKeys: '', totalTokens: 10 });
|
||||
const req = {};
|
||||
|
||||
await getRequestMemories({ req, userId: 'user-1', getFormattedMemories });
|
||||
await getRequestMemories({ req, userId: 'user-1', agentId: 'agent_a', getFormattedMemories });
|
||||
await getRequestMemories({ req, userId: 'user-1', agentId: 'agent_a', getFormattedMemories });
|
||||
/** Personal pool and agent partition are distinct cache entries. */
|
||||
expect(getFormattedMemories).toHaveBeenCalledTimes(2);
|
||||
expect(getFormattedMemories).toHaveBeenLastCalledWith({
|
||||
userId: 'user-1',
|
||||
agentId: 'agent_a',
|
||||
});
|
||||
|
||||
/** Invalidating one partition leaves the other cached. */
|
||||
invalidateRequestMemories(req, 'agent_a');
|
||||
await getRequestMemories({ req, userId: 'user-1', getFormattedMemories });
|
||||
expect(getFormattedMemories).toHaveBeenCalledTimes(2);
|
||||
await getRequestMemories({ req, userId: 'user-1', agentId: 'agent_a', getFormattedMemories });
|
||||
expect(getFormattedMemories).toHaveBeenCalledTimes(3);
|
||||
});
|
||||
});
|
||||
|
||||
describe('getMemoryAgentId', () => {
|
||||
it('resolves the agent partition only for memory_scope "agent"', () => {
|
||||
expect(getMemoryAgentId({ id: 'agent_a', memory_scope: MemoryScope.agent })).toBe('agent_a');
|
||||
expect(getMemoryAgentId({ id: 'agent_a', memory_scope: MemoryScope.user })).toBeUndefined();
|
||||
expect(getMemoryAgentId({ id: 'agent_a' })).toBeUndefined();
|
||||
expect(getMemoryAgentId({ memory_scope: MemoryScope.agent })).toBeUndefined();
|
||||
expect(getMemoryAgentId(null)).toBeUndefined();
|
||||
});
|
||||
|
||||
it('strips runtime id suffixes so added-conversation runs share the persisted partition', () => {
|
||||
expect(getMemoryAgentId({ id: 'agent_a____1', memory_scope: MemoryScope.agent })).toBe(
|
||||
'agent_a',
|
||||
);
|
||||
});
|
||||
});
|
||||
|
|
|
|||
|
|
@ -6,10 +6,12 @@ import { Run, Providers, GraphEvents } from '@librechat/agents';
|
|||
import { HumanMessage } from '@librechat/agents/langchain/messages';
|
||||
import {
|
||||
Tools,
|
||||
MemoryScope,
|
||||
Permissions,
|
||||
EModelEndpoint,
|
||||
PermissionTypes,
|
||||
AgentCapabilities,
|
||||
stripAgentIdSuffix,
|
||||
} from 'librechat-data-provider';
|
||||
import type {
|
||||
OpenAIClientOptions,
|
||||
|
|
@ -117,6 +119,7 @@ type MemoryArtifactRecord = Record<Tools.memory, MemoryArtifact>;
|
|||
*/
|
||||
export const createMemoryTool = ({
|
||||
userId,
|
||||
agentId,
|
||||
setMemory,
|
||||
validKeys,
|
||||
charLimit,
|
||||
|
|
@ -125,6 +128,8 @@ export const createMemoryTool = ({
|
|||
onWrite,
|
||||
}: {
|
||||
userId: string | ObjectId;
|
||||
/** Agent partition to write to; omit for the shared personal pool */
|
||||
agentId?: string;
|
||||
setMemory: MemoryMethods['setMemory'];
|
||||
validKeys?: string[];
|
||||
charLimit?: number;
|
||||
|
|
@ -220,10 +225,11 @@ export const createMemoryTool = ({
|
|||
value,
|
||||
tokenCount,
|
||||
type: 'update',
|
||||
...(agentId ? { agentId } : {}),
|
||||
},
|
||||
};
|
||||
|
||||
const result = await setMemory({ userId, key, value, tokenCount });
|
||||
const result = await setMemory({ userId, key, value, tokenCount, agentId });
|
||||
if (result.ok) {
|
||||
if (tokenLimit) {
|
||||
currentTotalTokens = newTotalTokens;
|
||||
|
|
@ -274,11 +280,14 @@ export const createMemoryTool = ({
|
|||
*/
|
||||
export const createDeleteMemoryTool = ({
|
||||
userId,
|
||||
agentId,
|
||||
deleteMemory,
|
||||
validKeys,
|
||||
onWrite,
|
||||
}: {
|
||||
userId: string | ObjectId;
|
||||
/** Agent partition to delete from; omit for the shared personal pool */
|
||||
agentId?: string;
|
||||
deleteMemory: MemoryMethods['deleteMemory'];
|
||||
validKeys?: string[];
|
||||
onWrite?: () => void;
|
||||
|
|
@ -299,10 +308,11 @@ export const createDeleteMemoryTool = ({
|
|||
[Tools.memory]: {
|
||||
key,
|
||||
type: 'delete',
|
||||
...(agentId ? { agentId } : {}),
|
||||
},
|
||||
};
|
||||
|
||||
const result = await deleteMemory({ userId, key });
|
||||
const result = await deleteMemory({ userId, key, agentId });
|
||||
if (result.ok) {
|
||||
onWrite?.();
|
||||
logger.debug(`Memory deleted for key "${key}" for user "${userId}"`);
|
||||
|
|
@ -433,7 +443,25 @@ type GetRoleByName = (
|
|||
fieldsToSelect?: string | string[],
|
||||
) => Promise<IRole | null>;
|
||||
|
||||
type InlineMemoryAgent = { tools?: unknown[]; memoryToolsRegistered?: boolean } | null | undefined;
|
||||
type InlineMemoryAgent =
|
||||
| { id?: string; tools?: unknown[]; memoryToolsRegistered?: boolean; memory_scope?: MemoryScope }
|
||||
| null
|
||||
| undefined;
|
||||
|
||||
/**
|
||||
* Resolves the memory partition an agent reads/writes: its persisted id when
|
||||
* the agent opted into isolated memory (`memory_scope: 'agent'`), otherwise
|
||||
* `undefined` for the shared personal pool. Runtime ids carry a `____N`
|
||||
* suffix in added-conversation paths, so the suffix is stripped to keep the
|
||||
* partition stable across single- and multi-agent runs. Ephemeral agents
|
||||
* never carry `memory_scope`, so they always resolve to the shared pool.
|
||||
*/
|
||||
export function getMemoryAgentId(agent: InlineMemoryAgent): string | undefined {
|
||||
if (agent?.memory_scope === MemoryScope.agent && typeof agent.id === 'string' && agent.id) {
|
||||
return stripAgentIdSuffix(agent.id);
|
||||
}
|
||||
return undefined;
|
||||
}
|
||||
|
||||
/**
|
||||
* Whether an agent carries the inline memory tools. Prefers the LibreChat-only
|
||||
|
|
@ -463,35 +491,44 @@ export function agentHasInlineMemoryTools(agent: InlineMemoryAgent): boolean {
|
|||
/**
|
||||
* Request-scoped cache so that multiple memory-enabled agents in one run (and
|
||||
* the run's memory context load) share a single `getFormattedMemories` call
|
||||
* instead of each re-fetching the same user's memories.
|
||||
* per partition instead of each re-fetching the same memories.
|
||||
*/
|
||||
const requestMemoriesCache = new WeakMap<object, Promise<FormattedMemoriesResult>>();
|
||||
const requestMemoriesCache = new WeakMap<object, Map<string, Promise<FormattedMemoriesResult>>>();
|
||||
|
||||
export function getRequestMemories({
|
||||
req,
|
||||
userId,
|
||||
agentId,
|
||||
getFormattedMemories,
|
||||
}: {
|
||||
req: object;
|
||||
userId: string | ObjectId;
|
||||
/** Agent partition; omit for the shared personal pool */
|
||||
agentId?: string;
|
||||
getFormattedMemories: MemoryMethods['getFormattedMemories'];
|
||||
}): Promise<FormattedMemoriesResult> {
|
||||
let cached = requestMemoriesCache.get(req);
|
||||
let partitions = requestMemoriesCache.get(req);
|
||||
if (!partitions) {
|
||||
partitions = new Map();
|
||||
requestMemoriesCache.set(req, partitions);
|
||||
}
|
||||
const partitionKey = agentId ?? '';
|
||||
let cached = partitions.get(partitionKey);
|
||||
if (!cached) {
|
||||
cached = getFormattedMemories({ userId });
|
||||
requestMemoriesCache.set(req, cached);
|
||||
cached = getFormattedMemories({ userId, agentId });
|
||||
partitions.set(partitionKey, cached);
|
||||
}
|
||||
return cached;
|
||||
}
|
||||
|
||||
/**
|
||||
* Drops the cached memories for a request so the next {@link getRequestMemories}
|
||||
* re-fetches. Inline `set_memory`/`delete_memory` writes call this on success so
|
||||
* a later tool round in the same response is seeded with the post-write usage
|
||||
* total instead of a stale pre-write one.
|
||||
* Drops the cached memories for a request partition so the next
|
||||
* {@link getRequestMemories} re-fetches. Inline `set_memory`/`delete_memory`
|
||||
* writes call this on success so a later tool round in the same response is
|
||||
* seeded with the post-write usage total instead of a stale pre-write one.
|
||||
*/
|
||||
export function invalidateRequestMemories(req: object): void {
|
||||
requestMemoriesCache.delete(req);
|
||||
export function invalidateRequestMemories(req: object, agentId?: string): void {
|
||||
requestMemoriesCache.get(req)?.delete(agentId ?? '');
|
||||
}
|
||||
|
||||
/**
|
||||
|
|
@ -563,6 +600,7 @@ export async function buildInlineMemoryTool({
|
|||
|
||||
const memoryConfig = req?.config?.memory;
|
||||
const validKeys = memoryConfig?.validKeys as string[] | undefined;
|
||||
const memoryAgentId = getMemoryAgentId(agent);
|
||||
|
||||
if (toolName === DELETE_MEMORY_TOOL_NAME) {
|
||||
const allowed = await isMemoryToolAllowed({
|
||||
|
|
@ -575,9 +613,10 @@ export async function buildInlineMemoryTool({
|
|||
}
|
||||
return createDeleteMemoryTool({
|
||||
userId,
|
||||
agentId: memoryAgentId,
|
||||
deleteMemory: memoryMethods.deleteMemory,
|
||||
validKeys,
|
||||
onWrite: () => invalidateRequestMemories(req),
|
||||
onWrite: () => invalidateRequestMemories(req, memoryAgentId),
|
||||
});
|
||||
}
|
||||
|
||||
|
|
@ -598,6 +637,7 @@ export async function buildInlineMemoryTool({
|
|||
const formatted = await getRequestMemories({
|
||||
req,
|
||||
userId,
|
||||
agentId: memoryAgentId,
|
||||
getFormattedMemories: memoryMethods.getFormattedMemories,
|
||||
});
|
||||
totalTokens = formatted?.totalTokens ?? 0;
|
||||
|
|
@ -611,12 +651,13 @@ export async function buildInlineMemoryTool({
|
|||
|
||||
return createMemoryTool({
|
||||
userId,
|
||||
agentId: memoryAgentId,
|
||||
setMemory: memoryMethods.setMemory,
|
||||
validKeys,
|
||||
charLimit,
|
||||
tokenLimit,
|
||||
totalTokens,
|
||||
onWrite: () => invalidateRequestMemories(req),
|
||||
onWrite: () => invalidateRequestMemories(req, memoryAgentId),
|
||||
});
|
||||
}
|
||||
|
||||
|
|
@ -647,6 +688,7 @@ export class BasicToolEndHandler implements EventHandler {
|
|||
export async function processMemory({
|
||||
res,
|
||||
userId,
|
||||
agentId,
|
||||
setMemory,
|
||||
deleteMemory,
|
||||
messages,
|
||||
|
|
@ -665,6 +707,8 @@ export async function processMemory({
|
|||
setMemory: MemoryMethods['setMemory'];
|
||||
deleteMemory: MemoryMethods['deleteMemory'];
|
||||
userId: string | ObjectId;
|
||||
/** Agent partition; omit for the shared personal pool */
|
||||
agentId?: string;
|
||||
memory: string;
|
||||
messageId: string;
|
||||
conversationId: string;
|
||||
|
|
@ -680,6 +724,7 @@ export async function processMemory({
|
|||
try {
|
||||
const memoryTool = createMemoryTool({
|
||||
userId,
|
||||
agentId,
|
||||
tokenLimit,
|
||||
setMemory,
|
||||
validKeys,
|
||||
|
|
@ -687,6 +732,7 @@ export async function processMemory({
|
|||
});
|
||||
const deleteMemoryTool = createDeleteMemoryTool({
|
||||
userId,
|
||||
agentId,
|
||||
validKeys,
|
||||
deleteMemory,
|
||||
});
|
||||
|
|
@ -874,6 +920,7 @@ ${memory ?? 'No existing memories'}`;
|
|||
export async function createMemoryProcessor({
|
||||
res,
|
||||
userId,
|
||||
agentId,
|
||||
messageId,
|
||||
memoryMethods,
|
||||
conversationId,
|
||||
|
|
@ -885,6 +932,8 @@ export async function createMemoryProcessor({
|
|||
messageId: string;
|
||||
conversationId: string;
|
||||
userId: string | ObjectId;
|
||||
/** Agent partition; omit for the shared personal pool */
|
||||
agentId?: string;
|
||||
memoryMethods: RequiredMemoryMethods;
|
||||
config?: MemoryConfig;
|
||||
streamId?: string | null;
|
||||
|
|
@ -895,6 +944,7 @@ export async function createMemoryProcessor({
|
|||
|
||||
const { withKeys, withoutKeys, totalTokens } = await memoryMethods.getFormattedMemories({
|
||||
userId,
|
||||
agentId,
|
||||
});
|
||||
|
||||
return [
|
||||
|
|
@ -904,6 +954,7 @@ export async function createMemoryProcessor({
|
|||
return await processMemory({
|
||||
res,
|
||||
userId,
|
||||
agentId,
|
||||
messages,
|
||||
validKeys,
|
||||
llmConfig,
|
||||
|
|
|
|||
|
|
@ -1,5 +1,5 @@
|
|||
import { z } from 'zod';
|
||||
import { MAX_SUBAGENTS, ViolationTypes, ErrorTypes } from 'librechat-data-provider';
|
||||
import { MemoryScope, MAX_SUBAGENTS, ViolationTypes, ErrorTypes } from 'librechat-data-provider';
|
||||
import type { Agent, TModelsConfig } from 'librechat-data-provider';
|
||||
import type { Request, Response } from 'express';
|
||||
|
||||
|
|
@ -388,6 +388,7 @@ export const agentBaseSchema: z.ZodObject<
|
|||
tools: z.ZodOptional<z.ZodArray<z.ZodString, 'many'>>;
|
||||
skills: z.ZodOptional<z.ZodArray<z.ZodString, 'many'>>;
|
||||
skills_enabled: z.ZodOptional<z.ZodBoolean>;
|
||||
memory_scope: z.ZodOptional<z.ZodNativeEnum<typeof MemoryScope>>;
|
||||
/** @deprecated Use edges instead */
|
||||
agent_ids: z.ZodOptional<z.ZodArray<z.ZodString, 'many'>>;
|
||||
edges: z.ZodOptional<
|
||||
|
|
@ -688,6 +689,7 @@ export const agentBaseSchema: z.ZodObject<
|
|||
tools: z.array(z.string()).optional(),
|
||||
skills: z.array(z.string()).optional(),
|
||||
skills_enabled: z.boolean().optional(),
|
||||
memory_scope: z.nativeEnum(MemoryScope).optional(),
|
||||
/** @deprecated Use edges instead */
|
||||
agent_ids: z.array(z.string()).optional(),
|
||||
edges: z.array(graphEdgeSchema).optional(),
|
||||
|
|
@ -732,6 +734,7 @@ export const agentCreateSchema: z.ZodObject<
|
|||
model_parameters: z.ZodOptional<z.ZodRecord<z.ZodString, z.ZodUnknown>>;
|
||||
skills: z.ZodOptional<z.ZodArray<z.ZodString, 'many'>>;
|
||||
skills_enabled: z.ZodOptional<z.ZodBoolean>;
|
||||
memory_scope: z.ZodOptional<z.ZodNativeEnum<typeof MemoryScope>>;
|
||||
agent_ids: z.ZodOptional<z.ZodArray<z.ZodString, 'many'>>;
|
||||
edges: z.ZodOptional<
|
||||
z.ZodArray<
|
||||
|
|
@ -1042,6 +1045,7 @@ export const agentUpdateSchema: z.ZodObject<
|
|||
tools: z.ZodOptional<z.ZodArray<z.ZodString, 'many'>>;
|
||||
skills: z.ZodOptional<z.ZodArray<z.ZodString, 'many'>>;
|
||||
skills_enabled: z.ZodOptional<z.ZodBoolean>;
|
||||
memory_scope: z.ZodOptional<z.ZodNativeEnum<typeof MemoryScope>>;
|
||||
agent_ids: z.ZodOptional<z.ZodArray<z.ZodString, 'many'>>;
|
||||
edges: z.ZodOptional<
|
||||
z.ZodArray<
|
||||
|
|
|
|||
|
|
@ -490,7 +490,8 @@ export const verifyTwoFactorTemp = () => `${BASE_URL}/api/auth/2fa/verify-temp`;
|
|||
|
||||
/* Memories */
|
||||
export const memories = () => `${BASE_URL}/api/memories`;
|
||||
export const memory = (key: string) => `${memories()}/${encodeURIComponent(key)}`;
|
||||
export const memory = (key: string, agentId?: string) =>
|
||||
`${memories()}/${encodeURIComponent(key)}${agentId ? `?agentId=${encodeURIComponent(agentId)}` : ''}`;
|
||||
export const memoryPreferences = () => `${memories()}/preferences`;
|
||||
|
||||
export const searchPrincipals = (params: q.PrincipalSearchParams) => {
|
||||
|
|
|
|||
|
|
@ -1307,16 +1307,17 @@ export const getMemories = (): Promise<q.MemoriesResponse> => {
|
|||
return request.get(endpoints.memories());
|
||||
};
|
||||
|
||||
export const deleteMemory = (key: string): Promise<void> => {
|
||||
return request.delete(endpoints.memory(key));
|
||||
export const deleteMemory = (key: string, agentId?: string): Promise<void> => {
|
||||
return request.delete(endpoints.memory(key, agentId));
|
||||
};
|
||||
|
||||
export const updateMemory = (
|
||||
key: string,
|
||||
value: string,
|
||||
originalKey?: string,
|
||||
agentId?: string,
|
||||
): Promise<q.TUserMemory> => {
|
||||
return request.patch(endpoints.memory(originalKey || key), { key, value });
|
||||
return request.patch(endpoints.memory(originalKey || key, agentId), { key, value });
|
||||
};
|
||||
|
||||
export const updateMemoryPreferences = (preferences: {
|
||||
|
|
@ -1328,6 +1329,7 @@ export const updateMemoryPreferences = (preferences: {
|
|||
export const createMemory = (data: {
|
||||
key: string;
|
||||
value: string;
|
||||
agentId?: string;
|
||||
}): Promise<{ created: boolean; memory: q.TUserMemory }> => {
|
||||
return request.post(endpoints.memories(), data);
|
||||
};
|
||||
|
|
|
|||
|
|
@ -344,6 +344,8 @@ export const defaultAgentFormValues = {
|
|||
subagents: undefined as
|
||||
| { enabled?: boolean; allowSelf?: boolean; agent_ids?: string[] }
|
||||
| undefined,
|
||||
/** Memory partition: 'agent' isolates memories per (user, agent); default shared pool */
|
||||
memory_scope: undefined as MemoryScope | undefined,
|
||||
};
|
||||
|
||||
export const ImageVisionTool: FunctionTool = {
|
||||
|
|
@ -837,11 +839,23 @@ export const tMessageSchema = z.object({
|
|||
quotes: z.array(z.string()).optional(),
|
||||
});
|
||||
|
||||
/**
|
||||
* Which memory partition an agent reads/writes.
|
||||
* `user` = the shared personal pool (default); `agent` = a partition
|
||||
* isolated per (user, agent) so the agent only sees its own memories.
|
||||
*/
|
||||
export enum MemoryScope {
|
||||
user = 'user',
|
||||
agent = 'agent',
|
||||
}
|
||||
|
||||
export type MemoryArtifact = {
|
||||
key: string;
|
||||
value?: string;
|
||||
tokenCount?: number;
|
||||
type: 'update' | 'delete' | 'error';
|
||||
/** Agent partition the write targeted; absent = shared personal pool */
|
||||
agentId?: string;
|
||||
};
|
||||
|
||||
export type UIResource = {
|
||||
|
|
|
|||
|
|
@ -1,5 +1,5 @@
|
|||
import type { OpenAPIV3 } from 'openapi-types';
|
||||
import type { AssistantsEndpoint, AgentProvider } from 'src/schemas';
|
||||
import type { AssistantsEndpoint, AgentProvider, MemoryScope } from 'src/schemas';
|
||||
import type { Agents, GraphEdge } from './agents';
|
||||
import type { ContentTypes } from './runs';
|
||||
import type { TFile } from './files';
|
||||
|
|
@ -301,6 +301,8 @@ export type Agent = {
|
|||
skills_enabled?: boolean;
|
||||
/** Subagent spawning configuration — isolated-context child agents. */
|
||||
subagents?: AgentSubagentsConfig;
|
||||
/** Memory partition: `agent` isolates memories per (user, agent); default shared pool */
|
||||
memory_scope?: MemoryScope;
|
||||
};
|
||||
|
||||
export type TAgentsMap = Record<string, Agent | undefined>;
|
||||
|
|
@ -329,6 +331,7 @@ export type AgentCreateParams = {
|
|||
| 'skills'
|
||||
| 'skills_enabled'
|
||||
| 'subagents'
|
||||
| 'memory_scope'
|
||||
>;
|
||||
|
||||
export type AgentUpdateParams = {
|
||||
|
|
@ -356,6 +359,7 @@ export type AgentUpdateParams = {
|
|||
| 'skills'
|
||||
| 'skills_enabled'
|
||||
| 'subagents'
|
||||
| 'memory_scope'
|
||||
>;
|
||||
|
||||
export type AgentListParams = {
|
||||
|
|
|
|||
|
|
@ -150,6 +150,10 @@ export type TUserMemory = {
|
|||
value: string;
|
||||
updated_at: string;
|
||||
tokenCount?: number;
|
||||
/** Agent partition this memory belongs to; absent = shared personal pool */
|
||||
agentId?: string;
|
||||
/** Display name of the partition's agent, resolved server-side when available */
|
||||
agentName?: string;
|
||||
};
|
||||
|
||||
export type MemoriesResponse = {
|
||||
|
|
|
|||
156
packages/data-schemas/src/methods/memory.spec.ts
Normal file
156
packages/data-schemas/src/methods/memory.spec.ts
Normal file
|
|
@ -0,0 +1,156 @@
|
|||
import mongoose from 'mongoose';
|
||||
import { MongoMemoryServer } from 'mongodb-memory-server';
|
||||
import type { IMemoryEntry } from '~/types';
|
||||
import { createMemoryMethods, type MemoryMethods } from './memory';
|
||||
import { createModels } from '~/models';
|
||||
|
||||
jest.mock('~/config/winston', () => ({
|
||||
error: jest.fn(),
|
||||
warn: jest.fn(),
|
||||
info: jest.fn(),
|
||||
debug: jest.fn(),
|
||||
}));
|
||||
|
||||
let mongoServer: InstanceType<typeof MongoMemoryServer>;
|
||||
let MemoryEntry: mongoose.Model<IMemoryEntry>;
|
||||
let methods: MemoryMethods;
|
||||
|
||||
const userId = new mongoose.Types.ObjectId();
|
||||
const otherUserId = new mongoose.Types.ObjectId();
|
||||
const agentId = 'agent_partition_a';
|
||||
const otherAgentId = 'agent_partition_b';
|
||||
|
||||
beforeAll(async () => {
|
||||
mongoServer = await MongoMemoryServer.create();
|
||||
createModels(mongoose);
|
||||
MemoryEntry = mongoose.models.MemoryEntry as mongoose.Model<IMemoryEntry>;
|
||||
methods = createMemoryMethods(mongoose);
|
||||
await mongoose.connect(mongoServer.getUri());
|
||||
});
|
||||
|
||||
afterAll(async () => {
|
||||
await mongoose.disconnect();
|
||||
await mongoServer.stop();
|
||||
});
|
||||
|
||||
beforeEach(async () => {
|
||||
await MemoryEntry.deleteMany({});
|
||||
});
|
||||
|
||||
describe('memory partitions', () => {
|
||||
it('writes to the shared personal pool when agentId is omitted', async () => {
|
||||
await methods.setMemory({ userId, key: 'preference', value: 'likes tea', tokenCount: 3 });
|
||||
|
||||
const personal = await methods.getUserMemories({ userId });
|
||||
expect(personal).toHaveLength(1);
|
||||
expect(personal[0].agentId == null).toBe(true);
|
||||
|
||||
const agentScoped = await methods.getUserMemories({ userId, agentId });
|
||||
expect(agentScoped).toHaveLength(0);
|
||||
});
|
||||
|
||||
it('isolates the same key across partitions', async () => {
|
||||
await methods.setMemory({ userId, key: 'context', value: 'personal', tokenCount: 1 });
|
||||
await methods.setMemory({ userId, key: 'context', value: 'agent a', tokenCount: 2, agentId });
|
||||
await methods.setMemory({
|
||||
userId,
|
||||
key: 'context',
|
||||
value: 'agent b',
|
||||
tokenCount: 3,
|
||||
agentId: otherAgentId,
|
||||
});
|
||||
|
||||
const personal = await methods.getUserMemories({ userId });
|
||||
const partitionA = await methods.getUserMemories({ userId, agentId });
|
||||
const partitionB = await methods.getUserMemories({ userId, agentId: otherAgentId });
|
||||
|
||||
expect(personal.map((m) => m.value)).toEqual(['personal']);
|
||||
expect(partitionA.map((m) => m.value)).toEqual(['agent a']);
|
||||
expect(partitionB.map((m) => m.value)).toEqual(['agent b']);
|
||||
expect(await MemoryEntry.countDocuments({ userId })).toBe(3);
|
||||
});
|
||||
|
||||
it('upserts within a partition instead of duplicating', async () => {
|
||||
await methods.setMemory({ userId, key: 'context', value: 'first', tokenCount: 1, agentId });
|
||||
await methods.setMemory({ userId, key: 'context', value: 'second', tokenCount: 2, agentId });
|
||||
|
||||
const partitionA = await methods.getUserMemories({ userId, agentId });
|
||||
expect(partitionA).toHaveLength(1);
|
||||
expect(partitionA[0].value).toBe('second');
|
||||
});
|
||||
|
||||
it('treats legacy entries without an agentId field as the personal pool', async () => {
|
||||
await MemoryEntry.collection.insertOne({
|
||||
userId,
|
||||
key: 'legacy',
|
||||
value: 'pre-partition entry',
|
||||
tokenCount: 4,
|
||||
updated_at: new Date(),
|
||||
});
|
||||
|
||||
const personal = await methods.getUserMemories({ userId });
|
||||
expect(personal.map((m) => m.key)).toEqual(['legacy']);
|
||||
|
||||
await methods.setMemory({ userId, key: 'legacy', value: 'updated', tokenCount: 4 });
|
||||
expect(await MemoryEntry.countDocuments({ userId })).toBe(1);
|
||||
expect((await methods.getUserMemories({ userId }))[0].value).toBe('updated');
|
||||
});
|
||||
|
||||
it('deletes only within the targeted partition', async () => {
|
||||
await methods.setMemory({ userId, key: 'context', value: 'personal', tokenCount: 1 });
|
||||
await methods.setMemory({ userId, key: 'context', value: 'agent a', tokenCount: 2, agentId });
|
||||
|
||||
const agentDelete = await methods.deleteMemory({ userId, key: 'context', agentId });
|
||||
expect(agentDelete.ok).toBe(true);
|
||||
expect(await methods.getUserMemories({ userId })).toHaveLength(1);
|
||||
expect(await methods.getUserMemories({ userId, agentId })).toHaveLength(0);
|
||||
|
||||
const missingDelete = await methods.deleteMemory({ userId, key: 'context', agentId });
|
||||
expect(missingDelete.ok).toBe(false);
|
||||
});
|
||||
|
||||
it('scopes createMemory duplicate detection to the partition', async () => {
|
||||
await methods.createMemory({ userId, key: 'context', value: 'personal', tokenCount: 1 });
|
||||
|
||||
await expect(
|
||||
methods.createMemory({ userId, key: 'context', value: 'dupe', tokenCount: 1 }),
|
||||
).rejects.toThrow('already exists');
|
||||
|
||||
const created = await methods.createMemory({
|
||||
userId,
|
||||
key: 'context',
|
||||
value: 'agent a',
|
||||
tokenCount: 2,
|
||||
agentId,
|
||||
});
|
||||
expect(created.ok).toBe(true);
|
||||
});
|
||||
|
||||
it('formats memories per partition with partition-scoped totals', async () => {
|
||||
await methods.setMemory({ userId, key: 'one', value: 'personal one', tokenCount: 5 });
|
||||
await methods.setMemory({ userId, key: 'two', value: 'agent two', tokenCount: 7, agentId });
|
||||
|
||||
const personal = await methods.getFormattedMemories({ userId });
|
||||
expect(personal.totalTokens).toBe(5);
|
||||
expect(personal.withoutKeys).toContain('personal one');
|
||||
expect(personal.withoutKeys).not.toContain('agent two');
|
||||
|
||||
const partitionA = await methods.getFormattedMemories({ userId, agentId });
|
||||
expect(partitionA.totalTokens).toBe(7);
|
||||
expect(partitionA.withKeys).toContain('"key": "two"');
|
||||
expect(partitionA.withKeys).not.toContain('personal one');
|
||||
});
|
||||
|
||||
it('returns every partition from getAllUserMemories and wipes them all on deleteAllUserMemories', async () => {
|
||||
await methods.setMemory({ userId, key: 'one', value: 'personal', tokenCount: 1 });
|
||||
await methods.setMemory({ userId, key: 'two', value: 'agent a', tokenCount: 1, agentId });
|
||||
await methods.setMemory({ userId: otherUserId, key: 'three', value: 'other', tokenCount: 1 });
|
||||
|
||||
const all = await methods.getAllUserMemories(userId);
|
||||
expect(all).toHaveLength(2);
|
||||
|
||||
const deleted = await methods.deleteAllUserMemories(userId);
|
||||
expect(deleted).toBe(2);
|
||||
expect(await MemoryEntry.countDocuments({ userId: otherUserId })).toBe(1);
|
||||
});
|
||||
});
|
||||
|
|
@ -9,14 +9,32 @@ const formatDate = (date: Date): string => {
|
|||
return date.toISOString().split('T')[0];
|
||||
};
|
||||
|
||||
/** Partition filter: `agentId: null` matches both null and missing fields,
|
||||
* so pre-partition entries remain part of the shared personal pool. */
|
||||
const partitionFilter = (agentId?: string) => ({ agentId: agentId ?? null });
|
||||
|
||||
// Factory function that takes mongoose instance and returns the methods
|
||||
export function createMemoryMethods(mongoose: typeof import('mongoose')): {
|
||||
setMemory: ({ userId, key, value, tokenCount }: t.SetMemoryParams) => Promise<t.MemoryResult>;
|
||||
createMemory: ({ userId, key, value, tokenCount }: t.SetMemoryParams) => Promise<t.MemoryResult>;
|
||||
deleteMemory: ({ userId, key }: t.DeleteMemoryParams) => Promise<t.MemoryResult>;
|
||||
setMemory: ({
|
||||
userId,
|
||||
key,
|
||||
value,
|
||||
tokenCount,
|
||||
agentId,
|
||||
}: t.SetMemoryParams) => Promise<t.MemoryResult>;
|
||||
createMemory: ({
|
||||
userId,
|
||||
key,
|
||||
value,
|
||||
tokenCount,
|
||||
agentId,
|
||||
}: t.SetMemoryParams) => Promise<t.MemoryResult>;
|
||||
deleteMemory: ({ userId, key, agentId }: t.DeleteMemoryParams) => Promise<t.MemoryResult>;
|
||||
getAllUserMemories: (userId: string | Types.ObjectId) => Promise<t.IMemoryEntryLean[]>;
|
||||
getUserMemories: ({ userId, agentId }: t.GetUserMemoriesParams) => Promise<t.IMemoryEntryLean[]>;
|
||||
getFormattedMemories: ({
|
||||
userId,
|
||||
agentId,
|
||||
}: t.GetFormattedMemoriesParams) => Promise<t.FormattedMemoriesResult>;
|
||||
deleteAllUserMemories: (userId: string | Types.ObjectId) => Promise<number>;
|
||||
} {
|
||||
|
|
@ -29,6 +47,7 @@ export function createMemoryMethods(mongoose: typeof import('mongoose')): {
|
|||
key,
|
||||
value,
|
||||
tokenCount = 0,
|
||||
agentId,
|
||||
}: t.SetMemoryParams): Promise<t.MemoryResult> {
|
||||
try {
|
||||
if (key?.toLowerCase() === 'nothing') {
|
||||
|
|
@ -36,7 +55,11 @@ export function createMemoryMethods(mongoose: typeof import('mongoose')): {
|
|||
}
|
||||
|
||||
const MemoryEntry = mongoose.models.MemoryEntry;
|
||||
const existingMemory = await MemoryEntry.findOne({ userId, key });
|
||||
const existingMemory = await MemoryEntry.findOne({
|
||||
userId,
|
||||
key,
|
||||
...partitionFilter(agentId),
|
||||
});
|
||||
if (existingMemory) {
|
||||
throw new Error('Memory with this key already exists');
|
||||
}
|
||||
|
|
@ -46,6 +69,7 @@ export function createMemoryMethods(mongoose: typeof import('mongoose')): {
|
|||
key,
|
||||
value,
|
||||
tokenCount,
|
||||
...(agentId ? { agentId } : {}),
|
||||
updated_at: new Date(),
|
||||
});
|
||||
|
||||
|
|
@ -65,6 +89,7 @@ export function createMemoryMethods(mongoose: typeof import('mongoose')): {
|
|||
key,
|
||||
value,
|
||||
tokenCount = 0,
|
||||
agentId,
|
||||
}: t.SetMemoryParams): Promise<t.MemoryResult> {
|
||||
try {
|
||||
if (key?.toLowerCase() === 'nothing') {
|
||||
|
|
@ -73,10 +98,11 @@ export function createMemoryMethods(mongoose: typeof import('mongoose')): {
|
|||
|
||||
const MemoryEntry = mongoose.models.MemoryEntry;
|
||||
await MemoryEntry.findOneAndUpdate(
|
||||
{ userId, key },
|
||||
{ userId, key, ...partitionFilter(agentId) },
|
||||
{
|
||||
value,
|
||||
tokenCount,
|
||||
...(agentId ? { agentId } : {}),
|
||||
updated_at: new Date(),
|
||||
},
|
||||
{
|
||||
|
|
@ -96,10 +122,18 @@ export function createMemoryMethods(mongoose: typeof import('mongoose')): {
|
|||
/**
|
||||
* Deletes a specific memory entry for a user
|
||||
*/
|
||||
async function deleteMemory({ userId, key }: t.DeleteMemoryParams): Promise<t.MemoryResult> {
|
||||
async function deleteMemory({
|
||||
userId,
|
||||
key,
|
||||
agentId,
|
||||
}: t.DeleteMemoryParams): Promise<t.MemoryResult> {
|
||||
try {
|
||||
const MemoryEntry = mongoose.models.MemoryEntry;
|
||||
const result = await MemoryEntry.findOneAndDelete({ userId, key });
|
||||
const result = await MemoryEntry.findOneAndDelete({
|
||||
userId,
|
||||
key,
|
||||
...partitionFilter(agentId),
|
||||
});
|
||||
return { ok: !!result };
|
||||
} catch (error) {
|
||||
throw new Error(
|
||||
|
|
@ -109,7 +143,7 @@ export function createMemoryMethods(mongoose: typeof import('mongoose')): {
|
|||
}
|
||||
|
||||
/**
|
||||
* Gets all memory entries for a user
|
||||
* Gets all memory entries for a user across every partition
|
||||
*/
|
||||
async function getAllUserMemories(
|
||||
userId: string | Types.ObjectId,
|
||||
|
|
@ -125,13 +159,35 @@ export function createMemoryMethods(mongoose: typeof import('mongoose')): {
|
|||
}
|
||||
|
||||
/**
|
||||
* Gets and formats all memories for a user in two different formats
|
||||
* Gets a user's memory entries for a single partition
|
||||
* (the shared personal pool when `agentId` is omitted)
|
||||
*/
|
||||
async function getUserMemories({
|
||||
userId,
|
||||
agentId,
|
||||
}: t.GetUserMemoriesParams): Promise<t.IMemoryEntryLean[]> {
|
||||
try {
|
||||
const MemoryEntry = mongoose.models.MemoryEntry;
|
||||
return (await MemoryEntry.find({
|
||||
userId,
|
||||
...partitionFilter(agentId),
|
||||
}).lean()) as t.IMemoryEntryLean[];
|
||||
} catch (error) {
|
||||
throw new Error(
|
||||
`Failed to get memories: ${error instanceof Error ? error.message : 'Unknown error'}`,
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Gets and formats a partition's memories for a user in two different formats
|
||||
*/
|
||||
async function getFormattedMemories({
|
||||
userId,
|
||||
agentId,
|
||||
}: t.GetFormattedMemoriesParams): Promise<t.FormattedMemoriesResult> {
|
||||
try {
|
||||
const memories = await getAllUserMemories(userId);
|
||||
const memories = await getUserMemories({ userId, agentId });
|
||||
|
||||
if (!memories || memories.length === 0) {
|
||||
return { withKeys: '', withoutKeys: '', totalTokens: 0 };
|
||||
|
|
@ -187,6 +243,7 @@ export function createMemoryMethods(mongoose: typeof import('mongoose')): {
|
|||
createMemory,
|
||||
deleteMemory,
|
||||
getAllUserMemories,
|
||||
getUserMemories,
|
||||
getFormattedMemories,
|
||||
deleteAllUserMemories,
|
||||
};
|
||||
|
|
|
|||
|
|
@ -125,6 +125,12 @@ const agentSchema: Schema<IAgent> = new Schema<IAgent>(
|
|||
type: Schema.Types.Mixed,
|
||||
default: undefined,
|
||||
},
|
||||
/** Memory partition: 'agent' isolates memories per (user, agent); default shared pool */
|
||||
memory_scope: {
|
||||
type: String,
|
||||
enum: ['user', 'agent'],
|
||||
default: undefined,
|
||||
},
|
||||
tenantId: {
|
||||
type: String,
|
||||
index: true,
|
||||
|
|
|
|||
|
|
@ -20,6 +20,11 @@ const MemoryEntrySchema: Schema<IMemoryEntry> = new Schema({
|
|||
type: String,
|
||||
required: true,
|
||||
},
|
||||
/** Agent partition; null/absent = shared personal pool */
|
||||
agentId: {
|
||||
type: String,
|
||||
default: undefined,
|
||||
},
|
||||
tokenCount: {
|
||||
type: Number,
|
||||
default: 0,
|
||||
|
|
@ -34,4 +39,6 @@ const MemoryEntrySchema: Schema<IMemoryEntry> = new Schema({
|
|||
},
|
||||
});
|
||||
|
||||
MemoryEntrySchema.index({ userId: 1, agentId: 1, key: 1 });
|
||||
|
||||
export default MemoryEntrySchema;
|
||||
|
|
|
|||
|
|
@ -1,6 +1,7 @@
|
|||
import { Document, Types } from 'mongoose';
|
||||
import type {
|
||||
GraphEdge,
|
||||
MemoryScope,
|
||||
AgentToolOptions,
|
||||
AgentToolResources,
|
||||
AgentSubagentsConfig,
|
||||
|
|
@ -50,5 +51,7 @@ export interface IAgent extends Omit<Document, 'model'> {
|
|||
tool_options?: AgentToolOptions;
|
||||
/** Subagent spawning configuration — isolated-context child agents. */
|
||||
subagents?: AgentSubagentsConfig;
|
||||
/** Memory partition: 'agent' isolates memories per (user, agent); default shared pool */
|
||||
memory_scope?: MemoryScope;
|
||||
tenantId?: string;
|
||||
}
|
||||
|
|
|
|||
|
|
@ -5,6 +5,8 @@ export interface IMemoryEntry extends Document {
|
|||
userId: Types.ObjectId;
|
||||
key: string;
|
||||
value: string;
|
||||
/** Agent partition; null/absent = shared personal pool */
|
||||
agentId?: string;
|
||||
tokenCount?: number;
|
||||
updated_at?: Date;
|
||||
tenantId?: string;
|
||||
|
|
@ -15,6 +17,7 @@ export interface IMemoryEntryLean {
|
|||
userId: Types.ObjectId;
|
||||
key: string;
|
||||
value: string;
|
||||
agentId?: string;
|
||||
tokenCount?: number;
|
||||
updated_at?: Date;
|
||||
__v?: number;
|
||||
|
|
@ -26,15 +29,24 @@ export interface SetMemoryParams {
|
|||
key: string;
|
||||
value: string;
|
||||
tokenCount?: number;
|
||||
/** Agent partition; omit for the shared personal pool */
|
||||
agentId?: string;
|
||||
}
|
||||
|
||||
export interface DeleteMemoryParams {
|
||||
userId: string | Types.ObjectId;
|
||||
key: string;
|
||||
agentId?: string;
|
||||
}
|
||||
|
||||
export interface GetUserMemoriesParams {
|
||||
userId: string | Types.ObjectId;
|
||||
agentId?: string;
|
||||
}
|
||||
|
||||
export interface GetFormattedMemoriesParams {
|
||||
userId: string | Types.ObjectId;
|
||||
agentId?: string;
|
||||
}
|
||||
|
||||
// Result interfaces
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue