Ver Fonte

update library interface

Bruce MacDonald há 1 ano atrás
pai
commit
8a1885abc4
3 ficheiros alterados com 430 adições e 193 exclusões
  1. 239 102
      src/index.ts
  2. 145 73
      src/interfaces.ts
  3. 46 18
      src/utils.ts

+ 239 - 102
src/index.ts

@@ -1,23 +1,32 @@
 import * as utils from "./utils.js";
 import * as utils from "./utils.js";
+import { promises, createReadStream } from 'fs';
+import { join, resolve, dirname } from 'path';
+import { createHash } from 'crypto';
+import { homedir } from 'os';
 
 
 import type {
 import type {
 	Fetch,
 	Fetch,
 	Config,
 	Config,
-	TagsResponse,
-	Tag,
+    GenerateRequest,
+    PullRequest,
+    PushRequest,
+    CreateRequest,
+    EmbeddingsRequest,
 	GenerateResponse,
 	GenerateResponse,
-	GenerateResponseEnd,
-	GenerateResult,
-	CreateResponse,
-	CreateStatus,
-	PullResponse,
-	PullResult,
 	EmbeddingsResponse,
 	EmbeddingsResponse,
-	GenerateOptions,
-	GenerateRequest,
-	ModelParameters
+    ListResponse,
+    ProgressResponse,
+    ErrorResponse,
+    StatusResponse,
+    DeleteRequest,
+    CopyRequest,
+    ShowResponse,
+    ShowRequest,
+    ChatRequest,
+    ChatResponse,
 } from "./interfaces.js";
 } from "./interfaces.js";
 
 
