mirror of
https://github.com/badlogic/pi-mono.git
synced 2026-11-05 23:37:25 +00:00
feat: expose provider stream events to extensions (#9901)
Add onProviderStreamEvent through pi-ai and agent core, and surface provider_stream_event to coding-agent extensions. Closes #9784
This commit is contained in:
parent
8d897edaa6
commit
002fc83852
36 changed files with 729 additions and 36 deletions
|
|
@ -119,6 +119,7 @@ export interface AgentOptions {
|
|||
getApiKey?: (provider: string) => Promise<string | undefined> | string | undefined;
|
||||
onPayload?: SimpleStreamOptions["onPayload"];
|
||||
onResponse?: SimpleStreamOptions["onResponse"];
|
||||
onProviderStreamEvent?: SimpleStreamOptions["onProviderStreamEvent"];
|
||||
beforeToolCall?: (context: BeforeToolCallContext, signal?: AbortSignal) => Promise<BeforeToolCallResult | undefined>;
|
||||
afterToolCall?: (context: AfterToolCallContext, signal?: AbortSignal) => Promise<AfterToolCallResult | undefined>;
|
||||
finishTurn?: FinishTurn;
|
||||
|
|
@ -196,6 +197,7 @@ export class Agent {
|
|||
public getApiKey?: (provider: string) => Promise<string | undefined> | string | undefined;
|
||||
public onPayload?: SimpleStreamOptions["onPayload"];
|
||||
public onResponse?: SimpleStreamOptions["onResponse"];
|
||||
public onProviderStreamEvent?: SimpleStreamOptions["onProviderStreamEvent"];
|
||||
public beforeToolCall?: (
|
||||
context: BeforeToolCallContext,
|
||||
signal?: AbortSignal,
|
||||
|
|
@ -235,6 +237,7 @@ export class Agent {
|
|||
this.getApiKey = runtimeOptions.getApiKey;
|
||||
this.onPayload = runtimeOptions.onPayload;
|
||||
this.onResponse = runtimeOptions.onResponse;
|
||||
this.onProviderStreamEvent = runtimeOptions.onProviderStreamEvent;
|
||||
this.beforeToolCall = runtimeOptions.beforeToolCall;
|
||||
this.afterToolCall = runtimeOptions.afterToolCall;
|
||||
this.finishTurn = runtimeOptions.finishTurn;
|
||||
|
|
@ -469,6 +472,7 @@ export class Agent {
|
|||
sessionId: this.sessionId,
|
||||
onPayload: this.onPayload,
|
||||
onResponse: this.onResponse,
|
||||
onProviderStreamEvent: this.onProviderStreamEvent,
|
||||
transport: this.transport,
|
||||
thinkingBudgets: this.thinkingBudgets,
|
||||
maxRetryDelayMs: this.maxRetryDelayMs,
|
||||
|
|
|
|||
|
|
@ -1166,6 +1166,28 @@ describe("Agent", () => {
|
|||
expect(agent.peekQueuedMessages()).toEqual([followUp]);
|
||||
});
|
||||
|
||||
it("forwards provider stream event observers through AgentOptions", async () => {
|
||||
const providerEvents: unknown[] = [];
|
||||
const agent = new Agent({
|
||||
onProviderStreamEvent: (data) => {
|
||||
providerEvents.push(data);
|
||||
},
|
||||
streamFn: (model, _context, options) => {
|
||||
const stream = new MockAssistantStream();
|
||||
queueMicrotask(async () => {
|
||||
await options?.onProviderStreamEvent?.({ request_cost: 0.01 }, model);
|
||||
const message = createAssistantMessage("ok");
|
||||
stream.push({ type: "done", reason: "stop", message });
|
||||
});
|
||||
return stream;
|
||||
},
|
||||
});
|
||||
|
||||
await agent.prompt("hello");
|
||||
|
||||
expect(providerEvents).toEqual([{ request_cost: 0.01 }]);
|
||||
});
|
||||
|
||||
it("forwards sessionId to streamFunction options", async () => {
|
||||
let receivedSessionId: string | undefined;
|
||||
const agent = new Agent({
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue