mirror of
https://github.com/danny-avila/LibreChat.git
synced 2026-08-03 22:32:42 +00:00
🧷 fix: Bind Agent File Context to Current Turn (#13506)
* fix: Bind agent file context to current turn * fix: Avoid duplicating agent file context * fix: Export agent file context prepender * test: Use exported file context prepender * fix: Keep file context transient for memory and counts
This commit is contained in:
Danny Avila
GitHub
parent
a1bfa3b298
commit
dc42748813
4 changed files with 249 additions and 45 deletions
|
|
@ -30,6 +30,7 @@ const {
|
|||
createMultiAgentMapper,
|
||||
filterMalformedContentParts,
|
||||
countFormattedMessageTokens,
|
||||
prependFileContext,
|
||||
hydrateMissingIndexTokenCounts,
|
||||
injectSkillPrimes,
|
||||
isSkillPrimeMessage,
|
||||
|
|
@ -135,6 +136,8 @@ class AgentClient extends BaseClient {
|
|||
this.usage;
|
||||
/** @type {Record<string, number>} */
|
||||
this.indexTokenCountMap = {};
|
||||
/** @type {Array<Record<string, unknown>> | null} */
|
||||
this.memoryPayload = null;
|
||||
/** @type {(messages: BaseMessage[]) => Promise<void>} */
|
||||
this.processMemory;
|
||||
}
|
||||
|
|
@ -315,39 +318,51 @@ class AgentClient extends BaseClient {
|
|||
}
|
||||
|
||||
/** @type {Record<number, number>} */
|
||||
const canonicalTokenCountMap = {};
|
||||
const indexTokenCountMap = {};
|
||||
/** @type {Record<string, number>} */
|
||||
const tokenCountMap = {};
|
||||
const memoryPayload = [];
|
||||
let hasFileContext = false;
|
||||
let promptTokenTotal = 0;
|
||||
const encoding = this.getEncoding();
|
||||
const formattedMessages = orderedMessages.map((message, i) => {
|
||||
const formattedMessage = formatMessage({
|
||||
message,
|
||||
userName: this.options?.name,
|
||||
assistantName: this.options?.modelLabel,
|
||||
});
|
||||
const memoryFormattedMessage = formatMessage({
|
||||
message,
|
||||
userName: this.options?.name,
|
||||
assistantName: this.options?.modelLabel,
|
||||
});
|
||||
|
||||
/** For non-latest messages, prepend file context directly to message content */
|
||||
if (message.fileContext && i !== orderedMessages.length - 1) {
|
||||
if (typeof formattedMessage.content === 'string') {
|
||||
formattedMessage.content = message.fileContext + '\n' + formattedMessage.content;
|
||||
} else {
|
||||
const textPart = formattedMessage.content.find((part) => part.type === 'text');
|
||||
textPart
|
||||
? (textPart.text = message.fileContext + '\n' + textPart.text)
|
||||
: formattedMessage.content.unshift({ type: 'text', text: message.fileContext });
|
||||
}
|
||||
/**
|
||||
* Bind file context to the message it belongs to. Historical attachments
|
||||
* are resent inline, so the current turn's text attachment must be inline
|
||||
* too instead of living only in the dynamic system tail.
|
||||
*/
|
||||
if (message.fileContext) {
|
||||
hasFileContext = true;
|
||||
prependFileContext(formattedMessage, message.fileContext);
|
||||
}
|
||||
|
||||
const dbTokenCount = orderedMessages[i].tokenCount;
|
||||
const needsTokenCount = !dbTokenCount || message.fileContext;
|
||||
memoryPayload.push(memoryFormattedMessage);
|
||||
|
||||
if (needsTokenCount || (this.isVisionModel && (message.image_urls || message.files))) {
|
||||
orderedMessages[i].tokenCount = countFormattedMessageTokens(
|
||||
formattedMessage,
|
||||
this.getEncoding(),
|
||||
);
|
||||
const dbTokenCount = Number(orderedMessages[i].tokenCount);
|
||||
const hasDbTokenCount = Number.isFinite(dbTokenCount) && dbTokenCount > 0;
|
||||
const needsCanonicalTokenCount =
|
||||
!hasDbTokenCount || (this.isVisionModel && (message.image_urls || message.files));
|
||||
|
||||
let canonicalTokenCount = hasDbTokenCount ? dbTokenCount : 0;
|
||||
if (needsCanonicalTokenCount) {
|
||||
canonicalTokenCount = countFormattedMessageTokens(memoryFormattedMessage, encoding);
|
||||
}
|
||||
|
||||
const promptMessageTokenCount = message.fileContext
|
||||
? countFormattedMessageTokens(formattedMessage, encoding)
|
||||
: canonicalTokenCount;
|
||||
|
||||
/* If message has files, calculate image token cost */
|
||||
if (this.message_file_map && this.message_file_map[message.messageId]) {
|
||||
const attachments = this.message_file_map[message.messageId];
|
||||
|
|
@ -362,13 +377,19 @@ class AgentClient extends BaseClient {
|
|||
}
|
||||
}
|
||||
|
||||
const tokenCount = Number(orderedMessages[i].tokenCount);
|
||||
const normalizedTokenCount = Number.isFinite(tokenCount) && tokenCount > 0 ? tokenCount : 0;
|
||||
canonicalTokenCountMap[i] = normalizedTokenCount;
|
||||
promptTokenTotal += normalizedTokenCount;
|
||||
const normalizedCanonicalTokenCount =
|
||||
Number.isFinite(canonicalTokenCount) && canonicalTokenCount > 0 ? canonicalTokenCount : 0;
|
||||
const normalizedPromptTokenCount =
|
||||
Number.isFinite(promptMessageTokenCount) && promptMessageTokenCount > 0
|
||||
? promptMessageTokenCount
|
||||
: 0;
|
||||
|
||||
orderedMessages[i].tokenCount = normalizedCanonicalTokenCount;
|
||||
indexTokenCountMap[i] = normalizedPromptTokenCount;
|
||||
promptTokenTotal += normalizedPromptTokenCount;
|
||||
|
||||
if (message.messageId) {
|
||||
tokenCountMap[message.messageId] = normalizedTokenCount;
|
||||
tokenCountMap[message.messageId] = normalizedCanonicalTokenCount;
|
||||
}
|
||||
|
||||
if (isEnabled(process.env.AGENT_DEBUG_LOGGING)) {
|
||||
|
|
@ -377,9 +398,10 @@ class AgentClient extends BaseClient {
|
|||
Array.isArray(message.content) && message.content.some((p) => p && p.type === 'summary');
|
||||
const suffix = hasSummary ? '[S]' : '';
|
||||
const id = (message.messageId ?? message.id ?? '').slice(-8);
|
||||
const recalced = needsTokenCount ? orderedMessages[i].tokenCount : null;
|
||||
const recalced = needsCanonicalTokenCount ? normalizedCanonicalTokenCount : null;
|
||||
const promptRecalced = message.fileContext ? normalizedPromptTokenCount : null;
|
||||
logger.debug(
|
||||
`[AgentClient] msg[${i}] ${role}${suffix} id=…${id} db=${dbTokenCount} needsRecount=${needsTokenCount} recalced=${recalced} tokens=${normalizedTokenCount}`,
|
||||
`[AgentClient] msg[${i}] ${role}${suffix} id=…${id} db=${dbTokenCount} needsRecount=${needsCanonicalTokenCount} recalced=${recalced} promptRecalced=${promptRecalced} tokens=${normalizedPromptTokenCount}`,
|
||||
);
|
||||
}
|
||||
|
||||
|
|
@ -387,22 +409,18 @@ class AgentClient extends BaseClient {
|
|||
});
|
||||
|
||||
payload = formattedMessages;
|
||||
this.memoryPayload = hasFileContext ? memoryPayload : null;
|
||||
messages = orderedMessages;
|
||||
promptTokens = promptTokenTotal;
|
||||
|
||||
/**
|
||||
* Build shared run context - applies to ALL agents in the run.
|
||||
* This includes file context from the latest message and augmented prompt (RAG).
|
||||
* Request attachment file context is already bound inline to the latest
|
||||
* user message above; only side-channel context belongs here.
|
||||
* Memory context is handled separately and applied per-agent based on config.
|
||||
*/
|
||||
const sharedRunContextParts = [];
|
||||
|
||||
/** File context from the latest message (attachments) */
|
||||
const latestMessage = orderedMessages[orderedMessages.length - 1];
|
||||
if (latestMessage?.fileContext) {
|
||||
sharedRunContextParts.push(latestMessage.fileContext);
|
||||
}
|
||||
|
||||
/** Augmented prompt from RAG/context handlers */
|
||||
if (this.contextHandlers) {
|
||||
this.augmentedPrompt = await this.contextHandlers.createContext();
|
||||
|
|
@ -428,8 +446,8 @@ class AgentClient extends BaseClient {
|
|||
tokenCountFn: (text) => countTokens(text),
|
||||
});
|
||||
|
||||
/** Preserve canonical pre-format token counts for all history entering graph formatting */
|
||||
this.indexTokenCountMap = canonicalTokenCountMap;
|
||||
/** Preserve prompt token counts for graph formatting and pruning. */
|
||||
this.indexTokenCountMap = indexTokenCountMap;
|
||||
|
||||
/** Extract contextMeta from the parent response (second-to-last in ordered chain;
|
||||
* last is the current user message). Seeds the pruner's calibration EMA for this run. */
|
||||
|
|
@ -970,6 +988,17 @@ class AgentClient extends BaseClient {
|
|||
tokenCounter,
|
||||
});
|
||||
|
||||
const memoryMessages =
|
||||
this.processMemory && this.memoryPayload
|
||||
? formatAgentMessages(
|
||||
this.memoryPayload,
|
||||
undefined,
|
||||
toolSet,
|
||||
skillPrimeResult?.skills,
|
||||
formatOptions,
|
||||
).messages
|
||||
: initialMessages;
|
||||
|
||||
/**
|
||||
* @param {BaseMessage[]} messages
|
||||
*/
|
||||
|
|
@ -1010,7 +1039,7 @@ class AgentClient extends BaseClient {
|
|||
// }
|
||||
|
||||
if (this.processMemory) {
|
||||
memoryPromise = this.runMemory(messages);
|
||||
memoryPromise = this.runMemory(memoryMessages);
|
||||
}
|
||||
|
||||
/** Seed calibration state from previous run if encoding matches */
|
||||
|
|
|
|||
|
|
@ -1,5 +1,5 @@
|
|||
const { Providers } = require('@librechat/agents');
|
||||
const { Constants, EModelEndpoint } = require('librechat-data-provider');
|
||||
const { Constants, ContentTypes, EModelEndpoint } = require('librechat-data-provider');
|
||||
const AgentClient = require('./client');
|
||||
|
||||
jest.mock('@librechat/agents', () => ({
|
||||
|
|
@ -1500,6 +1500,7 @@ describe('AgentClient - titleConvo', () => {
|
|||
beforeEach(() => {
|
||||
jest.clearAllMocks();
|
||||
mockFormatInstructions.mockResolvedValue('');
|
||||
require('@librechat/api').countFormattedMessageTokens.mockImplementation(() => 42);
|
||||
|
||||
mockAgent = {
|
||||
id: 'primary-agent',
|
||||
|
|
@ -1544,7 +1545,7 @@ describe('AgentClient - titleConvo', () => {
|
|||
client.useMemory = jest.fn().mockResolvedValue(undefined);
|
||||
});
|
||||
|
||||
it("applies shared request context plus each agent's own context docs only", async () => {
|
||||
it('places request context inline and applies each agent context doc only once', async () => {
|
||||
const requestFile = makeTextFile('request-file', 'request.txt', 'Shared request context');
|
||||
const primaryContext = makeTextFile(
|
||||
'primary-context',
|
||||
|
|
@ -1574,7 +1575,7 @@ describe('AgentClient - titleConvo', () => {
|
|||
]);
|
||||
client.agentConfigs = new Map([['handoff-agent', handoffAgent]]);
|
||||
|
||||
await client.buildMessages(
|
||||
const result = await client.buildMessages(
|
||||
[
|
||||
{
|
||||
messageId: 'msg-1',
|
||||
|
|
@ -1588,15 +1589,99 @@ describe('AgentClient - titleConvo', () => {
|
|||
{},
|
||||
);
|
||||
|
||||
expect(mockAgent.additional_instructions).toContain('Shared request context');
|
||||
expect(result.prompt[0].content).toContain('Shared request context');
|
||||
|
||||
expect(mockAgent.additional_instructions).toContain('Primary private context');
|
||||
expect(mockAgent.additional_instructions).not.toContain('Shared request context');
|
||||
expect(mockAgent.additional_instructions).not.toContain('Handoff private context');
|
||||
|
||||
expect(handoffAgent.additional_instructions).toContain('Shared request context');
|
||||
expect(handoffAgent.additional_instructions).toContain('Handoff private context');
|
||||
expect(handoffAgent.additional_instructions).not.toContain('Shared request context');
|
||||
expect(handoffAgent.additional_instructions).not.toContain('Primary private context');
|
||||
});
|
||||
|
||||
it('places current request file context on the latest user message', async () => {
|
||||
const currentFile = makeTextFile('current-file', 'current.txt', 'Current turn file body');
|
||||
const previousFileContext =
|
||||
'Attached document(s):\n```md\n# "previous.txt"\nPrevious turn file body\n```';
|
||||
|
||||
client.options.attachments = [currentFile];
|
||||
|
||||
const result = await client.buildMessages(
|
||||
[
|
||||
{
|
||||
messageId: 'msg-1',
|
||||
parentMessageId: null,
|
||||
sender: 'User',
|
||||
text: 'What is written here?',
|
||||
isCreatedByUser: true,
|
||||
fileContext: previousFileContext,
|
||||
},
|
||||
{
|
||||
messageId: 'msg-2',
|
||||
parentMessageId: 'msg-1',
|
||||
sender: 'Assistant',
|
||||
text: 'It describes the previous file.',
|
||||
isCreatedByUser: false,
|
||||
},
|
||||
{
|
||||
messageId: 'msg-3',
|
||||
parentMessageId: 'msg-2',
|
||||
sender: 'User',
|
||||
text: 'What is written here?',
|
||||
isCreatedByUser: true,
|
||||
},
|
||||
],
|
||||
'msg-3',
|
||||
{},
|
||||
);
|
||||
|
||||
expect(result.prompt[0].content).toContain('Previous turn file body');
|
||||
expect(result.prompt[2].content).toContain('Current turn file body');
|
||||
expect(result.prompt[2].content).toContain('What is written here?');
|
||||
expect(result.prompt[2].content).not.toContain('Previous turn file body');
|
||||
expect(client.memoryPayload[2].content).toContain('What is written here?');
|
||||
expect(client.memoryPayload[2].content).not.toContain('Current turn file body');
|
||||
expect(mockAgent.additional_instructions ?? '').not.toContain('Current turn file body');
|
||||
expect(result.prompt[2].content.indexOf('Current turn file body')).toBeLessThan(
|
||||
result.prompt[2].content.indexOf('What is written here?'),
|
||||
);
|
||||
});
|
||||
|
||||
it('persists canonical token counts while counting request file context for the prompt', async () => {
|
||||
const { countFormattedMessageTokens } = require('@librechat/api');
|
||||
const currentFile = makeTextFile('current-file', 'current.txt', 'Current turn file body');
|
||||
|
||||
countFormattedMessageTokens.mockImplementation(({ content }) => {
|
||||
const text = Array.isArray(content)
|
||||
? content.map((part) => part.text ?? part[ContentTypes.TEXT] ?? '').join('\n')
|
||||
: String(content ?? '');
|
||||
return text.includes('Current turn file body') ? 200 : 20;
|
||||
});
|
||||
|
||||
client.options.attachments = [currentFile];
|
||||
|
||||
const result = await client.buildMessages(
|
||||
[
|
||||
{
|
||||
messageId: 'msg-1',
|
||||
parentMessageId: null,
|
||||
sender: 'User',
|
||||
text: 'What is written here?',
|
||||
isCreatedByUser: true,
|
||||
},
|
||||
],
|
||||
'msg-1',
|
||||
{},
|
||||
);
|
||||
|
||||
expect(result.prompt[0].content).toContain('Current turn file body');
|
||||
expect(result.tokenCountMap['msg-1']).toBe(20);
|
||||
expect(result.promptTokens).toBe(200);
|
||||
expect(client.indexTokenCountMap[0]).toBe(200);
|
||||
expect(client.memoryPayload[0].content).toBe('What is written here?');
|
||||
});
|
||||
|
||||
it('does not duplicate a file that is both request context and scoped context', async () => {
|
||||
const sharedFile = makeTextFile('shared-file', 'shared.txt', 'Shared duplicate context');
|
||||
|
||||
|
|
@ -1604,7 +1689,7 @@ describe('AgentClient - titleConvo', () => {
|
|||
client.options.agentContextAttachmentsByAgentId = new Map([['primary-agent', [sharedFile]]]);
|
||||
client.agentConfigs = new Map();
|
||||
|
||||
await client.buildMessages(
|
||||
const result = await client.buildMessages(
|
||||
[
|
||||
{
|
||||
messageId: 'msg-1',
|
||||
|
|
@ -1618,10 +1703,10 @@ describe('AgentClient - titleConvo', () => {
|
|||
{},
|
||||
);
|
||||
|
||||
const occurrences = (
|
||||
mockAgent.additional_instructions.match(/Shared duplicate context/g) ?? []
|
||||
).length;
|
||||
expect(occurrences).toBe(1);
|
||||
const inlineOccurrences = (result.prompt[0].content.match(/Shared duplicate context/g) ?? [])
|
||||
.length;
|
||||
expect(inlineOccurrences).toBe(1);
|
||||
expect(mockAgent.additional_instructions ?? '').not.toContain('Shared duplicate context');
|
||||
});
|
||||
|
||||
it('keeps direct chats with context-doc agents working without request attachments', async () => {
|
||||
|
|
|
|||
54
packages/api/src/agents/client.spec.ts
Normal file
54
packages/api/src/agents/client.spec.ts
Normal file
|
|
@ -0,0 +1,54 @@
|
|||
import { ContentTypes } from 'librechat-data-provider';
|
||||
import { prependFileContext, type FormattedMessageWithContent } from './client';
|
||||
|
||||
describe('prependFileContext', () => {
|
||||
it('prepends file context to string content', () => {
|
||||
const message: FormattedMessageWithContent = { content: 'Answer this question.' };
|
||||
|
||||
prependFileContext(message, 'Attached file text');
|
||||
|
||||
expect(message.content).toBe('Attached file text\nAnswer this question.');
|
||||
});
|
||||
|
||||
it('prepends file context to the first text content part', () => {
|
||||
const message: FormattedMessageWithContent = {
|
||||
content: [
|
||||
{ type: ContentTypes.IMAGE_URL, image_url: { url: 'data:image/png;base64,abc' } },
|
||||
{ type: ContentTypes.TEXT, text: 'Answer this question.' },
|
||||
],
|
||||
};
|
||||
|
||||
prependFileContext(message, 'Attached file text');
|
||||
|
||||
expect(Array.isArray(message.content)).toBe(true);
|
||||
if (!Array.isArray(message.content)) {
|
||||
throw new Error('Expected array content');
|
||||
}
|
||||
expect(message.content[1].text).toBe('Attached file text\nAnswer this question.');
|
||||
expect(message.content[0]).toEqual({
|
||||
type: ContentTypes.IMAGE_URL,
|
||||
image_url: { url: 'data:image/png;base64,abc' },
|
||||
});
|
||||
});
|
||||
|
||||
it('adds a text content part when an array has no text part', () => {
|
||||
const message: FormattedMessageWithContent = {
|
||||
content: [{ type: ContentTypes.IMAGE_URL, image_url: { url: 'data:image/png;base64,abc' } }],
|
||||
};
|
||||
|
||||
prependFileContext(message, 'Attached file text');
|
||||
|
||||
expect(message.content).toEqual([
|
||||
{ type: ContentTypes.TEXT, text: 'Attached file text' },
|
||||
{ type: ContentTypes.IMAGE_URL, image_url: { url: 'data:image/png;base64,abc' } },
|
||||
]);
|
||||
});
|
||||
|
||||
it('leaves content unchanged when file context is empty', () => {
|
||||
const message: FormattedMessageWithContent = { content: 'Answer this question.' };
|
||||
|
||||
prependFileContext(message, '');
|
||||
|
||||
expect(message.content).toBe('Answer this question.');
|
||||
});
|
||||
});
|
||||
|
|
@ -56,6 +56,42 @@ type ContentBlock = {
|
|||
tool_call?: { name?: string; args?: string; output?: string };
|
||||
};
|
||||
|
||||
export type FormattedMessageContentPart = {
|
||||
type?: string;
|
||||
text?: string;
|
||||
[key: string]: unknown;
|
||||
};
|
||||
|
||||
export type FormattedMessageWithContent = {
|
||||
content?: string | FormattedMessageContentPart[];
|
||||
};
|
||||
|
||||
export function prependFileContext(
|
||||
formattedMessage: FormattedMessageWithContent,
|
||||
fileContext?: string | null,
|
||||
): void {
|
||||
if (!fileContext) {
|
||||
return;
|
||||
}
|
||||
|
||||
if (typeof formattedMessage.content === 'string') {
|
||||
formattedMessage.content = `${fileContext}\n${formattedMessage.content}`;
|
||||
return;
|
||||
}
|
||||
|
||||
if (!Array.isArray(formattedMessage.content)) {
|
||||
return;
|
||||
}
|
||||
|
||||
const textPart = formattedMessage.content.find((part) => part.type === ContentTypes.TEXT);
|
||||
if (textPart != null && typeof textPart.text === 'string') {
|
||||
textPart.text = `${fileContext}\n${textPart.text}`;
|
||||
return;
|
||||
}
|
||||
|
||||
formattedMessage.content.unshift({ type: ContentTypes.TEXT, text: fileContext });
|
||||
}
|
||||
|
||||
function estimateImageDataTokens(data: string, isClaude: boolean): number {
|
||||
const dims = extractImageDimensions(data);
|
||||
if (dims == null) {
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue