fix(@effect/ai-amazon-bedrock): emit finish part after metadata so streaming usage is populated

[kn0ll]

Type

• Bug Fix

Description

In makeStreamResponse, the messageStop event was emitting the finish part with a reference to a shared mutable usage object. The metadata event (which carries actual token counts) arrives after messageStop in the Bedrock Converse stream, but by that point @effect/ai's LanguageModel.streamText has already run Schema.decode on the finish part, creating a new Usage class instance from the values at decode time (all undefined).

The fix defers emitting the finish part: messageStop now stores the finish reason in a closure variable, and the metadata handler emits the finish part after populating usage.

Related

• Closes @effect/ai-amazon-bedrock: streaming token usage is always empty #6237

[changeset-bot]

⚠️ No Changeset found

Latest commit: 613527b

Merging this PR will not cause a version bump for any packages. If these changes should not result in a new version, you're good to go. If these changes should result in a version bump, you need to add a changeset.

This PR includes no changesets

When changesets are added to this PR, you'll see the packages that this PR includes changesets for and the associated semver types

Click here to learn what changesets are, and how to add one.

Click here if you're a maintainer who wants to add a changeset to this PR

[IMax153]

@kn0ll before I approve this, can you please provide a minimal reproduction program I can run to test the behavior you are referring to.

[kn0ll]

@IMax153 of course!

i got too clever with my issue & PR organization.

import * as LanguageModel from "@effect/ai/LanguageModel"
import * as AmazonBedrockClient from "@effect/ai-amazon-bedrock/AmazonBedrockClient"
import * as AmazonBedrockLanguageModel from "@effect/ai-amazon-bedrock/AmazonBedrockLanguageModel"
import * as NodeHttpClient from "@effect/platform-node/NodeHttpClient"
import { Effect, Layer, Redacted, Stream } from "effect"

const MainLive = AmazonBedrockLanguageModel.layer({
  model: "anthropic.claude-3-haiku-20240307-v1:0"
}).pipe(
  Layer.provide(AmazonBedrockClient.layer({
    accessKeyId: process.env.AWS_ACCESS_KEY_ID!,
    secretAccessKey: Redacted.make(process.env.AWS_SECRET_ACCESS_KEY!),
    region: "us-east-1"
  })),
  Layer.provide(NodeHttpClient.layerUndici)
)

const program = LanguageModel.streamText({ prompt: "Say hello in one sentence." }).pipe(
  Stream.runFold(undefined as any, (_acc, part) => part.type === "finish" ? part : _acc),
  Effect.andThen((finish) => {
    console.log(`inputTokens: ${finish.usage.inputTokens}, outputTokens: ${finish.usage.outputTokens}`)
  }),
  Effect.scoped,
  Effect.provide(MainLive)
)

Effect.runPromise(program)

this script produces undefined for inputTokens and outputTokens. applying the patch returns the correct numbers! :)

thank you for all your hard work!