67 lines
1.9 KiB
TypeScript
67 lines
1.9 KiB
TypeScript
import { type JSONSchema4 } from "json-schema";
|
|
import { type JsonValue } from "type-fest";
|
|
import { z } from "zod";
|
|
|
|
const ZodSupportedProvider = z.union([
|
|
z.literal("openai/ChatCompletion"),
|
|
z.literal("replicate/llama2"),
|
|
]);
|
|
|
|
export type SupportedProvider = z.infer<typeof ZodSupportedProvider>;
|
|
|
|
export const ZodModel = z.object({
|
|
name: z.string(),
|
|
contextWindow: z.number(),
|
|
promptTokenPrice: z.number().optional(),
|
|
completionTokenPrice: z.number().optional(),
|
|
pricePerSecond: z.number().optional(),
|
|
speed: z.union([z.literal("fast"), z.literal("medium"), z.literal("slow")]),
|
|
provider: ZodSupportedProvider,
|
|
description: z.string().optional(),
|
|
learnMoreUrl: z.string().optional(),
|
|
});
|
|
|
|
export type Model = z.infer<typeof ZodModel>;
|
|
|
|
export type FrontendModelProvider<SupportedModels extends string, OutputSchema> = {
|
|
name: string;
|
|
models: Record<SupportedModels, Model>;
|
|
|
|
normalizeOutput: (output: OutputSchema) => NormalizedOutput;
|
|
};
|
|
|
|
export type CompletionResponse<T> =
|
|
| { type: "error"; message: string; autoRetry: boolean; statusCode?: number }
|
|
| {
|
|
type: "success";
|
|
value: T;
|
|
timeToComplete: number;
|
|
statusCode: number;
|
|
promptTokens?: number;
|
|
completionTokens?: number;
|
|
cost?: number;
|
|
};
|
|
|
|
export type ModelProvider<SupportedModels extends string, InputSchema, OutputSchema> = {
|
|
getModel: (input: InputSchema) => SupportedModels | null;
|
|
shouldStream: (input: InputSchema) => boolean;
|
|
inputSchema: JSONSchema4;
|
|
getCompletion: (
|
|
input: InputSchema,
|
|
onStream: ((partialOutput: OutputSchema) => void) | null,
|
|
) => Promise<CompletionResponse<OutputSchema>>;
|
|
|
|
// This is just a convenience for type inference, don't use it at runtime
|
|
_outputSchema?: OutputSchema | null;
|
|
} & FrontendModelProvider<SupportedModels, OutputSchema>;
|
|
|
|
export type NormalizedOutput =
|
|
| {
|
|
type: "text";
|
|
value: string;
|
|
}
|
|
| {
|
|
type: "json";
|
|
value: JsonValue;
|
|
};
|