Use Effect primitives for AI and response fanout

ts/packages/flow/src/__tests__/text-completion-common.test.ts

@@ -1,8 +1,10 @@
 import { describe, expect, it } from "@effect/vitest";
 import type { LlmChunk } from "@trustgraph/base";
-import { Effect, Stream } from "effect";
+import { Effect, Layer, ManagedRuntime, Stream } from "effect";
+import { AiError, LanguageModel } from "effect/unstable/ai";
 import {
   llmStreamPart,
+  makeLanguageModelProvider,
   providerRuntimeError,
   providerStatusError,
   streamTextCompletionChunks,
@@ -10,6 +12,36 @@ import {
   toAsyncGenerator,
 } from "../model/text-completion/common.js";
 
+const languageModelRuntime = ManagedRuntime.make(Layer.empty);
+
+const usage = (inputTokens: number, outputTokens: number) => ({
+  inputTokens: {
+    uncached: undefined,
+    total: inputTokens,
+    cacheRead: undefined,
+    cacheWrite: undefined,
+  },
+  outputTokens: {
+    total: outputTokens,
+    text: undefined,
+    reasoning: undefined,
+  },
+});
+
+const finishPart = (inputTokens: number, outputTokens: number) => ({
+  type: "finish",
+  reason: "stop",
+  usage: usage(inputTokens, outputTokens),
+  response: undefined,
+});
+
+const aiError = (reason: AiError.AiErrorReason) =>
+  new AiError.AiError({
+    module: "FakeLanguageModel",
+    method: "generateText",
+    reason,
+  });
+
 const emptyChunkIterator = (): AsyncIterable<LlmChunk> => ({
   [Symbol.asyncIterator]: () => ({
     next: () => Promise.resolve({ done: true, value: undefined }),
@@ -84,4 +116,107 @@ describe("text completion common helpers", () => {
     expect(textFromContent([{ text: "a" }, { text: "b" }, { wrong: "skip" }])).toBe("ab");
     expect(textFromContent([{ text: 1 }])).toBe("");
   });
+
+  it("adapts Effect LanguageModel generateText responses to LlmProvider results", async () => {
+    const provider = makeLanguageModelProvider({
+      provider: "FakeLanguageModel",
+      defaultModel: "fake-model",
+      defaultTemperature: 0.1,
+      runtime: languageModelRuntime,
+      makeLanguageModel: ({ model, temperature }) =>
+        LanguageModel.make({
+          generateText: () =>
+            Effect.succeed([
+              { type: "text", text: `model=${model};temperature=${temperature}` },
+              finishPart(11, 7),
+            ]),
+          streamText: () => Stream.empty,
+        }),
+    });
+
+    await expect(provider.generateContent("system", "prompt", "override-model", 0.4)).resolves.toEqual({
+      text: "model=override-model;temperature=0.4",
+      inToken: 11,
+      outToken: 7,
+      model: "override-model",
+    });
+  });
+
+  it("adapts Effect LanguageModel stream parts to TrustGraph chunks", async () => {
+    const provider = makeLanguageModelProvider({
+      provider: "FakeLanguageModel",
+      defaultModel: "fake-stream-model",
+      defaultTemperature: 0,
+      runtime: languageModelRuntime,
+      makeLanguageModel: () =>
+        LanguageModel.make({
+          generateText: () =>
+            Effect.succeed([
+              { type: "text", text: "unused" },
+              finishPart(1, 1),
+            ]),
+          streamText: () =>
+            Stream.fromArray([
+              { type: "text-delta", id: "part-1", delta: "hel" },
+              { type: "text-delta", id: "part-1", delta: "lo" },
+              finishPart(13, 8),
+            ]),
+        }),
+    });
+
+    const chunks: Array<LlmChunk> = [];
+    for await (const chunk of provider.generateContentStream("system", "prompt")) {
+      chunks.push(chunk);
+    }
+
+    expect(chunks).toEqual([
+      {
+        text: "hel",
+        inToken: null,
+        outToken: null,
+        model: "fake-stream-model",
+        isFinal: false,
+      },
+      {
+        text: "lo",
+        inToken: null,
+        outToken: null,
+        model: "fake-stream-model",
+        isFinal: false,
+      },
+      {
+        text: "",
+        inToken: 13,
+        outToken: 8,
+        model: "fake-stream-model",
+        isFinal: true,
+      },
+    ]);
+  });
+
+  it("maps Effect AI rate and quota failures to TrustGraph retry errors", async () => {
+    const reasons = [
+      new AiError.RateLimitError({}),
+      new AiError.QuotaExhaustedError({}),
+    ];
+
+    for (const reason of reasons) {
+      const provider = makeLanguageModelProvider({
+        provider: "FakeLanguageModel",
+        defaultModel: "fake-model",
+        defaultTemperature: 0,
+        runtime: languageModelRuntime,
+        makeLanguageModel: () =>
+          LanguageModel.make({
+            generateText: () => Effect.fail(aiError(reason)),
+            streamText: () => Stream.fail(aiError(reason)),
+          }),
+      });
+
+      await expect(provider.generateContent("system", "prompt")).rejects.toMatchObject({
+        _tag: "TooManyRequestsError",
+        message: "Rate limit exceeded",
+      });
+    }
+  });
 });