86 lines
2.4 KiB
TypeScript
86 lines
2.4 KiB
TypeScript
import type {
|
|
OllamaChatOptions,
|
|
OllamaGenerateOptions,
|
|
OllamaStreamChunk,
|
|
OllamaGenerateChunk,
|
|
} from './types.js';
|
|
import { parseNdjsonStream } from './ndjson.js';
|
|
|
|
/**
|
|
* Stream a chat completion from Ollama.
|
|
*
|
|
* Yields NDJSON chunks from `POST /api/chat` as an async generator.
|
|
*
|
|
* @param baseUrl - Ollama server base URL
|
|
* @param options - Chat options (model, messages, signal, etc.)
|
|
*/
|
|
export async function* streamChat(
|
|
baseUrl: string,
|
|
options: OllamaChatOptions
|
|
): AsyncGenerator<OllamaStreamChunk> {
|
|
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`, {
|
|
method: 'POST',
|
|
headers: { 'Content-Type': 'application/json' },
|
|
body: JSON.stringify(body),
|
|
signal,
|
|
});
|
|
|
|
if (!res.ok) {
|
|
const text = await res.text().catch(() => '');
|
|
throw new Error(`Ollama chat failed (${res.status}): ${text.slice(0, 200)}`);
|
|
}
|
|
|
|
if (!res.body) {
|
|
throw new Error('No response body from Ollama');
|
|
}
|
|
|
|
yield* parseNdjsonStream<OllamaStreamChunk>(res.body);
|
|
}
|
|
|
|
/**
|
|
* Stream a text generation from Ollama.
|
|
*
|
|
* Yields NDJSON chunks from `POST /api/generate` as an async generator.
|
|
*
|
|
* @param baseUrl - Ollama server base URL
|
|
* @param options - Generate options (model, prompt, signal, etc.)
|
|
*/
|
|
export async function* streamGenerate(
|
|
baseUrl: string,
|
|
options: OllamaGenerateOptions
|
|
): AsyncGenerator<OllamaGenerateChunk> {
|
|
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`, {
|
|
method: 'POST',
|
|
headers: { 'Content-Type': 'application/json' },
|
|
body: JSON.stringify(body),
|
|
signal,
|
|
});
|
|
|
|
if (!res.ok) {
|
|
const text = await res.text().catch(() => '');
|
|
throw new Error(`Ollama generate failed (${res.status}): ${text.slice(0, 200)}`);
|
|
}
|
|
|
|
if (!res.body) {
|
|
throw new Error('No response body from Ollama');
|
|
}
|
|
|
|
yield* parseNdjsonStream<OllamaGenerateChunk>(res.body);
|
|
}
|