mirror of
https://github.com/danny-avila/LibreChat.git
synced 2025-12-17 17:00:15 +01:00
🧠 feat: Enforce Token Limit for Memory Usage (#8401)
This commit is contained in:
Samuel Path
GitHub
parent
2e1874e596
commit
8e869f2274
10 changed files with 765 additions and 31 deletions
|
|
@ -13,14 +13,25 @@ export default function MemoryArtifacts({ attachments }: { attachments?: TAttach
|
|||
const [isAnimating, setIsAnimating] = useState(false);
|
||||
const prevShowInfoRef = useRef<boolean>(showInfo);
|
||||
|
||||
const memoryArtifacts = useMemo(() => {
|
||||
const { hasErrors, memoryArtifacts } = useMemo(() => {
|
||||
let hasErrors = false;
|
||||
const result: MemoryArtifact[] = [];
|
||||
for (const attachment of attachments ?? []) {
|
||||
|
||||
if (!attachments || attachments.length === 0) {
|
||||
return { hasErrors, memoryArtifacts: result };
|
||||
}
|
||||
|
||||
for (const attachment of attachments) {
|
||||
if (attachment?.[Tools.memory] != null) {
|
||||
result.push(attachment[Tools.memory]);
|
||||
|
||||
if (!hasErrors && attachment[Tools.memory].type === 'error') {
|
||||
hasErrors = true;
|
||||
}
|
||||
}
|
||||
}
|
||||
return result;
|
||||
|
||||
return { hasErrors, memoryArtifacts: result };
|
||||
}, [attachments]);
|
||||
|
||||
useLayoutEffect(() => {
|
||||
|
|
@ -75,7 +86,12 @@ export default function MemoryArtifacts({ attachments }: { attachments?: TAttach
|
|||
<div className="flex items-center">
|
||||
<div className="inline-block">
|
||||
<button
|
||||
className="outline-hidden my-1 flex items-center gap-1 text-sm font-semibold text-text-secondary-alt transition-colors hover:text-text-primary"
|
||||
className={cn(
|
||||
'outline-hidden my-1 flex items-center gap-1 text-sm font-semibold transition-colors',
|
||||
hasErrors
|
||||
? 'text-red-500 hover:text-red-600 dark:text-red-400 dark:hover:text-red-500'
|
||||
: 'text-text-secondary-alt hover:text-text-primary',
|
||||
)}
|
||||
type="button"
|
||||
onClick={() => setShowInfo((prev) => !prev)}
|
||||
aria-expanded={showInfo}
|
||||
|
|
@ -102,7 +118,7 @@ export default function MemoryArtifacts({ attachments }: { attachments?: TAttach
|
|||
fill="currentColor"
|
||||
/>
|
||||
</svg>
|
||||
{localize('com_ui_memory_updated')}
|
||||
{hasErrors ? localize('com_ui_memory_error') : localize('com_ui_memory_updated')}
|
||||
</button>
|
||||
</div>
|
||||
</div>
|
||||
|
|
|
|||
|
|
@ -1,17 +1,47 @@
|
|||
import type { MemoryArtifact } from 'librechat-data-provider';
|
||||
import { useMemo } from 'react';
|
||||
import { useLocalize } from '~/hooks';
|
||||
|
||||
export default function MemoryInfo({ memoryArtifacts }: { memoryArtifacts: MemoryArtifact[] }) {
|
||||
const localize = useLocalize();
|
||||
|
||||
const { updatedMemories, deletedMemories, errorMessages } = useMemo(() => {
|
||||
const updated = memoryArtifacts.filter((art) => art.type === 'update');
|
||||
const deleted = memoryArtifacts.filter((art) => art.type === 'delete');
|
||||
const errors = memoryArtifacts.filter((art) => art.type === 'error');
|
||||
|
||||
const messages = errors.map((artifact) => {
|
||||
try {
|
||||
const errorData = JSON.parse(artifact.value as string);
|
||||
|
||||
if (errorData.errorType === 'already_exceeded') {
|
||||
return localize('com_ui_memory_already_exceeded', {
|
||||
tokens: errorData.tokenCount,
|
||||
});
|
||||
} else if (errorData.errorType === 'would_exceed') {
|
||||
return localize('com_ui_memory_would_exceed', {
|
||||
tokens: errorData.tokenCount,
|
||||
});
|
||||
} else {
|
||||
return localize('com_ui_memory_error');
|
||||
}
|
||||
} catch {
|
||||
return localize('com_ui_memory_error');
|
||||
}
|
||||
});
|
||||
|
||||
return {
|
||||
updatedMemories: updated,
|
||||
deletedMemories: deleted,
|
||||
errorMessages: messages,
|
||||
};
|
||||
}, [memoryArtifacts, localize]);
|
||||
|
||||
if (memoryArtifacts.length === 0) {
|
||||
return null;
|
||||
}
|
||||
|
||||
// Group artifacts by type
|
||||
const updatedMemories = memoryArtifacts.filter((artifact) => artifact.type === 'update');
|
||||
const deletedMemories = memoryArtifacts.filter((artifact) => artifact.type === 'delete');
|
||||
|
||||
if (updatedMemories.length === 0 && deletedMemories.length === 0) {
|
||||
if (updatedMemories.length === 0 && deletedMemories.length === 0 && errorMessages.length === 0) {
|
||||
return null;
|
||||
}
|
||||
|
||||
|
|
@ -23,8 +53,8 @@ export default function MemoryInfo({ memoryArtifacts }: { memoryArtifacts: Memor
|
|||
{localize('com_ui_memory_updated_items')}
|
||||
</h4>
|
||||
<div className="space-y-2">
|
||||
{updatedMemories.map((artifact, index) => (
|
||||
<div key={`update-${index}`} className="rounded-lg p-3">
|
||||
{updatedMemories.map((artifact) => (
|
||||
<div key={`update-${artifact.key}`} className="rounded-lg p-3">
|
||||
<div className="mb-1 text-xs font-medium uppercase tracking-wide text-text-secondary">
|
||||
{artifact.key}
|
||||
</div>
|
||||
|
|
@ -43,8 +73,8 @@ export default function MemoryInfo({ memoryArtifacts }: { memoryArtifacts: Memor
|
|||
{localize('com_ui_memory_deleted_items')}
|
||||
</h4>
|
||||
<div className="space-y-2">
|
||||
{deletedMemories.map((artifact, index) => (
|
||||
<div key={`delete-${index}`} className="rounded-lg p-3 opacity-60">
|
||||
{deletedMemories.map((artifact) => (
|
||||
<div key={`delete-${artifact.key}`} className="rounded-lg p-3 opacity-60">
|
||||
<div className="mb-1 text-xs font-medium uppercase tracking-wide text-text-secondary">
|
||||
{artifact.key}
|
||||
</div>
|
||||
|
|
@ -56,6 +86,24 @@ export default function MemoryInfo({ memoryArtifacts }: { memoryArtifacts: Memor
|
|||
</div>
|
||||
</div>
|
||||
)}
|
||||
|
||||
{errorMessages.length > 0 && (
|
||||
<div>
|
||||
<h4 className="mb-2 text-sm font-semibold text-red-500">
|
||||
{localize('com_ui_memory_storage_full')}
|
||||
</h4>
|
||||
<div className="space-y-2">
|
||||
{errorMessages.map((errorMessage) => (
|
||||
<div
|
||||
key={errorMessage}
|
||||
className="rounded-md bg-red-50 p-3 text-sm text-red-800 dark:bg-red-900/20 dark:text-red-400"
|
||||
>
|
||||
{errorMessage}
|
||||
</div>
|
||||
))}
|
||||
</div>
|
||||
</div>
|
||||
)}
|
||||
</div>
|
||||
);
|
||||
}
|
||||
|
|
|
|||
|
|
@ -0,0 +1,197 @@
|
|||
import React from 'react';
|
||||
import { render, screen, fireEvent } from '@testing-library/react';
|
||||
import '@testing-library/jest-dom/extend-expect';
|
||||
import MemoryArtifacts from '../MemoryArtifacts';
|
||||
import type { TAttachment, MemoryArtifact } from 'librechat-data-provider';
|
||||
import { Tools } from 'librechat-data-provider';
|
||||
|
||||
// Mock the localize hook
|
||||
jest.mock('~/hooks', () => ({
|
||||
useLocalize: () => (key: string) => {
|
||||
const translations: Record<string, string> = {
|
||||
com_ui_memory_updated: 'Updated saved memory',
|
||||
com_ui_memory_error: 'Memory Error',
|
||||
};
|
||||
return translations[key] || key;
|
||||
},
|
||||
}));
|
||||
|
||||
// Mock the MemoryInfo component
|
||||
jest.mock('../MemoryInfo', () => ({
|
||||
__esModule: true,
|
||||
default: ({ memoryArtifacts }: { memoryArtifacts: MemoryArtifact[] }) => (
|
||||
<div data-testid="memory-info">
|
||||
{memoryArtifacts.map((artifact, index) => (
|
||||
<div key={index} data-testid={`memory-artifact-${artifact.type}`}>
|
||||
{artifact.type}: {artifact.key}
|
||||
</div>
|
||||
))}
|
||||
</div>
|
||||
),
|
||||
}));
|
||||
|
||||
describe('MemoryArtifacts', () => {
|
||||
const createMemoryAttachment = (type: 'update' | 'delete' | 'error', key: string): TAttachment =>
|
||||
({
|
||||
type: Tools.memory,
|
||||
[Tools.memory]: {
|
||||
type,
|
||||
key,
|
||||
value:
|
||||
type === 'error'
|
||||
? JSON.stringify({ errorType: 'exceeded', tokenCount: 100 })
|
||||
: 'test value',
|
||||
} as MemoryArtifact,
|
||||
}) as TAttachment;
|
||||
|
||||
describe('Error State Handling', () => {
|
||||
test('displays error styling when memory artifacts contain errors', () => {
|
||||
const attachments = [
|
||||
createMemoryAttachment('error', 'system'),
|
||||
createMemoryAttachment('update', 'memory1'),
|
||||
];
|
||||
|
||||
render(<MemoryArtifacts attachments={attachments} />);
|
||||
|
||||
const button = screen.getByRole('button');
|
||||
expect(button).toHaveClass('text-red-500');
|
||||
expect(button).toHaveClass('hover:text-red-600');
|
||||
expect(button).toHaveClass('dark:text-red-400');
|
||||
expect(button).toHaveClass('dark:hover:text-red-500');
|
||||
});
|
||||
|
||||
test('displays normal styling when no errors present', () => {
|
||||
const attachments = [
|
||||
createMemoryAttachment('update', 'memory1'),
|
||||
createMemoryAttachment('delete', 'memory2'),
|
||||
];
|
||||
|
||||
render(<MemoryArtifacts attachments={attachments} />);
|
||||
|
||||
const button = screen.getByRole('button');
|
||||
expect(button).toHaveClass('text-text-secondary-alt');
|
||||
expect(button).toHaveClass('hover:text-text-primary');
|
||||
expect(button).not.toHaveClass('text-red-500');
|
||||
});
|
||||
|
||||
test('displays error message when errors are present', () => {
|
||||
const attachments = [createMemoryAttachment('error', 'system')];
|
||||
|
||||
render(<MemoryArtifacts attachments={attachments} />);
|
||||
|
||||
expect(screen.getByText('Memory Error')).toBeInTheDocument();
|
||||
expect(screen.queryByText('Updated saved memory')).not.toBeInTheDocument();
|
||||
});
|
||||
|
||||
test('displays normal message when no errors are present', () => {
|
||||
const attachments = [createMemoryAttachment('update', 'memory1')];
|
||||
|
||||
render(<MemoryArtifacts attachments={attachments} />);
|
||||
|
||||
expect(screen.getByText('Updated saved memory')).toBeInTheDocument();
|
||||
expect(screen.queryByText('Memory Error')).not.toBeInTheDocument();
|
||||
});
|
||||
});
|
||||
|
||||
describe('Memory Artifacts Filtering', () => {
|
||||
test('filters and passes only memory-type attachments to MemoryInfo', () => {
|
||||
const attachments = [
|
||||
createMemoryAttachment('update', 'memory1'),
|
||||
{ type: 'file' } as TAttachment, // Non-memory attachment
|
||||
createMemoryAttachment('error', 'system'),
|
||||
];
|
||||
|
||||
render(<MemoryArtifacts attachments={attachments} />);
|
||||
|
||||
// Click to expand
|
||||
fireEvent.click(screen.getByRole('button'));
|
||||
|
||||
// Check that only memory artifacts are passed to MemoryInfo
|
||||
expect(screen.getByTestId('memory-artifact-update')).toBeInTheDocument();
|
||||
expect(screen.getByTestId('memory-artifact-error')).toBeInTheDocument();
|
||||
});
|
||||
|
||||
test('correctly identifies multiple error artifacts', () => {
|
||||
const attachments = [
|
||||
createMemoryAttachment('error', 'system1'),
|
||||
createMemoryAttachment('error', 'system2'),
|
||||
createMemoryAttachment('update', 'memory1'),
|
||||
];
|
||||
|
||||
render(<MemoryArtifacts attachments={attachments} />);
|
||||
|
||||
const button = screen.getByRole('button');
|
||||
expect(button).toHaveClass('text-red-500');
|
||||
expect(screen.getByText('Memory Error')).toBeInTheDocument();
|
||||
});
|
||||
});
|
||||
|
||||
describe('Collapse/Expand Functionality', () => {
|
||||
test('toggles memory info visibility on button click', () => {
|
||||
const attachments = [createMemoryAttachment('update', 'memory1')];
|
||||
|
||||
render(<MemoryArtifacts attachments={attachments} />);
|
||||
|
||||
// Initially collapsed
|
||||
expect(screen.queryByTestId('memory-info')).not.toBeInTheDocument();
|
||||
|
||||
// Click to expand
|
||||
fireEvent.click(screen.getByRole('button'));
|
||||
expect(screen.getByTestId('memory-info')).toBeInTheDocument();
|
||||
|
||||
// Click to collapse
|
||||
fireEvent.click(screen.getByRole('button'));
|
||||
expect(screen.queryByTestId('memory-info')).not.toBeInTheDocument();
|
||||
});
|
||||
|
||||
test('updates aria-expanded attribute correctly', () => {
|
||||
const attachments = [createMemoryAttachment('update', 'memory1')];
|
||||
|
||||
render(<MemoryArtifacts attachments={attachments} />);
|
||||
|
||||
const button = screen.getByRole('button');
|
||||
expect(button).toHaveAttribute('aria-expanded', 'false');
|
||||
|
||||
fireEvent.click(button);
|
||||
expect(button).toHaveAttribute('aria-expanded', 'true');
|
||||
});
|
||||
});
|
||||
|
||||
describe('Edge Cases', () => {
|
||||
test('handles empty attachments array', () => {
|
||||
render(<MemoryArtifacts attachments={[]} />);
|
||||
expect(screen.queryByRole('button')).not.toBeInTheDocument();
|
||||
});
|
||||
|
||||
test('handles undefined attachments', () => {
|
||||
render(<MemoryArtifacts />);
|
||||
expect(screen.queryByRole('button')).not.toBeInTheDocument();
|
||||
});
|
||||
|
||||
test('handles attachments with no memory artifacts', () => {
|
||||
const attachments = [{ type: 'file' } as TAttachment, { type: 'image' } as TAttachment];
|
||||
|
||||
render(<MemoryArtifacts attachments={attachments} />);
|
||||
expect(screen.queryByRole('button')).not.toBeInTheDocument();
|
||||
});
|
||||
|
||||
test('handles malformed memory artifacts gracefully', () => {
|
||||
const attachments = [
|
||||
{
|
||||
type: Tools.memory,
|
||||
[Tools.memory]: {
|
||||
type: 'error',
|
||||
key: 'system',
|
||||
// Missing value
|
||||
},
|
||||
} as TAttachment,
|
||||
];
|
||||
|
||||
render(<MemoryArtifacts attachments={attachments} />);
|
||||
|
||||
const button = screen.getByRole('button');
|
||||
expect(button).toHaveClass('text-red-500');
|
||||
expect(screen.getByText('Memory Error')).toBeInTheDocument();
|
||||
});
|
||||
});
|
||||
});
|
||||
|
|
@ -0,0 +1,267 @@
|
|||
import React from 'react';
|
||||
import { render, screen } from '@testing-library/react';
|
||||
import '@testing-library/jest-dom/extend-expect';
|
||||
import MemoryInfo from '../MemoryInfo';
|
||||
import type { MemoryArtifact } from 'librechat-data-provider';
|
||||
|
||||
// Mock the localize hook
|
||||
jest.mock('~/hooks', () => ({
|
||||
useLocalize: () => (key: string, params?: Record<string, any>) => {
|
||||
const translations: Record<string, string> = {
|
||||
com_ui_memory_updated_items: 'Updated Memories',
|
||||
com_ui_memory_deleted_items: 'Deleted Memories',
|
||||
com_ui_memory_already_exceeded: `Memory storage already full - exceeded by ${params?.tokens || 0} tokens. Delete existing memories before adding new ones.`,
|
||||
com_ui_memory_would_exceed: `Cannot save - would exceed limit by ${params?.tokens || 0} tokens. Delete existing memories to make space.`,
|
||||
com_ui_memory_deleted: 'This memory has been deleted',
|
||||
com_ui_memory_storage_full: 'Memory Storage Full',
|
||||
com_ui_memory_error: 'Memory Error',
|
||||
com_ui_updated_successfully: 'Updated successfully',
|
||||
com_ui_none_selected: 'None selected',
|
||||
};
|
||||
return translations[key] || key;
|
||||
},
|
||||
}));
|
||||
|
||||
describe('MemoryInfo', () => {
|
||||
const createMemoryArtifact = (
|
||||
type: 'update' | 'delete' | 'error',
|
||||
key: string,
|
||||
value?: string,
|
||||
): MemoryArtifact => ({
|
||||
type,
|
||||
key,
|
||||
value: value || `test value for ${key}`,
|
||||
});
|
||||
|
||||
describe('Error Memory Display', () => {
|
||||
test('displays error section when memory is already exceeded', () => {
|
||||
const memoryArtifacts: MemoryArtifact[] = [
|
||||
{
|
||||
type: 'error',
|
||||
key: 'system',
|
||||
value: JSON.stringify({ errorType: 'already_exceeded', tokenCount: 150 }),
|
||||
},
|
||||
];
|
||||
|
||||
render(<MemoryInfo memoryArtifacts={memoryArtifacts} />);
|
||||
|
||||
expect(screen.getByText('Memory Storage Full')).toBeInTheDocument();
|
||||
expect(
|
||||
screen.getByText(
|
||||
'Memory storage already full - exceeded by 150 tokens. Delete existing memories before adding new ones.',
|
||||
),
|
||||
).toBeInTheDocument();
|
||||
});
|
||||
|
||||
test('displays error when memory would exceed limit', () => {
|
||||
const memoryArtifacts: MemoryArtifact[] = [
|
||||
{
|
||||
type: 'error',
|
||||
key: 'system',
|
||||
value: JSON.stringify({ errorType: 'would_exceed', tokenCount: 50 }),
|
||||
},
|
||||
];
|
||||
|
||||
render(<MemoryInfo memoryArtifacts={memoryArtifacts} />);
|
||||
|
||||
expect(screen.getByText('Memory Storage Full')).toBeInTheDocument();
|
||||
expect(
|
||||
screen.getByText(
|
||||
'Cannot save - would exceed limit by 50 tokens. Delete existing memories to make space.',
|
||||
),
|
||||
).toBeInTheDocument();
|
||||
});
|
||||
|
||||
test('displays multiple error messages', () => {
|
||||
const memoryArtifacts: MemoryArtifact[] = [
|
||||
{
|
||||
type: 'error',
|
||||
key: 'system1',
|
||||
value: JSON.stringify({ errorType: 'already_exceeded', tokenCount: 100 }),
|
||||
},
|
||||
{
|
||||
type: 'error',
|
||||
key: 'system2',
|
||||
value: JSON.stringify({ errorType: 'would_exceed', tokenCount: 25 }),
|
||||
},
|
||||
];
|
||||
|
||||
render(<MemoryInfo memoryArtifacts={memoryArtifacts} />);
|
||||
|
||||
expect(
|
||||
screen.getByText(
|
||||
'Memory storage already full - exceeded by 100 tokens. Delete existing memories before adding new ones.',
|
||||
),
|
||||
).toBeInTheDocument();
|
||||
expect(
|
||||
screen.getByText(
|
||||
'Cannot save - would exceed limit by 25 tokens. Delete existing memories to make space.',
|
||||
),
|
||||
).toBeInTheDocument();
|
||||
});
|
||||
|
||||
test('applies correct styling to error messages', () => {
|
||||
const memoryArtifacts: MemoryArtifact[] = [
|
||||
{
|
||||
type: 'error',
|
||||
key: 'system',
|
||||
value: JSON.stringify({ errorType: 'would_exceed', tokenCount: 50 }),
|
||||
},
|
||||
];
|
||||
|
||||
render(<MemoryInfo memoryArtifacts={memoryArtifacts} />);
|
||||
|
||||
const errorMessage = screen.getByText(
|
||||
'Cannot save - would exceed limit by 50 tokens. Delete existing memories to make space.',
|
||||
);
|
||||
const errorContainer = errorMessage.closest('div');
|
||||
|
||||
expect(errorContainer).toHaveClass('rounded-md');
|
||||
expect(errorContainer).toHaveClass('bg-red-50');
|
||||
expect(errorContainer).toHaveClass('p-3');
|
||||
expect(errorContainer).toHaveClass('text-sm');
|
||||
expect(errorContainer).toHaveClass('text-red-800');
|
||||
expect(errorContainer).toHaveClass('dark:bg-red-900/20');
|
||||
expect(errorContainer).toHaveClass('dark:text-red-400');
|
||||
});
|
||||
});
|
||||
|
||||
describe('Mixed Memory Types', () => {
|
||||
test('displays all sections when different memory types are present', () => {
|
||||
const memoryArtifacts: MemoryArtifact[] = [
|
||||
createMemoryArtifact('update', 'memory1', 'Updated content'),
|
||||
createMemoryArtifact('delete', 'memory2'),
|
||||
{
|
||||
type: 'error',
|
||||
key: 'system',
|
||||
value: JSON.stringify({ errorType: 'would_exceed', tokenCount: 200 }),
|
||||
},
|
||||
];
|
||||
|
||||
render(<MemoryInfo memoryArtifacts={memoryArtifacts} />);
|
||||
|
||||
// Check all sections are present
|
||||
expect(screen.getByText('Updated Memories')).toBeInTheDocument();
|
||||
expect(screen.getByText('Deleted Memories')).toBeInTheDocument();
|
||||
expect(screen.getByText('Memory Storage Full')).toBeInTheDocument();
|
||||
|
||||
// Check content
|
||||
expect(screen.getByText('memory1')).toBeInTheDocument();
|
||||
expect(screen.getByText('Updated content')).toBeInTheDocument();
|
||||
expect(screen.getByText('memory2')).toBeInTheDocument();
|
||||
expect(
|
||||
screen.getByText(
|
||||
'Cannot save - would exceed limit by 200 tokens. Delete existing memories to make space.',
|
||||
),
|
||||
).toBeInTheDocument();
|
||||
});
|
||||
|
||||
test('only displays sections with content', () => {
|
||||
const memoryArtifacts: MemoryArtifact[] = [
|
||||
{
|
||||
type: 'error',
|
||||
key: 'system',
|
||||
value: JSON.stringify({ errorType: 'already_exceeded', tokenCount: 10 }),
|
||||
},
|
||||
];
|
||||
|
||||
render(<MemoryInfo memoryArtifacts={memoryArtifacts} />);
|
||||
|
||||
// Only error section should be present
|
||||
expect(screen.getByText('Memory Storage Full')).toBeInTheDocument();
|
||||
expect(screen.queryByText('Updated Memories')).not.toBeInTheDocument();
|
||||
expect(screen.queryByText('Deleted Memories')).not.toBeInTheDocument();
|
||||
});
|
||||
});
|
||||
|
||||
describe('Edge Cases', () => {
|
||||
test('handles empty memory artifacts array', () => {
|
||||
const { container } = render(<MemoryInfo memoryArtifacts={[]} />);
|
||||
expect(container.firstChild).toBeNull();
|
||||
});
|
||||
|
||||
test('handles malformed error data gracefully', () => {
|
||||
const memoryArtifacts: MemoryArtifact[] = [
|
||||
{
|
||||
type: 'error',
|
||||
key: 'system',
|
||||
value: 'invalid json',
|
||||
},
|
||||
];
|
||||
|
||||
render(<MemoryInfo memoryArtifacts={memoryArtifacts} />);
|
||||
|
||||
// Should render generic error message
|
||||
expect(screen.getByText('Memory Storage Full')).toBeInTheDocument();
|
||||
expect(screen.getByText('Memory Error')).toBeInTheDocument();
|
||||
});
|
||||
|
||||
test('handles missing value in error artifact', () => {
|
||||
const memoryArtifacts: MemoryArtifact[] = [
|
||||
{
|
||||
type: 'error',
|
||||
key: 'system',
|
||||
// value is undefined
|
||||
} as MemoryArtifact,
|
||||
];
|
||||
|
||||
render(<MemoryInfo memoryArtifacts={memoryArtifacts} />);
|
||||
|
||||
expect(screen.getByText('Memory Storage Full')).toBeInTheDocument();
|
||||
expect(screen.getByText('Memory Error')).toBeInTheDocument();
|
||||
});
|
||||
|
||||
test('handles unknown errorType gracefully', () => {
|
||||
const memoryArtifacts: MemoryArtifact[] = [
|
||||
{
|
||||
type: 'error',
|
||||
key: 'system',
|
||||
value: JSON.stringify({ errorType: 'unknown_type', tokenCount: 30 }),
|
||||
},
|
||||
];
|
||||
|
||||
render(<MemoryInfo memoryArtifacts={memoryArtifacts} />);
|
||||
|
||||
// Should show generic error message for unknown types
|
||||
expect(screen.getByText('Memory Storage Full')).toBeInTheDocument();
|
||||
expect(screen.getByText('Memory Error')).toBeInTheDocument();
|
||||
});
|
||||
|
||||
test('returns null when no memories of any type exist', () => {
|
||||
const memoryArtifacts: MemoryArtifact[] = [{ type: 'unknown' as any, key: 'test' }];
|
||||
|
||||
const { container } = render(<MemoryInfo memoryArtifacts={memoryArtifacts} />);
|
||||
expect(container.firstChild).toBeNull();
|
||||
});
|
||||
});
|
||||
|
||||
describe('Update and Delete Memory Display', () => {
|
||||
test('displays updated memories correctly', () => {
|
||||
const memoryArtifacts: MemoryArtifact[] = [
|
||||
createMemoryArtifact('update', 'preferences', 'User prefers dark mode'),
|
||||
createMemoryArtifact('update', 'location', 'Lives in San Francisco'),
|
||||
];
|
||||
|
||||
render(<MemoryInfo memoryArtifacts={memoryArtifacts} />);
|
||||
|
||||
expect(screen.getByText('Updated Memories')).toBeInTheDocument();
|
||||
expect(screen.getByText('preferences')).toBeInTheDocument();
|
||||
expect(screen.getByText('User prefers dark mode')).toBeInTheDocument();
|
||||
expect(screen.getByText('location')).toBeInTheDocument();
|
||||
expect(screen.getByText('Lives in San Francisco')).toBeInTheDocument();
|
||||
});
|
||||
|
||||
test('displays deleted memories correctly', () => {
|
||||
const memoryArtifacts: MemoryArtifact[] = [
|
||||
createMemoryArtifact('delete', 'old_preference'),
|
||||
createMemoryArtifact('delete', 'outdated_info'),
|
||||
];
|
||||
|
||||
render(<MemoryInfo memoryArtifacts={memoryArtifacts} />);
|
||||
|
||||
expect(screen.getByText('Deleted Memories')).toBeInTheDocument();
|
||||
expect(screen.getByText('old_preference')).toBeInTheDocument();
|
||||
expect(screen.getByText('outdated_info')).toBeInTheDocument();
|
||||
});
|
||||
});
|
||||
});
|
||||
|
|
@ -860,6 +860,10 @@
|
|||
"com_ui_memory_key_validation": "Memory key must only contain lowercase letters and underscores.",
|
||||
"com_ui_memory_updated": "Updated saved memory",
|
||||
"com_ui_memory_updated_items": "Updated Memories",
|
||||
"com_ui_memory_storage_full": "Memory Storage Full",
|
||||
"com_ui_memory_error": "Memory Error",
|
||||
"com_ui_memory_already_exceeded": "Memory storage already full - exceeded by {{tokens}} tokens. Delete existing memories before adding new ones.",
|
||||
"com_ui_memory_would_exceed": "Cannot save - would exceed limit by {{tokens}} tokens. Delete existing memories to make space.",
|
||||
"com_ui_mention": "Mention an endpoint, assistant, or preset to quickly switch to it",
|
||||
"com_ui_min_tags": "Cannot remove more values, a minimum of {{0}} are required.",
|
||||
"com_ui_misc": "Misc.",
|
||||
|
|
|
|||
|
|
@ -17,6 +17,9 @@ import '@testing-library/jest-dom/extend-expect';
|
|||
// 'react-lottie' uses canvas
|
||||
import 'jest-canvas-mock';
|
||||
|
||||
// Mock ResizeObserver
|
||||
import './resizeObserver.mock';
|
||||
|
||||
beforeEach(() => {
|
||||
jest.clearAllMocks();
|
||||
});
|
||||
|
|
|
|||
|
|
@ -1,6 +1,7 @@
|
|||
export default {
|
||||
collectCoverageFrom: ['src/**/*.{js,jsx,ts,tsx}', '!<rootDir>/node_modules/'],
|
||||
coveragePathIgnorePatterns: ['/node_modules/', '/dist/'],
|
||||
testPathIgnorePatterns: ['/node_modules/', '/dist/'],
|
||||
coverageReporters: ['text', 'cobertura'],
|
||||
testResultsProcessor: 'jest-junit',
|
||||
moduleNameMapper: {
|
||||
|
|
|
|||
165
packages/api/src/agents/__tests__/memory.test.ts
Normal file
165
packages/api/src/agents/__tests__/memory.test.ts
Normal file
|
|
@ -0,0 +1,165 @@
|
|||
import { Tools, type MemoryArtifact } from 'librechat-data-provider';
|
||||
import { createMemoryTool } from '../memory';
|
||||
|
||||
// Mock the logger
|
||||
jest.mock('winston', () => ({
|
||||
createLogger: jest.fn(() => ({
|
||||
debug: jest.fn(),
|
||||
warn: jest.fn(),
|
||||
error: jest.fn(),
|
||||
})),
|
||||
format: {
|
||||
combine: jest.fn(),
|
||||
colorize: jest.fn(),
|
||||
simple: jest.fn(),
|
||||
},
|
||||
transports: {
|
||||
Console: jest.fn(),
|
||||
},
|
||||
}));
|
||||
|
||||
// Mock the Tokenizer
|
||||
jest.mock('~/utils', () => ({
|
||||
Tokenizer: {
|
||||
getTokenCount: jest.fn((text: string) => text.length), // Simple mock: 1 char = 1 token
|
||||
},
|
||||
}));
|
||||
|
||||
describe('createMemoryTool', () => {
|
||||
let mockSetMemory: jest.Mock;
|
||||
|
||||
beforeEach(() => {
|
||||
jest.clearAllMocks();
|
||||
mockSetMemory = jest.fn().mockResolvedValue({ ok: true });
|
||||
});
|
||||
|
||||
// Memory overflow tests
|
||||
describe('overflow handling', () => {
|
||||
it('should return error artifact when memory is already overflowing', async () => {
|
||||
const tool = createMemoryTool({
|
||||
userId: 'test-user',
|
||||
setMemory: mockSetMemory,
|
||||
tokenLimit: 100,
|
||||
totalTokens: 150, // Already over limit
|
||||
});
|
||||
|
||||
// Call the underlying function directly since invoke() doesn't handle responseFormat in tests
|
||||
const result = await tool.func({ key: 'test', value: 'new memory' });
|
||||
expect(result).toHaveLength(2);
|
||||
expect(result[0]).toBe('Memory storage exceeded. Cannot save new memories.');
|
||||
|
||||
const artifacts = result[1] as Record<Tools.memory, MemoryArtifact>;
|
||||
expect(artifacts[Tools.memory]).toBeDefined();
|
||||
expect(artifacts[Tools.memory].type).toBe('error');
|
||||
expect(artifacts[Tools.memory].key).toBe('system');
|
||||
|
||||
const errorData = JSON.parse(artifacts[Tools.memory].value as string);
|
||||
expect(errorData).toEqual({
|
||||
errorType: 'already_exceeded',
|
||||
tokenCount: 50,
|
||||
totalTokens: 150,
|
||||
tokenLimit: 100,
|
||||
});
|
||||
|
||||
expect(mockSetMemory).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it('should return error artifact when new memory would exceed limit', async () => {
|
||||
const tool = createMemoryTool({
|
||||
userId: 'test-user',
|
||||
setMemory: mockSetMemory,
|
||||
tokenLimit: 100,
|
||||
totalTokens: 80,
|
||||
});
|
||||
|
||||
// This would put us at 101 tokens total, exceeding the limit
|
||||
const result = await tool.func({ key: 'test', value: 'This is a 20 char str' });
|
||||
expect(result).toHaveLength(2);
|
||||
expect(result[0]).toBe('Memory storage would exceed limit. Cannot save this memory.');
|
||||
|
||||
const artifacts = result[1] as Record<Tools.memory, MemoryArtifact>;
|
||||
expect(artifacts[Tools.memory]).toBeDefined();
|
||||
expect(artifacts[Tools.memory].type).toBe('error');
|
||||
expect(artifacts[Tools.memory].key).toBe('system');
|
||||
|
||||
const errorData = JSON.parse(artifacts[Tools.memory].value as string);
|
||||
expect(errorData).toEqual({
|
||||
errorType: 'would_exceed',
|
||||
tokenCount: 1, // Math.abs(-1)
|
||||
totalTokens: 101,
|
||||
tokenLimit: 100,
|
||||
});
|
||||
|
||||
expect(mockSetMemory).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it('should successfully save memory when below limit', async () => {
|
||||
const tool = createMemoryTool({
|
||||
userId: 'test-user',
|
||||
setMemory: mockSetMemory,
|
||||
tokenLimit: 100,
|
||||
totalTokens: 50,
|
||||
});
|
||||
|
||||
const result = await tool.func({ key: 'test', value: 'small memory' });
|
||||
expect(result).toHaveLength(2);
|
||||
expect(result[0]).toBe('Memory set for key "test" (12 tokens)');
|
||||
|
||||
const artifacts = result[1] as Record<Tools.memory, MemoryArtifact>;
|
||||
expect(artifacts[Tools.memory]).toBeDefined();
|
||||
expect(artifacts[Tools.memory].type).toBe('update');
|
||||
expect(artifacts[Tools.memory].key).toBe('test');
|
||||
expect(artifacts[Tools.memory].value).toBe('small memory');
|
||||
|
||||
expect(mockSetMemory).toHaveBeenCalledWith({
|
||||
userId: 'test-user',
|
||||
key: 'test',
|
||||
value: 'small memory',
|
||||
tokenCount: 12,
|
||||
});
|
||||
});
|
||||
});
|
||||
|
||||
// Basic functionality tests
|
||||
describe('basic functionality', () => {
|
||||
it('should validate keys when validKeys is provided', async () => {
|
||||
const tool = createMemoryTool({
|
||||
userId: 'test-user',
|
||||
setMemory: mockSetMemory,
|
||||
validKeys: ['allowed', 'keys'],
|
||||
});
|
||||
|
||||
const result = await tool.func({ key: 'invalid', value: 'some value' });
|
||||
expect(result).toHaveLength(2);
|
||||
expect(result[0]).toBe('Invalid key "invalid". Must be one of: allowed, keys');
|
||||
expect(result[1]).toBeUndefined();
|
||||
expect(mockSetMemory).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it('should handle setMemory failure', async () => {
|
||||
mockSetMemory.mockResolvedValue({ ok: false });
|
||||
const tool = createMemoryTool({
|
||||
userId: 'test-user',
|
||||
setMemory: mockSetMemory,
|
||||
});
|
||||
|
||||
const result = await tool.func({ key: 'test', value: 'some value' });
|
||||
expect(result).toHaveLength(2);
|
||||
expect(result[0]).toBe('Failed to set memory for key "test"');
|
||||
expect(result[1]).toBeUndefined();
|
||||
});
|
||||
|
||||
it('should handle exceptions', async () => {
|
||||
mockSetMemory.mockRejectedValue(new Error('DB error'));
|
||||
const tool = createMemoryTool({
|
||||
userId: 'test-user',
|
||||
setMemory: mockSetMemory,
|
||||
});
|
||||
|
||||
const result = await tool.func({ key: 'test', value: 'some value' });
|
||||
expect(result).toHaveLength(2);
|
||||
expect(result[0]).toBe('Error setting memory for key "test"');
|
||||
expect(result[1]).toBeUndefined();
|
||||
});
|
||||
});
|
||||
});
|
||||
|
|
@ -71,7 +71,7 @@ const getDefaultInstructions = (
|
|||
/**
|
||||
* Creates a memory tool instance with user context
|
||||
*/
|
||||
const createMemoryTool = ({
|
||||
export const createMemoryTool = ({
|
||||
userId,
|
||||
setMemory,
|
||||
validKeys,
|
||||
|
|
@ -84,6 +84,9 @@ const createMemoryTool = ({
|
|||
tokenLimit?: number;
|
||||
totalTokens?: number;
|
||||
}) => {
|
||||
const remainingTokens = tokenLimit ? tokenLimit - totalTokens : Infinity;
|
||||
const isOverflowing = tokenLimit ? remainingTokens <= 0 : false;
|
||||
|
||||
return tool(
|
||||
async ({ key, value }) => {
|
||||
try {
|
||||
|
|
@ -93,24 +96,48 @@ const createMemoryTool = ({
|
|||
', ',
|
||||
)}`,
|
||||
);
|
||||
return `Invalid key "${key}". Must be one of: ${validKeys.join(', ')}`;
|
||||
return [`Invalid key "${key}". Must be one of: ${validKeys.join(', ')}`, undefined];
|
||||
}
|
||||
|
||||
const tokenCount = Tokenizer.getTokenCount(value, 'o200k_base');
|
||||
|
||||
if (tokenLimit && tokenCount > tokenLimit) {
|
||||
logger.warn(
|
||||
`Memory Agent failed to set memory: Value exceeds token limit. Value has ${tokenCount} tokens, but limit is ${tokenLimit}`,
|
||||
);
|
||||
return `Memory value too large: ${tokenCount} tokens exceeds limit of ${tokenLimit}`;
|
||||
if (isOverflowing) {
|
||||
const errorArtifact: Record<Tools.memory, MemoryArtifact> = {
|
||||
[Tools.memory]: {
|
||||
key: 'system',
|
||||
type: 'error',
|
||||
value: JSON.stringify({
|
||||
errorType: 'already_exceeded',
|
||||
tokenCount: Math.abs(remainingTokens),
|
||||
totalTokens: totalTokens,
|
||||
tokenLimit: tokenLimit!,
|
||||
}),
|
||||
tokenCount: totalTokens,
|
||||
},
|
||||
};
|
||||
return [`Memory storage exceeded. Cannot save new memories.`, errorArtifact];
|
||||
}
|
||||
|
||||
if (tokenLimit && totalTokens + tokenCount > tokenLimit) {
|
||||
const remainingCapacity = tokenLimit - totalTokens;
|
||||
logger.warn(
|
||||
`Memory Agent failed to set memory: Would exceed total token limit. Current usage: ${totalTokens}, new memory: ${tokenCount} tokens, limit: ${tokenLimit}`,
|
||||
);
|
||||
return `Cannot add memory: would exceed token limit. Current usage: ${totalTokens}/${tokenLimit} tokens. This memory requires ${tokenCount} tokens, but only ${remainingCapacity} tokens available.`;
|
||||
if (tokenLimit) {
|
||||
const newTotalTokens = totalTokens + tokenCount;
|
||||
const newRemainingTokens = tokenLimit - newTotalTokens;
|
||||
|
||||
if (newRemainingTokens < 0) {
|
||||
const errorArtifact: Record<Tools.memory, MemoryArtifact> = {
|
||||
[Tools.memory]: {
|
||||
key: 'system',
|
||||
type: 'error',
|
||||
value: JSON.stringify({
|
||||
errorType: 'would_exceed',
|
||||
tokenCount: Math.abs(newRemainingTokens),
|
||||
totalTokens: newTotalTokens,
|
||||
tokenLimit,
|
||||
}),
|
||||
tokenCount: totalTokens,
|
||||
},
|
||||
};
|
||||
return [`Memory storage would exceed limit. Cannot save this memory.`, errorArtifact];
|
||||
}
|
||||
}
|
||||
|
||||
const artifact: Record<Tools.memory, MemoryArtifact> = {
|
||||
|
|
@ -177,7 +204,7 @@ const createDeleteMemoryTool = ({
|
|||
', ',
|
||||
)}`,
|
||||
);
|
||||
return `Invalid key "${key}". Must be one of: ${validKeys.join(', ')}`;
|
||||
return [`Invalid key "${key}". Must be one of: ${validKeys.join(', ')}`, undefined];
|
||||
}
|
||||
|
||||
const artifact: Record<Tools.memory, MemoryArtifact> = {
|
||||
|
|
@ -269,7 +296,13 @@ export async function processMemory({
|
|||
llmConfig?: Partial<LLMConfig>;
|
||||
}): Promise<(TAttachment | null)[] | undefined> {
|
||||
try {
|
||||
const memoryTool = createMemoryTool({ userId, tokenLimit, setMemory, validKeys, totalTokens });
|
||||
const memoryTool = createMemoryTool({
|
||||
userId,
|
||||
tokenLimit,
|
||||
setMemory,
|
||||
validKeys,
|
||||
totalTokens,
|
||||
});
|
||||
const deleteMemoryTool = createDeleteMemoryTool({
|
||||
userId,
|
||||
validKeys,
|
||||
|
|
|
|||
|
|
@ -534,7 +534,7 @@ export type MemoryArtifact = {
|
|||
key: string;
|
||||
value?: string;
|
||||
tokenCount?: number;
|
||||
type: 'update' | 'delete';
|
||||
type: 'update' | 'delete' | 'error';
|
||||
};
|
||||
|
||||
export type TAttachmentMetadata = {
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue