fix(ollama-client): remove timeout from pull(), guard formatBytes NaN, simplify stream body, use fetchJson for delete()

This commit is contained in:
saravanakumardb1 2026-03-29 12:50:49 -07:00
parent 26ac2a3dec
commit 2688f7e3ca
5 changed files with 29 additions and 38 deletions

View File

@ -124,7 +124,7 @@ describe('OllamaClient', () => {
describe('delete', () => { describe('delete', () => {
it('sends DELETE request', async () => { it('sends DELETE request', async () => {
globalThis.fetch = vi.fn().mockResolvedValue({ ok: true, status: 200 }); globalThis.fetch = mockFetch({});
await client.delete('llama3'); await client.delete('llama3');
expect(globalThis.fetch).toHaveBeenCalledWith( expect(globalThis.fetch).toHaveBeenCalledWith(
@ -137,12 +137,8 @@ describe('OllamaClient', () => {
}); });
it('throws on failure', async () => { it('throws on failure', async () => {
globalThis.fetch = vi.fn().mockResolvedValue({ globalThis.fetch = mockFetch('model not found', { ok: false, status: 404 });
ok: false, await expect(client.delete('nope')).rejects.toThrow('Ollama /api/delete failed (404)');
status: 404,
text: () => Promise.resolve('model not found'),
});
await expect(client.delete('nope')).rejects.toThrow('Ollama delete failed (404)');
}); });
}); });

View File

@ -6,6 +6,7 @@ import type {
OllamaPullProgress, OllamaPullProgress,
OllamaVersionResponse, OllamaVersionResponse,
} from './types.js'; } from './types.js';
import { parseNdjsonStream } from './ndjson.js';
/** /**
* Ollama API client for model management operations. * Ollama API client for model management operations.
@ -75,11 +76,13 @@ export class OllamaClient {
stream: boolean = false stream: boolean = false
): Promise<{ status: string } | AsyncGenerator<OllamaPullProgress>> { ): Promise<{ status: string } | AsyncGenerator<OllamaPullProgress>> {
if (!stream) { if (!stream) {
// Model pulls can download GBs — use 10 minute timeout instead of the default
const pullTimeoutMs = Math.max(this.timeoutMs, 600_000);
const res = await fetch(`${this.baseUrl}/api/pull`, { const res = await fetch(`${this.baseUrl}/api/pull`, {
method: 'POST', method: 'POST',
headers: { 'Content-Type': 'application/json' }, headers: { 'Content-Type': 'application/json' },
body: JSON.stringify({ name: model, stream: false }), body: JSON.stringify({ name: model, stream: false }),
signal: AbortSignal.timeout(this.timeoutMs), signal: AbortSignal.timeout(pullTimeoutMs),
}); });
if (!res.ok) { if (!res.ok) {
const text = await res.text().catch(() => ''); const text = await res.text().catch(() => '');
@ -88,9 +91,8 @@ export class OllamaClient {
return (await res.json()) as { status: string }; return (await res.json()) as { status: string };
} }
// Streaming pull — return async generator // Streaming pull — return async generator (no timeout, consumer controls lifetime)
const baseUrl = this.baseUrl; const res = await fetch(`${this.baseUrl}/api/pull`, {
const res = await fetch(`${baseUrl}/api/pull`, {
method: 'POST', method: 'POST',
headers: { 'Content-Type': 'application/json' }, headers: { 'Content-Type': 'application/json' },
body: JSON.stringify({ name: model, stream: true }), body: JSON.stringify({ name: model, stream: true }),
@ -101,7 +103,6 @@ export class OllamaClient {
} }
if (!res.body) throw new Error('No response body from Ollama pull'); if (!res.body) throw new Error('No response body from Ollama pull');
const { parseNdjsonStream } = await import('./ndjson.js');
return parseNdjsonStream<OllamaPullProgress>(res.body); return parseNdjsonStream<OllamaPullProgress>(res.body);
} }
@ -130,16 +131,10 @@ export class OllamaClient {
/** Delete a model (DELETE /api/delete). */ /** Delete a model (DELETE /api/delete). */
async delete(model: string): Promise<void> { async delete(model: string): Promise<void> {
const res = await fetch(`${this.baseUrl}/api/delete`, { await this.fetchJson('/api/delete', {
method: 'DELETE', method: 'DELETE',
headers: { 'Content-Type': 'application/json' },
body: JSON.stringify({ name: model }), body: JSON.stringify({ name: model }),
signal: AbortSignal.timeout(this.timeoutMs),
}); });
if (!res.ok) {
const text = await res.text().catch(() => '');
throw new Error(`Ollama delete failed (${res.status}): ${text.slice(0, 200)}`);
}
} }
/** Get Ollama server version (GET /api/version). */ /** Get Ollama server version (GET /api/version). */

View File

@ -21,6 +21,10 @@ describe('formatBytes', () => {
it('formats gigabytes', () => { it('formats gigabytes', () => {
expect(formatBytes(4294967296)).toBe('4 GB'); expect(formatBytes(4294967296)).toBe('4 GB');
}); });
it('returns 0 B for negative input', () => {
expect(formatBytes(-100)).toBe('0 B');
});
}); });
describe('estimateTokens', () => { describe('estimateTokens', () => {

View File

@ -5,7 +5,7 @@
* @example formatBytes(0) // '0 B' * @example formatBytes(0) // '0 B'
*/ */
export function formatBytes(bytes: number): string { export function formatBytes(bytes: number): string {
if (bytes === 0) return '0 B'; if (bytes <= 0) return '0 B';
const k = 1024; const k = 1024;
const sizes = ['B', 'KB', 'MB', 'GB', 'TB']; const sizes = ['B', 'KB', 'MB', 'GB', 'TB'];
const i = Math.floor(Math.log(bytes) / Math.log(k)); const i = Math.floor(Math.log(bytes) / Math.log(k));

View File

@ -20,17 +20,15 @@ export async function* streamChat(
): AsyncGenerator<OllamaStreamChunk> { ): AsyncGenerator<OllamaStreamChunk> {
const { model, messages, signal, ...rest } = options; const { model, messages, signal, ...rest } = options;
const body: Record<string, unknown> = { model, messages, stream: true };
if (rest.options) body.options = rest.options;
if (rest.format) body.format = rest.format;
if (rest.keep_alive) body.keep_alive = rest.keep_alive;
const res = await fetch(`${baseUrl}/api/chat`, { const res = await fetch(`${baseUrl}/api/chat`, {
method: 'POST', method: 'POST',
headers: { 'Content-Type': 'application/json' }, headers: { 'Content-Type': 'application/json' },
body: JSON.stringify({ body: JSON.stringify(body),
model,
messages,
stream: true,
...('options' in rest && rest.options ? { options: rest.options } : {}),
...('format' in rest && rest.format ? { format: rest.format } : {}),
...('keep_alive' in rest && rest.keep_alive ? { keep_alive: rest.keep_alive } : {}),
}),
signal, signal,
}); });
@ -60,19 +58,17 @@ export async function* streamGenerate(
): AsyncGenerator<OllamaGenerateChunk> { ): AsyncGenerator<OllamaGenerateChunk> {
const { model, prompt, signal, ...rest } = options; const { model, prompt, signal, ...rest } = options;
const body: Record<string, unknown> = { model, prompt, stream: true };
if (rest.system) body.system = rest.system;
if (rest.options) body.options = rest.options;
if (rest.format) body.format = rest.format;
if (rest.keep_alive) body.keep_alive = rest.keep_alive;
if (rest.context) body.context = rest.context;
const res = await fetch(`${baseUrl}/api/generate`, { const res = await fetch(`${baseUrl}/api/generate`, {
method: 'POST', method: 'POST',
headers: { 'Content-Type': 'application/json' }, headers: { 'Content-Type': 'application/json' },
body: JSON.stringify({ body: JSON.stringify(body),
model,
prompt,
stream: true,
...('system' in rest && rest.system ? { system: rest.system } : {}),
...('options' in rest && rest.options ? { options: rest.options } : {}),
...('format' in rest && rest.format ? { format: rest.format } : {}),
...('keep_alive' in rest && rest.keep_alive ? { keep_alive: rest.keep_alive } : {}),
...('context' in rest && rest.context ? { context: rest.context } : {}),
}),
signal, signal,
}); });