preserve tool call results in AiResponse.merge (#4866)

Author: tim-smart

.changeset/tricky-months-sit.md

@@ -0,0 +1,5 @@
+---
+"@effect/ai": patch
+---
+
+preserve tool call results in AiResponse.merge

packages/ai/ai/src/AiLanguageModel.ts

@@ -358,14 +358,14 @@ export const make: <Config>(
       if (Option.isNone(spanTransformer)) {
         return stream
       }
-      let lastResponse: AiResponse.AiResponse | undefined
+      let finalResponse = AiResponse.empty
       return stream.pipe(
         Stream.map((response) => {
-          lastResponse = response
+          finalResponse = AiResponse.merge(finalResponse, response)
           return response
         }),
         Stream.ensuring(Effect.sync(() => {
-          spanTransformer.value({ ...options, response: lastResponse! })
+          spanTransformer.value({ ...options, response: finalResponse })
         }))
       )
     })),
@@ -470,49 +470,49 @@ const getDescription = (ast: AST.AST): string => {
   return AST.DescriptionAnnotationId in annotations ? annotations[AST.DescriptionAnnotationId] as string : ""
 }
 
-const resolveParts = <Tools extends AiTool.Any>(options: {
+const resolveParts = Effect.fnUntraced(function*<Tools extends AiTool.Any>(options: {
   readonly response: AiResponse.AiResponse
   readonly toolkit: AiToolkit.ToHandler<Tools>
   readonly concurrency: Concurrency | undefined
   readonly method: string
-}) =>
-  Effect.gen(function*() {
-    const toolNames: Array<string> = []
-    const toolParts = options.response.parts.filter(
-      (part): part is AiResponse.ToolCallPart => {
-        if (part._tag === "ToolCallPart") {
-          toolNames.push(part.name)
-          return true
-        }
-        return false
+}) {
+  const toolNames: Array<string> = []
+  const toolParts = options.response.parts.filter(
+    (part): part is AiResponse.ToolCallPart => {
+      if (part._tag === "ToolCallPart") {
+        toolNames.push(part.name)
+        return true
       }
-    )
-    if (toolParts.length === 0) {
-      return new AiResponse.WithToolCallResults({
-        parts: options.response.parts,
-        results: constEmptyMap,
-        encodedResults: constEmptyMap
-      }, constDisableValidation)
+      return false
     }
-    yield* Effect.annotateCurrentSpan("toolCalls", toolNames)
-    const results = new Map<AiResponse.ToolCallId, AiTool.Success<Tools>>()
-    const encodedResults = new Map<AiResponse.ToolCallId, unknown>()
-    yield* Effect.forEach(toolParts, (part) => {
-      const id = part.id as AiResponse.ToolCallId
-      const name = part.name as AiTool.Name<Tools>
-      const params = part.params as AiTool.Parameters<Tools>
-      const toolCall = options.toolkit.handle(name, params)
-      return Effect.map(toolCall, ({ encodedResult, result }) => {
-        results.set(id, result)
-        encodedResults.set(id, encodedResult)
-      })
-    }, { concurrency: options.concurrency, discard: true })
+  )
+  if (toolParts.length === 0) {
     return new AiResponse.WithToolCallResults({
       parts: options.response.parts,
-      results,
-      encodedResults
+      results: constEmptyMap,
+      encodedResults: constEmptyMap
     }, constDisableValidation)
-  })
+  }
+  yield* Effect.annotateCurrentSpan("toolCalls", toolNames)
+  const results = new Map<AiResponse.ToolCallId, AiTool.Success<Tools>>()
+  const encodedResults = new Map<AiResponse.ToolCallId, unknown>()
+  const resolve = Effect.forEach(toolParts, (part) => {
+    const id = part.id as AiResponse.ToolCallId
+    const name = part.name as AiTool.Name<Tools>
+    const params = part.params as AiTool.Parameters<Tools>
+    const toolCall = options.toolkit.handle(name, params)
+    return Effect.map(toolCall, ({ encodedResult, result }) => {
+      results.set(id, result)
+      encodedResults.set(id, encodedResult)
+    })
+  }, { concurrency: options.concurrency, discard: true })
+  yield* resolve
+  return new AiResponse.WithToolCallResults({
+    parts: options.response.parts,
+    results,
+    encodedResults
+  }, constDisableValidation)
+})
 
 /**
  * Generate text using a large language model for the specified `prompt`.

packages/ai/ai/src/AiResponse.ts

@@ -757,15 +757,56 @@ export const merge: {
   (other: AiResponse) => (self: AiResponse) => AiResponse,
   (self: AiResponse, other: AiResponse) => AiResponse
 >(2, (self, other) => {
-  if (other.parts.length === 0) {
-    return new AiResponse({
-      parts: self.parts
+  const newParts = mergeParts(self, other)
+  if (hasToolCallResults(self) && hasToolCallResults(other)) {
+    return new WithToolCallResults({
+      results: new Map([...self.results, ...other.results]) as any,
+      encodedResults: new Map([...self.encodedResults, ...other.encodedResults]),
+      parts: newParts
+    }, constDisableValidation)
+  } else if (hasToolCallResults(self)) {
+    return new WithToolCallResults({
+      results: self.results as any,
+      encodedResults: self.encodedResults,
+      parts: newParts
+    }, constDisableValidation)
+  } else if (hasToolCallResults(other)) {
+    return new WithToolCallResults({
+      results: other.results as any,
+      encodedResults: other.encodedResults,
+      parts: newParts
+    }, constDisableValidation)
+  } else if (isStructured(self) && isStructured(other)) {
+    return new WithStructuredOutput({
+      id: self.id,
+      name: self.name,
+      value: other.value,
+      parts: newParts
+    }, constDisableValidation)
+  } else if (isStructured(self)) {
+    return new WithStructuredOutput({
+      id: self.id,
+      name: self.name,
+      value: self.value,
+      parts: newParts
+    }, constDisableValidation)
+  } else if (isStructured(other)) {
+    return new WithStructuredOutput({
+      id: other.id,
+      name: other.name,
+      value: other.value,
+      parts: newParts
     }, constDisableValidation)
   }
+  return new AiResponse({ parts: newParts }, constDisableValidation)
+})
+
+const mergeParts = (self: AiResponse, other: AiResponse): ReadonlyArray<Part> => {
+  if (other.parts.length === 0) {
+    return self.parts
+  }
   if (self.parts.length === 0) {
-    return new AiResponse({
-      parts: other.parts
-    }, constDisableValidation)
+    return other.parts
   }
   const lastPart = self.parts[self.parts.length - 1]
   const newParts: Array<Part> = []
@@ -778,10 +819,8 @@ export const merge: {
   if (text.length > 0) {
     newParts.push(new TextPart({ text }, constDisableValidation))
   }
-  return newParts.length === 0 ? self : new AiResponse({
-    parts: [...self.parts.slice(0, self.parts.length - 1), ...newParts]
-  }, constDisableValidation)
-})
+  return newParts.length === 0 ? self.parts : [...self.parts.slice(0, self.parts.length - 1), ...newParts]
+}
 
 /**
  * Adds the specified tool calls to the provided `AiResponse`.