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

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);
}