Something went wrong. Try again.
fork of prime agent with changes i want but are not upstreamable probqbly (and some fixes)
Something went wrong. Try again.
5.6 kB · 187 lines
TypeScript
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188import { Agent, type AgentMessage } from "@earendil-works/pi-agent-core";import type { AssistantMessage, UserMessage } from "@earendil-works/pi-ai";import { completeWithProviderRetry, DEFAULT_PROVIDER_RETRY_POLICY, type ProviderRetryPolicy,} from "./provider-retry.js";import { unwrapSemanticEdgeStreamFn } from "./semantic-edges.js";
export type SideQuestionStatus = "running" | "complete" | "cancelled" | "error";
export interface SideQuestionEvent { id: string; question: string; answer: string; status: SideQuestionStatus; errorMessage?: string;}
export interface SideQuestionTurn { question: string; answer: string;}
export interface SideQuestionRun { done: Promise<void>; abort(): void;}
const SIDE_QUESTION_INSTRUCTION = "Answer this side question using only the conversation context above. Do not use tools. The user may send follow-up side questions; none of this side conversation is added to the main session.";
function sideQuestionPrompt(question: string, isFirstTurn: boolean): string { const body = isFirstTurn ? `${SIDE_QUESTION_INSTRUCTION}\n\n${question}` : question; return `<side_question>\n${body}\n</side_question>`;}
function readAssistantText(message: AgentMessage): string { if (message.role !== "assistant") { return ""; } return message.content .filter((block) => block.type === "text") .map((block) => block.text) .join("");}
export function startSideQuestion( parent: Agent, id: string, question: string, onEvent: (event: SideQuestionEvent) => void | Promise<void>, previousTurns: SideQuestionTurn[] = [], retry: ProviderRetryPolicy = DEFAULT_PROVIDER_RETRY_POLICY,): SideQuestionRun { const model = parent.state.model; if (!model) { throw new Error("Select a model before asking a side question"); }
// Each turn re-clones the live main conversation, so follow-ups always see // the newest main-thread context; earlier side turns are replayed after it. const previousTurnMessages: AgentMessage[] = previousTurns.flatMap((turn, index) => [ { role: "user", content: [{ type: "text", text: sideQuestionPrompt(turn.question, index === 0) }], timestamp: Date.now(), } satisfies UserMessage, { role: "assistant", content: [{ type: "text", text: turn.answer }], api: model.api, provider: model.provider, model: model.id, usage: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, totalTokens: 0, cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, }, stopReason: "stop", timestamp: Date.now(), } satisfies AssistantMessage, ]);
const sideAgent = new Agent({ initialState: { model, systemPrompt: parent.state.systemPrompt, messages: [...structuredClone(parent.state.messages), ...previousTurnMessages], thinkingLevel: "off", serviceTier: parent.state.serviceTier, tools: [], }, convertToLlm: parent.convertToLlm, transformContext: parent.transformContext, // Side questions are excluded from session history; their calls carry no provenance. streamFn: unwrapSemanticEdgeStreamFn(parent.streamFn), getApiKey: parent.getApiKey, onPayload: parent.onPayload, onResponse: parent.onResponse, shouldStopAfterTurn: () => true, sessionId: parent.sessionId, thinkingBudgets: parent.thinkingBudgets, transport: "sse", toolExecution: parent.toolExecution, });
let answer = ""; let abortRequested = false; let started = false; const retryAbortController = new AbortController(); const emit = (status: SideQuestionStatus, errorMessage?: string) => onEvent({ id, question, answer, status, ...(errorMessage ? { errorMessage } : {}) });
const unsubscribe = sideAgent.subscribe(async (event) => { if (event.type !== "message_update" && event.type !== "message_end") { return; } const nextAnswer = readAssistantText(event.message); if (nextAnswer === answer) { return; } answer = nextAnswer; await emit("running"); });
const prompt = sideQuestionPrompt(question, previousTurns.length === 0); const done = Promise.resolve() .then(() => emit("running")) .then(async () => { if (abortRequested) { await emit("cancelled"); return; } started = true; // Standalone side agents bypass the session auto-retry loop; retry here instead. let promptedOnce = false; await completeWithProviderRetry( async () => { if (promptedOnce) { // Session-loop recovery: drop the failed assistant turn and re-run. sideAgent.state.messages = sideAgent.state.messages.slice(0, -1); await sideAgent.continue(); } else { promptedOnce = true; await sideAgent.prompt(prompt); } const last = sideAgent.state.messages.at(-1); if (last?.role !== "assistant") { throw new Error(sideAgent.state.errorMessage || "Side question produced no assistant message"); } return last as AssistantMessage; }, { policy: retry, signal: retryAbortController.signal }, ); if (abortRequested) { await emit("cancelled"); return; } if (sideAgent.state.errorMessage) { await emit("error", sideAgent.state.errorMessage); return; } await emit("complete"); }) .catch(async (error) => { const errorMessage = error instanceof Error ? error.message : String(error); await Promise.resolve( emit(abortRequested ? "cancelled" : "error", abortRequested ? undefined : errorMessage), ).catch(() => undefined); }) .finally(unsubscribe);
return { done, abort() { abortRequested = true; retryAbortController.abort(); if (started) { sideAgent.abort(); } }, };}