diff --git a/api/app/clients/BaseClient.js b/api/app/clients/BaseClient.js index 13abe8443b..91ad11ab20 100644 --- a/api/app/clients/BaseClient.js +++ b/api/app/clients/BaseClient.js @@ -855,6 +855,10 @@ class BaseClient { } } + if (this.contextMeta) { + responseMessage.contextMeta = this.contextMeta; + } + responseMessage.databasePromise = this.saveMessageToDatabase( responseMessage, saveOptions, diff --git a/api/server/controllers/agents/client.js b/api/server/controllers/agents/client.js index 5bb2680056..1a97fb8e57 100644 --- a/api/server/controllers/agents/client.js +++ b/api/server/controllers/agents/client.js @@ -334,6 +334,14 @@ class AgentClient extends BaseClient { /** Preserve canonical pre-format token counts for all history entering graph formatting */ this.indexTokenCountMap = canonicalTokenCountMap; + /** 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. */ + const parentResponse = + orderedMessages.length >= 2 ? orderedMessages[orderedMessages.length - 2] : undefined; + if (parentResponse?.contextMeta && !parentResponse.isCreatedByUser) { + this.contextMeta = parentResponse.contextMeta; + } + const result = { tokenCountMap, prompt: payload, @@ -807,11 +815,26 @@ class AgentClient extends BaseClient { memoryPromise = this.runMemory(messages); + /** Seed calibration ratio from previous run if encoding matches */ + const currentEncoding = this.getEncoding(); + const prevMeta = this.contextMeta; + const calibrationRatio = + prevMeta?.encoding === currentEncoding && prevMeta?.calibrationRatio > 0 + ? prevMeta.calibrationRatio + : undefined; + + if (prevMeta) { + logger.debug( + `[AgentClient] contextMeta from parent: ratio=${prevMeta.calibrationRatio}, encoding=${prevMeta.encoding}, current=${currentEncoding}, seeded=${calibrationRatio ?? 'none'}`, + ); + } + run = await createRun({ agents, messages, indexTokenCountMap, initialSummary, + calibrationRatio, runId: this.responseMessageId, signal: abortController.signal, customHandlers: this.options.eventHandlers, @@ -849,6 +872,21 @@ class AgentClient extends BaseClient { const hideSequentialOutputs = config.configurable.hide_sequential_outputs; await runAgents(initialMessages); + + /** Capture calibration ratio from the completed run for persistence on the response message */ + if (this.run) { + const ratio = this.run.getCalibrationRatio(); + if (ratio > 0 && ratio !== 1) { + this.contextMeta = { + calibrationRatio: Math.round(ratio * 1000) / 1000, + encoding: this.getEncoding(), + }; + logger.debug( + `[AgentClient] contextMeta to persist: ratio=${this.contextMeta.calibrationRatio}, encoding=${this.contextMeta.encoding}`, + ); + } + } + /** @deprecated Agent Chain */ if (hideSequentialOutputs) { this.contentParts = this.contentParts.filter((part, index) => { diff --git a/packages/api/src/agents/run.ts b/packages/api/src/agents/run.ts index 20064374a2..4c8af970db 100644 --- a/packages/api/src/agents/run.ts +++ b/packages/api/src/agents/run.ts @@ -209,6 +209,7 @@ export async function createRun({ indexTokenCountMap, summarizationConfig, initialSummary, + calibrationRatio, streaming = true, streamUsage = true, }: { @@ -224,6 +225,8 @@ export async function createRun({ summarizationConfig?: SummarizationConfig; /** Cross-run summary from formatAgentMessages, forwarded to AgentContext */ initialSummary?: { text: string; tokenCount: number }; + /** Calibration ratio from previous run's contextMeta, seeds the pruner EMA */ + calibrationRatio?: number; } & Pick): Promise< Run > { @@ -421,5 +424,6 @@ export async function createRun({ tokenCounter, customHandlers, indexTokenCountMap, + calibrationRatio, }); } diff --git a/packages/data-provider/src/schemas.ts b/packages/data-provider/src/schemas.ts index 7eb0482e9f..c75a8dd300 100644 --- a/packages/data-provider/src/schemas.ts +++ b/packages/data-provider/src/schemas.ts @@ -630,6 +630,12 @@ export const tMessageSchema = z.object({ feedback: feedbackSchema.optional(), /** metadata */ metadata: z.record(z.unknown()).optional(), + contextMeta: z + .object({ + calibrationRatio: z.number().optional(), + encoding: z.string().optional(), + }) + .optional(), }); export type MemoryArtifact = { diff --git a/packages/data-schemas/src/schema/message.ts b/packages/data-schemas/src/schema/message.ts index 610251443d..ff3468918e 100644 --- a/packages/data-schemas/src/schema/message.ts +++ b/packages/data-schemas/src/schema/message.ts @@ -114,6 +114,14 @@ const messageSchema: Schema = new Schema( type: String, }, metadata: { type: mongoose.Schema.Types.Mixed }, + contextMeta: { + type: { + calibrationRatio: { type: Number }, + encoding: { type: String }, + }, + _id: false, + default: undefined, + }, attachments: { type: [{ type: mongoose.Schema.Types.Mixed }], default: undefined }, /* attachments: { diff --git a/packages/data-schemas/src/types/message.ts b/packages/data-schemas/src/types/message.ts index c3f465e711..201e5650ef 100644 --- a/packages/data-schemas/src/types/message.ts +++ b/packages/data-schemas/src/types/message.ts @@ -39,6 +39,10 @@ export interface IMessage extends Document { iconURL?: string; addedConvo?: boolean; metadata?: Record; + contextMeta?: { + calibrationRatio?: number; + encoding?: string; + }; attachments?: unknown[]; expiredAt?: Date | null; createdAt?: Date;