rubber-duck-mcp/apps/eval/lib/ollama.mjs

const DEFAULT_URL = "http://localhost:11434";

export class Ollama {
  constructor({ url = DEFAULT_URL, timeoutMs = 600000, idleTimeoutMs = null, maxTotalMs = null } = {}) {
    this.url = url.replace(/\/$/, "");
    this.timeoutMs = timeoutMs;
    // Dynamic timeout: if set, the timer resets on every new token; the request
    // only fails if this much time passes with NO progress (no new token).
    this.idleTimeoutMs = idleTimeoutMs ?? timeoutMs;
    // Absolute hard cap on one request's lifetime, even if tokens keep flowing
    // (guards against a model looping/chatting forever). Defaults to 5 minutes.
    this.maxTotalMs = maxTotalMs ?? 5 * 60_000;
  }

  async ping({ timeoutMs = 30000 } = {}) {
    if (process.env.OLLAMA_SKIP_PING === "1") return { ok: true, model: 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();
    const models = (data.models ?? []).map((m) => m.name);
    return { ok: true, models };
  }

  /**
   * Streaming chat with a dynamic (progress-based) timeout.
   * The timeout only fires when `timeoutMs` passes with no new token/chunk.
   * Optionally calls onChunk({token, content, type, toolCalls}) as data arrives.
   * Returns the same shape as the old non-streaming chat().
   */
  async chatStream({ model, messages, tools, temperature = 0, numCtx = 8192, timeoutMs, maxTotalMs, onChunk }) {
    const body = {
      model,
      messages,
      stream: true,
      options: { temperature },
    };
    if (numCtx) body.options.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 = [];
    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`);
        err.name = "TimeoutError";
        controller.abort(err);
        return;
      }
      if (Date.now() - lastActivity > limitMs) {
        const err = new Error(`No progress from Ollama for ${limitMs}ms; aborting`);
        err.name = "TimeoutError";
        controller.abort(err);
      }
    }, Math.min(500, Math.max(100, Math.floor(Math.min(limitMs, totalCapMs) / 4))));

    try {
      for (;;) {
        let chunk;
        try {
          chunk = await reader.read();
        } catch (readErr) {
          // Abort / network drop mid-stream. Surface whatever we got so the
          // caller can tell "slowly progressing" from "hung".
          const e = new Error(
            `Aborted mid-stream after ${content.length} chars (toolCalls=${toolCalls.length}): ${readErr.message}`
          );
          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;
        while ((idx = buffer.indexOf("\n")) !== -1) {
          const line = buffer.slice(0, idx).trim();
          buffer = buffer.slice(idx + 1);
          if (!line) continue;
          let obj;
          try {
            obj = JSON.parse(line);
          } catch {
            continue;
          }
          if (obj.prompt_eval_count != null) promptEvalCount = obj.prompt_eval_count;
          if (obj.eval_count != null) evalCount = obj.eval_count;
          const msg = obj.message ?? {};
          if (msg.content) {
            content += msg.content;
            onChunk?.({ token: msg.content, content, type: "content" });
          }
          if (msg.tool_calls && msg.tool_calls.length) {
            toolCalls = msg.tool_calls;
            onChunk?.({ token: null, content, type: "tool_calls", toolCalls: msg.tool_calls });
          }
        }
      }
    } 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 }) {
    return this.chatStream({
      model,
      messages,
      tools,
      temperature,
      numCtx,
      timeoutMs,
      onChunk: null,
    });
  }

  async warmup({ model, text = "say OK", timeoutMs = 60000 }) {
    try {
      return await this.chatStream({
        model,
        messages: [{ role: "user", content: text }],
        timeoutMs,
      });
    } catch (err) {
      return { error: err.message };
    }
  }
}

export async function listModels() {
  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();
  return (data.models ?? []).map((m) => m.name);
}