freezmo/lib/ollama.ts

/**
 * Minimal server-side Ollama HTTP client.
 *
 * Works with both supported setups:
 *  1. a local daemon (http://localhost:11434) that proxies :cloud models
 *     after `ollama signin`
 *  2. the direct cloud API (https://ollama.com) with a Bearer API key
 *
 * This module must only ever be imported from server code (route handlers).
 */

import { DEFAULT_LOCALE, translate, type Locale } from "./i18n-dict";
import { getServerConfig, resolveModelName, type OllamaAuth } from "./config";

export type OllamaMessage = {
  role: "system" | "user" | "assistant";
  content: string;
  /** Base64 strings (no data: prefix) or full data URLs; Ollama accepts both. */
  images?: string[];
};

export type OllamaErrorKind =
  | "unreachable"
  | "timeout"
  | "auth"
  | "no_credits"
  /** The free monthly Ollama usage is used up — only the user can fix this. */
  | "usage_limit"
  /** Ollama is throttling short-term requests; retrying soon usually works. */
  | "throttled"
  | "model_missing"
  | "http"
  | "bad_response";

/**
 * Ollama reports two very different things with HTTP 429:
 *  - the free monthly usage is exhausted ("upgrade", "quota", "monthly …")
 *  - a short-term rate limit ("too many requests", "try again …")
 * The body is the only signal, so classify it (kept pure for testing).
 */
const USAGE_LIMIT_PATTERN =
  /(monthly|per month|this month|quota|usage limit|limit reached|free (plan|tier|usage)|upgrade|billing|subscribe|exceeded your)/i;

export function classifyOllamaFailure(status: number, body: string): "usage_limit" | "throttled" | null {
  if (status !== 429) return null;
  return USAGE_LIMIT_PATTERN.test(body) ? "usage_limit" : "throttled";
}

export class OllamaError extends Error {
  kind: OllamaErrorKind;
  status?: number;
  detail?: string;

  constructor(kind: OllamaErrorKind, message: string, options?: { status?: number; detail?: string }) {
    super(message);
    this.name = "OllamaError";
    this.kind = kind;
    this.status = options?.status;
    this.detail = options?.detail;
  }
}

export function friendlyErrorMessage(
  error: unknown,
  baseUrl?: string,
  locale: Locale = DEFAULT_LOCALE,
): string {
  if (error instanceof OllamaError) {
    // Provider failures (lib/openai-compat) carry their own copy.
    if ((error as { provider?: string }).provider) return error.message;
    switch (error.kind) {
      case "unreachable":
        return translate(locale, "server.ollamaUnreachable", {
          url: baseUrl ?? "the configured URL",
        });
      case "timeout":
        return translate(locale, "server.ollamaTimeout");
      case "auth":
        return translate(locale, "server.ollamaAuth");
      case "usage_limit":
        return translate(locale, "server.usageLimit");
      case "throttled":
        return translate(locale, "server.ollamaThrottled");
      case "no_credits":
        return error.message;
      case "model_missing":
        return error.message;
      default:
        return error.message || translate(locale, "server.ollamaUnexpected");
    }
  }
  if (error instanceof Error) return error.message;
  return translate(locale, "server.ollamaGeneric");
}

type ChatOptions = {
  model: string;
  messages: OllamaMessage[];
  /** Credentials for this request (browser key, server key or local daemon). */
  auth: OllamaAuth;
  /** Ask Ollama to constrain the output to valid JSON. */
  json?: boolean;
  signal?: AbortSignal;
  timeoutMs?: number;
  temperature?: number;
};

export type ChatResult = {
  content: string;
  model: string;
  raw: unknown;
};

async function fetchWithTimeout(
  url: string,
  init: RequestInit,
  timeoutMs: number,
  external?: AbortSignal,
): Promise<Response> {
  const controller = new AbortController();
  const timer = setTimeout(() => controller.abort(), timeoutMs);
  const onAbort = () => controller.abort();
  if (external) {
    if (external.aborted) controller.abort();
    else external.addEventListener("abort", onAbort, { once: true });
  }

  try {
    return await fetch(url, { ...init, signal: controller.signal });
  } finally {
    clearTimeout(timer);
    external?.removeEventListener("abort", onAbort);
  }
}

export async function ollamaChat(options: ChatOptions): Promise<ChatResult> {
  const config = getServerConfig();
  const { baseUrl, apiKey, directCloud } = options.auth;
  const model = resolveModelName(options.model, directCloud);

  const headers: Record<string, string> = { "Content-Type": "application/json" };
  if (apiKey) headers.Authorization = `Bearer ${apiKey}`;

  const body: Record<string, unknown> = {
    model,
    messages: options.messages,
    stream: false,
  };
  if (options.json) body.format = "json";
  if (typeof options.temperature === "number") {
    body.options = { temperature: options.temperature };
  }
  if (config.think !== null) body.think = config.think;

  let response: Response;
  try {
    response = await fetchWithTimeout(
      `${baseUrl}/api/chat`,
      { method: "POST", headers, body: JSON.stringify(body), cache: "no-store" },
      options.timeoutMs ?? config.timeoutMs,
      options.signal,
    );
  } catch (error) {
    if (error instanceof Error && error.name === "AbortError") {
      throw new OllamaError("timeout", "The request to Ollama timed out.");
    }
    throw new OllamaError(
      "unreachable",
      `Could not connect to Ollama at ${baseUrl}.`,
      { detail: error instanceof Error ? error.message : String(error) },
    );
  }

  if (!response.ok) {
    const text = await response.text().catch(() => "");
    throwForStatus(response.status, text, model);
  }

  let payload: unknown;
  try {
    payload = await response.json();
  } catch {
    throw new OllamaError("bad_response", "Ollama sent a response that wasn't JSON.");
  }

  const content = extractContent(payload);
  if (content === null) {
    throw new OllamaError("bad_response", "Ollama sent a response with no message content.", {
      detail: JSON.stringify(payload).slice(0, 500),
    });
  }

  return { content, model, raw: payload };
}

/**
 * Maps a non-2xx Ollama response to the right `OllamaError`. Shared by the
 * buffered and streaming chat calls so both classify 429s identically.
 */
function throwForStatus(status: number, text: string, model: string): never {
  if (status === 401 || status === 403) {
    throw new OllamaError("auth", "Ollama refused the request (unauthorised).", {
      status,
      detail: text.slice(0, 500),
    });
  }
  if (status === 402) {
    throw new OllamaError(
      "no_credits",
      `Your Ollama account can't use "${model}": this model is not included in your plan. Add usage credits at ollama.com/settings, or switch to a free-tier model in .env.local.`,
      { status, detail: text.slice(0, 500) },
    );
  }
  if (status === 404) {
    throw new OllamaError(
      "model_missing",
      `Ollama could not find the model "${model}". Run \`ollama pull ${model}\` or pick another model in .env.local.`,
      { status, detail: text.slice(0, 500) },
    );
  }
  if (status === 429) {
    if (classifyOllamaFailure(status, text) === "usage_limit") {
      throw new OllamaError(
        "usage_limit",
        "Your free monthly Ollama usage has run out. Upgrade your Ollama account, or create a new one and use the new API key.",
        { status, detail: text.slice(0, 500) },
      );
    }
    throw new OllamaError(
      "throttled",
      "Ollama is temporarily rate limiting requests. Try again in a moment.",
      { status, detail: text.slice(0, 500) },
    );
  }
  throw new OllamaError("http", `Ollama responded with HTTP ${status}.`, {
    status,
    detail: text.slice(0, 500),
  });
}

/* ------------------------------------------------------------------ */
/* Streaming chat                                                      */
/* ------------------------------------------------------------------ */

export type ChatStreamResult = {
  content: string;
  thinking: string;
  model: string;
};

type StreamOptions = ChatOptions & {
  onToken?: (text: string) => void;
  onThinking?: (text: string) => void;
};

/**
 * Streaming counterpart of `ollamaChat`: POSTs with `stream: true` and hands
 * each delta to the callbacks. Reasoning models deliver `message.thinking`
 * separately — it is returned apart from the answer and never mixed in.
 */
export async function ollamaChatStream(options: StreamOptions): Promise<ChatStreamResult> {
  const config = getServerConfig();
  const { baseUrl, apiKey, directCloud } = options.auth;
  const model = resolveModelName(options.model, directCloud);

  const headers: Record<string, string> = { "Content-Type": "application/json" };
  if (apiKey) headers.Authorization = `Bearer ${apiKey}`;

  const body: Record<string, unknown> = {
    model,
    messages: options.messages,
    stream: true,
  };
  if (typeof options.temperature === "number") {
    body.options = { temperature: options.temperature };
  }
  if (config.think !== null) body.think = config.think;

  const controller = new AbortController();
  const timer = setTimeout(() => controller.abort(), options.timeoutMs ?? config.timeoutMs);
  const onAbort = () => controller.abort();
  if (options.signal) {
    if (options.signal.aborted) controller.abort();
    else options.signal.addEventListener("abort", onAbort, { once: true });
  }

  try {
    let response: Response;
    try {
      response = await fetch(`${baseUrl}/api/chat`, {
        method: "POST",
        headers,
        body: JSON.stringify(body),
        cache: "no-store",
        signal: controller.signal,
      });
    } catch (error) {
      if (error instanceof Error && error.name === "AbortError") {
        throw new OllamaError("timeout", "The request to Ollama timed out.");
      }
      throw new OllamaError("unreachable", `Could not connect to Ollama at ${baseUrl}.`, {
        detail: error instanceof Error ? error.message : String(error),
      });
    }

    if (!response.ok) {
      const text = await response.text().catch(() => "");
      throwForStatus(response.status, text, model);
    }
    if (!response.body) {
      throw new OllamaError("bad_response", "Ollama sent an empty stream.");
    }

    const reader = response.body.getReader();
    const decoder = new TextDecoder();
    let buffer = "";
    let content = "";
    let thinking = "";

    const handleLine = (line: string) => {
      const trimmed = line.trim();
      if (!trimmed) return;
      let event: unknown;
      try {
        event = JSON.parse(trimmed);
      } catch {
        return; // a half-written line must never break the stream
      }
      if (!event || typeof event !== "object") return;
      const record = event as Record<string, unknown>;
      if (typeof record.error === "string" && record.error) {
        throw new OllamaError("http", record.error.slice(0, 300));
      }
      const message = record.message as Record<string, unknown> | undefined;
      const delta = typeof message?.content === "string" ? message.content : "";
      const reason = typeof message?.thinking === "string" ? message.thinking : "";
      if (reason) {
        thinking += reason;
        options.onThinking?.(reason);
      }
      if (delta) {
        content += delta;
        options.onToken?.(delta);
      }
    };

    try {
      while (true) {
        const { done, value } = await reader.read();
        if (done) break;
        buffer += decoder.decode(value, { stream: true });
        let newline = buffer.indexOf("\n");
        while (newline !== -1) {
          handleLine(buffer.slice(0, newline));
          buffer = buffer.slice(newline + 1);
          newline = buffer.indexOf("\n");
        }
      }
      handleLine(buffer);
    } catch (error) {
      void reader.cancel().catch(() => {});
      if (error instanceof Error && error.name === "AbortError") {
        throw new OllamaError("timeout", "The request to Ollama timed out.");
      }
      throw error;
    }

    if (!content.trim()) {
      throw new OllamaError("bad_response", "Ollama finished the stream without any answer text.");
    }

    return { content, thinking, model };
  } finally {
    clearTimeout(timer);
    options.signal?.removeEventListener("abort", onAbort);
  }
}

function extractContent(payload: unknown): string | null {
  if (!payload || typeof payload !== "object") return null;
  const record = payload as Record<string, unknown>;
  const message = record.message;
  if (message && typeof message === "object") {
    const content = (message as Record<string, unknown>).content;
    if (typeof content === "string") return content;
  }
  if (typeof record.response === "string") return record.response;
  return null;
}

export type OllamaModelInfo = { name: string; model: string; size?: number };

/** List the models known to the given Ollama instance (or the cloud). */
export async function ollamaListModels(auth: OllamaAuth, timeoutMs?: number): Promise<OllamaModelInfo[]> {
  const config = getServerConfig();
  const { baseUrl, apiKey } = auth;
  const headers: Record<string, string> = {};
  if (apiKey) headers.Authorization = `Bearer ${apiKey}`;

  let response: Response;
  try {
    response = await fetchWithTimeout(
      `${baseUrl}/api/tags`,
      { headers, cache: "no-store" },
      timeoutMs ?? config.healthTimeoutMs,
    );
  } catch (error) {
    if (error instanceof Error && error.name === "AbortError") {
      throw new OllamaError("timeout", "Timed out while listing Ollama models.");
    }
    throw new OllamaError("unreachable", `Could not connect to Ollama at ${baseUrl}.`, {
      detail: error instanceof Error ? error.message : String(error),
    });
  }

  if (!response.ok) {
    if (response.status === 401 || response.status === 403) {
      throw new OllamaError("auth", "Ollama refused the request (unauthorised).", {
        status: response.status,
      });
    }
    throw new OllamaError("http", `Ollama responded with HTTP ${response.status} while listing models.`, {
      status: response.status,
    });
  }

  const payload = (await response.json().catch(() => null)) as
    | { models?: Array<{ name?: string; model?: string; size?: number }> }
    | null;

  return (payload?.models ?? []).map((entry) => ({
    name: entry.name ?? entry.model ?? "unknown",
    model: entry.model ?? entry.name ?? "unknown",
    size: entry.size,
  }));
}

/* ------------------------------------------------------------------ */
/* API key verification                                                */
/* ------------------------------------------------------------------ */

export type KeyProbe = {
  /** true = accepted, false = rejected, null = could not tell. */
  valid: boolean | null;
  reason?: "invalid" | "no_credits" | "usage_limit" | "throttled" | "model_missing" | "unreachable";
  detail?: string;
};

type CachedProbe = KeyProbe & { at: number };

const probeCacheStore = globalThis as unknown as { __freezmoKeyProbes?: Map<string, CachedProbe> };
const probeCache: Map<string, CachedProbe> = (probeCacheStore.__freezmoKeyProbes ??= new Map());
const PROBE_TTL_MS = 5 * 60_000;

/**
 * Model listing is public on ollama.com, so the only trustworthy way to check
 * a key is a real (tiny) chat request: 1 output token, cached for a few minutes.
 * Interpreted statuses:
 *   200        key works
 *   401 / 403  key rejected
 *   402        key works but the plan has no access to that model
 *   404        model unknown — the key isn't the problem
 */
export async function ollamaVerifyKey(auth: OllamaAuth, model: string, timeoutMs = 20_000): Promise<KeyProbe> {
  if (!auth.apiKey) return { valid: null };
  if (process.env.OLLAMA_VERIFY_KEY?.trim().toLowerCase() === "false") return { valid: null };

  const cacheKey = `${auth.baseUrl}|${model}|${hashKey(auth.apiKey)}|${auth.authMode}`;
  const cached = probeCache.get(cacheKey);
  if (cached && Date.now() - cached.at < PROBE_TTL_MS) {
    return { valid: cached.valid, reason: cached.reason, detail: cached.detail };
  }

  const resolved = resolveModelName(model, auth.directCloud);
  let response: Response;
  try {
    response = await fetchWithTimeout(
      `${auth.baseUrl}/api/chat`,
      {
        method: "POST",
        headers: {
          "Content-Type": "application/json",
          Authorization: `Bearer ${auth.apiKey}`,
        },
        cache: "no-store",
        body: JSON.stringify({
          model: resolved,
          messages: [{ role: "user", content: "ping" }],
          stream: false,
          options: { num_predict: 1 },
        }),
      },
      timeoutMs,
    );
  } catch {
    return { valid: null, reason: "unreachable" };
  }

  const detail = (await response.text().catch(() => "")).slice(0, 300);
  let probe: KeyProbe;
  if (response.ok) probe = { valid: true };
  else if (response.status === 401 || response.status === 403) probe = { valid: false, reason: "invalid", detail };
  else if (response.status === 402) probe = { valid: true, reason: "no_credits", detail };
  else if (response.status === 404) probe = { valid: null, reason: "model_missing", detail };
  else if (response.status === 429) {
    // The key itself is fine; the account is either out of monthly usage or
    // being throttled for a moment.
    probe =
      classifyOllamaFailure(response.status, detail) === "usage_limit"
        ? { valid: true, reason: "usage_limit", detail }
        : { valid: null, reason: "throttled", detail };
  } else probe = { valid: null, detail: `HTTP ${response.status}` };

  if (probeCache.size >= 500) {
    const now = Date.now();
    for (const [k, entry] of probeCache) {
      if (now - entry.at >= PROBE_TTL_MS) probeCache.delete(k);
    }
    if (probeCache.size >= 500) {
      const oldest = probeCache.keys().next().value;
      if (oldest) probeCache.delete(oldest);
    }
  }
  probeCache.set(cacheKey, { ...probe, at: Date.now() });
  return probe;
}

/** The key itself is never kept — only a hash, so the cache can't leak it. */
function hashKey(key: string): string {
  // Small non-cryptographic hash is enough to bucket cache entries.
  let hash = 2166136261;
  for (let i = 0; i < key.length; i++) {
    hash ^= key.charCodeAt(i);
    hash = Math.imul(hash, 16777619);
  }
  return (hash >>> 0).toString(36);
}