diff --git a/src/bundled/pi/host.live.test.ts b/src/bundled/pi/host.live.test.ts index 8723f10a01..0ce5faaa7a 100644 --- a/src/bundled/pi/host.live.test.ts +++ b/src/bundled/pi/host.live.test.ts @@ -73,7 +73,7 @@ out.a = { ...(await pi.rpc("wait", { agent, submission: a.submission })), calls: faux.setResponses([fauxAssistantMessage("Plain answer")]); const b = await pi.rpc("deliver", turn("B", "second")); out.b = { ...(await pi.rpc("wait", { agent, submission: b.submission })), conversation: b.conversation }; -// What a transcript reads: the places, and the entries without schemas or usage. +// What a transcript and usage read: the places, and the entries without schemas. out.places = (await pi.rpc("places", { agent })).places; out.entries = (await pi.rpc("entries", { agent, conversation: a.conversation })).entries.map((entry) => [entry.kind, ...(entry.model ?? []).map((message) => message.role + (message.usage ? "+usage" : ""))]); // A request while the run works steers it. @@ -195,12 +195,13 @@ describe.skipIf(!DEPS)("the Pi host on Pi durable", () => { expect(out.places).toEqual([ { key: "chan/root", conversation: out.a.conversation, seen: 100 }, ]); - // Newest first; a system entry's tool schemas are left out. + // Newest first; a system entry's tool schemas are left out, an answer's + // usage kept. expect(out.entries).toEqual([ - ["pi.assistant", "assistant"], + ["pi.assistant", "assistant+usage"], ["pi.user", "user"], ["pi.tool-result", "toolResult"], - ["pi.assistant", "assistant"], + ["pi.assistant", "assistant+usage"], ["pi.system"], ["pi.user", "user"], ]); diff --git a/src/bundled/pi/host.mjs b/src/bundled/pi/host.mjs index 9edfbba691..a76c582958 100644 --- a/src/bundled/pi/host.mjs +++ b/src/bundled/pi/host.mjs @@ -179,8 +179,9 @@ async function piDefaults(dir) { } } -/** A message as a transcript shows it: no usage, image data or signatures. */ -const slim = ({ usage, ...message }) => +/** A message as a transcript and usage show it: no image data or + * signatures. */ +const slim = (message) => Array.isArray(message.content) ? { ...message, diff --git a/src/bundled/pi/index.ts b/src/bundled/pi/index.ts index 5265220f66..7fe53d1460 100644 --- a/src/bundled/pi/index.ts +++ b/src/bundled/pi/index.ts @@ -69,6 +69,10 @@ export const apply: PluginModule["apply"] = (ctx) => { const { model, provider } = config(agent.config); return [provider, model].filter(Boolean).join(" ยท ") || "Pi"; }, + ...(host.spawn && { + usage: (agent, { since, signal }) => + runtime.usage(agent.pubkey, since, signal), + }), sessions: { list: (pubkey) => runtime.sessions(pubkey), subscribe: (listener) => runtime.subscribe(listener), diff --git a/src/bundled/pi/runtime.ts b/src/bundled/pi/runtime.ts index 46d6b30fcd..7396a69732 100644 --- a/src/bundled/pi/runtime.ts +++ b/src/bundled/pi/runtime.ts @@ -54,6 +54,7 @@ import { } from "./pi"; import { DEFAULT_THINKING, PiSetup, runsAt } from "./setup"; import { readTranscript } from "./transcript"; +import { readUsage } from "./usage"; export type Config = Readonly<{ /** A Pi provider; empty uses the one Pi's settings name. */ @@ -274,6 +275,12 @@ export class PiRuntime { }; }); } + /** The usage of the agent's Buzz conversations since `since`. */ + async usage(pubkey: string, since: number, signal: AbortSignal) { + // Reading opens its storage, which only an agent running here may hold. + if (!this.running.has(pubkey)) return []; + return readUsage(this.pi, pubkey, await this.places(pubkey), since, signal); + } /** The conversations a Buzz request of the agent's is working in. */ working(pubkey: string) { let view = this.views.get(pubkey); diff --git a/src/bundled/pi/transcript.ts b/src/bundled/pi/transcript.ts index 607c5cd372..93059b9d0b 100644 --- a/src/bundled/pi/transcript.ts +++ b/src/bundled/pi/transcript.ts @@ -37,6 +37,9 @@ type Message = Readonly<{ stopReason?: string; errorMessage?: string; model?: string; + responseId?: string; + /** pi-ai's token counts and cost of an answer. */ + usage?: unknown; }>; /** A Pi entry as the host reads it: its messages, as pi-ai has them. */ export type PiEntry = Readonly<{ diff --git a/src/bundled/pi/usage.test.ts b/src/bundled/pi/usage.test.ts new file mode 100644 index 0000000000..added27ca7 --- /dev/null +++ b/src/bundled/pi/usage.test.ts @@ -0,0 +1,138 @@ +// @vitest-environment node +import { expect, it, vi } from "vitest"; +import type { PiEntry } from "./transcript"; +import { conversationUsage, readUsage } from "./usage"; + +const PLACE = { conversation: 2, channelId: "chan", root: "root" }; +const ask = (id: number, text: string): PiEntry => ({ + id, + kind: "pi.user", + model: [ + { + role: "user", + content: `\nChannel: general (#chan)\n\n\n\nFrom: Sal (npub1x)\nContent: ${text}\nTags: []\n`, + timestamp: id * 1000, + }, + ], +}); +const answer = (id: number, usage: Record, more = {}) => ({ + id, + kind: "pi.assistant", + model: [ + { + role: "assistant", + content: [{ type: "text", text: "Done." }], + timestamp: id * 1000, + model: "claude-x", + responseId: `resp-${id}`, + usage, + ...more, + }, + ], +}); +const usage = (cost: number, more = {}) => ({ + input: 100, + output: 10, + cacheRead: 50, + cacheWrite: 5, + totalTokens: 165, + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: cost }, + ...more, +}); + +it("reports each answer's tokens and cost, where and what was asked", () => { + expect( + conversationUsage( + [ + ask(1, "fix the build"), + answer(2, usage(0.25, { cacheWrite1h: 5 })), + { + id: 3, + kind: "pi.tool-result", + model: [{ role: "toolResult", timestamp: 3000 }], + }, + // A model pi-ai has no rates for is priced at 0: the cost is unknown. + answer(4, usage(0), { model: "local-y", responseId: undefined }), + ], + PLACE, + ), + ).toEqual({ + id: "2", + channelId: "chan", + root: "root", + title: "fix the build", + startedAt: 1000, + lastAt: 4000, + gaps: [], + observations: [ + { + requestId: "resp-2", + at: 2000, + model: "claude-x", + tokens: { + input: 100, + output: 10, + cacheRead: 50, + cacheWrite: 5, + cacheWrite1h: 5, + }, + cost: { usd: 0.25, source: "pi" }, + }, + { + at: 4000, + model: "local-y", + tokens: { input: 100, output: 10, cacheRead: 50, cacheWrite: 5 }, + }, + ], + }); +}); + +it("leaves out counts pi-ai did not record, and a conversation with no messages", () => { + const session = conversationUsage( + [ask(1, "hi"), answer(2, { output: 3, input: -1, cacheRead: "x" })], + { conversation: 5, channelId: "chan" }, + ); + expect(session?.root).toBeUndefined(); + expect(session?.observations).toEqual([ + { requestId: "resp-2", at: 2000, model: "claude-x", tokens: { output: 3 } }, + ]); + expect(conversationUsage([], PLACE)).toBeUndefined(); +}); + +it("reads every page of the conversations used since `since`, newest first", async () => { + const pages: Record = { + "2:": { entries: [answer(4, usage(1)), ask(3, "second")], next: "c" }, + "2:c": { entries: [answer(2, usage(1)), ask(1, "first")] }, + "7:": { entries: [answer(9, usage(1)), ask(8, "later")] }, + }; + const request = vi.fn( + async ( + _method: string, + params: { conversation: number; cursor?: string }, + ) => pages[`${params.conversation}:${params.cursor ?? ""}`], + ); + const sessions = await readUsage( + { request } as never, + "agent", + [ + { conversation: 2, channelId: "chan", root: "root", at: 3000 }, + { conversation: 7, channelId: "other", at: 8000 }, + // Shown nothing since: not read. + { conversation: 9, channelId: "old", at: 500 }, + ], + 1500, + new AbortController().signal, + ); + expect(request).toHaveBeenCalledTimes(3); + expect(request).toHaveBeenCalledWith("entries", { + agent: "agent", + conversation: 2, + cursor: "c", + }); + expect(sessions.map((session) => [session.id, session.title])).toEqual([ + ["7", "later"], + ["2", "first"], + ]); + expect(sessions[1]?.observations).toHaveLength(2); + expect(sessions[0]?.revision).toMatch(/^[0-9a-f]{64}$/); +}); diff --git a/src/bundled/pi/usage.ts b/src/bundled/pi/usage.ts new file mode 100644 index 0000000000..a7dd6b39f9 --- /dev/null +++ b/src/bundled/pi/usage.ts @@ -0,0 +1,139 @@ +// The token use of an agent's Pi conversations, as pi-ai records it on each +// of their answers, read from the agent's Pi durable storage through the host. +// Pi owns the entries; nothing here writes or keeps them. +import { promptSections } from "../../features/agents/activity-transcript"; +import type { + SessionUsage, + UsageObservation, +} from "../../features/agents/service"; +import { withRevision } from "../../features/agents/session-usage"; +import type { Pi } from "./pi"; +import type { PiEntry } from "./transcript"; + +type Json = Record; +/** A Buzz conversation's Pi conversation, as `places` lists it. */ +type Place = Readonly<{ + conversation: number; + channelId: string; + root?: string; + at: number; +}>; + +/** pi-ai's counts, named as Buzz names them. */ +const TOKENS = [ + "input", + "output", + "cacheRead", + "cacheWrite", + "cacheWrite1h", +] as const; + +const object = (value: unknown): Json | undefined => + value !== null && typeof value === "object" && !Array.isArray(value) + ? (value as Json) + : undefined; +const str = (value: unknown) => (typeof value === "string" ? value : ""); +const count = (value: unknown) => + typeof value === "number" && Number.isFinite(value) && value >= 0 + ? value + : undefined; +const text = (content: unknown) => + typeof content === "string" + ? content + : (Array.isArray(content) ? content : []) + .map((part) => str(object(part)?.text)) + .join("\n"); + +/** What was asked, from the conversation's first prompt. */ +function title(entries: readonly PiEntry[]) { + const first = entries.find((entry) => entry.kind === "pi.user")?.model?.[0]; + return promptSections(text(first?.content)) + .find((section) => section.tag === "buzz-event") + ?.body.match(/^Content: (.*)$/m)?.[1] + ?.trim(); +} + +/** A conversation's model requests, its entries oldest first. pi-ai counts a + * request's input without its cached part, and prices the models it knows. */ +export function conversationUsage( + entries: readonly PiEntry[], + place: Omit, +): Omit | undefined { + const observations: UsageObservation[] = []; + const times: number[] = []; + for (const entry of entries) + for (const message of entry.model ?? []) { + if (typeof message.timestamp === "number") times.push(message.timestamp); + const usage = object(message.usage); + if (message.role !== "assistant" || !usage) continue; + const cost = count(object(usage.cost)?.total); + observations.push({ + ...(message.responseId ? { requestId: message.responseId } : {}), + at: message.timestamp ?? 0, + model: message.model || "pi", + tokens: Object.fromEntries( + TOKENS.flatMap((key) => { + const value = count(usage[key]); + return value === undefined ? [] : [[key, value]]; + }), + ), + // pi-ai prices a model it has no rates for at 0: that is unknown. + ...(cost ? { cost: { usd: cost, source: "pi" } } : {}), + }); + } + if (!times.length) return undefined; + const asked = title(entries); + return { + id: String(place.conversation), + channelId: place.channelId, + ...(place.root ? { root: place.root } : {}), + ...(asked ? { title: asked } : {}), + startedAt: Math.min(...times), + lastAt: Math.max(...times), + gaps: [], + observations, + }; +} + +/** All of a conversation's entries, oldest first. */ +async function entries( + pi: Pick, + agent: string, + conversation: number, + signal: AbortSignal, +) { + const found: PiEntry[] = []; + let cursor: unknown; + do { + const page = await pi.request<{ entries: PiEntry[]; next?: unknown }>( + "entries", + { agent, conversation, ...(cursor ? { cursor } : {}) }, + ); + signal.throwIfAborted(); + found.push(...page.entries); + cursor = page.next; + } while (cursor); + return found.reverse(); +} + +/** The agent's Buzz conversations used since `since`. A conversation's + * requests follow the newest message it was shown. */ +export async function readUsage( + pi: Pick, + agent: string, + places: readonly Place[], + since: number, + signal: AbortSignal, +): Promise { + const found: SessionUsage[] = []; + for (const place of places) { + if (place.at < since) continue; + const session = conversationUsage( + await entries(pi, agent, place.conversation, signal), + place, + ); + if (session && session.lastAt >= since) + found.push(await withRevision(session)); + } + return found.sort((a, b) => b.lastAt - a.lastAt); +}