Tools can now emit "done" event and end chat early gracefully

package.json

@@ -1,6 +1,6 @@
 {
 	"name": "@ztimson/ai-utils",
-	"version": "1.2.11",
+	"version": "1.2.12",
 	"description": "AI Utility library",
 	"author": "Zak Timson",
 	"license": "MIT",

src/antrhopic.ts

@@ -21,10 +21,10 @@ export class Anthropic extends LLMProvider {
 				messages.push(<any>{timestamp, ...h});
 			} else {
 				const textContent = h.content?.filter((c: any) => c.type == 'text').map((c: any) => c.text).join('\n\n');
-				if(textContent) messages.push({timestamp, role: h.role, content: textContent});
+				if(textContent) messages.push({role: h.role, content: textContent, timestamp: timestamp});
 				h.content.forEach((c: any) => {
 					if(c.type == 'tool_use') {
-						messages.push({timestamp, role: 'tool', id: c.id, name: c.name, args: c.input, content: undefined});
+						messages.push({role: 'tool', id: c.id, name: c.name, args: c.input, timestamp: c.timestamp, content: undefined});
 					} else if(c.type == 'tool_result') {
 						const m: any = messages.findLast(m => (<any>m).id == c.tool_use_id);
 						if(m) m[c.is_error ? 'error' : 'content'] = c.content;
@@ -46,7 +46,7 @@ export class Anthropic extends LLMProvider {
 				i++;
 			}
 		}
-		return history.map(({timestamp, ...h}) => h);
+		return history;
 	}
 
 	ask(message: string, options: LLMRequest = {}): AbortablePromise<string | any> {
@@ -83,8 +83,9 @@ export class Anthropic extends LLMProvider {
 				};
 			}
 
-			let resp: any, isFirstMessage = true;
+			let resp: any, isFirstMessage = true, terminal = false;
 			do {
+				requestParams.messages = history.map(({timestamp, ...m}) => m);
 				resp = await this.client.messages.create(requestParams).catch(err => {
 					err.message += `\n\nMessages:\n${JSON.stringify(history, null, 2)}`;
 					throw err;
@@ -113,7 +114,7 @@ export class Anthropic extends LLMProvider {
 							}
 						} else if(chunk.type === 'content_block_stop') {
 							const last = resp.content.at(-1);
-							if(last.input != null) last.input = last.input ? JSONAttemptParse(last.input, {}) : {};
+							if(last?.input != null) last.input = last.input ? JSONAttemptParse(last.input, {}) : {};
 						} else if(chunk.type === 'message_stop') {
 							break;
 						}
@@ -123,31 +124,35 @@ export class Anthropic extends LLMProvider {
 				// Run tools
 				const toolCalls = resp.content.filter((c: any) => c.type === 'tool_use');
 				if(toolCalls.length && !controller.signal.aborted) {
-					history.push({role: 'assistant', content: resp.content});
+					history.push({role: 'assistant', content: resp.content, timestamp: Date.now()});
 					const results = await Promise.all(toolCalls.map(async (toolCall: any) => {
 						const tool = tools.find(findByProp('name', toolCall.name));
 						if(options.stream) options.stream({tool: toolCall.name});
 						if(!tool) return {tool_use_id: toolCall.id, is_error: true, content: 'Tool not found'};
 						try {
-							const result = await tool.fn(toolCall.input, options?.stream, this.ai);
+							// Wrap stream so a tool's `done` ends turn gracefully
+							const toolStream = options.stream && ((chunk: any) => {
+								if(chunk.done) { terminal = true; return; }
+								options.stream!(chunk);
+							});
+							const result = await tool.fn(toolCall.input, toolStream, this.ai);
 							return {type: 'tool_result', tool_use_id: toolCall.id, content: typeof result == 'object' ? JSONSanitize(result) : result};
 						} catch (err: any) {
 							return {type: 'tool_result', tool_use_id: toolCall.id, is_error: true, content: err?.message || err?.toString() || 'Unknown'};
 						}
 					}));
-					history.push({role: 'user', content: results});
+					history.push({role: 'user', content: results, timestamp: Date.now()});
 					requestParams.messages = history;
 				}
-			} while (!controller.signal.aborted && resp.content.some((c: any) => c.type === 'tool_use'));
+			} while (!terminal && !controller.signal.aborted && resp.content.some((c: any) => c.type === 'tool_use'));
 
-			const textContent = resp.content.filter((c: any) => c.type == 'text').map((c: any) => c.text).join('\n\n');
-			history.push({role: 'assistant', content: textContent});
+			if(!terminal) {
+				const textContent = resp.content.filter((c: any) => c.type == 'text').map((c: any) => c.text).join('\n\n');
+				history.push({role: 'assistant', content: textContent, timestamp: Date.now()});
+			}
 			history = this.toStandard(history);
-
 			if(options.stream) options.stream({done: true});
 			if(options.history) options.history.splice(0, options.history.length, ...history);
-
-			// Return parsed JSON if schema provided
 			const finalContent = history.at(-1)?.content;
 			res(options.schema ? JSONAttemptParse(finalContent, finalContent) : finalContent);
 		}), {abort: () => controller.abort()});

src/open-ai.ts

@@ -51,15 +51,16 @@ export class OpenAi extends LLMProvider {
 					content: null,
 					tool_calls: [{ id: h.id, type: 'function', function: { name: h.name, arguments: JSON.stringify(h.args) } }],
 					refusal: null,
-					annotations: []
+					annotations: [],
+					timestamp: h.timestamp,
 				}, {
 					role: 'tool',
 					tool_call_id: h.id,
-					content: h.error || h.content
+					content: h.error || h.content,
+					timestamp: h.timestamp,
 				});
 			} else {
-				const {timestamp, ...rest} = h;
-				result.push(rest);
+				result.push(h);
 			}
 			return result;
 		}, [] as any[]);
@@ -106,8 +107,9 @@ export class OpenAi extends LLMProvider {
 				};
 			}
 
-			let resp: any, isFirstMessage = true;
+			let resp: any, isFirstMessage = true, terminal = false;
 			do {
+				requestParams.messages = history.map(({timestamp, ...m}) => m);
 				resp = await this.client.chat.completions.create(requestParams).catch(err => {
 					err.message += `\n\nMessages:\n${JSON.stringify(history, null, 2)}`;
 					throw err;
@@ -116,7 +118,7 @@ export class OpenAi extends LLMProvider {
 				if(options.stream) {
 					if(!isFirstMessage) options.stream({text: '\n\n'});
 					else isFirstMessage = false;
-					resp.choices = [{message: {role: 'assistant', content: '', tool_calls: []}}];
+					resp.choices = [{message: {role: 'assistant', content: '', tool_calls: [], timestamp: Date.now()}}];
 					for await (const chunk of resp) {
 						if(controller.signal.aborted) break;
 						if(chunk.choices[0].delta.content) {
@@ -158,28 +160,32 @@ export class OpenAi extends LLMProvider {
 					const results = await Promise.all(toolCalls.map(async (toolCall: any) => {
 						const tool = tools?.find(findByProp('name', toolCall.function.name));
 						if(options.stream) options.stream({tool: toolCall.function.name});
-						if(!tool) return {role: 'tool', tool_call_id: toolCall.id, content: '{"error": "Tool not found"}'};
+						if(!tool) return {role: 'tool', tool_call_id: toolCall.id, content: '{"error": "Tool not found"}', timestamp: Date.now()};
 						try {
 							const args = JSONAttemptParse(toolCall.function.arguments, {});
-							const result = await tool.fn(args, options.stream, this.ai);
-							return {role: 'tool', tool_call_id: toolCall.id, content: typeof result == 'object' ? JSONSanitize(result) : result};
+							// Wrap stream so a tool's `done` ends turn gracefully
+							const toolStream = options.stream && ((chunk: any) => {
+								if(chunk.done) { terminal = true; return; }
+								options.stream!(chunk);
+							});
+							const result = await tool.fn(args, toolStream, this.ai);
+							return {role: 'tool', tool_call_id: toolCall.id, content: typeof result == 'object' ? JSONSanitize(result) : result, timestamp: Date.now()};
 						} catch (err: any) {
-							return {role: 'tool', tool_call_id: toolCall.id, content: JSONSanitize({error: err?.message || err?.toString() || 'Unknown'})};
+							return {role: 'tool', tool_call_id: toolCall.id, content: JSONSanitize({error: err?.message || err?.toString() || 'Unknown'}), timestamp: Date.now()};
 						}
 					}));
 					history.push(...results);
 					requestParams.messages = history;
 				}
-			} while (!controller.signal.aborted && resp.choices?.[0]?.message?.tool_calls?.length);
+			} while (!terminal && !controller.signal.aborted && resp.choices?.[0]?.message?.tool_calls?.length);
 
-			const textContent = resp.choices[0].message.content?.trim() || '';
-			history.push({role: 'assistant', content: textContent});
+			if(!terminal) {
+				const textContent = resp.choices[0].message.content?.trim() || '';
+				history.push({role: 'assistant', content: textContent, timestamp: Date.now()});
+			}
 			history = this.toStandard(history);
-
 			if(options.stream) options.stream({done: true});
 			if(options.history) options.history.splice(0, options.history.length, ...history);
-
-			// Return parsed JSON if schema provided
 			const finalContent = history.at(-1)?.content;
 			res(options.schema ? JSONAttemptParse(finalContent, finalContent) : finalContent);
 		}), {abort: () => controller.abort()});