feat: move runner to typescript, refactor
This commit is contained in:
@@ -0,0 +1,244 @@
|
||||
interface OllamaConstructorOpts {
|
||||
url?: string;
|
||||
timeoutMs?: number;
|
||||
idleTimeoutMs?: number | null;
|
||||
maxTotalMs?: number | null;
|
||||
}
|
||||
|
||||
interface PingResult {
|
||||
ok: boolean;
|
||||
models: string[];
|
||||
}
|
||||
|
||||
export type OllamaToolCall = {
|
||||
function: { name: string; arguments: string };
|
||||
};
|
||||
|
||||
export type OllamaToolDefinition = {
|
||||
type: "function";
|
||||
function: {
|
||||
name: string;
|
||||
description: string;
|
||||
parameters: {
|
||||
type: "object";
|
||||
properties?: Record<string, unknown>;
|
||||
[key: string]: unknown;
|
||||
};
|
||||
};
|
||||
};
|
||||
|
||||
export type OllamaMessage = {
|
||||
role: string;
|
||||
content: string;
|
||||
tool_calls?: OllamaToolCall[];
|
||||
};
|
||||
|
||||
export type OllamaChunk = {
|
||||
token: string | null;
|
||||
content: string;
|
||||
type: "content" | "tool_calls";
|
||||
toolCalls?: OllamaToolCall[];
|
||||
};
|
||||
|
||||
export interface ChatStreamOpts {
|
||||
model: string;
|
||||
messages: OllamaMessage[];
|
||||
tools?: unknown[];
|
||||
temperature?: number;
|
||||
numCtx?: number;
|
||||
timeoutMs?: number;
|
||||
maxTotalMs?: number;
|
||||
onChunk?: ((chunk: OllamaChunk) => void) | null;
|
||||
}
|
||||
|
||||
export interface OllamaChatResult {
|
||||
role: "assistant";
|
||||
content: string;
|
||||
toolCalls: OllamaToolCall[];
|
||||
promptEvalCount: number;
|
||||
evalCount: number;
|
||||
raw: { stream: boolean };
|
||||
}
|
||||
|
||||
export interface OllamaError extends Error {
|
||||
partialContent?: string;
|
||||
partialToolCalls?: OllamaToolCall[];
|
||||
partialPrompt?: number;
|
||||
partialGen?: number;
|
||||
}
|
||||
|
||||
const DEFAULT_URL = "http://localhost:11434";
|
||||
|
||||
export class Ollama {
|
||||
private url: string;
|
||||
private timeoutMs: number;
|
||||
private idleTimeoutMs: number;
|
||||
private maxTotalMs: number;
|
||||
|
||||
constructor({ url = DEFAULT_URL, timeoutMs = 600000, idleTimeoutMs = null, maxTotalMs = null }: OllamaConstructorOpts = {}) {
|
||||
this.url = url.replace(/\/$/, "");
|
||||
this.timeoutMs = timeoutMs;
|
||||
this.idleTimeoutMs = idleTimeoutMs ?? timeoutMs;
|
||||
this.maxTotalMs = maxTotalMs ?? 5 * 60_000;
|
||||
}
|
||||
|
||||
async ping({ timeoutMs = 30000 } = {}): Promise<PingResult> {
|
||||
if (process.env.OLLAMA_SKIP_PING === "1") return { ok: true, models: [process.env.OLLAMA_MODEL ?? ""] };
|
||||
const res = await fetch(`${this.url}/api/tags`, { signal: AbortSignal.timeout(timeoutMs) });
|
||||
if (!res.ok) throw new Error(`Ollama not reachable (HTTP ${res.status}) at ${this.url}`);
|
||||
const data = await res.json() as { models?: Array<{ name: string }> };
|
||||
const models = (data.models ?? []).map((m) => m.name);
|
||||
return { ok: true, models };
|
||||
}
|
||||
|
||||
async chatStream({
|
||||
model,
|
||||
messages,
|
||||
tools,
|
||||
temperature = 0,
|
||||
numCtx = 8192,
|
||||
timeoutMs,
|
||||
maxTotalMs,
|
||||
onChunk,
|
||||
}: ChatStreamOpts): Promise<OllamaChatResult> {
|
||||
const body: Record<string, unknown> = {
|
||||
model,
|
||||
messages,
|
||||
stream: true,
|
||||
options: { temperature },
|
||||
};
|
||||
if (numCtx) (body.options as Record<string, unknown>).num_ctx = numCtx;
|
||||
if (tools && tools.length) body.tools = tools;
|
||||
|
||||
const limitMs = timeoutMs ?? this.idleTimeoutMs;
|
||||
const totalCapMs = maxTotalMs ?? this.maxTotalMs;
|
||||
const controller = new AbortController();
|
||||
const startedAt = Date.now();
|
||||
|
||||
const res = await fetch(`${this.url}/api/chat`, {
|
||||
method: "POST",
|
||||
headers: { "Content-Type": "application/json" },
|
||||
body: JSON.stringify(body),
|
||||
signal: controller.signal,
|
||||
});
|
||||
if (!res.ok) {
|
||||
const t = await res.text();
|
||||
throw new Error(`Ollama chat HTTP ${res.status}: ${t.slice(0, 300)}`);
|
||||
}
|
||||
if (!res.body || !res.body.getReader) {
|
||||
throw new Error("Ollama streaming response has no body reader");
|
||||
}
|
||||
|
||||
const reader = res.body.getReader();
|
||||
const decoder = new TextDecoder();
|
||||
let buffer = "";
|
||||
let content = "";
|
||||
let toolCalls: OllamaToolCall[] = [];
|
||||
let promptEvalCount = 0;
|
||||
let evalCount = 0;
|
||||
let lastActivity = Date.now();
|
||||
|
||||
const guardInterval = setInterval(() => {
|
||||
if (Date.now() - startedAt > totalCapMs) {
|
||||
const err = new Error(`Request exceeded hard cap of ${totalCapMs}ms; aborting`) as OllamaError;
|
||||
err.name = "TimeoutError";
|
||||
controller.abort(err);
|
||||
return;
|
||||
}
|
||||
if (Date.now() - lastActivity > limitMs) {
|
||||
const err = new Error(`No progress from Ollama for ${limitMs}ms; aborting`) as OllamaError;
|
||||
err.name = "TimeoutError";
|
||||
controller.abort(err);
|
||||
}
|
||||
}, Math.min(500, Math.max(100, Math.floor(Math.min(limitMs, totalCapMs) / 4))));
|
||||
|
||||
try {
|
||||
for (;;) {
|
||||
let chunk: ReadableStreamReadResult<Uint8Array>;
|
||||
try {
|
||||
chunk = await reader.read();
|
||||
} catch (readErr) {
|
||||
const e = new Error(
|
||||
`Aborted mid-stream after ${content.length} chars (toolCalls=${toolCalls.length}): ${(readErr as Error).message}`
|
||||
) as OllamaError;
|
||||
e.name = "TimeoutError";
|
||||
e.partialContent = content;
|
||||
e.partialToolCalls = toolCalls;
|
||||
e.partialPrompt = promptEvalCount;
|
||||
e.partialGen = evalCount;
|
||||
throw e;
|
||||
}
|
||||
const { done, value } = chunk;
|
||||
if (done) break;
|
||||
lastActivity = Date.now();
|
||||
buffer += decoder.decode(value, { stream: true });
|
||||
let idx: number;
|
||||
while ((idx = buffer.indexOf("\n")) !== -1) {
|
||||
const line = buffer.slice(0, idx).trim();
|
||||
buffer = buffer.slice(idx + 1);
|
||||
if (!line) continue;
|
||||
let obj: Record<string, unknown>;
|
||||
try {
|
||||
obj = JSON.parse(line);
|
||||
} catch {
|
||||
continue;
|
||||
}
|
||||
if (obj.prompt_eval_count != null) promptEvalCount = obj.prompt_eval_count as number;
|
||||
if (obj.eval_count != null) evalCount = obj.eval_count as number;
|
||||
const msg = (obj.message ?? {}) as Record<string, unknown>;
|
||||
if (msg.content) {
|
||||
content += msg.content as string;
|
||||
onChunk?.({ token: msg.content as string, content, type: "content" });
|
||||
}
|
||||
if (msg.tool_calls && (msg.tool_calls as unknown[]).length) {
|
||||
toolCalls = msg.tool_calls as OllamaToolCall[];
|
||||
onChunk?.({ token: null, content, type: "tool_calls", toolCalls });
|
||||
}
|
||||
}
|
||||
}
|
||||
} finally {
|
||||
clearInterval(guardInterval);
|
||||
}
|
||||
|
||||
return {
|
||||
role: "assistant",
|
||||
content,
|
||||
toolCalls: Array.isArray(toolCalls) ? toolCalls : [],
|
||||
promptEvalCount,
|
||||
evalCount,
|
||||
raw: { stream: true },
|
||||
};
|
||||
}
|
||||
|
||||
async chat({ model, messages, tools, temperature = 0, numCtx = 8192, timeoutMs }: ChatStreamOpts): Promise<OllamaChatResult> {
|
||||
return this.chatStream({
|
||||
model,
|
||||
messages,
|
||||
tools,
|
||||
temperature,
|
||||
numCtx,
|
||||
timeoutMs,
|
||||
onChunk: null,
|
||||
});
|
||||
}
|
||||
|
||||
async warmup({ model, text = "say OK", timeoutMs = 60000 }: { model: string; text?: string; timeoutMs?: number }): Promise<Partial<OllamaChatResult> & { error?: string }> {
|
||||
try {
|
||||
return await this.chatStream({
|
||||
model,
|
||||
messages: [{ role: "user", content: text }],
|
||||
timeoutMs,
|
||||
});
|
||||
} catch (err) {
|
||||
return { error: (err as Error).message };
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
export async function listModels(): Promise<string[]> {
|
||||
const o = new Ollama();
|
||||
const res = await fetch(`${o["url"]}/api/tags`);
|
||||
if (!res.ok) throw new Error(`Ollama not reachable (HTTP ${res.status})`);
|
||||
const data = await res.json() as { models?: Array<{ name: string }> };
|
||||
return (data.models ?? []).map((m) => m.name);
|
||||
}
|
||||
Reference in New Issue
Block a user