File size: 1,375 Bytes
3a01622 41f8b74 3a01622 41f8b74 3a01622 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 |
import { z } from "zod";
import {
embeddingEndpointTei,
embeddingEndpointTeiParametersSchema,
} from "./tei/embeddingEndpoints";
import {
embeddingEndpointTransformersJS,
embeddingEndpointTransformersJSParametersSchema,
} from "./transformersjs/embeddingEndpoints";
// parameters passed when generating text
interface EmbeddingEndpointParameters {
inputs: string[];
}
export type Embedding = number[];
// type signature for the endpoint
export type EmbeddingEndpoint = (params: EmbeddingEndpointParameters) => Promise<Embedding[]>;
export const embeddingEndpointSchema = z.discriminatedUnion("type", [
embeddingEndpointTeiParametersSchema,
embeddingEndpointTransformersJSParametersSchema,
]);
type EmbeddingEndpointTypeOptions = z.infer<typeof embeddingEndpointSchema>["type"];
// generator function that takes in type discrimantor value for defining the endpoint and return the endpoint
export type EmbeddingEndpointGenerator<T extends EmbeddingEndpointTypeOptions> = (
inputs: Extract<z.infer<typeof embeddingEndpointSchema>, { type: T }>
) => EmbeddingEndpoint | Promise<EmbeddingEndpoint>;
// list of all endpoint generators
export const embeddingEndpoints: {
[Key in EmbeddingEndpointTypeOptions]: EmbeddingEndpointGenerator<Key>;
} = {
tei: embeddingEndpointTei,
transformersjs: embeddingEndpointTransformersJS,
};
export default embeddingEndpoints;
|