diff --git a/package.json b/package.json index 03e570c..3a325e3 100644 --- a/package.json +++ b/package.json @@ -1,6 +1,6 @@ { "name": "@ztimson/ai-utils", - "version": "1.6.5", + "version": "1.6.6", "description": "AI Utility library", "author": "Zak Timson", "license": "MIT", diff --git a/src/ai.ts b/src/ai.ts index 4f1ee35..9609d5b 100644 --- a/src/ai.ts +++ b/src/ai.ts @@ -4,7 +4,7 @@ import { Audio } from './audio.ts'; import {Vision} from './vision.ts'; export type AbortablePromise = Promise & { - abort: () => any + abort: (keep?: boolean) => any }; export type AiOptions = { diff --git a/src/llm.ts b/src/llm.ts index 9255dba..2669aef 100644 --- a/src/llm.ts +++ b/src/llm.ts @@ -265,7 +265,7 @@ class LLM { }; } - private setupAgent(agents: Agent[] = [], allAgents: Agent[], history: LLMMessage[], aborts: (() => void)[], depth = 0, delegateState: {resp: string | null}): AiTool[] { + private setupAgent(agents: Agent[] = [], allAgents: Agent[], history: LLMMessage[], aborts: ((keep?: boolean) => void)[], depth = 0, delegateState: {resp: string | null}): AiTool[] { return agents.map(a => { const toolName = `${a.delegate ? '' : 'sub'}agent_${snakeCase(a.name)}`; return { @@ -397,11 +397,13 @@ ${a.system}`, if(!this.models[m]) throw new Error(`Model does not exist: ${m}`); let request: AbortablePromise | null = null; let aborted = false; - const nestedAborts: (() => void)[] = []; - const abort = () => { + let keepOnAbort = true; + const nestedAborts: ((keep?: boolean) => void)[] = []; + const abort = (keep = true) => { aborted = true; - request?.abort?.(); - nestedAborts.forEach(a => a()); + keepOnAbort = keep; + request?.abort?.(keep); + nestedAborts.forEach(a => a(keep)); }; let promise: any; @@ -411,9 +413,25 @@ ${a.system}`, let tools: AiTool[] = options.tools || this.ai.options.llm?.tools || []; const prompts: string[] = []; let history = options.history || []; + const historyStart = history.length; const files = options.files || []; if(message || files.length) history.push({role: 'user', content: message || '', timestamp: Date.now()}); + // Accumulate streamed text so it can be committed to history if aborted mid-generation + let partialText = ''; + const onStream = options.stream; + const stream = (chunk: {text?: string, tool?: string, done?: true}) => { + if(chunk.text) partialText += chunk.text; + return onStream?.(chunk); + }; + + /** Commit (keep) or discard this turn's progress on abort, then throw */ + const abortNow = (): never => { + if(keepOnAbort) { if(partialText) history.push({role: 'assistant', content: partialText, timestamp: Date.now()}); } + else history.splice(historyStart, history.length - historyStart); + throw Object.assign(new Error('Aborted'), {name: 'AbortError'}); + }; + // MCP const mcp = options.mcp || this.ai.options?.llm?.mcp; if(mcp?.length) { @@ -441,8 +459,8 @@ ${a.system}`, const mems = mem.memory instanceof MemoryCache ? mem.memory.memories : mem.memory; if(mems.length) { if(mem.inject) { - const pool = 15; // candidates considered, cheap since only refs are listed - const budget = mem.maxTokens ?? 2000; // actual content injected + const pool = 15; + const budget = mem.maxTokens ?? 2000; const relevant = await this.memoryManager.recollect(message, mem.memory, pool); let used = 0; @@ -481,7 +499,7 @@ Linked: ${makeUnique([...r.links, ...r.backlinks]).join(', ')} } } - if(aborted) throw Object.assign(new Error('Aborted'), {name: 'AbortError'}); + if(aborted) abortNow(); const lastMsg = history[history.length - 1]; if(files.length && lastMsg?.role === 'user') lastMsg.files = files; @@ -500,11 +518,17 @@ Linked: ${makeUnique([...r.links, ...r.backlinks]).join(', ')} const toolTimings = new Map(); tools = this.wrapToolTiming(tools, toolTimings); - if(aborted) throw Object.assign(new Error('Aborted'), {name: 'AbortError'}); + if(aborted) abortNow(); prompts.unshift(options.system || this.ai.options.llm?.system || ''); - request = this.models[m].ask('', {...options, tools, system: prompts.filter(Boolean).join('\n\n')}); - let resp = await request; + request = this.models[m].ask('', {...options, tools, stream, system: prompts.filter(Boolean).join('\n\n')}); + let resp: string; + try { + resp = await request; + } catch(err: any) { + if(aborted) return abortNow(); + throw err; + } // Strip the file injection shim restores.forEach(({msg, content}) => msg.content = content);