@effect-uai/core 0.5.0 → 0.5.2

package/src/language-model/LanguageModel.ts

@@ -1,4 +1,4 @@
-import { Array as Arr, Context, Data, Effect, Option, type Schedule, Stream } from "effect"
+import { Array as Arr, Context, Effect, Option, 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"
@@ -38,6 +38,16 @@ export type CommonRequest = {
 
 export type LanguageModelService = {
   readonly streamTurn: (request: CommonRequest) => Stream.Stream<TurnEvent, AiError.AiError>
+  /**
+   * Drain a single turn and return the assembled `Turn` from the
+   * terminal `TurnComplete` event. Fails with `IncompleteTurn` if the
+   * stream ends without one.
+   *
+   * Most providers derive this from `streamTurn` via
+   * {@link turnFromStream}; providers with a native non-streaming
+   * endpoint may override with a cheaper direct call.
+   */
+  readonly turn: (request: CommonRequest) => Effect.Effect<Turn, AiError.AiError>
 }
 
 export class LanguageModel extends Context.Service<LanguageModel, LanguageModelService>()(
@@ -53,64 +63,31 @@ export const streamTurn = (
   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.
+ * Drain a single turn and return the assembled `Turn`. Delegates to the
+ * service's `turn` method — providers with a native complete endpoint
+ * can override; the rest get the default streamTurn-drain via
+ * {@link turnFromStream}.
  */
-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 }
+export const turn = (request: CommonRequest): Effect.Effect<Turn, AiError.AiError, LanguageModel> =>
+  Effect.flatMap(LanguageModel.asEffect(), (m) => m.turn(request))
 
 /**
- * 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.
+ * Build a `turn` implementation from a `streamTurn` implementation.
+ * Providers without a native non-streaming endpoint use this to
+ * populate the service's `turn` field. Generic over the request type so
+ * provider-typed services (with their narrowed request) can reuse it.
  */
-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),
+export const turnFromStream =
+  <Req>(streamTurn: (request: Req) => Stream.Stream<TurnEvent, AiError.AiError>) =>
+  (request: Req): Effect.Effect<Turn, AiError.AiError> =>
+    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),
+          }),
+        ),
       ),
     )

package/src/retry/Retry.test.ts

