packages/ai/src/loop.ts
Rama dev · réplica de lectura del repositorio original.
import { Effect, Stream } from "effect"
import {
AIError,
AssistantMessage,
InvalidProviderOutputError,
LLMResponse,
ResponseFold,
Message,
Usage,
LLMEvent,
ToolCallPart,
ToolChoice,
type Finish,
type FinishReasonDetails,
type LLMRequest,
type ToolInputError,
} from "./schema/index.js"
import type { DispatchResult, PendingCall } from "./tool-runtime.js"
/**
* Client calls to answer before the next step. A call whose input was not valid JSON carries the error the model
* gets back, so it can retry. A step is final when it calls `answerTool`, even alongside other tools, or when it hits
* the output token limit, which may have cut its calls off.
*/
export const pendingCalls = (response: LLMResponse, answerTool?: string): ReadonlyArray<PendingCall> => {
if (response.finishReason.normalized === "length") return []
if (answerTool !== undefined && response.toolCalls.some((call) => call.name === answerTool)) return []
return [
...response.toolCalls.filter((call) => call.providerExecuted !== true).map((call) => ({ call })),
...response.events.filter(LLMEvent.is.toolInputError).map((event) => ({
call: invalidCall(event),
error: `Invalid tool arguments: the input is not valid JSON. Received: ${event.raw}`,
})),
]
}
/** The call to replay for input that was not valid JSON. Providers require an object input, so it is empty. */
const invalidCall = (event: ToolInputError) =>
ToolCallPart.make({ id: event.id, name: event.name, namespace: event.namespace, input: {} })
/** Add the replayed calls for invalid input to the assistant message, so each tool result has its call. */
const withInvalidCalls = (message: AssistantMessage, calls: ReadonlyArray<ToolCallPart>) =>
calls.length === 0 ? message : new AssistantMessage({ ...message, content: [...message.content, ...calls] })
/** The step's assistant message followed by one tool message per call, in call order. */
export const stepMessages = (
response: LLMResponse,
calls: ReadonlyArray<PendingCall>,
results: ReadonlyArray<DispatchResult>,
): ReadonlyArray<Message> => [
withInvalidCalls(
response.message,
calls.flatMap((item) => (item.error === undefined ? [] : [item.call])),
),
...calls.map((item, i) =>
Message.tool({
id: item.call.id,
name: item.call.name,
namespace: item.call.namespace,
result: results[i].result,
providerMetadata: item.call.providerMetadata,
}),
),
]
/**
* The request for the step after `response`, with its messages and tool results appended. A caller's forced
* `toolChoice` applies to the first step only; with `keepForcedChoice` (the answer tool's) it stays.
*/
export const nextRequest = (
request: LLMRequest,
response: LLMResponse,
calls: ReadonlyArray<PendingCall>,
results: ReadonlyArray<DispatchResult>,
keepForcedChoice: boolean,
): LLMRequest => ({
...request,
toolChoice:
!keepForcedChoice && (request.toolChoice?.type === "required" || request.toolChoice?.type === "tool")
? new ToolChoice({ type: "auto", disableParallelToolUse: request.toolChoice.disableParallelToolUse })
: request.toolChoice,
messages: [...request.messages, ...stepMessages(response, calls, results)],
})
/** Sum usage across steps. */
export const addUsage = (total: Usage | undefined, step: Usage | undefined) => {
if (total === undefined || step === undefined) return total ?? step
const sum = (
key:
| "inputTokens"
| "outputTokens"
| "nonCachedInputTokens"
| "cacheReadInputTokens"
| "cacheWriteInputTokens"
| "reasoningTokens"
| "totalTokens"
| "cost",
) => (total[key] === undefined && step[key] === undefined ? undefined : (total[key] ?? 0) + (step[key] ?? 0))
return new Usage({
// Context size describes the latest call, not a running total.
contextTokens: step.contextTokens ?? total.contextTokens,
inputTokens: sum("inputTokens"),
outputTokens: sum("outputTokens"),
nonCachedInputTokens: sum("nonCachedInputTokens"),
cacheReadInputTokens: sum("cacheReadInputTokens"),
cacheWriteInputTokens: sum("cacheWriteInputTokens"),
reasoningTokens: sum("reasoningTokens"),
totalTokens: sum("totalTokens"),
cost: sum("cost"),
})
}
/** The result of `collect`: every step of a stream from `stream` or `loop`, folded into one value. */
export interface Collected {
/** Text of the final step. */
readonly text: string
/** One response per model call, in order. */
readonly steps: ReadonlyArray<LLMResponse>
/** Each step's assistant message followed by its tool results, ready to append to the request's messages. */
readonly messages: ReadonlyArray<Message>
/** Usage from the terminal `finish`, summed across steps for a loop. */
readonly usage?: Usage
readonly finishReason: FinishReasonDetails
}
type FinishOf<Event> = Extract<Event, { readonly type: "finish" }>
/** `Collected`, plus `object` when the stream was started with `schema`. */
export type CollectResult<Event> = "object" extends keyof FinishOf<Event>
? Collected & Pick<FinishOf<Event>, "object">
: Collected
/** Run a stream from `stream` or `loop` to the end and fold its events into the final result. */
export function collect<Event extends LLMEvent, E, R>(
events: Stream.Stream<Event, E, R>,
): Effect.Effect<CollectResult<Event>, E | AIError, R>
export function collect<E, R>(
events: Stream.Stream<LLMEvent, E, R>,
): Effect.Effect<Collected & { readonly object?: unknown }, E | AIError, R> {
return Effect.suspend(() => {
const steps: LLMResponse[] = []
const messages: Message[] = []
let step = ResponseFold.empty()
// Between a `step-finish` and the next `step-start`, tool results come from the loop running tools.
let betweenSteps = false
// The last step's assistant message, and its calls with invalid input that the loop may answer.
let assistant = -1
let invalid: ReadonlyArray<ToolCallPart> = []
let finish: Finish | undefined
const close = (reason: FinishReasonDetails, usage: Usage | undefined) => {
if (step.events.length === 0) return
const response = new LLMResponse({
message: step.message,
events: step.events,
usage: usage ?? step.usage,
finishReason: reason,
})
steps.push(response)
assistant = messages.push(response.message) - 1
invalid = step.events.filter(LLMEvent.is.toolInputError).map(invalidCall)
step = ResponseFold.empty()
}
return Stream.runForEach(events, (event) =>
Effect.sync(() => {
if (event.type === "finish") {
close(event.reason, event.usage)
finish = event
return
}
if (event.type === "tool-result" && betweenSteps) {
const call = invalid.find((item) => item.id === event.id)
const previous = messages[assistant]
if (call && previous.role === "assistant") messages[assistant] = withInvalidCalls(previous, [call])
messages.push(
Message.tool({
id: event.id,
name: event.name,
namespace: event.namespace,
result: event.result,
providerMetadata: event.providerMetadata,
}),
)
return
}
if (event.type === "step-start") betweenSteps = false
step = ResponseFold.reduce(step, event)
if (event.type !== "step-finish") return
close(event.reason, event.usage)
betweenSteps = true
}),
).pipe(
Effect.flatMap(() => {
if (finish === undefined)
return Effect.fail(
new AIError({
reason: new InvalidProviderOutputError({
message: "The stream ended without a finish event.",
classification: "incomplete-stream",
}),
}),
)
const result = {
text: steps.at(-1)?.text ?? "",
steps,
messages,
usage: finish.usage,
finishReason: finish.reason,
...("object" in finish ? { object: finish.object } : {}),
}
return Effect.succeed(result)
}),
)
})
}