Use Match for Effect AI stream parts

This commit is contained in:
elpresidank 2026-06-04 05:38:08 -05:00
parent 3a256096f8
commit 664aef44a7
3 changed files with 33 additions and 18 deletions

View file

@ -1,7 +1,7 @@
import { describe, expect, it } from "@effect/vitest";
import type { LlmChunk } from "@trustgraph/base";
import { Effect, Layer, ManagedRuntime, Stream } from "effect";
import { AiError, LanguageModel } from "effect/unstable/ai";
import { AiError, LanguageModel, Response } from "effect/unstable/ai";
import {
llmStreamPart,
makeLanguageModelProvider,
@ -157,6 +157,7 @@ describe("text completion common helpers", () => {
]),
streamText: () =>
Stream.fromArray([
Response.makePart("text-start", { id: "part-1" }),
{ type: "text-delta", id: "part-1", delta: "hel" },
{ type: "text-delta", id: "part-1", delta: "lo" },
finishPart(13, 8),

View file

@ -7,7 +7,7 @@ import {
type LlmResult,
type LlmProvider,
} from "@trustgraph/base";
import { Config, Effect, Layer, ManagedRuntime, Ref, Result, Stream } from "effect";
import { Config, Effect, Layer, ManagedRuntime, Match, Ref, Result, Stream } from "effect";
import * as O from "effect/Option";
import * as Predicate from "effect/Predicate";
import * as S from "effect/Schema";
@ -258,29 +258,26 @@ const languageModelStreamChunk = (
provider: string,
model: string,
part: Response.StreamPart<{}>,
): Effect.Effect<Result.Result<LlmChunk, undefined>, TextCompletionRuntimeError> => {
switch (part.type) {
case "text-delta":
return Effect.succeed(
): Effect.Effect<Result.Result<LlmChunk, undefined>, TextCompletionRuntimeError> =>
Match.value(part).pipe(
Match.discriminators("type")({
"text-delta": (part) => Effect.succeed(
part.delta.length > 0
? Result.succeed(textChunk(model, part.delta))
: Result.fail(undefined),
);
case "finish":
return Effect.succeed(
),
finish: (part) => Effect.succeed(
Result.succeed(
finalChunk(model, {
inToken: usageInputTokens(part.usage),
outToken: usageOutputTokens(part.usage),
}),
),
);
case "error":
return Effect.fail(effectAiProviderError(provider, part.error));
default:
return Effect.succeed(Result.fail(undefined));
}
};
),
error: (part) => Effect.fail(effectAiProviderError(provider, part.error)),
}),
Match.orElse(() => Effect.succeed(Result.fail(undefined))),
);
const runLanguageModelStream = <RuntimeRequirements, StreamRequirements extends RuntimeRequirements>(
runtime: ManagedRuntime.ManagedRuntime<RuntimeRequirements, TextCompletionRuntimeError>,