Something went wrong. Try again.
forked niri
Something went wrong. Try again.
15 kB · 375 lines
TypeScript
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376import { stream } from "../llm"import type { AssistantMessage, AssistantMessageEvent, TextContent, ThinkingContent, Usage } from "../llm"import type { OpenAICompletionsOptions } from "../llm"import { logMessage } from "../db"import { buildCompletionMessages, rememberRecalledMemoryChunks } from "../memory"import { recordMetric } from "../metrics"import { emit } from "../stream"import type { LoopState } from "./types"import { API_BASE, ENABLE_THINKING, FALLBACK_API_KEY, FALLBACK_BASE, FALLBACK_MODEL, FALLBACK_TOOL_CHOICE, MODEL, PRIMARY_API_KEY, PRIMARY_TOOL_CHOICE, RESPECT_CACHE, SUMMARY_API_KEY, SUMMARY_MODEL, THINKING_EFFORT, USE_FALLBACK, buildContext, errorSummary, estimatePromptTokens, fallbackContextWindow, fallbackModel, findSummaryMessageIndex, isContentFilterError, isImageParseError, isPromptTooLargeError, loadAgentSummaryContext, primaryModel, retryDelayMs, sanitizeMessages, scrubImagesFromConversation, shouldFallback, summaryModel, summarizeConversationViaLLM,} from "./util"import type { CompletionTurnResult } from "./loop-shared"import type { Model } from "../llm"
/** * Resolves the configured summary model pair. * When respect_cache is true, returns null — no summaries are produced. */export function configuredSummaryProvider(): { model: Model | null; apiKey: string | undefined } { if (RESPECT_CACHE) return { model: null, apiKey: undefined } if (summaryModel && SUMMARY_MODEL) return { model: summaryModel, apiKey: SUMMARY_API_KEY } return { model: USE_FALLBACK ? fallbackModel : primaryModel, apiKey: USE_FALLBACK ? FALLBACK_API_KEY : PRIMARY_API_KEY, }}
/** * Appends an assistant message to state and persists it to conversation logs. */export function addAssistantMessage(convId: number, state: LoopState, msg: AssistantMessage): void { state.conversation.push(msg) const textContent = msg.content .filter((p): p is TextContent => p.type === "text") .map((p) => p.text) .join("") const toolCalls = msg.content .filter((p) => p.type === "toolCall") .map((p) => p) logMessage(convId, msg.role, textContent, toolCalls.length > 0 ? toolCalls : undefined)}
function recordPromptResponse( model: string, toolChoice: string | undefined, state: LoopState, result: CompletionTurnResult, promptMetricId: number | null,): void { recordMetric({ type: "prompt_response", promptMetricId: promptMetricId ?? undefined, model, toolChoice, messages: state.conversation as import("../types").Message[], response: result.message, usage: result.usage, })}
/** * Applies token usage from a completion response to loop state counters. */export function applyUsage( state: LoopState, usage: Usage, timing: Pick<CompletionTurnResult, "elapsedMs" | "tokensPerSecond"> = {},): void { state.tokenCount += usage.totalTokens if (usage.input) state.contextSize = usage.input + usage.cacheRead const rate = typeof timing.tokensPerSecond === "number" ? ` ${timing.tokensPerSecond.toFixed(1)} tok/s` : "" console.log( `[tokens] +${usage.totalTokens} total=${state.tokenCount}` + (usage.cacheRead ? ` cache_read=${usage.cacheRead}` : "") + (usage.cacheWrite ? ` cache_write=${usage.cacheWrite}` : "") + rate, ) emit({ type: "usage", promptTokens: usage.input, completionTokens: usage.output, totalTokens: usage.totalTokens, cachedPromptTokens: usage.cacheRead || undefined, cacheWriteTokens: usage.cacheWrite || undefined, elapsedMs: timing.elapsedMs, tokensPerSecond: timing.tokensPerSecond, }) recordMetric({ type: "usage", usage })}
/** * Emits model reasoning/thinking when present in the assistant message. */export function emitThinking(msg: AssistantMessage): void { const thinking = msg.content .filter((c): c is ThinkingContent => c.type === "thinking") .map((c) => c.thinking) .join("") .trim() if (thinking) emit({ type: "thinking", text: thinking })}
function sleep(ms: number): Promise<void> { return new Promise((resolve) => setTimeout(resolve, ms))}
/** * Streams one completion from a model. */async function createStreamedCompletion( model: Model, apiKey: string | undefined, toolChoice: "required" | "auto" | "none", state: LoopState,): Promise<CompletionTurnResult> { const startedAt = Date.now() const context = buildContext(state)
const resolvedToolChoice = ENABLE_THINKING && toolChoice === "required" ? "auto" : toolChoice const options: OpenAICompletionsOptions = { ...(apiKey ? { apiKey } : {}), toolChoice: resolvedToolChoice, ...(ENABLE_THINKING ? { reasoningEffort: THINKING_EFFORT } : {}), }
const promptMetricId = recordMetric({ type: "prompt", messages: state.conversation as import("../types").Message[] }) const eventStream = stream(model, context, options)
let emittedText = false let emittedThinking = false let bufferedThinking = "" let errorMessage: string | undefined
for await (const event of eventStream as AsyncIterable<AssistantMessageEvent>) { if (event.type === "text_delta") { emit({ type: "text", text: event.delta }) emittedText = true } else if (event.type === "thinking_start") { bufferedThinking = "" } else if (event.type === "thinking_delta") { if (ENABLE_THINKING) bufferedThinking += event.delta } else if (event.type === "thinking_end") { if (ENABLE_THINKING && bufferedThinking) { emit({ type: "thinking", text: bufferedThinking }) emittedThinking = true } } else if (event.type === "error") { errorMessage = event.error.errorMessage ?? "unknown API error" } }
const finalMessage = await (eventStream as { result(): Promise<AssistantMessage> }).result()
if (finalMessage.stopReason === "error" || errorMessage) { throw new Error(errorMessage ?? finalMessage.errorMessage ?? "unknown API error") }
const usage = finalMessage.usage const elapsedMs = Math.max(0, Date.now() - startedAt) const tokensPerSecond = usage.output && elapsedMs > 0 ? usage.output / (elapsedMs / 1000) : undefined
recordPromptResponse(model.id, toolChoice, state, { message: finalMessage, usage, emittedText, emittedThinking, bufferedThinking, elapsedMs, tokensPerSecond }, promptMetricId)
return { message: finalMessage, usage, emittedText, emittedThinking, bufferedThinking, elapsedMs, tokensPerSecond }}
async function createFallbackCompletion(state: LoopState): Promise<CompletionTurnResult> { return createStreamedCompletion(fallbackModel, FALLBACK_API_KEY, FALLBACK_TOOL_CHOICE, state)}
async function createPrimaryCompletion(state: LoopState): Promise<CompletionTurnResult> { return createStreamedCompletion(primaryModel!, PRIMARY_API_KEY, PRIMARY_TOOL_CHOICE, state)}
function logPromptSizeDebug(state: LoopState, errMsg: string, label: string): void { const messageCount = state.conversation.length const estimate = estimatePromptTokens(state.conversation) const charLength = JSON.stringify(state.conversation).length console.warn( `[api] ${label}: ${errMsg} - messages=${messageCount} est_tokens=${estimate} chars=${charLength} observedPromptTokens=${state.contextSize}`, )}
async function recoverFromPromptTooLarge(state: LoopState, attempt: number): Promise<boolean> { const beforeEstimate = estimatePromptTokens(state.conversation) const summaryProvider = configuredSummaryProvider() if (!summaryProvider.model) { console.warn(`[context] recovery: no summary model available; cannot llm-summarize`) return false }
console.warn(`[context] recovery: attempting llm summarization via ${summaryProvider.model.id} (attempt=${attempt + 1})`) const agentContext = await loadAgentSummaryContext() const summarized = await summarizeConversationViaLLM( state.conversation, summaryProvider.model, summaryProvider.apiKey, { agentContext }, ) if (!summarized) { console.warn(`[context] recovery: llm summarization returned no changes`) return false }
const afterEstimate = estimatePromptTokens(summarized) if (afterEstimate >= beforeEstimate) { console.warn(`[context] recovery: llm summary not smaller (${beforeEstimate} -> ${afterEstimate}); keeping original`) return false }
state.conversation = summarized state.contextSize = afterEstimate
const summaryIdx = findSummaryMessageIndex(state.conversation) const summary = summaryIdx >= 0 ? (state.conversation[summaryIdx]?.role === "user" ? (typeof (state.conversation[summaryIdx] as { content: unknown }).content === "string" ? (state.conversation[summaryIdx] as { content: string }).content : undefined) : undefined) : undefined
console.warn(`[context] recovery: llm-summarized conversation (${beforeEstimate} -> ${afterEstimate} tokens)`) recordMetric({ type: "compaction", before: beforeEstimate, after: afterEstimate, method: "force-llm", summary }) return true}
/** * Fetches the next assistant completion, including fallback and backoff behavior. */export async function fetchCompletion( state: LoopState, _baseState: LoopState = state,): Promise<CompletionTurnResult> { let promptTooLargeAttempts = 0 let imagesScrubbed = false
const recoverByScrubbingImages = (errMsg: string, label: string): boolean => { if (imagesScrubbed) return false if (!isContentFilterError(errMsg) && !isImageParseError(errMsg)) return false const reason = isContentFilterError(errMsg) ? "content-filter rejection" : "image parse/format rejection" const scrubbed = scrubImagesFromConversation(state.conversation) imagesScrubbed = true if (scrubbed > 0) { console.warn(`[api] ${label} ${reason}; scrubbed ${scrubbed} image attachment(s) and retrying`) return true } console.warn(`[api] ${label} ${reason} but no images found to scrub`) return false }
while (true) { state.conversation = sanitizeMessages(state.conversation)
const requestState = state.memoryRecallPending ? await (async () => { const ctx = await buildCompletionMessages(state.conversation, state.memoryRecallCooldowns, state.memoryRecallTurn) return { ...state, conversation: ctx.messages, _recalledChunkIds: ctx.recalledChunkIds } })() : { ...state, _recalledChunkIds: [] as number[] }
const recalledChunkIds = (requestState as typeof requestState & { _recalledChunkIds: number[] })._recalledChunkIds
if (USE_FALLBACK) { const fallbackWindow = fallbackContextWindow(requestState.conversation) if (fallbackWindow.nearLimit) { console.warn(`[fallback] prompt estimate ${fallbackWindow.estimate} nearing fallback limit ${fallbackWindow.softLimit} (${FALLBACK_MODEL})`) }
try { const completion = await createFallbackCompletion(requestState) state.memoryRecallCooldowns = rememberRecalledMemoryChunks(state.memoryRecallCooldowns, recalledChunkIds, state.memoryRecallTurn) state.memoryRecallPending = false return completion } catch (fallbackErr) { const errMsg = errorSummary(fallbackErr) if (isPromptTooLargeError(errMsg) && promptTooLargeAttempts < 2) { logPromptSizeDebug(state, errMsg, `fallback rejected prompt (attempt ${promptTooLargeAttempts + 1}/2)`) const recovered = await recoverFromPromptTooLarge(state, promptTooLargeAttempts) promptTooLargeAttempts++ if (recovered) continue } if (recoverByScrubbingImages(errMsg, "fallback")) continue if (shouldFallback(errMsg)) { console.warn(`[fallback] transient failure (${errMsg}); retrying after ${Math.ceil(retryDelayMs / 1000)}s`) await sleep(retryDelayMs) continue } console.error(`[api] fallback error: ${errMsg} model=${FALLBACK_MODEL} api=${FALLBACK_BASE}`) throw fallbackErr } }
try { const completion = await createPrimaryCompletion(requestState) state.memoryRecallCooldowns = rememberRecalledMemoryChunks(state.memoryRecallCooldowns, recalledChunkIds, state.memoryRecallTurn) state.memoryRecallPending = false return completion } catch (primaryErr) { const errMsg = errorSummary(primaryErr) if (isPromptTooLargeError(errMsg) && promptTooLargeAttempts < 2) { logPromptSizeDebug(state, errMsg, `primary rejected prompt (attempt ${promptTooLargeAttempts + 1}/2)`) const recovered = await recoverFromPromptTooLarge(state, promptTooLargeAttempts) promptTooLargeAttempts++ if (recovered) continue throw primaryErr }
if (recoverByScrubbingImages(errMsg, "primary")) continue
if (!shouldFallback(errMsg)) { console.error(`[api] primary error: ${errMsg} model=${MODEL} api=${API_BASE}`) throw primaryErr }
const fallbackWindow = fallbackContextWindow(requestState.conversation) if (fallbackWindow.skip) { console.warn(`[api] primary down (${errMsg}) and fallback context estimate ${fallbackWindow.estimate} exceeds hard limit; retrying primary after backoff`) await sleep(retryDelayMs) continue }
if (fallbackWindow.nearLimit) { console.warn(`[fallback] prompt estimate ${fallbackWindow.estimate} nearing fallback limit ${fallbackWindow.softLimit} (${FALLBACK_MODEL})`) }
console.warn(`[api] primary down (${errMsg}) - switching to fallback`) try { const completion = await createFallbackCompletion(requestState) state.memoryRecallCooldowns = rememberRecalledMemoryChunks(state.memoryRecallCooldowns, recalledChunkIds, state.memoryRecallTurn) state.memoryRecallPending = false return completion } catch (fallbackErr) { const fallbackErrMsg = errorSummary(fallbackErr) if (isPromptTooLargeError(fallbackErrMsg) && promptTooLargeAttempts < 2) { logPromptSizeDebug(state, fallbackErrMsg, `fallback rejected prompt during failover (attempt ${promptTooLargeAttempts + 1}/2)`) const recovered = await recoverFromPromptTooLarge(state, promptTooLargeAttempts) promptTooLargeAttempts++ if (recovered) continue } if (recoverByScrubbingImages(fallbackErrMsg, "fallback-failover")) continue console.warn(`[api] fallback failed (${fallbackErrMsg}) after primary failure (${errMsg}); retrying primary after backoff`) await sleep(retryDelayMs) } } }}