@effect-uai/core 0.3.0 → 0.5.0

package/src/language-model/LanguageModel.test.ts

@@ -0,0 +1,170 @@
+import { Cause, Effect, Exit, Layer, Option, Ref, Schedule, Stream } from "effect"
+import { describe, expect, it } from "vitest"
+import * as AiError from "../domain/AiError.js"
+import * as Items from "../domain/Items.js"
+import { type Turn, TurnEvent } from "../domain/Turn.js"
+import { LanguageModel, retry, turn } from "./LanguageModel.js"
+import * as MockProvider from "../testing/MockProvider.js"
+
+const oneTextTurn = (text: string): Turn => ({
+  items: [Items.assistantText(text)],
+  usage: { input_tokens: 1, output_tokens: 1 },
+  stop_reason: "stop",
+})
+
+describe("LanguageModel.turn", () => {
+  it("returns the assembled Turn from the terminal TurnComplete event", async () => {
+    const expected = oneTextTurn("hello world")
+    const program = turn({ history: [Items.userText("hi")], model: "mock" })
+
+    const result = await Effect.runPromise(
+      program.pipe(Effect.provide(MockProvider.layer([expected]))),
+    )
+
+    expect(result).toEqual(expected)
+  })
+
+  it("fails with IncompleteTurn when the stream ends without TurnComplete", async () => {
+    // Custom service whose stream emits a single TextDelta and then ends.
+    const broken = Layer.succeed(LanguageModel, {
+      streamTurn: () => Stream.fromIterable<TurnEvent>([TurnEvent.TextDelta({ text: "partial" })]),
+    })
+
+    const program = turn({ history: [Items.userText("hi")], model: "mock" })
+    const exit = await Effect.runPromise(Effect.exit(program.pipe(Effect.provide(broken))))
+
+    expect(Exit.isFailure(exit)).toBe(true)
+    if (Exit.isFailure(exit)) {
+      const failure = Cause.findErrorOption(exit.cause)
+      expect(Option.isSome(failure)).toBe(true)
+      if (Option.isSome(failure)) {
+        expect(failure.value).toBeInstanceOf(AiError.IncompleteTurn)
+      }
+    }
+  })
+
+  it("propagates an AiError raised by streamTurn", async () => {
+    const rateLimited = new AiError.RateLimited({ provider: "mock", raw: null })
+    const failing = Layer.succeed(LanguageModel, {
+      streamTurn: () => Stream.fail<AiError.AiError>(rateLimited),
+    })
+
+    const program = turn({ history: [], model: "mock" })
+    const exit = await Effect.runPromise(Effect.exit(program.pipe(Effect.provide(failing))))
+
+    expect(Exit.isFailure(exit)).toBe(true)
+    if (Exit.isFailure(exit)) {
+      expect(Cause.findErrorOption(exit.cause)).toEqual(Option.some(rateLimited))
+    }
+  })
+
+  it("returns the LAST TurnComplete when the stream contains multiple (defensive)", async () => {
+    // A misbehaving provider might emit two TurnComplete events; turn
+    // should pick the last one (the most recent assembled Turn).
+    const first = oneTextTurn("first")
+    const second = oneTextTurn("second")
+    const weird = Layer.succeed(LanguageModel, {
+      streamTurn: () =>
+        Stream.fromIterable<TurnEvent>([
+          TurnEvent.TurnComplete({ turn: first }),
+          TurnEvent.TurnComplete({ turn: second }),
+        ]),
+    })
+
+    const program = turn({ history: [], model: "mock" })
+    const result = await Effect.runPromise(program.pipe(Effect.provide(weird)))
+
+    expect(result).toEqual(second)
+  })
+})
+
+describe("LanguageModel.retry", () => {
+  const textDelta = (text: string): TurnEvent => TurnEvent.TextDelta({ text })
+  const textTurn = (text: string): Turn => ({
+    items: [Items.assistantText(text)],
+    usage: { input_tokens: 0, output_tokens: 0 },
+    stop_reason: "stop",
+  })
+  const completeEvent = (text: string): TurnEvent =>
+    TurnEvent.TurnComplete({ turn: textTurn(text) })
+
+  // Builds a stream that emits a failure or success based on attempt counter.
+  // Each call to the returned Effect produces a fresh attempt stream.
+  const attemptStream = (
+    attempts: Ref.Ref<number>,
+    plan: ReadonlyArray<Stream.Stream<TurnEvent, AiError.AiError>>,
+  ): Stream.Stream<TurnEvent, AiError.AiError> =>
+    Stream.unwrap(
+      Ref.getAndUpdate(attempts, (n) => n + 1).pipe(
+        Effect.map((n) => plan[Math.min(n, plan.length - 1)]!),
+      ),
+    )
+
+  it("retries on RateLimited and yields the success on the next attempt", async () => {
+    const program = Effect.gen(function* () {
+      const attempts = yield* Ref.make(0)
+      const stream = attemptStream(attempts, [
+        Stream.fail(new AiError.RateLimited({ provider: "mock", raw: null })),
+        Stream.fromIterable([textDelta("ok"), completeEvent("ok")]),
+      ]).pipe(retry(Schedule.recurs(3)))
+      const events = yield* Stream.runCollect(stream)
+      const count = yield* Ref.get(attempts)
+      return { events: Array.from(events), count }
+    })
+
+    const { events, count } = await Effect.runPromise(program)
+    expect(count).toBe(2)
+    expect(events.map((e) => e._tag)).toEqual(["TextDelta", "TurnComplete"])
+  })
+
+  it("surfaces the underlying retryable failure when retries are exhausted", async () => {
+    const cause = new AiError.Unavailable({ provider: "mock", raw: null })
+    const stream = Stream.fail<AiError.AiError>(cause).pipe(retry(Schedule.recurs(2)))
+
+    const exit = await Effect.runPromise(Effect.exit(Stream.runCollect(stream)))
+    expect(Exit.isFailure(exit)).toBe(true)
+    if (Exit.isFailure(exit)) {
+      expect(Cause.findErrorOption(exit.cause)).toEqual(Option.some(cause))
+    }
+  })
+
+  it("bypasses retry for non-retryable AiError (ContentFiltered)", async () => {
+    const program = Effect.gen(function* () {
+      const attempts = yield* Ref.make(0)
+      const cause = new AiError.ContentFiltered({ provider: "mock", raw: null })
+      const stream = attemptStream(attempts, [Stream.fail(cause)]).pipe(retry(Schedule.recurs(5)))
+      const exit = yield* Effect.exit(Stream.runCollect(stream))
+      const count = yield* Ref.get(attempts)
+      return { exit, count, cause }
+    })
+
+    const { exit, count, cause } = await Effect.runPromise(program)
+    expect(count).toBe(1) // no retry happened
+    expect(Exit.isFailure(exit)).toBe(true)
+    if (Exit.isFailure(exit)) {
+      expect(Cause.findErrorOption(exit.cause)).toEqual(Option.some(cause))
+    }
+  })
+
+  it("preserves deltas emitted before a retryable failure (and replays on retry)", async () => {
+    // Documents the 'replays on retry' caveat in the JSDoc — first attempt
+    // emits a delta then fails; second attempt is a clean success. Consumer
+    // sees the first delta twice (once from the failed attempt, once from
+    // the replay).
+    const program = Effect.gen(function* () {
+      const attempts = yield* Ref.make(0)
+      const stream = attemptStream(attempts, [
+        Stream.concat(
+          Stream.succeed<TurnEvent>(textDelta("partial")),
+          Stream.fail(new AiError.Timeout({ provider: "mock", raw: null })),
+        ),
+        Stream.fromIterable([textDelta("partial"), completeEvent("done")]),
+      ]).pipe(retry(Schedule.recurs(1)))
+      const events = yield* Stream.runCollect(stream)
+      return Array.from(events)
+    })
+
+    const events = await Effect.runPromise(program)
+    expect(events.map((e) => e._tag)).toEqual(["TextDelta", "TextDelta", "TurnComplete"])
+  })
+})

