inference-server 1.0.0-beta.19

package/src/types/index.ts

@@ -0,0 +1,488 @@
+import type { SomeJSONSchema } from 'ajv/dist/types/json-schema'
+import type { Sharp } from 'sharp'
+import type { BuiltInEngineName } from '#package/engines/index.js'
+import type { Logger } from '#package/lib/logger.js'
+import type { ModelPool } from '#package/pool.js'
+import type { ModelStore } from '#package/store.js'
+import {
+	AssistantMessage,
+	ChatMessage,
+	CompletionFinishReason,
+	TextCompletionParams,
+	ToolDefinition,
+} from '#package/types/completions.js'
+import type { ContextShiftStrategy } from '#package/engines/node-llama-cpp/types.js'
+import type {
+	StableDiffusionWeightType,
+	StableDiffusionSamplingMethod,
+  StableDiffusionSchedule,
+} from '#package/engines/stable-diffusion-cpp/types.js'
+import type {
+  TransformersJsModelClass,
+  TransformersJsTokenizerClass,
+  TransformersJsProcessorClass,
+  TransformersJsDataType,
+} from '#package/engines/transformers-js/types.js'
+export * from '#package/types/completions.js'
+
+export type ModelTaskType =
+	| 'text-completion'
+	| 'embedding'
+	| 'image-to-text'
+	| 'image-to-image'
+	| 'text-to-image'
+	| 'speech-to-text'
+
+export interface ModelOptionsBase {
+	engine: BuiltInEngineName | (string & {})
+	task: ModelTaskType | (string & {})
+	prepare?: 'blocking' | 'async' | 'on-demand'
+	minInstances?: number
+	maxInstances?: number
+	location?: string
+}
+
+export interface BuiltInModelOptionsBase extends ModelOptionsBase {
+	engine: BuiltInEngineName
+	task: ModelTaskType
+	url?: string
+	location?: string
+}
+
+export interface ModelConfigBase extends ModelOptionsBase {
+	id: string
+	minInstances: number
+	maxInstances: number
+	modelsCachePath: string
+}
+
+export interface ModelConfig extends ModelConfigBase {
+	url?: string
+	location?: string
+	task: ModelTaskType | (string & {})
+	engine: BuiltInEngineName | (string & {})
+	// minInstances: number
+	// maxInstances: number
+	ttl?: number
+	prefix?: string
+	initialMessages?: ChatMessage[]
+	device?: {
+		gpu?: boolean | 'auto' | (string & {})
+		// gpuLayers?: number
+		// cpuThreads?: number
+		// memLock?: boolean
+	}
+}
+
+// export interface ChatModelConfig extends ModelConfig {
+// 	initialMessages?: ChatMessage[]
+// }
+
+export interface CompletionChunk {
+	tokens: number[]
+	text: string
+}
+
+export interface ProcessingOptions {
+	timeout?: number
+	signal?: AbortSignal
+}
+
+export interface Image {
+	handle: Sharp
+	width: number
+	height: number
+	channels: 1 | 2 | 3 | 4
+}
+
+export interface CompletionProcessingOptions extends ProcessingOptions {
+	onChunk?: (chunk: CompletionChunk) => void
+}
+
+export interface SpeechToTextProcessingOptions extends ProcessingOptions {
+	onChunk?: (chunk: { text: string }) => void
+}
+
+export interface EngineContext<
+	TModelConfig = ModelConfig,
+	TModelMeta = unknown,
+> {
+	config: TModelConfig
+	meta?: TModelMeta
+	log: Logger
+}
+
+export interface TextCompletionRequestBase extends TextCompletionParams {
+	model: string
+	stream?: boolean
+}
+
+export interface TextCompletionRequest extends TextCompletionRequestBase {
+	prompt?: string
+}
+
+export interface ChatCompletionRequest extends TextCompletionRequestBase {
+	messages: ChatMessage[]
+	grammar?: string
+	tools?: Record<string, ToolDefinition>
+}
+
+export interface TextEmbeddingInput {
+	type: 'text'
+	content: string
+}
+
+export interface ImageEmbeddingInput {
+	type: 'image'
+	content: Image
+}
+
+export type EmbeddingInput = TextEmbeddingInput | ImageEmbeddingInput | string
+
+export interface EmbeddingRequest {
+	model: string
+	input: EmbeddingInput | EmbeddingInput[]
+	dimensions?: number
+	pooling?: 'cls' | 'mean'
+}
+
+export interface ImageToTextRequest {
+	model: string
+	image: Image
+	prompt?: string
+	maxTokens?: number
+}
+
+export interface StableDiffusionRequest {
+	negativePrompt?: string
+	guidance?: number
+	styleRatio?: number
+	strength?: number
+	sampleSteps?: number
+	batchCount?: number
+	samplingMethod?: StableDiffusionSamplingMethod
+	cfgScale?: number
+	controlStrength?: number
+}
+
+export interface TextToImageRequest extends StableDiffusionRequest {
+	model: string
+	prompt: string
+	width?: number
+	height?: number
+	seed?: number
+}
+
+export interface ImageToImageRequest extends StableDiffusionRequest {
+	model: string
+	image: Image
+	prompt: string
+	width?: number
+	height?: number
+	seed?: number
+}
+
+export interface SpeechToTextRequest {
+	model: string
+	url?: string
+	file?: string
+	language?: string
+	prompt?: string
+	maxTokens?: number
+}
+
+export interface ModelRequestMeta {
+	sequence: number
+	abortController: AbortController
+}
+export type IncomingRequest =
+	| TextCompletionRequest
+	| ChatCompletionRequest
+	| EmbeddingRequest
+	| ImageToTextRequest
+	| SpeechToTextRequest
+export type ModelInstanceRequest = ModelRequestMeta & IncomingRequest
+
+export interface EngineTextCompletionArgs<
+	TModelConfig = unknown,
+	TModelMeta = unknown,
+> extends EngineContext<TModelConfig, TModelMeta> {
+	onChunk?: (chunk: CompletionChunk) => void
+	resetContext?: boolean
+	request: TextCompletionRequest
+}
+
+export interface EngineChatCompletionArgs<
+	TModelConfig = unknown,
+	TModelMeta = unknown,
+> extends EngineContext<TModelConfig, TModelMeta> {
+	onChunk?: (chunk: CompletionChunk) => void
+	resetContext?: boolean
+	request: ChatCompletionRequest
+}
+
+export interface EngineEmbeddingArgs<
+	TModelConfig = unknown,
+	TModelMeta = unknown,
+> extends EngineContext<TModelConfig, TModelMeta> {
+	request: EmbeddingRequest
+}
+
+export interface EngineImageToTextArgs<
+	TModelConfig = unknown,
+	TModelMeta = unknown,
+> extends EngineContext<TModelConfig, TModelMeta> {
+	request: ImageToTextRequest
+}
+
+export interface EngineTextToImageArgs<
+	TModelConfig = unknown,
+	TModelMeta = unknown,
+> extends EngineContext<TModelConfig, TModelMeta> {
+	request: TextToImageRequest
+}
+
+export interface EngineImageToImageArgs<
+	TModelConfig = unknown,
+	TModelMeta = unknown,
+> extends EngineContext<TModelConfig, TModelMeta> {
+	request: ImageToImageRequest
+}
+
+export interface EngineSpeechToTextArgs<
+	TModelConfig = unknown,
+	TModelMeta = unknown,
+> extends EngineContext<TModelConfig, TModelMeta> {
+	request: SpeechToTextRequest
+	onChunk?: (chunk: { text: string }) => void
+}
+
+export interface FileDownloadProgress {
+	file: string
+	loadedBytes: number
+	totalBytes: number
+}
+
+export interface EngineStartContext {
+	pool: ModelPool
+	store: ModelStore
+}
+
+export interface ModelEngine<
+	TInstance = unknown,
+	TModelConfig extends ModelConfig = ModelConfig,
+	TModelMeta = unknown,
+> {
+	autoGpu?: boolean
+	start?: (ctx: EngineStartContext) => Promise<void>
+	prepareModel: (
+		ctx: EngineContext<TModelConfig, TModelMeta>,
+		onProgress?: (progress: FileDownloadProgress) => void,
+		signal?: AbortSignal,
+	) => Promise<TModelMeta>
+	createInstance: (
+		ctx: EngineContext<TModelConfig, TModelMeta>,
+		signal?: AbortSignal,
+	) => Promise<TInstance>
+	disposeInstance: (instance: TInstance) => Promise<void>
+	processChatCompletionTask?: (
+		args: EngineChatCompletionArgs<TModelConfig, TModelMeta>,
+		instance: TInstance,
+		signal?: AbortSignal,
+	) => Promise<EngineChatCompletionResult>
+	processTextCompletionTask?: (
+		args: EngineTextCompletionArgs<TModelConfig, TModelMeta>,
+		instance: TInstance,
+		signal?: AbortSignal,
+	) => Promise<EngineTextCompletionResult>
+	processEmbeddingTask?: (
+		args: EngineEmbeddingArgs<TModelConfig, TModelMeta>,
+		instance: TInstance,
+		signal?: AbortSignal,
+	) => Promise<EngineEmbeddingResult>
+	processImageToTextTask?: (
+		args: EngineImageToTextArgs<TModelConfig, TModelMeta>,
+		instance: TInstance,
+		signal?: AbortSignal,
+	) => Promise<EngineImageToTextResult>
+	processSpeechToTextTask?: (
+		args: EngineSpeechToTextArgs<TModelConfig, TModelMeta>,
+		instance: TInstance,
+		signal?: AbortSignal,
+	) => Promise<EngineSpeechToTextResult>
+	processTextToImageTask?: (
+		args: EngineTextToImageArgs<TModelConfig, TModelMeta>,
+		instance: TInstance,
+		signal?: AbortSignal,
+	) => Promise<EngineTextToImageResult>
+	processImageToImageTask?: (
+		args: EngineImageToImageArgs<TModelConfig, TModelMeta>,
+		instance: TInstance,
+		signal?: AbortSignal,
+	) => Promise<EngineImageToImageResult>
+}
+
+interface EmbeddingModelOptions {
+	task: 'embedding'
+}
+
+export type TextCompletionGrammar = string | SomeJSONSchema
+
+interface TextCompletionModelOptions {
+	task: 'text-completion'
+	contextSize?: number
+	grammars?: Record<string, TextCompletionGrammar>
+	completionDefaults?: TextCompletionParams
+	initialMessages?: ChatMessage[]
+	prefix?: string
+	batchSize?: number
+}
+
+interface LlamaCppModelOptionsBase extends BuiltInModelOptionsBase {
+	engine: 'node-llama-cpp'
+	task: 'text-completion' | 'embedding'
+	sha256?: string
+	batchSize?: number
+	contextShiftStrategy?: ContextShiftStrategy
+	tools?: {
+		definitions: Record<string, ToolDefinition>
+		includeParamsDocumentation?: boolean
+		parallelism?: number
+	}
+	device?: {
+		gpu?: boolean | 'auto' | (string & {})
+		gpuLayers?: number
+		cpuThreads?: number
+		memLock?: boolean
+	}
+}
+
+interface LlamaCppEmbeddingModelOptions
+	extends LlamaCppModelOptionsBase,
+		EmbeddingModelOptions {
+	task: 'embedding'
+}
+
+export interface LlamaCppTextCompletionModelOptions
+	extends LlamaCppModelOptionsBase,
+		TextCompletionModelOptions {
+	task: 'text-completion'
+}
+
+interface GPT4AllModelOptions extends BuiltInModelOptionsBase {
+	engine: 'gpt4all'
+	task: 'text-completion' | 'embedding'
+	md5?: string
+	device?: {
+		gpu?: boolean | 'auto' | (string & {})
+		gpuLayers?: number
+		cpuThreads?: number
+	}
+}
+
+type GPT4AllTextCompletionModelOptions = TextCompletionModelOptions &
+	GPT4AllModelOptions
+
+type GPT4AllEmbeddingModelOptions = GPT4AllModelOptions & EmbeddingModelOptions
+
+export interface TransformersJsModel {
+	processor?: {
+  	url?: string
+		file?: string
+	}
+	processorClass?: TransformersJsProcessorClass
+	tokenizerClass?: TransformersJsTokenizerClass
+	modelClass?: TransformersJsModelClass
+	dtype?: Record<string, TransformersJsDataType> | TransformersJsDataType
+}
+
+interface TransformersJsModelOptions extends BuiltInModelOptionsBase {
+	engine: 'transformers-js'
+	task: 'image-to-text' | 'speech-to-text' | 'text-completion' | 'embedding'
+	textModel?: TransformersJsModel
+	visionModel?: TransformersJsModel
+	speechModel?: TransformersJsModel
+	device?: {
+		gpu?: boolean | 'auto' | (string & {})
+	}
+}
+
+export interface ModelFileSource {
+	url?: string
+	file?: string
+	sha256?: string
+}
+
+interface StableDiffusionModelOptions extends BuiltInModelOptionsBase {
+	engine: 'stable-diffusion-cpp'
+	task: 'image-to-text' | 'text-to-image' | 'image-to-image'
+	sha256?: string
+	url?: string
+	diffusionModel?: boolean
+	vae?: ModelFileSource
+	clipL?: ModelFileSource
+	clipG?: ModelFileSource
+	t5xxl?: ModelFileSource
+	taesd?: ModelFileSource
+	controlNet?: ModelFileSource
+	samplingMethod?: StableDiffusionSamplingMethod
+	weightType?: StableDiffusionWeightType
+	schedule?: StableDiffusionSchedule
+	loras?: ModelFileSource[]
+}
+
+export interface CustomEngineModelOptions extends ModelOptionsBase {}
+
+export type BuiltInModelOptions =
+	| LlamaCppTextCompletionModelOptions
+	| LlamaCppEmbeddingModelOptions
+	| GPT4AllTextCompletionModelOptions
+	| GPT4AllEmbeddingModelOptions
+	| TransformersJsModelOptions
+	| StableDiffusionModelOptions
+
+export type ModelOptions = BuiltInModelOptions | CustomEngineModelOptions
+
+export interface EngineEmbeddingResult {
+	embeddings: Float32Array[]
+	inputTokens: number
+}
+
+export interface ChatCompletionResult extends EngineChatCompletionResult {
+	id: string
+	model: string
+}
+
+export interface EngineChatCompletionResult {
+	message: AssistantMessage
+	finishReason: CompletionFinishReason
+	promptTokens: number
+	completionTokens: number
+	contextTokens: number
+}
+
+export interface EngineTextCompletionResult {
+	text: string
+	finishReason?: CompletionFinishReason
+	promptTokens: number
+	completionTokens: number
+	contextTokens: number
+}
+
+export interface EngineImageToTextResult {
+	text: string
+}
+
+export interface EngineTextToImageResult {
+	images: Image[]
+	seed: number
+}
+
+export interface EngineImageToImageResult {
+	images: Image[]
+	seed: number
+}
+
+export interface EngineSpeechToTextResult {
+	text: string
+}

