Skip to content

fix: retain bounded tool images across later tool batches #609

New issue

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

Sign up for GitHub

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

Already on GitHub? Sign in to your account

Draft
ldavis9000aws wants to merge 1 commit into
base: main
Choose a base branch
from
Draft

fix: retain bounded tool images across later tool batches #609

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

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
52 changes: 46 additions & 6 deletions src/graphs/Graph.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3885,7 +3885,8 @@ export class StandardGraph extends Graph<t.BaseGraphState, t.GraphNode> {
candidate: BaseMessage[],
provider: t.ProviderName,
clientOptions: t.ClientOptions | undefined,
maxChars: number
maxChars: number,
replayImages = true
): BaseMessage[] => {
if (!(candidate[candidate.length - 1] instanceof ToolMessage)) {
return candidate;
Expand All @@ -3902,7 +3903,7 @@ export class StandardGraph extends Graph<t.BaseGraphState, t.GraphNode> {
) {
return trackProviderMessageOrigins(
candidate,
projectArtifactPayload(candidate, maxChars)
projectArtifactPayload(candidate, maxChars, replayImages)
);
}
return candidate;
Expand Down Expand Up @@ -3952,15 +3953,25 @@ export class StandardGraph extends Graph<t.BaseGraphState, t.GraphNode> {

let artifactBaseMessages: BaseMessage[] | undefined;
if (lastMessageY instanceof ToolMessage) {
const artifactCandidate = projectServingArtifacts(
let artifactCandidate = projectServingArtifacts(
finalMessages,
agentContext.provider,
agentContext.clientOptions,
maxProviderToolResultChars
);

if (artifactCandidate !== finalMessages) {
const projection = measureProviderPayload(artifactCandidate);
let projection = measureProviderPayload(artifactCandidate);
if (!projection.fits) {
artifactCandidate = projectServingArtifacts(
finalMessages,
agentContext.provider,
agentContext.clientOptions,
maxProviderToolResultChars,
false
);
projection = measureProviderPayload(artifactCandidate);
}
if (projection.fits) {
artifactBaseMessages = finalMessages;
finalMessages = artifactCandidate;
Expand All @@ -3987,14 +3998,28 @@ export class StandardGraph extends Graph<t.BaseGraphState, t.GraphNode> {
if (artifactBaseMessages != null) {
if (!finalProjection.fits) {
finalMessages = projectProviderReferences(
applyProviderMessageTransforms(artifactBaseMessages)
applyProviderMessageTransforms(
projectServingArtifacts(
artifactBaseMessages,
agentContext.provider,
agentContext.clientOptions,
maxProviderToolResultChars,
false
)
)
);
finalProjection = measureProviderPayload(finalMessages);
if (!finalProjection.fits) {
finalMessages = projectProviderReferences(
applyProviderMessageTransforms(artifactBaseMessages)
);
finalProjection = measureProviderPayload(finalMessages);
}
emitAgentLog(
config,
'warn',
'graph',
'Artifact payload omitted after final provider formatting exceeded the remaining context budget',
'Earlier tool images or artifact payload omitted after final provider formatting exceeded the remaining context budget',
{
projectedMessageTokens: finalProjection.projectedMessageTokens,
availableMessageTokens: finalProjection.availableMessageTokens,
Expand Down Expand Up @@ -4746,6 +4771,21 @@ export class StandardGraph extends Graph<t.BaseGraphState, t.GraphNode> {
servingFallbackMessages
);
let projection = preparedFallbackRequest.projection;
if (!projection.fits && artifactSource !== boundedSource) {
servingFallbackMessages = transformFallback(
projectServingArtifacts(
boundedSource,
fallbackProvider,
fallbackClientOptions,
fallbackToolResultChars,
false
)
);
preparedFallbackRequest = prepareFallback(
servingFallbackMessages
);
projection = preparedFallbackRequest.projection;
}
if (!projection.fits && artifactSource !== boundedSource) {
servingFallbackMessages = transformFallback(boundedSource);
preparedFallbackRequest = prepareFallback(
Expand Down
65 changes: 64 additions & 1 deletion src/graphs/__tests__/Graph.contextOverflow.test.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -12,8 +12,8 @@ import {
import type { BaseMessage } from '@langchain/core/messages';
import type * as t from '@/types';
import { OVERFLOW_SIGNATURES } from '@/utils/__tests__/fixtures/contextOverflowSignatures';
import { ContentTypes, GraphEvents, Providers } from '@/common';
import { getProviderSourceMessageIds } from '@/messages/provenance';
import { ContentTypes, GraphEvents, Providers } from '@/common';
import * as init from '@/llm/init';
import { Run } from '@/run';

Expand Down Expand Up @@ -194,6 +194,69 @@ async function captureFallback(
}

describe('context overflow recovery', () => {
it.each([false, true])(
'replays earlier tool images with context pressure=%s',
async (pressure) => {
const imageUrl = 'data:image/png;base64,QUFBQUFBQUFBQUFB';
const currentImageUrl = 'data:image/png;base64,QkJCQkJCQkJCQkJC';
const messages: BaseMessage[] = [
new HumanMessage('inspect the image and then read metadata'),
new AIMessage({
content: '',
tool_calls: [{ id: 'image', name: 'read', args: {} }],
}),
new ToolMessage({
content: 'image read',
tool_call_id: 'image',
artifact: {
content: [{ type: 'image_url', image_url: { url: imageUrl } }],
},
}),
new AIMessage({
content: '',
tool_calls: [{ id: 'metadata', name: 'read', args: {} }],
}),
new ToolMessage({
content: 'metadata read',
tool_call_id: 'metadata',
artifact: {
content: [
{ type: 'image_url', image_url: { url: currentImageUrl } },
],
},
}),
];
const imageTokens = pressure ? 20_000 : 100;
const imageCounter: t.TokenCounter = (message) =>
JSON.stringify(message.content).includes(imageUrl) ? imageTokens : 10;
const run = await createRun({
runId: `tool-image-replay-${pressure}`,
maxContextTokens: 10_000,
provider: Providers.OPENAI,
tokenCounter: imageCounter,
indexTokenCountMap: Object.fromEntries(
messages.map((message, index) => [index, imageCounter(message)])
),
});
if (!run.Graph) throw new Error('Expected graph');
const model = new OverflowThenSucceedModel({}, 0);
run.Graph.overrideModel = model;
await run.processStream({ messages }, streamConfig);
expect(model.calls).toHaveLength(1);
expect(
JSON.stringify(
model.calls[0].map((message) => message.content)
).includes(imageUrl)
).toBe(!pressure);
expect(
JSON.stringify(model.calls[0].map((message) => message.content))
).toContain(currentImageUrl);
expect(
(messages[2] as ToolMessage).artifact.content[0].image_url.url
).toBe(imageUrl);
}
);

it.each([Providers.ANTHROPIC, Providers.OPENAI, Providers.GOOGLE])(
'preserves fallback artifacts for %s',
async (provider) => {
Expand Down
122 changes: 110 additions & 12 deletions src/messages/core.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -9,13 +9,13 @@ import {
} from '@langchain/core/messages';
import type { ContentBlock as LangChainContentBlock } from '@langchain/core/messages';
import type { ToolCall, ToolCallChunk } from '@langchain/core/messages/tool';
import type { ProviderMessageProvenancePart } from './provenance';
import type * as t from '@/types';
import type {
ResponsesReplayPosition,
ToolHistoryPreparation,
OrderedToolHistoryProjection,
} from './toolHistoryProjection';
import type { ProviderMessageProvenancePart } from './provenance';
import type * as t from '@/types';
import {
cloneToolMessageWithContent,
compactToolContent,
Expand All @@ -27,6 +27,13 @@ import {
isComputerCallOutputMessage,
serializeToolContentBounded,
} from '@/utils/toolContent';
import {
OPENAI_RESPONSES_REPLAY_POSITIONS_KEY,
createToolHistoryPreparation,
getGeneratedImageMimeType,
getResponsesHistorySource,
isResponsesReplayPosition,
} from './toolHistoryProjection';
import {
getProviderMessageProvenance,
getProviderSourceMessageIds,
Expand All @@ -35,17 +42,10 @@ import {
stampSyntheticProviderMessage,
} from './provenance';
import { HARD_MAX_TOOL_RESULT_CHARS } from '@/utils/truncation';
import { stripAnthropicCacheControl } from './cache';
import { isReasoningContentBlock } from './reasoningTypes';
import { stripAnthropicCacheControl } from './cache';
import { ContentTypes, Providers } from '@/common';
import { toLangChainContent } from './langchain';
import {
OPENAI_RESPONSES_REPLAY_POSITIONS_KEY,
createToolHistoryPreparation,
getGeneratedImageMimeType,
getResponsesHistorySource,
isResponsesReplayPosition,
} from './toolHistoryProjection';

type ReasoningSummary = { summary?: Array<{ text?: string }> };
type ReasoningDetail = { type?: string; text?: string };
Expand Down Expand Up @@ -2619,7 +2619,8 @@ export function formatAnthropicArtifactContent(messages: BaseMessage[]): void {

export function projectArtifactPayload(
messages: BaseMessage[],
maxChars = HARD_MAX_TOOL_RESULT_CHARS
maxChars = HARD_MAX_TOOL_RESULT_CHARS,
replayImages = true
): BaseMessage[] {
const lastMessageY = messages[messages.length - 1];
if (!(lastMessageY instanceof ToolMessage)) return messages;
Expand Down Expand Up @@ -2717,7 +2718,104 @@ export function projectArtifactPayload(
setProviderMessageProvenance(artifactPayload, artifactProvenanceParts);
formattedMessages?.push(artifactPayload);
}
return formattedMessages ?? messages;
const projected = formattedMessages ?? messages;
return replayImages
? projectEarlierToolImages(
messages,
projected,
latestAIParentIndex,
aggregate.remainingChars
)
: projected;
}

/** Replay only images from completed tool calls in the current user turn.
* Newest images consume the remaining artifact budget first. The projection
* never writes image bytes back into graph state or persisted messages. */
function projectEarlierToolImages(
source: BaseMessage[],
projected: BaseMessage[],
end: number,
maxChars: number
): BaseMessage[] {
if (maxChars <= 0) {
return projected;
}
const eligible = new Set<number>();
let pending = new Set<string>();
for (let i = 0; i < end; i++) {
const message = source[i];
if (message instanceof HumanMessage) {
eligible.clear();
pending.clear();
} else if (
message instanceof AIMessage ||
message instanceof AIMessageChunk
) {
pending = new Set(
message.tool_calls?.flatMap((call) =>
call.id == null ? [] : [call.id]
)
);
} else if (
message instanceof ToolMessage &&
pending.delete(message.tool_call_id)
) {
eligible.add(i);
}
}
if (eligible.size === 0) {
return projected;
}
const aggregate = createBoundedContentAccumulator(maxChars);
const payloads: HumanMessage[] = [];
for (let i = end - 1; i >= 0 && aggregate.remainingChars > 0; i--) {
const message = source[i];
if (
!eligible.has(i) ||
!(message instanceof ToolMessage) ||
isComputerCallOutputMessage(message) ||
!Array.isArray(message.artifact?.content)
) {
continue;
}
const images: t.MessageContentComplex[] = [];
for (const block of message.artifact.content) {
const image = getComputerCallOutputScreenshot([block]);
if (image == null || !('image_url' in image)) {
continue;
}
const normalized = {
type: 'image_url',
image_url: {
url: image.image_url,
...('detail' in image &&
image.detail != null && { detail: image.detail }),
},
};
const chars = getToolContentCharLength([normalized]) + 1;
if (chars > aggregate.remainingChars) {
continue;
}
images.push(normalized);
aggregate.remainingChars -= chars;
}
if (images.length === 0) {
continue;
}
const payload = new HumanMessage({ content: toLangChainContent(images) });
setProviderMessageProvenance(payload, [
projectUnindexedProviderMessageAttribution(message, 'tool'),
]);
payloads.push(payload);
}
return payloads.length === 0
? projected
: [
...projected.slice(0, source.length),
...payloads.reverse(),
...projected.slice(source.length),
];
}

/**
Expand Down
Loading