package/src/language-model/LanguageModel.ts

@@ -1,4 +1,4 @@
-import { Context, Effect, Stream } from "effect"
+import { Array as Arr, Context, Data, Effect, Option, type Schedule, Stream } from "effect"
 import * as AiError from "../domain/AiError.js"
 import type { Item } from "../domain/Items.js"
 import type * as StructuredFormat from "../structured-format/StructuredFormat.js"
@@ -51,3 +51,66 @@ export const streamTurn = (
   request: CommonRequest,
 ): Stream.Stream<TurnEvent, AiError.AiError, LanguageModel> =>
   Stream.unwrap(Effect.map(LanguageModel.asEffect(), (m) => m.streamTurn(request)))
+
+/**
+ * Drain `streamTurn` and return the assembled `Turn` from the terminal
+ * `TurnComplete` event. Fails with `IncompleteTurn` if the stream ends
+ * without one. Derived from `streamTurn`; providers get it for free.
+ */
+export const turn = (
+  request: CommonRequest,
+): Effect.Effect<Turn, AiError.AiError | AiError.IncompleteTurn, LanguageModel> =>
+  streamTurn(request).pipe(
+    Stream.runCollect,
+    Effect.flatMap((events) =>
+      Arr.findLast(events, isTurnComplete).pipe(
+        Option.match({
+          onNone: () => Effect.fail(new AiError.IncompleteTurn({})),
+          onSome: (e) => Effect.succeed(e.turn),
+        }),
+      ),
+    ),
+  )
+
+// ---------------------------------------------------------------------------
+// retry — retry the retryable subset of AiError, let other failures escape
+// ---------------------------------------------------------------------------
+
+/** Internal wrapper around the retryable subset of `AiError`. */
+export class Retryable extends Data.TaggedError("RetryableAi")<{
+  readonly cause: AiError.RateLimited | AiError.Unavailable | AiError.Timeout
+}> {}
+
+const isRetryable = (
+  e: AiError.AiError,
+): e is AiError.RateLimited | AiError.Unavailable | AiError.Timeout =>
+  e._tag === "RateLimited" || e._tag === "Unavailable" || e._tag === "Timeout"
+
+// Lift events to Items, non-retryable failures to Terminal values (escape
+// retry), retryable failures to wrapped errors (only thing retry sees).
+type Lifted<A> =
+  | { readonly _tag: "Item"; readonly value: A }
+  | { readonly _tag: "Terminal"; readonly cause: AiError.AiError }
+
+/**
+ * Retry a stream of `AiError` on the retryable subset
+ * (`RateLimited | Unavailable | Timeout`). Other failures bypass the
+ * schedule and propagate unchanged. Like all `Stream.retry`, the entire
+ * stream re-runs — deltas before the failure replay on the next attempt.
+ */
+export const retry =
+  <Out>(schedule: Schedule.Schedule<Out, Retryable>) =>
+  <A, R>(stream: Stream.Stream<A, AiError.AiError, R>): Stream.Stream<A, AiError.AiError, R> =>
+    stream.pipe(
+      Stream.map((value): Lifted<A> => ({ _tag: "Item", value })),
+      Stream.catchIf(
+        isRetryable,
+        (cause) => Stream.fail(new Retryable({ cause })),
+        (cause) => Stream.succeed<Lifted<A>>({ _tag: "Terminal", cause }),
+      ),
+      Stream.retry(schedule),
+      Stream.catchTag("RetryableAi", (e) => Stream.fail<AiError.AiError>(e.cause)),
+      Stream.flatMap((item) =>
+        item._tag === "Item" ? Stream.succeed(item.value) : Stream.fail(item.cause),
+      ),
+    )

