Share text completion stream helpers

This commit is contained in:
elpresidank 2026-06-02 04:33:48 -05:00
parent 9d3f745fb0
commit 32fc7ea32d
9 changed files with 297 additions and 297 deletions

View file

@ -24,9 +24,11 @@ import {
} from "@trustgraph/base";
import { Effect, Layer, ManagedRuntime, Stream } from "effect";
import {
llmStreamPart,
optionalStringConfig,
providerStatusError,
requiredString,
streamTextCompletionChunks,
toAsyncGenerator,
type TextCompletionRuntimeError,
} from "./common.ts";
@ -147,53 +149,18 @@ export function makeOpenAICompatibleProvider(
catch: mapOpenAICompatibleError,
}),
).pipe(
Stream.flatMap((openAIStream) => {
const iterator = openAIStream[Symbol.asyncIterator]();
let totalInputTokens = 0;
let totalOutputTokens = 0;
return Stream.unfold<"pulling" | "done", LlmChunk, TextCompletionRuntimeError, never>(
"pulling",
(state) => {
if (state === "done") return Effect.as(Effect.void, undefined);
return Effect.gen(function* () {
while (true) {
const next = yield* Effect.tryPromise({
try: () => iterator.next(),
catch: mapOpenAICompatibleError,
});
if (next.done === true) {
return [{
text: "",
inToken: totalInputTokens,
outToken: totalOutputTokens,
model: modelName,
isFinal: true,
}, "done"] as const;
}
const chunk = next.value;
const content = chunk.choices[0]?.delta?.content;
if (chunk.usage !== null && chunk.usage !== undefined) {
totalInputTokens = chunk.usage.prompt_tokens;
totalOutputTokens = chunk.usage.completion_tokens;
}
if (content !== null && content !== undefined && content.length > 0) {
return [{
text: content,
inToken: null,
outToken: null,
model: modelName,
isFinal: false,
}, "pulling"] as const;
}
}
});
},
);
}),
Stream.flatMap((openAIStream) =>
streamTextCompletionChunks(openAIStream, {
model: modelName,
mapError: mapOpenAICompatibleError,
extract: (chunk) =>
llmStreamPart({
text: chunk.choices[0]?.delta?.content,
inToken: chunk.usage?.prompt_tokens,
outToken: chunk.usage?.completion_tokens,
}),
})
),
);
return toAsyncGenerator(Stream.toAsyncIterable(stream), mapOpenAICompatibleError);