@@ -0,0 +1,163 @@
+import { Cause, Effect, Exit, Option, Ref, Schedule, Stream } from "effect"
+import { describe, expect, it } from "vitest"
+import * as AiError from "../domain/AiError.js"
+import * as Retry from "./Retry.js"
+
+// Each call pulls the next plan entry; the last entry repeats forever
+// (so "succeed on attempt N" plans naturally re-yield success after).
+const planned = <A, E>(
+  attempts: Ref.Ref<number>,
+  plan: ReadonlyArray<Stream.Stream<A, E>>,
+): Stream.Stream<A, E> =>
+  Stream.unwrap(
+    Ref.getAndUpdate(attempts, (n) => n + 1).pipe(
+      Effect.map((n) => plan[Math.min(n, plan.length - 1)]!),
+    ),
+  )
+
+const plannedEffect = <A, E>(
+  attempts: Ref.Ref<number>,
+  plan: ReadonlyArray<Effect.Effect<A, E>>,
+): Effect.Effect<A, E> =>
+  Ref.getAndUpdate(attempts, (n) => n + 1).pipe(
+    Effect.flatMap((n) => plan[Math.min(n, plan.length - 1)]!),
+  )
+
+describe("Retry.stream", () => {
+  it("retries on RateLimited and yields success on the next attempt", async () => {
+    const program = Effect.gen(function* () {
+      const attempts = yield* Ref.make(0)
+      const s = planned(attempts, [
+        Stream.fail<AiError.AiError>(new AiError.RateLimited({ provider: "mock", raw: null })),
+        Stream.fromIterable(["a", "b"]),
+      ]).pipe(Retry.stream(Schedule.recurs(3)))
+      const out = yield* Stream.runCollect(s)
+      const count = yield* Ref.get(attempts)
+      return { out: Array.from(out), count }
+    })
+
+    const { out, count } = await Effect.runPromise(program)
+    expect(count).toBe(2)
+    expect(out).toEqual(["a", "b"])
+  })
+
+  it("surfaces the underlying retryable failure when retries are exhausted", async () => {
+    const cause = new AiError.Unavailable({ provider: "mock", raw: null })
+    const s = Stream.fail<AiError.AiError>(cause).pipe(Retry.stream(Schedule.recurs(2)))
+
+    const exit = await Effect.runPromise(Effect.exit(Stream.runCollect(s)))
+    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 s = planned(attempts, [Stream.fail<AiError.AiError>(cause)]).pipe(
+        Retry.stream(Schedule.recurs(5)),
+      )
+      const exit = yield* Effect.exit(Stream.runCollect(s))
+      const count = yield* Ref.get(attempts)
+      return { exit, count, cause }
+    })
+
+    const { exit, count, cause } = await Effect.runPromise(program)
+    expect(count).toBe(1)
+    expect(Exit.isFailure(exit)).toBe(true)
+    if (Exit.isFailure(exit)) {
+      expect(Cause.findErrorOption(exit.cause)).toEqual(Option.some(cause))
+    }
+  })
+
+  it("preserves items emitted before a retryable failure (and replays on retry)", async () => {
+    // Documents the 'replays on retry' caveat: the entire stream re-runs.
+    const program = Effect.gen(function* () {
+      const attempts = yield* Ref.make(0)
+      const s = planned(attempts, [
+        Stream.concat(
+          Stream.succeed("partial"),
+          Stream.fail<AiError.AiError>(new AiError.Timeout({ provider: "mock", raw: null })),
+        ),
+        Stream.fromIterable(["partial", "done"]),
+      ]).pipe(Retry.stream(Schedule.recurs(1)))
+      const out = yield* Stream.runCollect(s)
+      return Array.from(out)
+    })
+
+    const out = await Effect.runPromise(program)
+    expect(out).toEqual(["partial", "partial", "done"])
+  })
+})
+
+describe("Retry.effect", () => {
+  it("retries on RateLimited and yields success on the next attempt", async () => {
+    const program = Effect.gen(function* () {
+      const attempts = yield* Ref.make(0)
+      const value = yield* plannedEffect(attempts, [
+        Effect.fail<AiError.AiError>(new AiError.RateLimited({ provider: "mock", raw: null })),
+        Effect.succeed("ok"),
+      ]).pipe(Retry.effect(Schedule.recurs(3)))
+      const count = yield* Ref.get(attempts)
+      return { value, count }
+    })
+
+    const { value, count } = await Effect.runPromise(program)
+    expect(count).toBe(2)
+    expect(value).toBe("ok")
+  })
+
+  it("surfaces the underlying retryable failure when retries are exhausted", async () => {
+    const cause = new AiError.Unavailable({ provider: "mock", raw: null })
+    const eff = Effect.fail<AiError.AiError>(cause).pipe(Retry.effect(Schedule.recurs(2)))
+
+    const exit = await Effect.runPromise(Effect.exit(eff))
+    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 exit = yield* Effect.exit(
+        plannedEffect(attempts, [Effect.fail<AiError.AiError>(cause)]).pipe(
+          Retry.effect(Schedule.recurs(5)),
+        ),
+      )
+      const count = yield* Ref.get(attempts)
+      return { exit, count, cause }
+    })
+
+    const { exit, count, cause } = await Effect.runPromise(program)
+    expect(count).toBe(1)
+    expect(Exit.isFailure(exit)).toBe(true)
+    if (Exit.isFailure(exit)) {
+      expect(Cause.findErrorOption(exit.cause)).toEqual(Option.some(cause))
+    }
+  })
+
+  it("does NOT replay successful prior work — at-most-once semantics", async () => {
+    // Contrast with Retry.stream: an Effect.retry only re-invokes the
+    // failing computation. No replay of work done before the failure
+    // (there's nothing to replay — Effects are single-shot).
+    const program = Effect.gen(function* () {
+      const attempts = yield* Ref.make(0)
+      const value = yield* plannedEffect(attempts, [
+        Effect.fail<AiError.AiError>(new AiError.Timeout({ provider: "mock", raw: null })),
+        Effect.fail<AiError.AiError>(new AiError.Timeout({ provider: "mock", raw: null })),
+        Effect.succeed("done"),
+      ]).pipe(Retry.effect(Schedule.recurs(5)))
+      const count = yield* Ref.get(attempts)
+      return { value, count }
+    })
+
+    const { value, count } = await Effect.runPromise(program)
+    expect(count).toBe(3)
+    expect(value).toBe("done")
+  })
+})

package/src/retry/Retry.ts