package/src/loop/Loop.test.ts

@@ -1,13 +1,19 @@
-import { Deferred, Effect, Fiber, Latch, Ref, Stream, SubscriptionRef } from "effect"
-import { describe, expect, it } from "vitest"
+import { Deferred, Effect, Fiber, Latch, pipe, Ref, Stream, SubscriptionRef } from "effect"
+import { describe, expect, expectTypeOf, it } from "vitest"
+import * as AiError from "../domain/AiError.js"
+import { TurnEvent } from "../domain/Turn.js"
 import {
   type Event,
   loop,
+  loopFrom,
   loopWithState,
   next,
   nextAfter,
-  stopEvent,
+  onTurnComplete,
+  stop,
   stopAfter,
+  stopEvent,
+  stopWith,
   value,
 } from "./Loop.js"
 
@@ -101,6 +107,80 @@ describe("Loop.loop", () => {
     expect(result).toEqual([0, 1, 2])
   })
 
+  it("type: data-last (pipe) form preserves the body's E channel", () => {
+    // Regression for the prior inference bug: when used as
+    // `pipe(initial, loop(body))`, the body's E must propagate to the outer
+    // stream instead of collapsing to `never`. Generics live on the outer
+    // return of each overload, so neither calling form can erase them.
+    const result = pipe(
+      { count: 0 },
+      loop((_state) => Stream.fail(new AiError.RateLimited({ provider: "test", raw: null }))),
+    )
+    type E = typeof result extends Stream.Stream<unknown, infer X, unknown> ? X : never
+    expectTypeOf<E>().toEqualTypeOf<AiError.RateLimited>()
+  })
+
+  it("type: data-first form preserves the body's E channel", () => {
+    const result = loop({ count: 0 }, (_state) =>
+      Stream.fail(new AiError.RateLimited({ provider: "test", raw: null })),
+    )
+    type E = typeof result extends Stream.Stream<unknown, infer X, unknown> ? X : never
+    expectTypeOf<E>().toEqualTypeOf<AiError.RateLimited>()
+  })
+
+  it("type: onTurnComplete inside loop infers S and A from the handler without annotation", () => {
+    // Regression: when piped through loop, onTurnComplete's handler return
+    // type (Stream<Event<A, S>>) is the single source of truth for the loop
+    // body's element type. The previous workaround required explicit
+    // <S, A> at the call site. Now the loop's outer-return generics pull
+    // them through automatically.
+    type LoopState = { readonly turns: number }
+    type ToolEvent = { readonly _tag: "tool"; readonly name: string }
+
+    const result = pipe(
+      { turns: 0 } as LoopState,
+      loop((state) =>
+        Effect.gen(function* () {
+          const deltas: Stream.Stream<TurnEvent> = Stream.empty
+          return deltas.pipe(
+            onTurnComplete(() =>
+              Effect.sync(() =>
+                state.turns >= 1
+                  ? stop
+                  : nextAfter(Stream.succeed<ToolEvent>({ _tag: "tool", name: "x" }), {
+                      turns: state.turns + 1,
+                    }),
+              ),
+            ),
+          )
+        }),
+      ),
+    )
+
+    type Element = typeof result extends Stream.Stream<infer X, unknown, unknown> ? X : never
+    expectTypeOf<Element>().toEqualTypeOf<TurnEvent | ToolEvent>()
+  })
+
+  it("onTurnComplete: data-first form (Function.dual) works at runtime", async () => {
+    // Pin both calling forms: deltas.pipe(onTurnComplete(handler)) and
+    // onTurnComplete(deltas, handler). Same dispatch as loop's dual.
+    const turnComplete: TurnEvent = TurnEvent.TurnComplete({
+      turn: { items: [], usage: { input_tokens: 0, output_tokens: 0 }, stop_reason: "stop" },
+    })
+    const textDelta: TurnEvent = TurnEvent.TextDelta({ text: "hi" })
+    const deltas: Stream.Stream<TurnEvent> = Stream.fromIterable([textDelta, turnComplete])
+
+    const dataFirst = onTurnComplete(deltas, () => Effect.sync(() => stop))
+    const dataLast = deltas.pipe(onTurnComplete(() => Effect.sync(() => stop)))
+
+    const a = await Effect.runPromise(Stream.runCollect(dataFirst))
+    const b = await Effect.runPromise(Stream.runCollect(dataLast))
+
+    // Two value(delta) wraps + one stop sentinel from the handler.
+    expect(a.length).toBe(3)
+    expect(b.length).toBe(3)
+  })
+
   it("is stack-safe and linear-time across many iterations", async () => {
     // 100k iterations far exceeds V8's typical stack depth (~10–15k frames).
     const N = 100_000
@@ -522,3 +602,176 @@ describe("Loop.loopWithState", () => {
     expect(await Effect.runPromise(program)).toEqual([0, 0.5, 1, 1.5, 2, 2.5, 3])
   })
 })
