alexion
4def357d09
Imported as a linear vendor commit so forge rebases do not replay the upstream root tree. git-subtree-dir: third_party/pi git-subtree-split: d981de1229ef899957bbe968bc8dcda02a21f477
99 lines
3.1 KiB
TypeScript
99 lines
3.1 KiB
TypeScript
import type { Api, AssistantMessage, AssistantMessageEvent, Model, ProviderStreams } from "../types.ts";
|
|
import { AssistantMessageEventStream } from "../utils/event-stream.ts";
|
|
|
|
function createSetupErrorMessage(model: Model<Api>, error: unknown): AssistantMessage {
|
|
return {
|
|
role: "assistant",
|
|
content: [],
|
|
api: model.api,
|
|
provider: model.provider,
|
|
model: model.id,
|
|
usage: {
|
|
input: 0,
|
|
output: 0,
|
|
cacheRead: 0,
|
|
cacheWrite: 0,
|
|
totalTokens: 0,
|
|
cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
|
|
},
|
|
stopReason: "error",
|
|
errorMessage: error instanceof Error ? error.message : String(error),
|
|
timestamp: Date.now(),
|
|
};
|
|
}
|
|
|
|
function hasResult(
|
|
source: AsyncIterable<AssistantMessageEvent>,
|
|
): source is AsyncIterable<AssistantMessageEvent> & { result(): Promise<AssistantMessage> } {
|
|
return typeof (source as { result?: unknown }).result === "function";
|
|
}
|
|
|
|
async function forwardStream(
|
|
target: AssistantMessageEventStream,
|
|
source: AsyncIterable<AssistantMessageEvent>,
|
|
): Promise<void> {
|
|
for await (const event of source) {
|
|
target.push(event);
|
|
}
|
|
target.end(hasResult(source) ? await source.result() : undefined);
|
|
}
|
|
|
|
/**
|
|
* Returns a stream synchronously while running async setup (auth resolution,
|
|
* lazy module loading) behind it. Setup failures terminate the stream with an
|
|
* error event.
|
|
*/
|
|
export function lazyStream(
|
|
model: Model<Api>,
|
|
setup: () => Promise<AsyncIterable<AssistantMessageEvent>>,
|
|
): AssistantMessageEventStream {
|
|
const outer = new AssistantMessageEventStream();
|
|
|
|
setup()
|
|
.then((inner) => forwardStream(outer, inner))
|
|
.catch((error) => {
|
|
const message = createSetupErrorMessage(model, error);
|
|
outer.push({ type: "error", reason: "error", error: message });
|
|
outer.end(message);
|
|
});
|
|
|
|
return outer;
|
|
}
|
|
|
|
/**
|
|
* Wraps a dynamically imported API implementation module as `ProviderStreams`.
|
|
* The module loads on first stream call; the host's import cache deduplicates
|
|
* loads. Load failures terminate the returned stream with an error event.
|
|
*/
|
|
export interface LazyApiCapabilities {
|
|
fetchDeferred?: boolean;
|
|
cancelDeferred?: boolean;
|
|
}
|
|
|
|
export function lazyApi(load: () => Promise<ProviderStreams>, capabilities?: LazyApiCapabilities): ProviderStreams {
|
|
const api: ProviderStreams = {
|
|
stream: (model, context, options) =>
|
|
lazyStream(model, async () => (await load()).stream(model, context, options)),
|
|
streamSimple: (model, context, options) =>
|
|
lazyStream(model, async () => (await load()).streamSimple(model, context, options)),
|
|
};
|
|
|
|
if (capabilities?.fetchDeferred) {
|
|
api.fetchDeferred = (model, handle, options) =>
|
|
lazyStream(model, async () => {
|
|
const implementation = await load();
|
|
if (!implementation.fetchDeferred) throw new Error("API does not support deferred responses");
|
|
return implementation.fetchDeferred(model, handle, options);
|
|
});
|
|
}
|
|
if (capabilities?.cancelDeferred) {
|
|
api.cancelDeferred = async (model, handle, options) => {
|
|
const implementation = await load();
|
|
if (!implementation.cancelDeferred) throw new Error("API cannot cancel deferred responses");
|
|
await implementation.cancelDeferred(model, handle, options);
|
|
};
|
|
}
|
|
|
|
return api;
|
|
}
|