+
 export class Ollama {
 export class Ollama {
 	private readonly config: Config;
 	private readonly config: Config;
 	private readonly fetch: Fetch;
 	private readonly fetch: Fetch;
@@ -44,109 +53,237 @@ export class Ollama {
 		this.fetch = f;
 		this.fetch = f;
 	}
 	}
 
 
-	async tags (): Promise<Tag[]> {
-		const response = await utils.get(this.fetch, `${this.config.address}/api/tags`);
-		const json = await response.json() as TagsResponse;
-
-		return json.models.map(m => ({
-			name: m.name,
-			modifiedAt: new Date(m.modified_at),
-			size: m.size
-		}));
-	}
-
-	async * generate (model: string, prompt: string, options?: Partial<GenerateOptions>): AsyncGenerator<string, GenerateResult> {
-		const parameters = options?.parameters;
-
-		delete options?.parameters;
-
-		const request: GenerateRequest = { model, prompt, ...options };
-
-		if (parameters != null) {
-			request.options = parameters;
+    private async processStreamableRequest<T extends object>(endpoint: string, request: { stream?: boolean } & Record<string, any>): Promise<T | AsyncGenerator<T>> {
+        request.stream = request.stream ?? false;
+        const response = await utils.post(this.fetch, `${this.config.address}/api/${endpoint}`, { ...request });
+    
+        if (!response.body) {
+            throw new Error("Missing body");
+        }
+    
+        const itr = utils.parseJSON<T | ErrorResponse>(response.body);
+    
+        if (request.stream) {
+            return (async function* () {
+                for await (const message of itr) {
+                    if ('error' in message) {
+                        throw new Error(message.error);
+                    }
+                    yield message;
+                    // message will be done in the case of chat and generate
+                    // message will be success in the case of a progress response (pull, push, create)
+                    if ((message as any).done || (message as any).status === "success") {
+                        return;
+                    }
+                }
+                throw new Error("Did not receive done or success response in stream.");
+            })();
+        } else {
+            const message = await itr.next();
+            if (!message.value.done && (message.value as any).status !== "success") {
+                throw new Error("Expected a completed response.");
+            }
+            return message.value;
+        }
+    }
+
+	private async encodeImage(image: Uint8Array | Buffer | string): Promise<string> {
+		if (typeof image !== 'string') {
+			// image is Uint8Array or Buffer, convert it to base64
+			const result = Buffer.from(image).toString('base64');
+			return result;
 		}
 		}
-
-		const response = await utils.post(this.fetch, `${this.config.address}/api/generate`, { ...request });
-
-		if (!response.body) {
-			throw new Error("Missing body");
-		}
-
-		const itr = utils.parseJSON<GenerateResponse | GenerateResponseEnd>(response.body);
-
-		for await (const message of itr) {
-			if (message.done) {
-				return {
-					model: message.model,
-					createdAt: new Date(message.created_at),
-					context: message.context,
-					totalDuration: message.total_duration,
-					loadDuration: message.load_duration,
-					promptEvalCount: message.prompt_eval_count,
-					evalCount: message.eval_count,
-					evalDuration: message.eval_duration
-				};
-			}
-
-			yield message.response;
-		}
-
-		throw new Error("Did not recieve done response in stream.");
+		const base64Pattern = /^[A-Za-z0-9+/]+={1,2}$/; // detect by checking for equals signs at the end
+		if (base64Pattern.test(image)) {
+			// the string is already base64 encoded
+			return image;
+		} 
+		// this is a filepath, read the file and convert it to base64
+		const fileBuffer = await promises.readFile(resolve(image));
+		return Buffer.from(fileBuffer).toString('base64');
+	}	
+
+    private async parseModelfile(modelfile: string, mfDir: string = process.cwd()): Promise<string> {
+        const out: string[] = [];
+        const lines = modelfile.split('\n');
+        for (const line of lines) {
+            const [command, args] = line.split(' ', 2);
+            if (['FROM', 'ADAPTER'].includes(command.toUpperCase())) {
+                const path = this.resolvePath(args.trim(), mfDir);
+                if (await this.fileExists(path)) {
+                    out.push(`${command} @${await this.createBlob(path)}`);
+                } else {
+                    out.push(`${command} ${args}`);
+                }
+            } else {
+                out.push(line);
+            }
+        }
+        return out.join('\n');
+    }
+
+    private resolvePath(inputPath, mfDir) {
+        if (inputPath.startsWith('~')) {
+            return join(homedir(), inputPath.slice(1));
+        }
+        return resolve(mfDir, inputPath);
+    }
+
+    private async fileExists(path: string): Promise<boolean> {
+        try {
+            await promises.access(path);
+            return true;
+        } catch {
+            return false;
+        }
+    }
+
+    private async createBlob(path: string): Promise<string> {
+        if (typeof ReadableStream === 'undefined') {
+            // Not all fetch implementations support streaming
+            // TODO: support non-streaming uploads
+            throw new Error("Streaming uploads are not supported in this environment.");
+        }
+
+        // Create a stream for reading the file
+        const fileStream = createReadStream(path);
+
+        // Compute the SHA256 digest
+        const sha256sum = await new Promise<string>((resolve, reject) => {
+            const hash = createHash('sha256');
+            fileStream.on('data', data => hash.update(data));
+            fileStream.on('end', () => resolve(hash.digest('hex')));
+            fileStream.on('error', reject);
+        });
+
+        const digest = `sha256:${sha256sum}`;
+
+        try {
+            await utils.head(this.fetch, `${this.config.address}/api/blobs/${digest}`);
+        } catch (e) {
+            if (e instanceof Error && e.message.includes('404')) {
+                // Create a new readable stream for the fetch request
+                const readableStream = new ReadableStream({
+                    start(controller) {
+                        fileStream.on('data', chunk => {
+                            controller.enqueue(chunk);  // Enqueue the chunk directly
+                        });
+                
+                        fileStream.on('end', () => {
+                            controller.close();  // Close the stream when the file ends
+                        });
+                
+                        fileStream.on('error', err => {
+                            controller.error(err);  // Propagate errors to the stream
+                        });
+                    }
+                });
+
+                await utils.post(this.fetch, `${this.config.address}/api/blobs/${digest}`, readableStream);
+            } else {
+                throw e;
+            }
+        }
+
+        return digest;
+    }
+
+    generate(request: GenerateRequest & { stream: true }): Promise<AsyncGenerator<GenerateResponse>>;
+    generate(request: GenerateRequest & { stream?: false }): Promise<GenerateResponse>;
+
+    async generate(request: GenerateRequest): Promise<GenerateResponse | AsyncGenerator<GenerateResponse>> {
+        if (request.images) {
+            request.images = await Promise.all(request.images.map(this.encodeImage.bind(this)));
+        }
+        return this.processStreamableRequest<GenerateResponse>('generate', request);
+    }
+
+    chat(request: ChatRequest & { stream: true }): Promise<AsyncGenerator<ChatResponse>>;
+    chat(request: ChatRequest & { stream?: false }): Promise<ChatResponse>;
+
+    async chat(request: ChatRequest): Promise<ChatResponse | AsyncGenerator<ChatResponse>> {
+        if (request.messages) {
+            for (const message of request.messages) {
+                if (message.images) {
+                    message.images = await Promise.all(message.images.map(this.encodeImage.bind(this)));
+                }
+            }
+        }
+        return this.processStreamableRequest<ChatResponse>('chat', request);
+    }
+
+    pull(request: PullRequest & { stream: true }): Promise<AsyncGenerator<ProgressResponse>>;
+    pull(request: PullRequest & { stream?: false }): Promise<ProgressResponse>;
+
+    async pull (request: PullRequest):  Promise<ProgressResponse | AsyncGenerator<ProgressResponse>> {
+        return this.processStreamableRequest<ProgressResponse>('pull', {
+			name: request.model,
+			stream: request.stream,
+			insecure: request.insecure,
+			username: request.username,
+			password: request.password,
+		});
 	}
 	}
 
 
-	async * create (name: string, path: string): AsyncGenerator<CreateStatus> {
-		const response = await utils.post(this.fetch, `${this.config.address}/api/create`, { name, path });
+    push(request: PushRequest & { stream: true }): Promise<AsyncGenerator<ProgressResponse>>;
+    push(request: PushRequest & { stream?: false }): Promise<ProgressResponse>;
 
 
-		if (!response.body) {
-			throw new Error("Missing body");
-		}
-
-		const itr = utils.parseJSON<CreateResponse>(response.body);
-
-		for await (const message of itr) {
-			yield message.status;
-		}
-	}
-
-	async copy (source: string, destination: string): Promise<void> {
-		await utils.post(this.fetch, `${this.config.address}/api/copy`, {
-			source,
-			destination
+    async push (request: PushRequest):  Promise<ProgressResponse | AsyncGenerator<ProgressResponse>> {
+        return this.processStreamableRequest<ProgressResponse>('push', {
+			name: request.model,
+			stream: request.stream,
+			insecure: request.insecure,
+			username: request.username,
+			password: request.password,
 		});
 		});
 	}
 	}
 
 
-	async delete (name: string): Promise<void> {
-		await utils.del(this.fetch, `${this.config.address}/api/delete`, { name });
+    create(request: CreateRequest & { stream: true }): Promise<AsyncGenerator<ProgressResponse>>;
+    create(request: CreateRequest & { stream?: false }): Promise<ProgressResponse>;
+
+	async create (request: CreateRequest): Promise<ProgressResponse | AsyncGenerator<ProgressResponse>> {
+        let modelfileContent = '';
+        if (request.path) {
+            modelfileContent = await promises.readFile(request.path, { encoding: 'utf8' });
+            modelfileContent = await this.parseModelfile(modelfileContent, dirname(request.path));
+        } else if (request.modelfile) {
+            modelfileContent = await this.parseModelfile(request.modelfile);
+        } else {
+            throw new Error('Must provide either path or modelfile to create a model');
+        }
+
+        return this.processStreamableRequest<ProgressResponse>('create', {
+            name: request.model,
+            stream: request.stream,
+            modelfile: modelfileContent,
+        });
 	}
 	}
 
 
-	async * pull (name: string): AsyncGenerator<PullResult> {
-		const response = await utils.post(this.fetch, `${this.config.address}/api/pull`, { name });
-
-		if (!response.body) {
-			throw new Error("Missing body");
-		}
-
-		const itr = utils.parseJSON<PullResponse>(response.body);
-
-		for await (const message of itr) {
-			yield {
-				status: message.status,
-				digest: message["digest"] ?? "",
-				total: message["total"] ?? 0,
-				completed: message["completed"] ?? 0
-			};
-		}
+    async delete (request: DeleteRequest): Promise<StatusResponse> {
+		await utils.del(this.fetch, `${this.config.address}/api/delete`, { name: request.model });
+        return { status: "success" };
+	}
+    
+    async copy (request: CopyRequest): Promise<StatusResponse> {
+		await utils.post(this.fetch, `${this.config.address}/api/copy`, { ...request });
+        return { status: "success" };
 	}
 	}
 
 
-	async embeddings (model: string, prompt: string, parameters?: Partial<ModelParameters>): Promise<number[]> {
-		const response = await utils.post(this.fetch, `${this.config.address}/api/embeddings`, {
-			model,
-			prompt,
-			options: parameters ?? {}
-		});
+    async list (): Promise<ListResponse> {
+		const response = await utils.get(this.fetch, `${this.config.address}/api/tags`);
+		const listResponse = await response.json() as ListResponse;
+		return listResponse;
+	}
 
 
-		const json = await response.json() as EmbeddingsResponse;
+    async show (request: ShowRequest): Promise<ShowResponse> {
+        const response = await utils.post(this.fetch, `${this.config.address}/api/show`, { ...request });
+        const showResponse = await response.json() as ShowResponse;
+        return showResponse;
+    }
 
 
-		return json.embedding;
+	async embeddings (request: EmbeddingsRequest): Promise<EmbeddingsResponse> {
+		const response = await utils.post(this.fetch, `${this.config.address}/api/embeddings`, { request });
+		const embeddingsResponse = await response.json() as EmbeddingsResponse;
+		return embeddingsResponse;
 	}
 	}
 }
 }

+ 145 - 73
src/interfaces.ts

@@ -5,118 +5,190 @@ export interface Config {
 	fetch?: Fetch
 	fetch?: Fetch
 }
 }
 
 
-export interface ModelParameters {
-	mirostat: number
-	mirostat_eta: number
-	mirostat_tau: number
-	num_ctx: number
-	num_gqa: number
-	num_thread: number
-	repeat_last_n: number
-	repeat_penalty: number
-	temperature: number
-	stop: string
-	tfs_z: number
-	top_k: number
-	top_p: number
-}
-
-export interface GenerateOptions {
-	parameters: Partial<ModelParameters>
-	context: number[]
-	template: string
-	system: string
+// request types
+
+export interface Options {
+    numa: boolean;
+    num_ctx: number;
+    num_batch: number;
+    main_gpu: number;
+    low_vram: boolean;
+    f16_kv: boolean;
+    logits_all: boolean;
+    vocab_only: boolean;
+    use_mmap: boolean;
+    use_mlock: boolean;
+    embedding_only: boolean;
+    num_thread: number;
+
+    // Runtime options
+    num_keep: number;
+    seed: number;
+    num_predict: number;
+    top_k: number;
+    top_p: number;
+    tfs_z: number;
+    typical_p: number;
+    repeat_last_n: number;
+    temperature: number;
+    repeat_penalty: number;
+    presence_penalty: number;
+    frequency_penalty: number;
+    mirostat: number;
+    mirostat_tau: number;
+    mirostat_eta: number;
+    penalize_newline: boolean;
+    stop: string[];
 }
 }
 
 