+
+describe("Loop.loopFrom", () => {
+  it("runs a multi-turn inner loop per input until the body emits stop", async () => {
+    // Per input: emit (input + turnsSoFar) twice, then stop. State counts
+    // total turns ACROSS inputs. Demonstrates that `next(s)` continues with
+    // the SAME input, multiple times per input — not one body call per item.
+    const result = await Effect.runPromise(
+      Stream.fromIterable(["a", "b"]).pipe(
+        loopFrom(0, (turns: number, input: string) => {
+          if (turns >= 2 * (input === "a" ? 1 : 2)) return Stream.fromIterable([stopEvent])
+          return Stream.fromIterable([value(`${input}:${turns}`), next(turns + 1)])
+        }),
+        Stream.runCollect,
+      ),
+    )
+
+    // input="a": turns 0,1 → emit "a:0","a:1"; turns=2 → stop. State threads.
+    // input="b": turns 2,3 → emit "b:2","b:3"; turns=4 → stop.
+    expect(result).toEqual(["a:0", "a:1", "b:2", "b:3"])
+  })
+
+  it("threads state across inputs (audio-pipeline shape)", async () => {
+    // History accumulates across inputs. Each input emits its joined view of
+    // history+input, then `stopWith` ends the inner loop AND carries the
+    // updated history to the next input.
+    const result = await Effect.runPromise(
+      Stream.fromIterable(["x", "y", "z"]).pipe(
+        loopFrom([] as ReadonlyArray<string>, (history: ReadonlyArray<string>, input: string) =>
+          Stream.fromIterable([
+            value([...history, input].join(",")),
+            stopWith([...history, input]),
+          ]),
+        ),
+        Stream.runCollect,
+      ),
+    )
+
+    expect(result).toEqual(["x", "x,y", "x,y,z"])
+  })
+
+  it("simulates a stream of documents with multi-turn tool calls per document", async () => {
+    // Document arrives → model "thinks" (one text turn) → calls a tool
+    // (one tool turn) → emits final text (one text turn) → done.
+    // Three turns per document, two documents.
+    type Turn =
+      | { readonly kind: "text"; readonly doc: string; readonly text: string }
+      | { readonly kind: "tool"; readonly doc: string; readonly tool: string }
+    type State = { readonly turn: number; readonly totalTurns: number }
+
+    const result = await Effect.runPromise(
+      Stream.fromIterable(["doc1", "doc2"]).pipe(
+        loopFrom({ turn: 0, totalTurns: 0 } as State, (state, doc: string) => {
+          // Each document runs three turns then stops.
+          if (state.turn === 0) {
+            return Stream.fromIterable([
+              value<Turn>({ kind: "text", doc, text: "thinking" }),
+              next({ turn: 1, totalTurns: state.totalTurns + 1 }),
+            ])
+          }
+          if (state.turn === 1) {
+            return Stream.fromIterable([
+              value<Turn>({ kind: "tool", doc, tool: "search" }),
+              next({ turn: 2, totalTurns: state.totalTurns + 1 }),
+            ])
+          }
+          // Final turn — `stopWith` emits the final value, advances state
+          // (reset turn to 0 for the next document, bump totalTurns), and
+          // ends this document's inner loop in one shot.
+          return Stream.fromIterable([
+            value<Turn>({ kind: "text", doc, text: "final" }),
+            stopWith({ turn: 0, totalTurns: state.totalTurns + 1 }),
+          ])
+        }),
+        Stream.runCollect,
+      ),
+    )
+
+    expect(result).toEqual([
+      { kind: "text", doc: "doc1", text: "thinking" },
+      { kind: "tool", doc: "doc1", tool: "search" },
+      { kind: "text", doc: "doc1", text: "final" },
+      { kind: "text", doc: "doc2", text: "thinking" },
+      { kind: "tool", doc: "doc2", tool: "search" },
+      { kind: "text", doc: "doc2", text: "final" },
+    ])
+  })
+
+  it("ends cleanly when the input stream ends mid-conversation", async () => {
+    // Single-input case: body advances via `next` then stops cleanly.
+    const result = await Effect.runPromise(
+      Stream.fromIterable(["only"]).pipe(
+        loopFrom(0, (turns: number, input: string) =>
+          turns >= 2
+            ? Stream.fromIterable([stopEvent])
+            : Stream.fromIterable([value(`${input}:${turns}`), next(turns + 1)]),
+        ),
+        Stream.runCollect,
+      ),
+    )
+
+    expect(result).toEqual(["only:0", "only:1"])
+  })
+
+  it("body's `stop` advances to the next input (does NOT halt the whole stream)", async () => {
+    // Three inputs, body always stops on its first emission. All three
+    // are processed — `stop` is per-input, not global. To halt the whole
+    // stream, end the INPUT stream upstream.
+    const result = await Effect.runPromise(
+      Stream.fromIterable([1, 2, 3]).pipe(
+        loopFrom(0, (_state: number, input: number) =>
+          Stream.fromIterable([value(input * 10), stopEvent]),
+        ),
+        Stream.runCollect,
+      ),
+    )
+
+    expect(result).toEqual([10, 20, 30])
+  })
+
+  it("data-first form (Function.dual) runs identically to data-last", async () => {
+    const inputs = Stream.fromIterable([1, 2])
+    const result = await Effect.runPromise(
+      Stream.runCollect(
+        loopFrom(inputs, 0, (state: number, input: number) =>
+          state >= input
+            ? Stream.fromIterable([stopEvent])
+            : Stream.fromIterable([value(state + input), next(state + 1)]),
+        ),
+      ),
+    )
+
+    // input=1: state=0 → emit 1, state→1; state=1≥1 → stop.
+    // input=2: state=1 → emit 3, state→2; state=2≥2 → stop.
+    expect(result).toEqual([1, 3])
+  })
+
+  it("supports Effect-returning bodies (parity with loop)", async () => {
+    const result = await Effect.runPromise(
+      Stream.fromIterable(["a"]).pipe(
+        loopFrom(0, (turns: number, input: string) =>
+          Effect.gen(function* () {
+            const cur = yield* Effect.succeed(turns)
+            if (cur >= 2) return Stream.fromIterable([stopEvent])
+            return Stream.fromIterable([value(`${input}:${cur}`), next(cur + 1)])
+          }),
+        ),
+        Stream.runCollect,
+      ),
+    )
+
+    expect(result).toEqual(["a:0", "a:1"])
+  })
+
+  it("type: data-last (pipe) form preserves the body's E channel", () => {
+    const result = pipe(
+      Stream.fromIterable([1]),
+      loopFrom(0, (_state: number, _input: number) =>
+        Stream.fail(new AiError.RateLimited({ provider: "test", raw: null })),
+      ),
+    )
+    type E = typeof result extends Stream.Stream<unknown, infer X, unknown> ? X : never
+    expectTypeOf<E>().toEqualTypeOf<AiError.RateLimited>()
+  })
+
+  it("type: data-first form preserves the body's E channel and unifies with input's E", () => {
+    const input: Stream.Stream<number, Error> = Stream.fail(new Error("boom"))
+    const result = loopFrom(input, 0, (_state: number, _i: number) =>
+      Stream.fail(new AiError.RateLimited({ provider: "test", raw: null })),
+    )
+    type E = typeof result extends Stream.Stream<unknown, infer X, unknown> ? X : never
+    expectTypeOf<E>().toEqualTypeOf<AiError.RateLimited | Error>()
+  })
+})