File size: 1,724 Bytes
3a01622
 
 
 
41f8b74
3a01622
 
 
41f8b74
f7db219
 
 
 
04b77a5
3a01622
 
 
 
 
 
 
 
 
 
 
 
 
 
f7db219
04b77a5
3a01622
 
 
 
 
 
 
 
 
 
 
 
 
 
 
f7db219
04b77a5
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
43
44
45
46
47
48
49
50
51
import { z } from "zod";
import {
	embeddingEndpointTei,
	embeddingEndpointTeiParametersSchema,
} from "./tei/embeddingEndpoints";
import {
	embeddingEndpointTransformersJS,
	embeddingEndpointTransformersJSParametersSchema,
} from "./transformersjs/embeddingEndpoints";
import {
	embeddingEndpointOpenAI,
	embeddingEndpointOpenAIParametersSchema,
} from "./openai/embeddingEndpoints";
import { embeddingEndpointHfApi, embeddingEndpointHfApiSchema } from "./hfApi/embeddingHfApi";

// 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,
	embeddingEndpointOpenAIParametersSchema,
	embeddingEndpointHfApiSchema,
]);

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,
	openai: embeddingEndpointOpenAI,
	hfapi: embeddingEndpointHfApi,
};

export default embeddingEndpoints;