-export interface GenerateResult {
+export interface GenerateRequest {
 	model: string
 	model: string
-	createdAt: Date
-	context: number[]
-	totalDuration: number
-	loadDuration: number
-	promptEvalCount: number
-	evalCount: number
-	evalDuration: number
+	prompt: string
+	system?: string
+	template?: string
+	context?: number[]
+	stream?: boolean
+	raw?: boolean
+	format?: string
+	images?: Uint8Array[] | string[]
+
+	options?: Partial<Options>
 }
 }
 
 
-export interface Tag {
-	name: string
-	modifiedAt: Date
-	size: number
+export interface Message {
+	role: string
+	content: string
+	images?: Uint8Array[] | string[]
 }
 }
 
 
-export interface PullResult {
-	status: PullStatus
-	digest: string
-	total: number
-	completed: number
+export interface ChatRequest {
+	model: string
+	messages?: Message[]
+	stream?: boolean
+	format?: string
+
+	options?: Partial<Options>
 }
 }
 
 
-// Responses:
-export interface ErrorResponse {
-	error: string
+export interface PullRequest {
+	model: string
+	insecure?: boolean
+	username?: string
+	password?: string
+	stream?: boolean
 }
 }
 
 
-export interface TagsResponse {
-	models: {
-		name: string
-		modified_at: string
-		size: number
-	}[]
+export interface PushRequest {
+	model: string
+	insecure?: boolean
+	username?: string
+	password?: string
+	stream?: boolean
 }
 }
 
 
-export interface GenerateRequest {
+export interface CreateRequest {
+	model: string
+	path?: string
+	modelfile?: string
+	stream?: boolean
+}
+
+export interface DeleteRequest {
+	model: string
+}
+
+export interface CopyRequest {
+	source: string
+	destination: string
+}
+
+export interface ShowRequest {
 	model: string
 	model: string
-	prompt: string
-	options?: Partial<ModelParameters>
 	system?: string
 	system?: string
 	template?: string
 	template?: string
-	context?: number[]
+	options?: Partial<Options>
 }
 }
 
 
-export interface GenerateResponse {
+export interface EmbeddingsRequest {
 	model: string
 	model: string
-	created_at: string
-	response: string
-	done: false
+	prompt: string
+
+	options?: Partial<Options>
 }
 }
 
 
-export interface GenerateResponseEnd {
+// response types
+
+export interface GenerateResponse {
 	model: string
 	model: string
-	created_at: string
-	done: true
+	created_at: Date
+	response: string
+	done: boolean
 	context: number[]
 	context: number[]
 	total_duration: number
 	total_duration: number
 	load_duration: number
 	load_duration: number
 	prompt_eval_count: number
 	prompt_eval_count: number
+	prompt_eval_duration: number
 	eval_count: number
 	eval_count: number
 	eval_duration: number
 	eval_duration: number
 }
 }
 
 
-export type CreateStatus = "parsing modelfile" | "looking for model" | "creating model layer" | "creating model template layer" | "creating model system layer" | "creating parameter layer" | "creating config layer" | `writing layer ${string}` | `using already created layer ${string}` | "writing manifest" | "removing any unused layers" | "success"
-
-export interface CreateResponse {
-	status: CreateStatus
+export interface ChatResponse {
+	model: string
+	created_at: Date
+	message: Message
+	done: boolean
+	total_duration: number
+	load_duration: number
+	prompt_eval_count: number
+	prompt_eval_duration: number
+	eval_count: number
+	eval_duration: number
 }
 }
 
 
-export type PullStatus = "" | "pulling manifest" | "verifying sha256 digest" | "writing manifest" | "removing any unused layers" | "success" | `downloading ${string}`
-
-interface PullResponseStatus {
-	status: PullStatus
+export interface EmbeddingsResponse {
+	embedding: number[]
 }
 }
 
 
-interface PullResponseDownloadStart {
-	status: `downloading ${string}`
+export interface ProgressResponse {
+	status: string
 	digest: string
 	digest: string
 	total: number
 	total: number
+	completed: number
 }
 }
 
 
