feat(context-engine): add toolsOverride field to AssembleResult
Allow context-engine plugins to filter tool schemas sent to the model. When assemble() returns toolsOverride, only those tool schemas are included in model requests, reducing token usage while keeping all tool executors available.
This commit is contained in:
parent
843e3c1efb
commit
6fc2a0af60
@ -127,6 +127,7 @@ import { dropThinkingBlocks } from "../thinking.js";
|
||||
import { collectAllowedToolNames } from "../tool-name-allowlist.js";
|
||||
import { installToolResultContextGuard } from "../tool-result-context-guard.js";
|
||||
import { splitSdkTools } from "../tool-split.js";
|
||||
import { wrapStreamFnWithToolsOverride } from "../stream-payload-utils.js";
|
||||
import { describeUnknownError, mapThinkingLevel } from "../utils.js";
|
||||
import { flushPendingToolResultsAfterIdle } from "../wait-for-idle-before-flush.js";
|
||||
import { waitForCompactionRetryWithAggregateTimeout } from "./compaction-retry-aggregate-timeout.js";
|
||||
@ -2123,6 +2124,15 @@ export async function runEmbeddedAttempt(
|
||||
`context engine: prepended system prompt addition (${assembled.systemPromptAddition.length} chars)`,
|
||||
);
|
||||
}
|
||||
if (assembled.toolsOverride && assembled.toolsOverride.length > 0) {
|
||||
activeSession.agent.streamFn = wrapStreamFnWithToolsOverride(
|
||||
activeSession.agent.streamFn,
|
||||
assembled.toolsOverride,
|
||||
);
|
||||
log.debug(
|
||||
`context engine: applied tools override (${assembled.toolsOverride.length} tools)`,
|
||||
);
|
||||
}
|
||||
} catch (assembleErr) {
|
||||
log.warn(
|
||||
`context engine assemble failed, using pipeline messages: ${String(assembleErr)}`,
|
||||
|
||||
@ -1,4 +1,28 @@
|
||||
import type { StreamFn } from "@mariozechner/pi-agent-core";
|
||||
import type { AnyAgentTool } from "../pi-tools.types.js";
|
||||
|
||||
/**
|
||||
* Creates a streamFn wrapper that filters the tools context to only include
|
||||
* tools whose names are in the override set. This reduces token usage by
|
||||
* sending fewer tool schemas to the model while keeping all executors available.
|
||||
*/
|
||||
export function wrapStreamFnWithToolsOverride(
|
||||
underlying: StreamFn,
|
||||
toolsOverride: AnyAgentTool[],
|
||||
): StreamFn {
|
||||
const allowedNames = new Set(toolsOverride.map((tool) => tool.name));
|
||||
return (model, context, options) => {
|
||||
const ctx = context as { tools?: unknown[] };
|
||||
if (!Array.isArray(ctx.tools)) {
|
||||
return underlying(model, context, options);
|
||||
}
|
||||
const filteredTools = ctx.tools.filter((tool) => {
|
||||
const name = (tool as { name?: string })?.name;
|
||||
return typeof name === "string" && allowedNames.has(name);
|
||||
});
|
||||
return underlying(model, { ...context, tools: filteredTools }, options);
|
||||
};
|
||||
}
|
||||
|
||||
export function streamWithPayloadPatch(
|
||||
underlying: StreamFn,
|
||||
|
||||
@ -1,4 +1,5 @@
|
||||
import type { AgentMessage } from "@mariozechner/pi-agent-core";
|
||||
import type { AnyAgentTool } from "../agents/pi-tools.types.js";
|
||||
|
||||
// Result types
|
||||
|
||||
@ -9,6 +10,12 @@ export type AssembleResult = {
|
||||
estimatedTokens: number;
|
||||
/** Optional context-engine-provided instructions prepended to the runtime system prompt */
|
||||
systemPromptAddition?: string;
|
||||
/**
|
||||
* Optional filtered list of tools to send to the model instead of the default full set.
|
||||
* When present, only the schemas for these tools are included in the model request,
|
||||
* reducing token usage. Tool executors for all registered tools remain available.
|
||||
*/
|
||||
toolsOverride?: AnyAgentTool[];
|
||||
};
|
||||
|
||||
export type CompactResult = {
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user