Ensure tool call parts are emitted immediately when streaming (#5634)

Author: IMax153

.changeset/social-walls-share.md

@@ -0,0 +1,5 @@
+---
+"@effect/ai": patch
+---
+
+Ensure that tool calls are emitted as soon as possible when streaming

packages/ai/ai/src/LanguageModel.ts

@@ -51,10 +51,12 @@
 import * as Chunk from "effect/Chunk"
 import * as Context from "effect/Context"
 import * as Effect from "effect/Effect"
+import * as Mailbox from "effect/Mailbox"
 import * as Option from "effect/Option"
 import * as ParseResult from "effect/ParseResult"
 import * as Predicate from "effect/Predicate"
 import * as Schema from "effect/Schema"
+import type * as Scope from "effect/Scope"
 import * as Stream from "effect/Stream"
 import type { Span } from "effect/Tracer"
 import type { Concurrency, Mutable, NoExcessProperties } from "effect/Types"
@@ -797,7 +799,9 @@ export const make: (params: ConstructorParams) => Effect.Effect<Service> = Effec
     >(options: Options & GenerateTextOptions<Tools>, providerOptions: Mutable<ProviderOptions>) => Effect.Effect<
       Stream.Stream<Response.StreamPart<Tools>, AiError.AiError | ParseResult.ParseError, IdGenerator>,
       Options extends { readonly toolkit: Effect.Effect<Toolkit.WithHandler<Tools>, infer _E, infer _R> } ? _E : never,
-      Options extends { readonly toolkit: Effect.Effect<Toolkit.WithHandler<Tools>, infer _E, infer _R> } ? _R : never
+      | (Options extends { readonly toolkit: Effect.Effect<Toolkit.WithHandler<Tools>, infer _E, infer _R> } ? _R
+        : never)
+      | Scope.Scope
     > = Effect.fnUntraced(
       function*<
         Tools extends Record<string, Tool.Any>,
@@ -842,16 +846,21 @@ export const make: (params: ConstructorParams) => Effect.Effect<Service> = Effec
           ) as Stream.Stream<Response.StreamPart<Tools>, AiError.AiError | ParseResult.ParseError, IdGenerator>
         }
 
-        const ResponseSchema = Schema.Chunk(Response.StreamPart(toolkit))
+        const mailbox = yield* Mailbox.make<Response.StreamPart<Tools>, AiError.AiError | ParseResult.ParseError>()
+        const ResponseSchema = Schema.Array(Response.StreamPart(toolkit))
         const decode = Schema.decode(ResponseSchema)
-        return params.streamText(providerOptions).pipe(
-          Stream.mapChunksEffect(Effect.fnUntraced(function*(chunk) {
+        yield* params.streamText(providerOptions).pipe(
+          Stream.runForEachChunk(Effect.fnUntraced(function*(chunk) {
             const rawContent = Chunk.toArray(chunk)
-            const toolResults = yield* resolveToolCalls(rawContent, toolkit, options.concurrency)
             const content = yield* decode(rawContent)
-            return Chunk.unsafeFromArray([...content, ...toolResults])
-          }))
-        ) as Stream.Stream<Response.StreamPart<Tools>, AiError.AiError | ParseResult.ParseError, IdGenerator>
+            yield* mailbox.offerAll(content)
+            const toolResults = yield* resolveToolCalls(rawContent, toolkit, options.concurrency)
+            yield* mailbox.offerAll(toolResults as any)
+          })),
+          Mailbox.into(mailbox),
+          Effect.forkScoped
+        )
+        return Mailbox.toStream(mailbox)
       }
     )
 

packages/ai/ai/test/LanguageModel.test.ts

@@ -0,0 +1,86 @@
+import * as LanguageModel from "@effect/ai/LanguageModel"
+import * as Response from "@effect/ai/Response"
+import * as Tool from "@effect/ai/Tool"
+import * as Toolkit from "@effect/ai/Toolkit"
+import { assert, describe, it } from "@effect/vitest"
+import * as Effect from "effect/Effect"
+import * as Schema from "effect/Schema"
+import * as Stream from "effect/Stream"
+import * as TestClock from "effect/TestClock"
+import * as TestUtils from "./utilities.js"
+
+const MyTool = Tool.make("MyTool", {
+  parameters: { testParam: Schema.String },
+  success: Schema.Struct({ testSuccess: Schema.String })
+})
+
+const MyToolkit = Toolkit.make(MyTool)
+
+const MyToolkitLayer = MyToolkit.toLayer({
+  MyTool: () =>
+    Effect.succeed({ testSuccess: "test-success" }).pipe(
+      Effect.delay("10 seconds")
+    )
+})
+
+describe("LanguageModel", () => {
+  describe("streamText", () => {
+    it.effect("should emit tool calls before executing tool handlers", () =>
+      Effect.gen(function*() {
+        const parts: Array<Response.StreamPart<Toolkit.Tools<typeof MyToolkit>>> = []
+        const latch = yield* Effect.makeLatch()
+
+        const toolCallId = "tool-abc123"
+        const toolName = "MyTool"
+        const toolParams = { testParam: "test-param" }
+        const toolResult = { testSuccess: "test-success" }
+
+        yield* LanguageModel.streamText({
+          prompt: [],
+          toolkit: MyToolkit
+        }).pipe(
+          Stream.runForEach((part) =>
+            Effect.andThen(latch.open, () => {
+              parts.push(part)
+            })
+          ),
+          TestUtils.withLanguageModel({
+            streamText: [
+              {
+                type: "tool-call",
+                id: toolCallId,
+                name: toolName,
+                params: toolParams
+              }
+            ]
+          }),
+          Effect.provide(MyToolkitLayer),
+          Effect.fork
+        )
+
+        yield* latch.await
+
+        const toolCallPart = Response.makePart("tool-call", {
+          id: toolCallId,
+          name: toolName,
+          params: toolParams,
+          providerExecuted: false
+        })
+
+        const toolResultPart = Response.toolResultPart({
+          id: toolCallId,
+          name: toolName,
+          result: toolResult,
+          encodedResult: toolResult,
+          isFailure: false,
+          providerExecuted: false
+        })
+
+        assert.deepStrictEqual(parts, [toolCallPart])
+
+        yield* TestClock.adjust("10 seconds")
+
+        assert.deepStrictEqual(parts, [toolCallPart, toolResultPart])
+      }))
+  })
+})