@@ -0,0 +1,84 @@
+import { Data, Effect, type Schedule, Stream } from "effect"
+import type * as AiError from "../domain/AiError.js"
+
+/**
+ * Subset-aware retry helpers for `AiError`.
+ *
+ * Effect's own `Effect.retry` / `Stream.retry` apply the schedule to
+ * every failure. For AI calls that's the wrong policy — replaying a
+ * `ContentFiltered` or `InvalidRequest` won't change the outcome and
+ * just wastes quota. These helpers gate the schedule on the
+ * {@link Retryable} subset (rate limits, transport hiccups, timeouts);
+ * everything else propagates unchanged.
+ *
+ * Two carriers:
+ * - {@link stream} — for `Stream<A, AiError, R>` (e.g. `streamTurn`,
+ *   `streamSynthesis`, `streamTranscriptionFrom`). Caveat: the entire
+ *   stream re-runs on retry, so deltas before the failure replay on
+ *   the next attempt.
+ * - {@link effect} — for `Effect<A, AiError, R>` (e.g. `turn`, `embed`,
+ *   `embedMany`, `synthesize`, `transcribe`). At-most-once.
+ *
+ * Composes with `Stream.retry` / `Effect.retry` directly (no naming
+ * collision since these live under the `Retry` namespace).
+ */
+
+/**
+ * Wrapper around the retryable subset of `AiError`. Internal to the
+ * lift/retry/unlift dance — exported so callers can build their own
+ * variants if needed (e.g. a custom schedule that observes `cause`).
+ */
+export class Retryable extends Data.TaggedError("RetryableAi")<{
+  readonly cause: AiError.RateLimited | AiError.Unavailable | AiError.Timeout
+}> {}
+
+/** Type-narrowing predicate for the retryable subset. */
+export 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<A, AiError, R>` on the retryable subset. Other
+ * failures bypass the schedule and propagate unchanged. The whole
+ * stream replays on each attempt — deltas emitted before the failure
+ * are re-emitted on the next try.
+ */
+export const stream =
+  <Out>(schedule: Schedule.Schedule<Out, Retryable>) =>
+  <A, R>(s: Stream.Stream<A, AiError.AiError, R>): Stream.Stream<A, AiError.AiError, R> =>
+    s.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),
+      ),
+    )
+
+/**
+ * Retry an `Effect<A, AiError, R>` on the retryable subset. Other
+ * failures bypass the schedule and propagate unchanged. At-most-once
+ * — the underlying call only re-runs if the previous attempt failed
+ * with a retryable tag.
+ *
+ * Skips the lift/unlift dance the Stream variant needs: `Effect.retry`
+ * exposes a `while:` predicate so the error channel stays a plain
+ * `AiError` throughout.
+ */
+export const effect =
+  <Out>(schedule: Schedule.Schedule<Out, AiError.AiError>) =>
+  <A, R>(eff: Effect.Effect<A, AiError.AiError, R>): Effect.Effect<A, AiError.AiError, R> =>
+    Effect.retry(eff, { schedule, while: isRetryable })

package/src/testing/MockProvider.ts

@@ -1,7 +1,11 @@
 import { Array as Arr, Duration, Effect, Layer, Match, Option, Ref, Schedule, Stream } from "effect"
 import * as AiError from "../domain/AiError.js"
 import { type Item, isOutputText } from "../domain/Items.js"
-import { LanguageModel, type LanguageModelService } from "../language-model/LanguageModel.js"
+import {
+  LanguageModel,
+  type LanguageModelService,
+  turnFromStream,
+} from "../language-model/LanguageModel.js"
 import { type Turn, TurnEvent } from "../domain/Turn.js"
 
 export type MockOptions = {
@@ -81,8 +85,8 @@ const buildService = (
   options: MockOptions | undefined,
   cursor: Ref.Ref<number>,
   record: (call: Call) => Effect.Effect<void>,
-): LanguageModelService => ({
-  streamTurn: (request) =>
+): LanguageModelService => {
+  const streamTurn: LanguageModelService["streamTurn"] = (request) =>
     Stream.unwrap(
       Ref.getAndUpdate(cursor, (n) => n + 1).pipe(
         Effect.flatMap(
@@ -96,8 +100,9 @@ const buildService = (
             }),
         ),
       ),
-    ),
-})
+    )
+  return { streamTurn, turn: turnFromStream(streamTurn) }
+}
 
 // ---------------------------------------------------------------------------
 // Recorder handle. Unsafe Ref is local: it backs both the `record` write