mirror of
https://github.com/trustgraph-ai/trustgraph.git
synced 2026-07-06 20:02:11 +02:00
Share text completion stream helpers
This commit is contained in:
parent
9d3f745fb0
commit
32fc7ea32d
9 changed files with 297 additions and 297 deletions
|
|
@ -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);
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue