add usage info to branch summary, compaction and tool result entries (#6671)
* add usage info to branch summary entries * add usage to compaction entries * allow custom tools to report llm usage in tool results * allow observing and patching usage in tool_result hooks * agent-harness: save usage in entries for compaction, branch summaries and tool results
This commit is contained in:
@@ -737,6 +737,7 @@ async function finalizeExecutedToolCall(
|
||||
...result,
|
||||
content: afterResult.content ?? result.content,
|
||||
details: afterResult.details ?? result.details,
|
||||
usage: afterResult.usage ?? result.usage,
|
||||
terminate: afterResult.terminate ?? result.terminate,
|
||||
};
|
||||
isError = afterResult.isError ?? isError;
|
||||
@@ -780,6 +781,7 @@ function createToolResultMessage(finalized: FinalizedToolCallOutcome): ToolResul
|
||||
// so the null never enters session history or provider payloads.
|
||||
content: finalized.result.content ?? [],
|
||||
details: finalized.result.details,
|
||||
usage: finalized.result.usage,
|
||||
...(finalized.result.addedToolNames?.length ? { addedToolNames: finalized.result.addedToolNames } : {}),
|
||||
isError: finalized.isError,
|
||||
timestamp: Date.now(),
|
||||
|
||||
@@ -32,6 +32,7 @@ import type {
|
||||
AgentHarnessResources,
|
||||
AgentHarnessStreamOptions,
|
||||
AgentHarnessStreamOptionsPatch,
|
||||
CompactResult,
|
||||
ExecutionEnv,
|
||||
NavigateTreeResult,
|
||||
PendingSessionWrite,
|
||||
@@ -434,9 +435,16 @@ export class AgentHarness<
|
||||
content: result.content,
|
||||
details: result.details,
|
||||
isError,
|
||||
usage: result.usage,
|
||||
});
|
||||
return patch
|
||||
? { content: patch.content, details: patch.details, isError: patch.isError, terminate: patch.terminate }
|
||||
? {
|
||||
content: patch.content,
|
||||
details: patch.details,
|
||||
isError: patch.isError,
|
||||
usage: patch.usage,
|
||||
terminate: patch.terminate,
|
||||
}
|
||||
: undefined;
|
||||
},
|
||||
prepareNextTurn: async () => {
|
||||
@@ -690,9 +698,7 @@ export class AgentHarness<
|
||||
}
|
||||
}
|
||||
|
||||
async compact(
|
||||
customInstructions?: string,
|
||||
): Promise<{ summary: string; firstKeptEntryId: string; tokensBefore: number; details?: unknown }> {
|
||||
async compact(customInstructions?: string): Promise<CompactResult> {
|
||||
if (this.phase !== "idle") throw new AgentHarnessError("busy", "compact() requires idle harness");
|
||||
this.phase = "compaction";
|
||||
try {
|
||||
@@ -723,6 +729,7 @@ export class AgentHarness<
|
||||
result.tokensBefore,
|
||||
result.details,
|
||||
provided !== undefined,
|
||||
result.usage,
|
||||
);
|
||||
const entry = await this.session.getEntry(entryId);
|
||||
if (entry?.type === "compaction") {
|
||||
@@ -764,6 +771,7 @@ export class AgentHarness<
|
||||
let summaryEntry: NavigateTreeResult["summaryEntry"];
|
||||
let summaryText: string | undefined = hookResult?.summary?.summary;
|
||||
let summaryDetails: unknown = hookResult?.summary?.details;
|
||||
let summaryUsage = hookResult?.summary?.usage;
|
||||
if (!summaryText && options?.summarize && entries.length > 0) {
|
||||
const model = this.model;
|
||||
if (!model) throw new AgentHarnessError("invalid_state", "No model set for branch summary");
|
||||
@@ -779,6 +787,7 @@ export class AgentHarness<
|
||||
throw new AgentHarnessError("branch_summary", branchSummary.error.message, branchSummary.error);
|
||||
}
|
||||
summaryText = branchSummary.value.summary;
|
||||
summaryUsage = branchSummary.value.usage;
|
||||
summaryDetails = {
|
||||
readFiles: branchSummary.value.readFiles,
|
||||
modifiedFiles: branchSummary.value.modifiedFiles,
|
||||
@@ -798,7 +807,12 @@ export class AgentHarness<
|
||||
const summaryId = await this.session.moveTo(
|
||||
newLeafId,
|
||||
summaryText
|
||||
? { summary: summaryText, details: summaryDetails, fromHook: hookResult?.summary !== undefined }
|
||||
? {
|
||||
summary: summaryText,
|
||||
details: summaryDetails,
|
||||
usage: summaryUsage,
|
||||
fromHook: hookResult?.summary !== undefined,
|
||||
}
|
||||
: undefined,
|
||||
);
|
||||
if (summaryId) {
|
||||
|
||||
@@ -252,6 +252,7 @@ export async function generateBranchSummary(
|
||||
|
||||
return ok({
|
||||
summary: summary || "No summary generated",
|
||||
usage: response.usage,
|
||||
readFiles,
|
||||
modifiedFiles,
|
||||
});
|
||||
|
||||
@@ -101,10 +101,35 @@ export interface CompactionResult<T = unknown> {
|
||||
firstKeptEntryId: string;
|
||||
/** Estimated context tokens before compaction. */
|
||||
tokensBefore: number;
|
||||
/** Usage from the LLM call(s) that generated this summary, if available. */
|
||||
usage?: Usage;
|
||||
/** Optional implementation-specific details stored with the compaction entry. */
|
||||
details?: T;
|
||||
}
|
||||
|
||||
function combineUsage(first: Usage, second: Usage): Usage {
|
||||
return {
|
||||
input: first.input + second.input,
|
||||
output: first.output + second.output,
|
||||
cacheRead: first.cacheRead + second.cacheRead,
|
||||
cacheWrite: first.cacheWrite + second.cacheWrite,
|
||||
...(first.cacheWrite1h !== undefined || second.cacheWrite1h !== undefined
|
||||
? { cacheWrite1h: (first.cacheWrite1h ?? 0) + (second.cacheWrite1h ?? 0) }
|
||||
: {}),
|
||||
...(first.reasoning !== undefined || second.reasoning !== undefined
|
||||
? { reasoning: (first.reasoning ?? 0) + (second.reasoning ?? 0) }
|
||||
: {}),
|
||||
totalTokens: first.totalTokens + second.totalTokens,
|
||||
cost: {
|
||||
input: first.cost.input + second.cost.input,
|
||||
output: first.cost.output + second.cost.output,
|
||||
cacheRead: first.cost.cacheRead + second.cost.cacheRead,
|
||||
cacheWrite: first.cost.cacheWrite + second.cost.cacheWrite,
|
||||
total: first.cost.total + second.cost.total,
|
||||
},
|
||||
};
|
||||
}
|
||||
|
||||
/** Compaction thresholds and retention settings. */
|
||||
export interface CompactionSettings {
|
||||
/** Enable automatic compaction decisions. */
|
||||
@@ -474,7 +499,7 @@ export async function generateSummary(
|
||||
customInstructions?: string,
|
||||
previousSummary?: string,
|
||||
thinkingLevel?: ThinkingLevel,
|
||||
): Promise<Result<string, CompactionError>> {
|
||||
): Promise<Result<{ text: string; usage: Usage }, CompactionError>> {
|
||||
const maxTokens = Math.min(
|
||||
Math.floor(0.8 * reserveTokens),
|
||||
model.maxTokens > 0 ? model.maxTokens : Number.POSITIVE_INFINITY,
|
||||
@@ -523,7 +548,7 @@ export async function generateSummary(
|
||||
|
||||
const textContent = contentText(response.content);
|
||||
|
||||
return ok(textContent);
|
||||
return ok({ text: textContent, usage: response.usage });
|
||||
}
|
||||
|
||||
/** Prepared inputs for a compaction run. */
|
||||
@@ -656,22 +681,26 @@ export async function compact(
|
||||
}
|
||||
|
||||
let summary: string;
|
||||
let summaryUsage: Usage;
|
||||
|
||||
if (isSplitTurn && turnPrefixMessages.length > 0) {
|
||||
const historyResult =
|
||||
messagesToSummarize.length > 0
|
||||
? await generateSummary(
|
||||
messagesToSummarize,
|
||||
models,
|
||||
model,
|
||||
settings.reserveTokens,
|
||||
signal,
|
||||
customInstructions,
|
||||
previousSummary,
|
||||
thinkingLevel,
|
||||
)
|
||||
: ok<string, CompactionError>("No prior history.");
|
||||
if (!historyResult.ok) return err(historyResult.error);
|
||||
let historyText = "No prior history.";
|
||||
let historyUsage: Usage | undefined;
|
||||
if (messagesToSummarize.length > 0) {
|
||||
const historyResult = await generateSummary(
|
||||
messagesToSummarize,
|
||||
models,
|
||||
model,
|
||||
settings.reserveTokens,
|
||||
signal,
|
||||
customInstructions,
|
||||
previousSummary,
|
||||
thinkingLevel,
|
||||
);
|
||||
if (!historyResult.ok) return err(historyResult.error);
|
||||
historyText = historyResult.value.text;
|
||||
historyUsage = historyResult.value.usage;
|
||||
}
|
||||
const turnPrefixResult = await generateTurnPrefixSummary(
|
||||
turnPrefixMessages,
|
||||
models,
|
||||
@@ -681,7 +710,10 @@ export async function compact(
|
||||
thinkingLevel,
|
||||
);
|
||||
if (!turnPrefixResult.ok) return err(turnPrefixResult.error);
|
||||
summary = `${historyResult.value}\n\n---\n\n**Turn Context (split turn):**\n\n${turnPrefixResult.value}`;
|
||||
summary = `${historyText}\n\n---\n\n**Turn Context (split turn):**\n\n${turnPrefixResult.value.text}`;
|
||||
summaryUsage = historyUsage
|
||||
? combineUsage(historyUsage, turnPrefixResult.value.usage)
|
||||
: turnPrefixResult.value.usage;
|
||||
} else {
|
||||
const summaryResult = await generateSummary(
|
||||
messagesToSummarize,
|
||||
@@ -694,7 +726,8 @@ export async function compact(
|
||||
thinkingLevel,
|
||||
);
|
||||
if (!summaryResult.ok) return err(summaryResult.error);
|
||||
summary = summaryResult.value;
|
||||
summary = summaryResult.value.text;
|
||||
summaryUsage = summaryResult.value.usage;
|
||||
}
|
||||
|
||||
const { readFiles, modifiedFiles } = computeFileLists(fileOps);
|
||||
@@ -704,6 +737,7 @@ export async function compact(
|
||||
summary,
|
||||
firstKeptEntryId,
|
||||
tokensBefore,
|
||||
usage: summaryUsage,
|
||||
details: { readFiles, modifiedFiles } as CompactionDetails,
|
||||
});
|
||||
}
|
||||
@@ -714,7 +748,7 @@ async function generateTurnPrefixSummary(
|
||||
reserveTokens: number,
|
||||
signal?: AbortSignal,
|
||||
thinkingLevel?: ThinkingLevel,
|
||||
): Promise<Result<string, CompactionError>> {
|
||||
): Promise<Result<{ text: string; usage: Usage }, CompactionError>> {
|
||||
const maxTokens = Math.min(
|
||||
Math.floor(0.5 * reserveTokens),
|
||||
model.maxTokens > 0 ? model.maxTokens : Number.POSITIVE_INFINITY,
|
||||
@@ -749,5 +783,8 @@ async function generateTurnPrefixSummary(
|
||||
);
|
||||
}
|
||||
|
||||
return ok(contentText(response.content));
|
||||
return ok({
|
||||
text: contentText(response.content),
|
||||
usage: response.usage,
|
||||
});
|
||||
}
|
||||
|
||||
@@ -1,4 +1,4 @@
|
||||
import type { ImageContent, TextContent } from "@earendil-works/pi-ai";
|
||||
import type { ImageContent, TextContent, Usage } from "@earendil-works/pi-ai";
|
||||
import type { AgentMessage } from "../../types.ts";
|
||||
import { createBranchSummaryMessage, createCompactionSummaryMessage, createCustomMessage } from "../messages.ts";
|
||||
import type {
|
||||
@@ -247,6 +247,7 @@ export class Session<TMetadata extends SessionMetadata = SessionMetadata> {
|
||||
tokensBefore: number,
|
||||
details?: T,
|
||||
fromHook?: boolean,
|
||||
usage?: Usage,
|
||||
): Promise<string> {
|
||||
return this.appendTypedEntry({
|
||||
type: "compaction",
|
||||
@@ -257,6 +258,7 @@ export class Session<TMetadata extends SessionMetadata = SessionMetadata> {
|
||||
firstKeptEntryId,
|
||||
tokensBefore,
|
||||
details,
|
||||
usage,
|
||||
fromHook,
|
||||
} satisfies CompactionEntry<T>);
|
||||
}
|
||||
@@ -317,7 +319,7 @@ export class Session<TMetadata extends SessionMetadata = SessionMetadata> {
|
||||
|
||||
async moveTo(
|
||||
entryId: string | null,
|
||||
summary?: { summary: string; details?: unknown; fromHook?: boolean },
|
||||
summary?: { summary: string; details?: unknown; usage?: Usage; fromHook?: boolean },
|
||||
): Promise<string | undefined> {
|
||||
if (entryId !== null && !(await this.storage.getEntry(entryId))) {
|
||||
throw new SessionError("not_found", `Entry ${entryId} not found`);
|
||||
@@ -332,6 +334,7 @@ export class Session<TMetadata extends SessionMetadata = SessionMetadata> {
|
||||
fromId: entryId ?? "root",
|
||||
summary: summary.summary,
|
||||
details: summary.details,
|
||||
usage: summary.usage,
|
||||
fromHook: summary.fromHook,
|
||||
} satisfies BranchSummaryEntry);
|
||||
}
|
||||
|
||||
@@ -1,4 +1,12 @@
|
||||
import type { ImageContent, Model, Models, SimpleStreamOptions, TextContent, Transport } from "@earendil-works/pi-ai";
|
||||
import type {
|
||||
ImageContent,
|
||||
Model,
|
||||
Models,
|
||||
SimpleStreamOptions,
|
||||
TextContent,
|
||||
Transport,
|
||||
Usage,
|
||||
} from "@earendil-works/pi-ai";
|
||||
import type { AgentEvent, AgentMessage, AgentTool, QueueMode, ThinkingLevel } from "../index.ts";
|
||||
import type { Session } from "./session/session.ts";
|
||||
|
||||
@@ -365,6 +373,7 @@ export interface CompactionEntry<T = unknown> extends SessionTreeEntryBase {
|
||||
firstKeptEntryId: string;
|
||||
tokensBefore: number;
|
||||
details?: T;
|
||||
usage?: Usage;
|
||||
fromHook?: boolean;
|
||||
}
|
||||
|
||||
@@ -373,6 +382,7 @@ export interface BranchSummaryEntry<T = unknown> extends SessionTreeEntryBase {
|
||||
fromId: string;
|
||||
summary: string;
|
||||
details?: T;
|
||||
usage?: Usage;
|
||||
fromHook?: boolean;
|
||||
}
|
||||
|
||||
@@ -572,6 +582,7 @@ export interface ToolResultEvent {
|
||||
content: Array<TextContent | ImageContent>;
|
||||
details: unknown;
|
||||
isError: boolean;
|
||||
usage?: Usage;
|
||||
}
|
||||
|
||||
export interface SessionBeforeCompactEvent {
|
||||
@@ -687,6 +698,7 @@ export interface ToolResultPatch {
|
||||
content?: Array<TextContent | ImageContent>;
|
||||
details?: unknown;
|
||||
isError?: boolean;
|
||||
usage?: Usage;
|
||||
terminate?: boolean;
|
||||
}
|
||||
|
||||
@@ -697,7 +709,12 @@ export interface SessionBeforeCompactResult {
|
||||
|
||||
export interface SessionBeforeTreeResult {
|
||||
cancel?: boolean;
|
||||
summary?: { summary: string; details?: unknown };
|
||||
summary?: {
|
||||
summary: string;
|
||||
details?: unknown;
|
||||
/** Usage from the LLM call that generated this summary, if available. */
|
||||
usage?: Usage;
|
||||
};
|
||||
customInstructions?: string;
|
||||
replaceInstructions?: boolean;
|
||||
label?: string;
|
||||
@@ -738,6 +755,8 @@ export interface CompactResult {
|
||||
summary: string;
|
||||
firstKeptEntryId: string;
|
||||
tokensBefore: number;
|
||||
/** Usage from the LLM call(s) that generated this summary, if available. */
|
||||
usage?: Usage;
|
||||
details?: unknown;
|
||||
}
|
||||
|
||||
@@ -793,6 +812,7 @@ export interface GenerateBranchSummaryOptions {
|
||||
|
||||
export interface BranchSummaryResult {
|
||||
summary: string;
|
||||
usage?: Usage;
|
||||
readFiles: string[];
|
||||
modifiedFiles: string[];
|
||||
}
|
||||
|
||||
@@ -11,6 +11,7 @@ import type {
|
||||
TextContent,
|
||||
Tool,
|
||||
ToolResultMessage,
|
||||
Usage,
|
||||
} from "@earendil-works/pi-ai";
|
||||
import type { Static, TSchema } from "typebox";
|
||||
|
||||
@@ -69,15 +70,18 @@ export interface BeforeToolCallResult {
|
||||
* - `content`: if provided, replaces the tool result content array in full
|
||||
* - `details`: if provided, replaces the tool result details value in full
|
||||
* - `isError`: if provided, replaces the tool result error flag
|
||||
* - `usage`: if provided, replaces the tool result usage
|
||||
* - `terminate`: if provided, replaces the early-termination hint
|
||||
*
|
||||
* Omitted fields keep the original executed tool result values.
|
||||
* There is no deep merge for `content` or `details`.
|
||||
* There is no deep merge for `content`, `details`, or `usage`.
|
||||
*/
|
||||
export interface AfterToolCallResult {
|
||||
content?: (TextContent | ImageContent)[];
|
||||
details?: unknown;
|
||||
isError?: boolean;
|
||||
/** Usage from the final tool execution itself, if available. Not used for main LLM context accounting. */
|
||||
usage?: Usage;
|
||||
/**
|
||||
* Hint that the agent should stop after the current tool batch.
|
||||
* Early termination only happens when every finalized tool result in the batch sets this to true.
|
||||
@@ -273,6 +277,7 @@ export interface AgentLoopConfig extends SimpleStreamOptions {
|
||||
* - `content` replaces the full content array
|
||||
* - `details` replaces the full details payload
|
||||
* - `isError` replaces the error flag
|
||||
* - `usage` replaces the tool result usage
|
||||
* - `terminate` replaces the early-termination hint
|
||||
*
|
||||
* Any omitted fields keep their original values. No deep merge is performed.
|
||||
@@ -352,6 +357,8 @@ export interface AgentToolResult<T> {
|
||||
content: (TextContent | ImageContent)[];
|
||||
/** Arbitrary structured details for logs or UI rendering. */
|
||||
details: T;
|
||||
/** Usage from the final tool execution itself, if available. Not used for main LLM context accounting. */
|
||||
usage?: Usage;
|
||||
/** Names of tools introduced by this result and available from this transcript point onward. */
|
||||
addedToolNames?: string[];
|
||||
/**
|
||||
|
||||
Reference in New Issue
Block a user