-interface PullResponseDownloadUpdate extends PullResponseDownloadStart {
-	completed: number
+export interface ModelResponse {
+	name: string
+	modified_at: Date
+	size: number
+	digest: string
+	format: string
+	family: string
+	families: string[]
+	parameter_size: string
+	quatization_level: number
+}
+
+export interface ShowResponse {
+	license: string
+	modelfile: string
+	parameters: string
+	template: string
+	system: string
+	format: string
+	family: string
+	families: string[]
+	parameter_size: string
+	quatization_level: number
 }
 }
 
 
-export type PullResponse = PullResponseStatus | PullResponseDownloadStart | PullResponseDownloadUpdate
+export interface ListResponse {
+	models: ModelResponse[]
+}
 
 
-export interface EmbeddingsResponse {
-	embedding: number[]
+export interface ErrorResponse {
+	error: string
+}
+
+export interface StatusResponse {
+	status: string
 }
 }

+ 46 - 18
src/utils.ts

@@ -13,17 +13,28 @@ export const formatAddress = (address: string): string => {
 };
 };
 
 
 const checkOk = async (response: Response): Promise<void> => {
 const checkOk = async (response: Response): Promise<void> => {
-	if (!response.ok) {
-		let message = await response.text();
-
-		try {
-			message = (JSON.parse(message) as ErrorResponse).error;
-		} catch(error) {
-			// Do nothing.
-		}
-
-		throw new Error(message);
-	}
+    if (!response.ok) {
+        let message = `Error ${response.status}: ${response.statusText}`;
+
+        if (response.headers.get('content-type')?.includes('application/json')) {
+            try {
+                const errorResponse = await response.json() as ErrorResponse;
+                message = errorResponse.error || message;
+            } catch(error) {
+                console.log("Failed to parse error response as JSON");
+            }
+        } else {
+            try {
+                console.log("Getting text from response");
+                const textResponse = await response.text();
+                message = textResponse || message;
+            } catch (error) {
+                console.log("Failed to get text from error response");
+            }
+        }
+
+        throw new Error(message);
+    }
 };
 };
 
 
 export const get = async (fetch: Fetch, address: string): Promise<Response> => {
 export const get = async (fetch: Fetch, address: string): Promise<Response> => {
@@ -34,17 +45,34 @@ export const get = async (fetch: Fetch, address: string): Promise<Response> => {
 	return response;
 	return response;
 };
 };
 
 
-export const post = async (fetch: Fetch, address: string, data?: Record<string, unknown>): Promise<Response> => {
-	const response = await fetch(formatAddress(address), {
-		method: "POST",
-		body: JSON.stringify(data)
-	});
+export const head = async (fetch: Fetch, address: string): Promise<Response> => {
+    const response = await fetch(formatAddress(address), {
+        method: "HEAD"
+    });
 
 
-	await checkOk(response);
+    await checkOk(response);
 
 
-	return response;
+    return response;
+};
+
+export const post = async (fetch: Fetch, address: string, data?: Record<string, unknown> | BodyInit): Promise<Response> => {
+    const isRecord = (input: any): input is Record<string, unknown> => {
+        return input !== null && typeof input === 'object' && !Array.isArray(input);
+    };
+
+    const formattedData = isRecord(data) ? JSON.stringify(data) : data;
+
+    const response = await fetch(formatAddress(address), {
+        method: "POST",
+        body: formattedData
+    });
+
+    await checkOk(response);
+
+    return response;
 };
 };
 
 
+
 export const del = async (fetch: Fetch, address: string, data?: Record<string, unknown>): Promise<Response> => {
 export const del = async (fetch: Fetch, address: string, data?: Record<string, unknown>): Promise<Response> => {
 	const response = await fetch(formatAddress(address), {
 	const response = await fetch(formatAddress(address), {
 		method: "DELETE",
 		method: "DELETE",