package/tsconfig.json

@@ -0,0 +1,29 @@
+{
+	"compilerOptions": {
+		"target": "es2022",
+		"module": "es2022",
+		"sourceMap": true,
+		"declaration": true,
+		"stripInternal": true,
+		"strict": true,
+		"noImplicitAny": true,
+		"noImplicitReturns": true,
+		"noImplicitThis": true,
+		"noImplicitOverride": true,
+		"strictNullChecks": true,
+		"noFallthroughCasesInSwitch": true,
+		"removeComments": false,
+		"isolatedModules": true,
+		"moduleResolution": "node",
+		"lib": ["es2022"],
+		"resolveJsonModule": false,
+		"esModuleInterop": true,
+		"forceConsistentCasingInFileNames": true,
+		"allowSyntheticDefaultImports": true,
+		"skipLibCheck": true,
+		"paths": {
+			"#package/*": ["./src/*"]
+		}
+	},
+	"include": ["src/**/*"]
+}

package/tsconfig.release.json

@@ -0,0 +1,11 @@
+{
+	"extends": "./tsconfig.json",
+	"compilerOptions": {
+		"rootDir": "src",
+		"outDir": "dist",
+		"noEmit": false
+	},
+	"include": [
+		"src/**/*"
+	]
+}

package/vitest.config.ts

@@ -0,0 +1,18 @@
+import { defineConfig } from 'vitest/config'
+
+export default defineConfig({
+	test: {
+		pool: 'forks',
+		maxWorkers: 1,
+		minWorkers: 1,
+		maxConcurrency: 1,
+		testTimeout: 60000,
+		hookTimeout: 120000,
+		poolOptions: {
+			threads: {
+				minThreads: 1,
+				maxThreads: 1,
+			},
+		},
+	},
+})