packages/ai/test/loop.test.ts
Rama dev · réplica de lectura del repositorio original.
import { describe, expect } from "bun:test"
import { Effect, Schema, Stream } from "effect"
import {
ai,
AIError,
LanguageModel,
type LLM,
LLMEvent,
ProviderInternalError,
tool,
ToolFailure,
} from "../src/effect.js"
import { AI } from "../src/index.js"
import { OpenAIChat } from "../src/protocols/openai-chat.js"
import { TestLLM } from "../src/testing.js"
import { testEffect } from "./lib/effect.js"
const model = LanguageModel.make({ id: "fictional-model", provider: "fixture", route: OpenAIChat.route })
const it = testEffect(TestLLM.testLayer())
const read_file = tool({
name: "read_file",
description: "Read a file.",
input: Schema.Struct({ path: Schema.String }),
success: Schema.String,
execute: ({ path }) =>
path === "missing.txt" ? new ToolFailure({ message: `No such file: ${path}` }) : Effect.succeed("hello"),
})
const collect = (until?: LLM.LLMLoopInput["until"]) =>
ai.loop({ model, prompt: "Read the file.", tools: [read_file], until }).pipe(
Stream.runCollect,
Effect.map((events) => Array.from(events)),
)
describe("ai.loop", () => {
it.effect("runs tool calls and sends the full assistant turn back to the model", () =>
Effect.gen(function* () {
const client = yield* TestLLM.Test
yield* client.push(
TestLLM.toolCalls(
LLMEvent.textStart({ id: "text-1" }),
LLMEvent.textDelta({ id: "text-1", text: "Reading." }),
LLMEvent.textEnd({ id: "text-1" }),
LLMEvent.toolCall({ id: "call-1", name: "read_file", input: { path: "a.txt" } }),
),
TestLLM.text("The file says hello.", "text-2"),
)
const events = yield* collect()
const requests = yield* client.requests()
expect(requests[0].tools.map((tool) => tool.name)).toEqual(["read_file"])
expect(requests[1].messages.map((message) => message.role)).toEqual(["user", "assistant", "tool"])
expect(requests[1].messages[1].content).toEqual([
expect.objectContaining({ type: "text", text: "Reading." }),
expect.objectContaining({ type: "tool-call", id: "call-1", name: "read_file", input: { path: "a.txt" } }),
])
expect(requests[1].messages[2].content).toEqual([
expect.objectContaining({ type: "tool-result", id: "call-1", result: { type: "text", value: "hello" } }),
])
expect(events.filter(LLMEvent.is.toolResult)).toEqual([
expect.objectContaining({ id: "call-1", result: { type: "text", value: "hello" } }),
])
expect(events.filter(LLMEvent.is.stepStart).map((event) => event.index)).toEqual([0, 1])
expect(events.filter(LLMEvent.is.stepFinish).map((event) => event.index)).toEqual([0, 1])
expect(events.filter(LLMEvent.is.finish)).toHaveLength(1)
expect(events.at(-1)?.type).toBe("finish")
}),
)
it.effect("ends with one finish carrying the final reason and usage summed across steps", () =>
Effect.gen(function* () {
const client = yield* TestLLM.Test
yield* client.push(
TestLLM.complete(
{
reason: { normalized: "tool-calls" },
usage: { inputTokens: 10, outputTokens: 4, contextTokens: 10, cost: 0.5 },
},
LLMEvent.toolCall({ id: "call-1", name: "read_file", input: { path: "a.txt" } }),
),
TestLLM.complete(
{ reason: { normalized: "stop" }, usage: { inputTokens: 20, outputTokens: 6, contextTokens: 20 } },
LLMEvent.textDelta({ id: "text-1", text: "Done." }),
),
)
const events = yield* collect()
expect(events.filter(LLMEvent.is.stepFinish).map((event) => event.usage?.inputTokens)).toEqual([10, 20])
const finishes = events.filter(LLMEvent.is.finish)
expect(finishes).toHaveLength(1)
expect(events.at(-1)).toBe(finishes[0])
expect(finishes[0].reason.normalized).toBe("stop")
expect(finishes[0].usage).toMatchObject({ inputTokens: 30, outputTokens: 10, contextTokens: 20, cost: 0.5 })
}),
)
it.effect("sends a failing tool back to the model as an error result", () =>
Effect.gen(function* () {
const client = yield* TestLLM.Test
yield* client.push(
TestLLM.toolCall("call-1", "read_file", { path: "missing.txt" }),
TestLLM.text("That file does not exist.", "text-1"),
)
const events = yield* collect()
const requests = yield* client.requests()
expect(events.filter(LLMEvent.is.toolResult)).toEqual([
expect.objectContaining({ id: "call-1", result: { type: "error", value: "No such file: missing.txt" } }),
])
expect(events.filter((event) => event.type.startsWith("tool-") && event.type !== "tool-call")).toHaveLength(1)
expect(requests[1].messages[2].content).toEqual([
expect.objectContaining({ result: { type: "error", value: "No such file: missing.txt" } }),
])
expect(events.filter(LLMEvent.is.textDelta).map((event) => event.text)).toEqual(["That file does not exist."])
}),
)
it.effect("runs Promise tools, validating success and projecting toModelOutput", () =>
Effect.gen(function* () {
const client = yield* TestLLM.Test
yield* client.push(
TestLLM.toolCalls(
LLMEvent.toolCall({ id: "call-1", name: "add", input: { a: 2, b: 3 } }),
LLMEvent.toolCall({ id: "call-2", name: "add", input: { a: -1, b: 0 } }),
),
TestLLM.text("Done.", "text-1"),
)
const add = AI.tool({
name: "add",
description: "Add two numbers.",
input: Schema.Struct({ a: Schema.Number, b: Schema.Number }),
success: Schema.Number,
execute: async ({ a, b }) => (a < 0 ? "negative" : a + b) as number,
toModelOutput: ({ output }) => [{ type: "text", text: `sum=${output}` }],
})
const events = yield* ai.loop({ model, prompt: "Add.", tools: [add] }).pipe(Stream.runCollect)
expect(Array.from(events).filter(LLMEvent.is.toolResult)).toEqual([
expect.objectContaining({ id: "call-1", result: { type: "text", value: "sum=5" } }),
expect.objectContaining({
result: { type: "error", value: expect.stringContaining("invalid value for its success schema") },
}),
])
}),
)
it.effect("stops when until returns true without running that step's tools", () =>
Effect.gen(function* () {
const client = yield* TestLLM.Test
yield* client.always(TestLLM.toolCall("call-1", "read_file", { path: "a.txt" }))
const seen: number[] = []
const events = yield* collect(({ steps }) => {
seen.push(steps.length)
expect(steps.at(-1)?.toolCalls.map((call) => call.id)).toEqual(["call-1"])
return steps.length >= 2
})
expect(seen).toEqual([1, 2])
expect(yield* client.requests()).toHaveLength(2)
expect(events.filter(LLMEvent.is.toolResult)).toHaveLength(1)
expect(events.at(-1)?.type).toBe("finish")
}),
)
it.effect("sends unparseable tool input back to the model as an error result and continues", () =>
Effect.gen(function* () {
const client = yield* TestLLM.Test
yield* client.push(
TestLLM.toolCalls(LLMEvent.toolInputError({ id: "call-1", name: "read_file", raw: '{"path": "a.txt"' })),
TestLLM.toolCall("call-2", "read_file", { path: "a.txt" }),
TestLLM.text("The file says hello.", "text-1"),
)
const events = yield* collect()
const requests = yield* client.requests()
const error = {
type: "error",
value: 'Invalid tool arguments: the input is not valid JSON. Received: {"path": "a.txt"',
} as const
expect(requests).toHaveLength(3)
expect(requests[1].messages[1].content).toEqual([
expect.objectContaining({ type: "tool-call", id: "call-1", name: "read_file", input: {} }),
])
expect(requests[1].messages[2].content).toEqual([
expect.objectContaining({ type: "tool-result", id: "call-1", result: error }),
])
expect(events.filter(LLMEvent.is.toolResult).map((event) => event.result)).toEqual([
error,
{ type: "text", value: "hello" },
])
}),
)
it.effect("stops at the output token limit without running the step's tool calls", () =>
Effect.gen(function* () {
const client = yield* TestLLM.Test
yield* client.push(
TestLLM.complete(
{ reason: { normalized: "length" } },
LLMEvent.toolCall({ id: "call-1", name: "read_file", input: { path: "a.txt" } }),
LLMEvent.toolInputError({ id: "call-2", name: "read_file", raw: '{"path": "b' }),
),
)
const events = yield* collect()
expect(yield* client.requests()).toHaveLength(1)
expect(events.filter(LLMEvent.is.toolResult)).toHaveLength(0)
expect(events.at(-1)).toEqual(expect.objectContaining({ type: "finish", reason: { normalized: "length" } }))
}),
)
it.effect("forces the caller's toolChoice on the first step only", () =>
Effect.gen(function* () {
const client = yield* TestLLM.Test
yield* client.push(
TestLLM.toolCall("call-1", "read_file", { path: "a.txt" }),
TestLLM.text("The file says hello.", "text-1"),
)
yield* ai
.loop({
model,
prompt: "Read the file.",
tools: [read_file],
toolChoice: { type: "required", disableParallelToolUse: true },
})
.pipe(Stream.runDrain)
const requests = yield* client.requests()
expect(requests.map((request) => ({ ...request.toolChoice }))).toEqual([
{ type: "required", disableParallelToolUse: true },
{ type: "auto", disableParallelToolUse: true },
])
}),
)
it.effect("keeps the answer tool forced on every step with schema", () =>
Effect.gen(function* () {
const client = yield* TestLLM.Test
yield* client.push(
TestLLM.toolCall("call-1", "read_file", { path: "a.txt" }),
TestLLM.toolCall("call-2", "generate_object", { text: "hello" }),
)
yield* ai
.loop({
model,
prompt: "Read the file.",
tools: [read_file],
toolChoice: "read_file",
schema: Schema.Struct({ text: Schema.String }),
})
.pipe(Stream.runDrain)
expect((yield* client.requests()).map((request) => request.toolChoice?.type)).toEqual(["required", "required"])
}),
)
it.effect("fails without running tools when a step's stream fails", () =>
Effect.gen(function* () {
const client = yield* TestLLM.Test
yield* client.push(
TestLLM.failAfter(
new AIError({ reason: new ProviderInternalError({ message: "Overloaded" }) }),
LLMEvent.stepStart({ index: 0 }),
LLMEvent.toolCall({ id: "call-1", name: "read_file", input: { path: "a.txt" } }),
),
)
const error = yield* collect().pipe(Effect.flip)
expect(yield* client.requests()).toHaveLength(1)
expect(error.reason.message).toBe("Overloaded")
expect(error.partial?.toolCalls.map((call) => call.id)).toEqual(["call-1"])
}),
)
it.effect("sends tools in array order, after wire entries given before them", () =>
Effect.gen(function* () {
const client = yield* TestLLM.Test
yield* client.push(TestLLM.text("Done.", "text-1"))
const search = tool({ name: "search", description: "Search.", input: Schema.String })
yield* ai.generate({
model,
prompt: "Hi.",
tools: [{ name: "hosted", description: "Hosted.", inputSchema: {} }, read_file, search],
})
expect((yield* client.requests())[0].tools.map((tool) => tool.name)).toEqual(["hosted", "read_file", "search"])
}),
)
it.effect("fails the request when two tools share a name", () =>
Effect.gen(function* () {
const client = yield* TestLLM.Test
const copy = tool({ name: "read_file", description: "Read again.", input: Schema.String })
const loop = yield* ai
.loop({ model, prompt: "Hi.", tools: [read_file, read_file] })
.pipe(Stream.runDrain, Effect.flip)
const generate = yield* ai.generate({ model, prompt: "Hi.", tools: [read_file, copy] }).pipe(Effect.flip)
expect(loop.reason).toMatchObject({ _tag: "InvalidRequest", message: "Duplicate tool name: read_file" })
expect(generate.reason).toMatchObject({ _tag: "InvalidRequest", message: "Duplicate tool name: read_file" })
expect(yield* client.requests()).toHaveLength(0)
}),
)
it.effect("fails untyped input that has neither or both of prompt and messages", () =>
Effect.gen(function* () {
const client = yield* TestLLM.Test
const untyped = (input: object) => input as LLM.LLMRequestInput
const message = "Pass exactly one of `prompt` or `messages`, not both or neither"
const neither = yield* ai.generate(untyped({ model })).pipe(Effect.flip)
const both = yield* ai
.loop({ ...untyped({ model, prompt: "Hi.", messages: [] }), tools: [read_file] })
.pipe(Stream.runDrain, Effect.flip)
expect(neither.reason).toMatchObject({ _tag: "InvalidRequest", message })
expect(both.reason).toMatchObject({ _tag: "InvalidRequest", message })
expect(() => ai.request(untyped({ model, system: "Be brief." }))).toThrow(message)
expect(yield* client.requests()).toHaveLength(0)
}),
)
})
describe("ai.collect", () => {
it.effect("folds a stream into one step whose finish carries the message and usage", () =>
Effect.gen(function* () {
const client = yield* TestLLM.Test
yield* client.push(TestLLM.textWithUsage("Hello.", "text-1", 7))
const events = Array.from(yield* ai.stream({ model, prompt: "Say hello." }).pipe(Stream.runCollect))
const finish = events.find(LLMEvent.is.finish)
expect(finish?.message?.content).toEqual([expect.objectContaining({ type: "text", text: "Hello." })])
expect(finish?.usage?.inputTokens).toBe(7)
yield* client.push(TestLLM.textWithUsage("Hello.", "text-1", 7))
const result = yield* ai.collect(ai.stream({ model, prompt: "Say hello." }))
expect(result.text).toBe("Hello.")
expect(result.steps).toHaveLength(1)
expect(result.messages).toEqual([result.steps[0].message])
expect(finish?.message).toEqual(result.steps[0].message)
expect(result.usage?.inputTokens).toBe(7)
expect(result.finishReason.normalized).toBe("stop")
}),
)
it.effect("folds a loop into its steps, one message per step and tool result, and the total usage", () =>
Effect.gen(function* () {
const client = yield* TestLLM.Test
yield* client.push(
TestLLM.complete(
{ reason: { normalized: "tool-calls" }, usage: { inputTokens: 10 } },
LLMEvent.textStart({ id: "text-1" }),
LLMEvent.textDelta({ id: "text-1", text: "Reading." }),
LLMEvent.textEnd({ id: "text-1" }),
LLMEvent.toolCall({ id: "call-1", name: "read_file", input: { path: "a.txt" } }),
),
TestLLM.complete(
{ reason: { normalized: "stop" }, usage: { inputTokens: 20 } },
LLMEvent.textDelta({ id: "text-2", text: "The file says hello." }),
),
)
const result = yield* ai.collect(ai.loop({ model, prompt: "Read the file.", tools: [read_file] }))
const requests = yield* client.requests()
expect(result.text).toBe("The file says hello.")
expect(result.steps.map((step) => step.text)).toEqual(["Reading.", "The file says hello."])
expect(result.steps.map((step) => step.usage?.inputTokens)).toEqual([10, 20])
expect(result.usage?.inputTokens).toBe(30)
expect(result.messages.map((message) => message.role)).toEqual(["assistant", "tool", "assistant"])
// The first two messages are exactly what the loop sent back for the second step.
expect(result.messages.slice(0, 2)).toEqual(requests[1].messages.slice(1))
}),
)
it.effect("gives each error result for unparseable tool input its call, as the loop sent it", () =>
Effect.gen(function* () {
const client = yield* TestLLM.Test
yield* client.push(
TestLLM.toolCalls(LLMEvent.toolInputError({ id: "call-1", name: "read_file", raw: '{"path": "a.txt"' })),
TestLLM.text("Sorry.", "text-1"),
)
const result = yield* ai.collect(ai.loop({ model, prompt: "Read the file.", tools: [read_file] }))
const requests = yield* client.requests()
expect(result.messages.map((message) => message.role)).toEqual(["assistant", "tool", "assistant"])
expect(result.messages[0].content).toEqual([
expect.objectContaining({ type: "tool-call", id: "call-1", name: "read_file", input: {} }),
])
expect(result.messages.slice(0, 2)).toEqual(requests[1].messages.slice(1))
}),
)
it.effect("carries the decoded object when the stream has a schema", () =>
Effect.gen(function* () {
const client = yield* TestLLM.Test
yield* client.push(TestLLM.toolCall("call-1", "generate_object", { path: "a.txt" }))
const result = yield* ai.collect(
ai.stream({ model, prompt: "Pick a file.", schema: Schema.Struct({ path: Schema.String }) }),
)
const path: string = result.object.path
expect(path).toBe("a.txt")
}),
)
it.effect("fails when the stream ends without a finish event", () =>
Effect.gen(function* () {
const error = yield* ai.collect(Stream.make(LLMEvent.stepStart({ index: 0 }))).pipe(Effect.flip)
expect(error.reason).toMatchObject({ _tag: "InvalidProviderOutput", classification: "incomplete-stream" })
}),
)
})