mirror of
https://github.com/badlogic/pi-mono.git
synced 2026-11-05 23:37:25 +00:00
feat(durable): beforeProviderRequest generation hook (#11619)
Adds a `beforeProviderRequest` generation hook so extensions can replace the provider request body before it is sent. It is the durable counterpart of coding-agent's `before_provider_request`. Durable calls `streamSimple()`, so apps can't set provider-specific fields such as Anthropic's `thinking.display: "updates"` (#10804) without wrapping `Models`. Now it's an extension: ```ts hook(GenerationTask, { beforeProviderRequest: ({ model, payload }) => { if (model.api !== "anthropic-messages") return undefined; // change payload return { payload }; }, }) ```
This commit is contained in:
parent
3bbb06aad7
commit
d280e1ffa2
6 changed files with 109 additions and 2 deletions
|
|
@ -152,6 +152,7 @@
|
|||
- Run-end tasks: a `GenerationHooks.afterRun(run, api)` hook returns `{ tasks }`, built with the typed `runEndTask(task, input)`, after a run answered (a final answer or a `terminate` round; not when an `onYield` continuation is taken, not after a failure or an abort). The answer is committed and the run's inputs settle `done` first, so `Submission.wait()` resolves at the answer; the generation then creates the tasks as its children and waits for them, keeping `pi.live.run` (with `run.afterRun`) and the conversation busy, so idle waits include them, Stop aborts them, and a crash resumes them. Input queued meanwhile, steers included, is placed when they are terminal and starts the next run; a task admits a model-visible write or a fix-run follow-up in its own commit with `submit(tx, draft)`. A crash before the tasks exist asks the hooks again; a broken request (an uninstalled or built-in task, a non-JSON input, a throwing `initial()`) is reported and skipped. `run` names the run's entries: `first` (`pi.live.run.first`, carried into `runFailure.first` so a continue covers a failed run's history), `answer`, and `last` ([#11141](https://github.com/earendil-works/pi/issues/11141))
|
||||
|
||||
- `registryReader(value, context)`: a `RegistryReader` over a Chord v2 registry value whose composition is a registry snapshot (`RegistryValue<RegistrySnapshot>`): `snapshot()` reads the current composition, `subscribe` observes its compositions synchronously inside Chord's batch, so a Harness opened in a Chord object takes its extensions from a Chord registry.
|
||||
- `GenerationHooks.beforeProviderRequest({ model, payload }, api)` replaces the provider request body of a generation request, for that request only, e.g. to set a provider-specific field. The body is in the provider API's format; the hook returns `{ payload }` or `undefined`. Hooks chain in extension order and run each time pi-ai builds a body; a throw is reported and the chain continues. Compaction requests do not use it ([#10804](https://github.com/earendil-works/pi/issues/10804))
|
||||
|
||||
### Changed
|
||||
|
||||
|
|
|
|||
|
|
@ -659,7 +659,7 @@ const Guard = defineExtension({
|
|||
});
|
||||
```
|
||||
|
||||
- **Generation:** `beforeRequest` (replace the messages of one request), `afterResponse`, `onYield` (continue the run with another user message), `afterTools` (runs once a round's tools are done), and `afterRun` (tasks to run before the run ends, below).
|
||||
- **Generation:** `beforeRequest` (replace the messages of one request), `beforeProviderRequest` (replace the provider request body of one request, e.g. to set a provider-specific field), `afterResponse`, `onYield` (continue the run with another user message), `afterTools` (runs once a round's tools are done), and `afterRun` (tasks to run before the run ends, below).
|
||||
- **Tools:** `beforeTool` (block or rewrite arguments) and `afterTool` (replace the result). `afterTool` sees every result a tool task writes, also the error results the Harness writes itself (unavailable, invalid, blocked, interrupted, restart limit, aborted), told apart by their diagnostic code; past `settings.restartLimit` it no longer runs. On an abort it runs inside the abort handler, so a slow hook delays `abort()`.
|
||||
|
||||
To limit a hook to some conversations, select its extension only there, for example with `configure({ extensions: { add: [Guard] } })`.
|
||||
|
|
|
|||
|
|
@ -3339,6 +3339,7 @@ consuming commit may rerun them. Abort errors always propagate.
|
|||
| hook | composition | ordinary throw |
|
||||
|---|---|---|
|
||||
| `beforeRequest` | replacement chain | report, continue |
|
||||
| `beforeProviderRequest` | replacement chain | report, continue |
|
||||
| `afterResponse` | all observers | report, continue |
|
||||
| `onYield` | first continuation wins | report, continue |
|
||||
| `beforeTool` | argument replacement chain; first block wins | block tool with error text |
|
||||
|
|
@ -3380,6 +3381,12 @@ interface GenerationHooks {
|
|||
api: HookApi,
|
||||
context: Context,
|
||||
): HookResult<{ readonly messages: readonly Message[] }>;
|
||||
/** After pi-ai builds the provider request body, before it is sent; replaces the body for that request only. */
|
||||
beforeProviderRequest(
|
||||
request: { readonly model: Model<Api>; readonly payload: unknown },
|
||||
api: HookApi,
|
||||
context: Context,
|
||||
): HookResult<{ readonly payload: unknown }>;
|
||||
/** Every terminal provider message, before classification. */
|
||||
afterResponse(message: AssistantMessage, api: HookApi, context: Context): void | Promise<void>;
|
||||
/** A final answer; `continue` appends a user message and continues the run. */
|
||||
|
|
@ -4900,7 +4907,9 @@ The run's inputs live in `pi.live.run`, not in the task input.
|
|||
(section 4.1), so Storage gets them at the next commit that reads `pi.live`
|
||||
or within `settings.progress.persistWithinMs`.
|
||||
Before streaming, the `beforeRequest` chain may
|
||||
replace the messages for this request only. After a close, recovery resends
|
||||
replace the messages for this request only. Once pi-ai has built the provider
|
||||
request body, the `beforeProviderRequest` chain may replace it, again for this
|
||||
request only; it runs each time pi-ai builds a body. After a close, recovery resends
|
||||
the same committed messages with the same pinned model, thinking level, and
|
||||
stream options, and reruns `beforeRequest`; after a crash the attempt is
|
||||
charged instead, as above.
|
||||
|
|
|
|||
|
|
@ -292,6 +292,7 @@ export const GenerationTask = defineTask<GenerationInput, GenerationCheckpoint,
|
|||
telemetry: context.telemetry,
|
||||
sessionId: await ensureProviderSessionId(runtime, context),
|
||||
...(thinkingLevel === "off" ? {} : { reasoning: thinkingLevel }),
|
||||
onPayload: (payload) => providerPayload(runtime, model, payload, context),
|
||||
};
|
||||
const message = await streamResponse(runtime, model, messages, options, attempt, context);
|
||||
const request = { attempt, compacted, model: ref, thinkingLevel, cutoff, messages: view.messages };
|
||||
|
|
@ -664,6 +665,21 @@ export async function convertPartial(tx: Tx, live: Draft<LiveState>, conversatio
|
|||
await appendAssistant(tx, conversationId, { ...message, stopReason: "aborted" });
|
||||
}
|
||||
|
||||
/** The `beforeProviderRequest` replacement chain over one provider request body. */
|
||||
async function providerPayload(
|
||||
runtime: Runtime,
|
||||
model: Model<Api>,
|
||||
payload: unknown,
|
||||
context: Context,
|
||||
): Promise<unknown> {
|
||||
let current = payload;
|
||||
await runtime.hooks.each("beforeProviderRequest", async (hook) => {
|
||||
const replaced = await hook({ model, payload: current }, runtime, context);
|
||||
if (replaced !== undefined) current = replaced.payload;
|
||||
});
|
||||
return current;
|
||||
}
|
||||
|
||||
/**
|
||||
* Stream one request and return the terminal message. Partials commit as trailing writes at most every
|
||||
* `progress.partialIntervalMs` (default 100 ms) with one commit in flight, deferred (`PROGRESS_COMMIT`); `finally` stops
|
||||
|
|
|
|||
|
|
@ -1,9 +1,11 @@
|
|||
import type { Context, JsonValue, WatchHandle } from "@earendil-works/chord";
|
||||
import type {
|
||||
Api,
|
||||
AssistantMessage,
|
||||
CacheRetention,
|
||||
ClientIdentity,
|
||||
Message,
|
||||
Model,
|
||||
Models,
|
||||
ModelThinkingLevel,
|
||||
Static,
|
||||
|
|
@ -970,6 +972,16 @@ export interface GenerationHooks {
|
|||
api: HookApi,
|
||||
context: Context,
|
||||
): HookResult<{ readonly messages: readonly Message[] }>;
|
||||
/**
|
||||
* After pi-ai builds the provider request body, before it is sent; the result replaces the body for that request only.
|
||||
* Runs again whenever pi-ai rebuilds the body, such as a retry within one attempt. The body is in the provider API's
|
||||
* own format (`request.model.api`), so a hook checks the API before changing it.
|
||||
*/
|
||||
beforeProviderRequest(
|
||||
request: { readonly model: Model<Api>; readonly payload: unknown },
|
||||
api: HookApi,
|
||||
context: Context,
|
||||
): HookResult<{ readonly payload: unknown }>;
|
||||
/** Every terminal provider message, before classification. */
|
||||
afterResponse(message: AssistantMessage, api: HookApi, context: Context): void | Promise<void>;
|
||||
/** A final answer; the first `continue` appends a user message and continues the run. */
|
||||
|
|
|
|||
|
|
@ -752,6 +752,75 @@ describe("generation hooks", () => {
|
|||
await harness.close(context);
|
||||
});
|
||||
|
||||
it("chains beforeProviderRequest over every provider request body", async () => {
|
||||
const setup = chatSetup();
|
||||
const seen: { hook: string; model: string; payload: unknown }[] = [];
|
||||
addHooks(setup.registry, GenerationTask, {
|
||||
beforeProviderRequest: ({ model, payload }) => {
|
||||
seen.push({ hook: "first", model: `${model.provider}/${model.id}`, payload });
|
||||
return { payload: { ...(payload as object), first: true } };
|
||||
},
|
||||
});
|
||||
addHooks(setup.registry, GenerationTask, {
|
||||
beforeProviderRequest: ({ model, payload }) => {
|
||||
seen.push({ hook: "keep", model: `${model.provider}/${model.id}`, payload });
|
||||
return undefined;
|
||||
},
|
||||
});
|
||||
addHooks(setup.registry, GenerationTask, {
|
||||
beforeProviderRequest: ({ payload }) => ({ payload: { ...(payload as object), last: true } }),
|
||||
});
|
||||
const sent: unknown[] = [];
|
||||
// pi-ai builds a body per provider call, and again when it resends within one attempt.
|
||||
const respond: FauxResponseStep = async (_context, options, _state, model) => {
|
||||
sent.push(await options!.onPayload!({ body: 1 }, model));
|
||||
sent.push(await options!.onPayload!({ body: 2 }, model));
|
||||
return DONE;
|
||||
};
|
||||
const { harness, status } = await run(setup, [respond]);
|
||||
expect(status).toBe("done");
|
||||
expect(sent).toEqual([
|
||||
{ body: 1, first: true, last: true },
|
||||
{ body: 2, first: true, last: true },
|
||||
]);
|
||||
expect(seen).toEqual([
|
||||
{ hook: "first", model: "faux/faux-1", payload: { body: 1 } },
|
||||
{ hook: "keep", model: "faux/faux-1", payload: { body: 1, first: true } },
|
||||
{ hook: "first", model: "faux/faux-1", payload: { body: 2 } },
|
||||
{ hook: "keep", model: "faux/faux-1", payload: { body: 2, first: true } },
|
||||
]);
|
||||
await harness.close(context);
|
||||
});
|
||||
|
||||
it("reports a throwing beforeProviderRequest hook and continues the chain", async () => {
|
||||
const setup = chatSetup();
|
||||
addHooks(
|
||||
setup.registry,
|
||||
GenerationTask,
|
||||
{
|
||||
beforeProviderRequest: () => {
|
||||
throw new Error("broken hook");
|
||||
},
|
||||
},
|
||||
"broken",
|
||||
);
|
||||
addHooks(setup.registry, GenerationTask, {
|
||||
beforeProviderRequest: ({ payload }) => ({ payload: { ...(payload as object), after: true } }),
|
||||
});
|
||||
let sent: unknown;
|
||||
const respond: FauxResponseStep = async (_context, options, _state, model) => {
|
||||
sent = await options!.onPayload!({ body: 1 }, model);
|
||||
return DONE;
|
||||
};
|
||||
const { harness, status } = await run(setup, [respond]);
|
||||
expect(status).toBe("done");
|
||||
expect(sent).toEqual({ body: 1, after: true });
|
||||
expect(setup.infos.filter((info) => info.kind === "extension_error")).toEqual([
|
||||
expect.objectContaining({ extension: "broken" }),
|
||||
]);
|
||||
await harness.close(context);
|
||||
});
|
||||
|
||||
it("keeps the run's input open across an onYield continuation and answers it with the final answer", async () => {
|
||||
const setup = chatSetup();
|
||||
let yields = 0;
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue