From 9b9c73a8c90f581d66c90a6b4f70c6d1c171ac73 Mon Sep 17 00:00:00 2001
From: iseki <admin@iseki.space>
Date: Sun, 12 Jan 2025 13:49:09 +0800
Subject: [PATCH] feat(generate-text): add toolCalls, toolResults (#22)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

* feat(generate-text): add toolCalls and toolResults

* Update packages/generate-object/src/index.ts

Co-authored-by: 藍+85CD <50108258+kwaa@users.noreply.github.com>
---
 packages/generate-object/src/index.ts |  3 +-
 packages/generate-text/src/index.ts   | 55 ++++++++++++++++++++++++---
 2 files changed, 51 insertions(+), 7 deletions(-)

diff --git a/packages/generate-object/src/index.ts b/packages/generate-object/src/index.ts
index 135ba6c..6b29fe3 100644
--- a/packages/generate-object/src/index.ts
+++ b/packages/generate-object/src/index.ts
@@ -8,7 +8,8 @@ export interface GenerateObjectOptions<T extends Schema> extends GenerateTextOpt
   schemaName?: string
 }
 
-export interface GenerateObjectResult<T extends Schema> extends Omit<GenerateTextResult, 'text'> {
+// TODO: toolCalls, toolResults
+export interface GenerateObjectResult<T extends Schema> extends Omit<GenerateTextResult, 'text' | 'toolCalls' | 'toolResults'> {
   object: Infer<T>
 }
 
diff --git a/packages/generate-text/src/index.ts b/packages/generate-text/src/index.ts
index 67dfd89..ad0679b 100644
--- a/packages/generate-text/src/index.ts
+++ b/packages/generate-text/src/index.ts
@@ -34,10 +34,26 @@ export interface GenerateTextResponseUsage {
   total_tokens: number
 }
 
+export interface ToolCall {
+  args: string
+  toolCallId: string
+  toolCallType: 'function'
+  toolName: string
+}
+
+export interface ToolResult {
+  args: Record<string, unknown>
+  result: string
+  toolCallId: string
+  toolName: string
+}
+
 export interface GenerateTextResult {
   finishReason: FinishReason
   steps: StepResult[]
   text?: string
+  toolCalls: ToolCall[]
+  toolResults: ToolResult[]
   usage: GenerateTextResponseUsage
 }
 
@@ -45,8 +61,8 @@ export interface StepResult {
   text?: string
   // TODO: step type
   // type: 'continue' | 'initial' | 'tool-result'
-  // TODO: toolCalls
-  // TODO: toolResults
+  toolCalls: ToolCall[]
+  toolResults: ToolResult[]
   usage: GenerateTextResponseUsage
 }
 
@@ -63,7 +79,8 @@ export const generateText = async (options: GenerateTextOptions): Promise<Genera
 
   const steps: StepResult[] = []
   const messages: Message[] = options.messages
-
+  const toolCalls: ToolCall[] = []
+  const toolResults: ToolResult[] = []
   while (currentStep < (options.maxSteps ?? 1)) {
     currentStep += 1
 
@@ -80,11 +97,13 @@ export const generateText = async (options: GenerateTextOptions): Promise<Genera
     text = message.content
     usage = data.usage
 
-    steps.push({
+    const stepResult: StepResult = {
       text: message.content,
+      toolCalls: [],
+      toolResults: [],
       // type: 'initial',
       usage,
-    })
+    }
 
     // TODO: fix types
     messages.push({ ...message, content: message.content! })
@@ -93,7 +112,8 @@ export const generateText = async (options: GenerateTextOptions): Promise<Genera
       // execute tools
       for (const toolCall of message.tool_calls ?? []) {
         const tool = (options.tools as Tool[]).find(tool => tool.function.name === toolCall.function.name)!
-        const toolResult = await tool.execute(JSON.parse(toolCall.function.arguments))
+        const parsedArgs: Record<string, any> = JSON.parse(toolCall.function.arguments)
+        const toolResult = await tool.execute(parsedArgs)
         const toolMessage = {
           content: toolResult,
           role: 'tool',
@@ -101,13 +121,34 @@ export const generateText = async (options: GenerateTextOptions): Promise<Genera
         } satisfies Message
 
         messages.push(toolMessage)
+
+        const toolCallData = {
+          args: toolCall.function.arguments,
+          toolCallId: toolCall.id,
+          toolCallType: toolCall.type,
+          toolName: toolCall.function.name,
+        }
+        toolCalls.push(toolCallData)
+        stepResult.toolCalls.push(toolCallData)
+        const toolResultData = {
+          args: parsedArgs,
+          result: toolResult,
+          toolCallId: toolCall.id,
+          toolName: toolCall.function.name,
+        }
+        toolResults.push(toolResultData)
+        stepResult.toolResults.push(toolResultData)
       }
+      steps.push(stepResult)
     }
     else {
+      steps.push(stepResult)
       return {
         finishReason: finish_reason,
         steps,
         text: message.content,
+        toolCalls,
+        toolResults,
         usage,
       }
     }
@@ -117,6 +158,8 @@ export const generateText = async (options: GenerateTextOptions): Promise<Genera
     finishReason,
     steps,
     text,
+    toolCalls,
+    toolResults,
     usage,
   }
 }