mirror of
https://github.com/danny-avila/LibreChat.git
synced 2026-08-29 05:20:49 +00:00
fix(api): preserve bounded phase partitions
This commit is contained in:
Danny Avila
parent
c090648916
commit
909c804c53
2 changed files with 144 additions and 40 deletions
|
|
@ -1307,7 +1307,7 @@ describe('createActivityPhaseWiring', () => {
|
|||
expect(parts[3]).toBeUndefined();
|
||||
});
|
||||
|
||||
it('bounds persisted evidence while preserving the full activity count', async () => {
|
||||
it('bounds persisted evidence and the tracked activity window', async () => {
|
||||
const parts: LooseContentPart[] = [];
|
||||
const wiring = createActivityPhaseWiring({
|
||||
getContentParts: () => parts,
|
||||
|
|
@ -1324,7 +1324,7 @@ describe('createActivityPhaseWiring', () => {
|
|||
|
||||
const snapshot = wiring.snapshot();
|
||||
expect(snapshot.version).toBe(2);
|
||||
expect(snapshot.activityCount).toBe(100);
|
||||
expect(snapshot.activityCount).toBe(77);
|
||||
expect(snapshot.activities).toHaveLength(13);
|
||||
expect(snapshot.overflowActivityStartIndex).toBe(99);
|
||||
expect(snapshot.overflowToolCallIds).toHaveLength(64);
|
||||
|
|
@ -2758,6 +2758,76 @@ describe('createActivityPhaseWiring', () => {
|
|||
expect(payloads[0].assistantContext).toBeUndefined();
|
||||
expect(payloads[1].assistantContext).toEqual(['Context for the later work.']);
|
||||
});
|
||||
|
||||
it('retains a completed tool batch when any call crosses a substantial boundary', async () => {
|
||||
const parts: LooseContentPart[] = [
|
||||
{ type: ContentTypes.TOOL_CALL, tool_call: { id: 'tool-1' } },
|
||||
{ type: ContentTypes.TOOL_CALL, tool_call: { id: 'tool-2' } },
|
||||
{ type: ContentTypes.TOOL_CALL, tool_call: { id: 'parallel-1' } },
|
||||
{ type: ContentTypes.TEXT, text: '' },
|
||||
{ type: ContentTypes.TOOL_CALL, tool_call: { id: 'parallel-2' } },
|
||||
];
|
||||
const payloads: GenerateActivityPhasePayload[] = [];
|
||||
const wiring = createActivityPhaseWiring({
|
||||
getContentParts: () => parts,
|
||||
getStepIndex: (stepId) => (stepId === 'boundary' ? 3 : undefined),
|
||||
bumpIndexOffset: jest.fn(),
|
||||
emitLabelEvent: jest.fn(async () => undefined),
|
||||
trackPendingFill: jest.fn(),
|
||||
generatePhase: jest.fn(async (payload: GenerateActivityPhasePayload) => {
|
||||
payloads.push(payload);
|
||||
return { label: 'Completed one phase' };
|
||||
}),
|
||||
});
|
||||
await wiring.hook(batch('tool-1'), new AbortController().signal);
|
||||
await wiring.hook(batch('tool-2'), new AbortController().signal);
|
||||
const parallelBatch = batch('parallel-1');
|
||||
parallelBatch.entries.push({
|
||||
...parallelBatch.entries[0],
|
||||
toolUseId: 'parallel-2',
|
||||
toolInput: { query: 'parallel-2' },
|
||||
});
|
||||
await wiring.hook(parallelBatch, new AbortController().signal);
|
||||
const handlers = wiring.handlers({
|
||||
[GraphEvents.ON_RUN_STEP]: { handle: jest.fn() },
|
||||
[GraphEvents.ON_MESSAGE_DELTA]: {
|
||||
handle: (_event, data) => {
|
||||
const delta = data as { delta?: { content?: { text?: string } } };
|
||||
parts[3] = { type: ContentTypes.TEXT, text: delta.delta?.content?.text ?? '' };
|
||||
},
|
||||
},
|
||||
});
|
||||
handlers?.[GraphEvents.ON_RUN_STEP]?.handle(
|
||||
GraphEvents.ON_RUN_STEP,
|
||||
{
|
||||
id: 'boundary',
|
||||
stepDetails: {
|
||||
type: StepTypes.MESSAGE_CREATION,
|
||||
message_creation: { message_id: 'm', content_type: 'text' },
|
||||
},
|
||||
},
|
||||
undefined,
|
||||
undefined,
|
||||
);
|
||||
handlers?.[GraphEvents.ON_MESSAGE_DELTA]?.handle(
|
||||
GraphEvents.ON_MESSAGE_DELTA,
|
||||
{
|
||||
id: 'boundary',
|
||||
delta: { content: { type: ContentTypes.TEXT, text: substantialText('Boundary result.') } },
|
||||
},
|
||||
undefined,
|
||||
undefined,
|
||||
);
|
||||
parts[6] = { type: ContentTypes.TOOL_CALL, tool_call: { id: 'tool-3' } };
|
||||
await wiring.hook(batch('tool-3'), new AbortController().signal);
|
||||
|
||||
wiring.complete();
|
||||
await flushDetached();
|
||||
|
||||
expect(payloads.map(({ totalActivityCount }) => totalActivityCount)).toEqual([2, 2]);
|
||||
expect(parts[5]).toMatchObject({ activity_end_index: 3, activity_count: 2 });
|
||||
expect(parts[7]).toMatchObject({ activity_start_index: 4, activity_count: 2 });
|
||||
});
|
||||
});
|
||||
|
||||
describe('createAssistantPhaseStampingHandlers', () => {
|
||||
|
|
|
|||
|
|
@ -26,6 +26,8 @@ export interface ActivityPhaseEntry {
|
|||
|
||||
export type TrackedActivity = ActivityPhaseEntry & {
|
||||
startIndex: number;
|
||||
/** A prior boundary can retain only the materialized tail of a straddling batch. */
|
||||
partitionStartIndex?: number;
|
||||
childLabelIndex?: number;
|
||||
/** Stable anchors survive content filtering and prepends across HITL resume. */
|
||||
toolCallIds?: string[];
|
||||
|
|
@ -216,30 +218,14 @@ function findLastPartIndex(
|
|||
return Math.max(0, parts.length - 1);
|
||||
}
|
||||
|
||||
function findBatchStart(
|
||||
parts: ReadonlyArray<LooseContentPart | null | undefined>,
|
||||
toolCallIds: Set<string>,
|
||||
): number {
|
||||
let first = -1;
|
||||
for (const index of definedPartIndices(parts)) {
|
||||
const part = parts[index];
|
||||
if (
|
||||
part?.type === ContentTypes.TOOL_CALL &&
|
||||
typeof part.tool_call?.id === 'string' &&
|
||||
toolCallIds.has(part.tool_call.id)
|
||||
) {
|
||||
first = first < 0 ? index : Math.min(first, index);
|
||||
}
|
||||
}
|
||||
return first >= 0 ? first : Math.max(0, parts.length - 1);
|
||||
}
|
||||
|
||||
function findTrackedStart(
|
||||
parts: ReadonlyArray<LooseContentPart | null | undefined>,
|
||||
activity: TrackedActivity,
|
||||
): number {
|
||||
const materializedStart = findMaterializedActivityStart(parts, activity);
|
||||
return materializedStart ?? Math.min(activity.startIndex, Math.max(0, parts.length - 1));
|
||||
const startIndex =
|
||||
materializedStart ?? Math.min(activity.startIndex, Math.max(0, parts.length - 1));
|
||||
return Math.max(startIndex, activity.partitionStartIndex ?? 0);
|
||||
}
|
||||
|
||||
function findMaterializedActivityStart(
|
||||
|
|
@ -312,16 +298,25 @@ function findTrackedToolStart(
|
|||
parts: ReadonlyArray<LooseContentPart | null | undefined>,
|
||||
activity: TrackedActivity,
|
||||
): number | undefined {
|
||||
if (activity.toolCallIds != null && activity.toolCallIds.length > 0) {
|
||||
const toolStart = findBatchStart(parts, new Set(activity.toolCallIds));
|
||||
if (
|
||||
parts[toolStart]?.type === ContentTypes.TOOL_CALL &&
|
||||
activity.toolCallIds.includes(String(parts[toolStart]?.tool_call?.id ?? ''))
|
||||
) {
|
||||
return toolStart;
|
||||
}
|
||||
return findTrackedToolIndices(parts, activity)[0];
|
||||
}
|
||||
|
||||
function findTrackedToolIndices(
|
||||
parts: ReadonlyArray<LooseContentPart | null | undefined>,
|
||||
activity: TrackedActivity,
|
||||
): number[] {
|
||||
if (activity.toolCallIds == null || activity.toolCallIds.length === 0) {
|
||||
return [];
|
||||
}
|
||||
return undefined;
|
||||
const ids = new Set(activity.toolCallIds);
|
||||
return definedPartIndices(parts).filter((index) => {
|
||||
const part = parts[index];
|
||||
return (
|
||||
part?.type === ContentTypes.TOOL_CALL &&
|
||||
typeof part.tool_call?.id === 'string' &&
|
||||
ids.has(part.tool_call.id)
|
||||
);
|
||||
});
|
||||
}
|
||||
|
||||
function findReasoningStart(
|
||||
|
|
@ -715,6 +710,9 @@ export function createActivityPhaseWiring(deps: ActivityPhaseHostDeps): Activity
|
|||
} else {
|
||||
overflowActivities.push({
|
||||
startIndex: activity.startIndex,
|
||||
...(activity.partitionStartIndex != null && {
|
||||
partitionStartIndex: activity.partitionStartIndex,
|
||||
}),
|
||||
...(activity.toolCallIds != null && {
|
||||
toolCallIds: activity.toolCallIds.slice(-MAX_RETAINED_TOOL_ENTRIES),
|
||||
}),
|
||||
|
|
@ -727,7 +725,13 @@ export function createActivityPhaseWiring(deps: ActivityPhaseHostDeps): Activity
|
|||
status: activity.status,
|
||||
});
|
||||
if (overflowActivities.length > MAX_OVERFLOW_ACTIVITY_ANCHORS) {
|
||||
overflowActivities.shift();
|
||||
const removed = overflowActivities.shift();
|
||||
activityCount -= 1;
|
||||
if (removed?.status === 'error') {
|
||||
failedActivityCount -= 1;
|
||||
} else if (removed?.status === 'partial') {
|
||||
partialActivityCount -= 1;
|
||||
}
|
||||
}
|
||||
if (overflowActivityStartIndex == null || activity.startIndex > overflowActivityStartIndex) {
|
||||
overflowActivityStartIndex = activity.startIndex;
|
||||
|
|
@ -828,6 +832,7 @@ export function createActivityPhaseWiring(deps: ActivityPhaseHostDeps): Activity
|
|||
childLabelIndex,
|
||||
toolCallIds,
|
||||
startIndex: _startIndex,
|
||||
partitionStartIndex: _partitionStartIndex,
|
||||
unresolvedToolStartIndex: _unresolvedToolStartIndex,
|
||||
...activity
|
||||
}) => {
|
||||
|
|
@ -882,7 +887,9 @@ export function createActivityPhaseWiring(deps: ActivityPhaseHostDeps): Activity
|
|||
* answer content. */
|
||||
const resolvedActivities = activities.map((activity) => {
|
||||
const toolStart = findTrackedToolStart(currentParts, activity);
|
||||
return toolStart != null ? { ...activity, startIndex: toolStart } : activity;
|
||||
return toolStart != null
|
||||
? { ...activity, startIndex: Math.max(toolStart, activity.partitionStartIndex ?? 0) }
|
||||
: activity;
|
||||
});
|
||||
const closesBeforeBoundary = (activity: TrackedActivity): boolean => {
|
||||
if (requestedEndIndex == null) {
|
||||
|
|
@ -895,12 +902,30 @@ export function createActivityPhaseWiring(deps: ActivityPhaseHostDeps): Activity
|
|||
return false;
|
||||
}
|
||||
const materializedStart = findMaterializedActivityStart(currentParts, activity);
|
||||
return materializedStart == null || materializedStart < requestedEndIndex;
|
||||
const materializedToolIndices = findTrackedToolIndices(currentParts, activity);
|
||||
return materializedToolIndices.length > 0
|
||||
? materializedToolIndices.every((index) => index < requestedEndIndex)
|
||||
: materializedStart == null || materializedStart < requestedEndIndex;
|
||||
};
|
||||
const snapshot = resolvedActivities.filter(closesBeforeBoundary);
|
||||
const retainedActivities = resolvedActivities.filter(
|
||||
(activity) => !closesBeforeBoundary(activity),
|
||||
);
|
||||
const reanchorRetainedActivity = (activity: TrackedActivity): TrackedActivity => {
|
||||
if (requestedEndIndex == null) {
|
||||
return activity;
|
||||
}
|
||||
const firstRetainedToolIndex = findTrackedToolIndices(currentParts, activity).find(
|
||||
(index) => index >= requestedEndIndex,
|
||||
);
|
||||
return firstRetainedToolIndex != null
|
||||
? {
|
||||
...activity,
|
||||
startIndex: firstRetainedToolIndex,
|
||||
partitionStartIndex: firstRetainedToolIndex,
|
||||
}
|
||||
: activity;
|
||||
};
|
||||
const retainedActivities = resolvedActivities
|
||||
.filter((activity) => !closesBeforeBoundary(activity))
|
||||
.map(reanchorRetainedActivity);
|
||||
const overflowCount = Math.max(0, activityCount - resolvedActivities.length);
|
||||
const overflowFailedCount = Math.max(
|
||||
0,
|
||||
|
|
@ -922,7 +947,9 @@ export function createActivityPhaseWiring(deps: ActivityPhaseHostDeps): Activity
|
|||
? resolvedOverflowActivities.filter(closesBeforeBoundary)
|
||||
: [];
|
||||
const retainedOverflowActivities = hasExactOverflowActivities
|
||||
? resolvedOverflowActivities.filter((activity) => !closesBeforeBoundary(activity))
|
||||
? resolvedOverflowActivities
|
||||
.filter((activity) => !closesBeforeBoundary(activity))
|
||||
.map(reanchorRetainedActivity)
|
||||
: [];
|
||||
const overflowCloses = hasExactOverflowActivities
|
||||
? closingOverflowActivities.length > 0
|
||||
|
|
@ -1176,10 +1203,17 @@ export function createActivityPhaseWiring(deps: ActivityPhaseHostDeps): Activity
|
|||
}
|
||||
}
|
||||
}
|
||||
const batchToolIndices = definedPartIndices(parts).filter((index) => {
|
||||
const part = parts[index];
|
||||
return (
|
||||
part?.type === ContentTypes.TOOL_CALL &&
|
||||
typeof part.tool_call?.id === 'string' &&
|
||||
ids.has(part.tool_call.id)
|
||||
);
|
||||
});
|
||||
if (
|
||||
batchStartIndex != null &&
|
||||
emittedContentRanges.some(
|
||||
({ start, end }) => batchStartIndex >= start && batchStartIndex < end,
|
||||
batchToolIndices.some((toolIndex) =>
|
||||
emittedContentRanges.some(({ start, end }) => toolIndex >= start && toolIndex < end),
|
||||
)
|
||||
) {
|
||||
pendingReasoning.delete(reasoningKey);
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue