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) {
    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();
  external?.addEventListener("abort", onAbort);

  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(() => "");
    if (response.status === 401 || response.status === 403) {
      throw new OllamaError("auth", "Ollama refused the request (unauthorised).", {
        status: response.status,
        detail: text.slice(0, 500),
      });
    }
    if (response.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: response.status, detail: text.slice(0, 500) },
      );
    }
    if (response.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: response.status, detail: text.slice(0, 500) },
      );
    }
    if (response.status === 429) {
      if (classifyOllamaFailure(response.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: response.status, detail: text.slice(0, 500) },
        );
      }
      throw new OllamaError(
        "throttled",
        "Ollama is temporarily rate limiting requests. Try again in a moment.",
        { status: response.status, detail: text.slice(0, 500) },
      );
    }
    throw new OllamaError(
      "http",
      `Ollama responded with HTTP ${response.status}.`,
      { status: response.status, detail: text.slice(0, 500) },
    );
  }

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

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

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