freezmo/lib/chat-models.ts

"use client";

/**
 * The chat model list and the remembered choice, shared by the Chat screen and
 * Teach it back so both offer exactly the same models and remember the same
 * pick. The server's health check lists the Ollama models; every enabled cloud
 * provider contributes its own (live listing first, built-in seeds after).
 */

import { useCallback, useEffect, useMemo, useState } from "react";

import { readStoredChatModel, readStoredChatProvider, storeChatModel, storeChatProvider, type ChatModelChoice } from "./chat";
import { PROVIDERS } from "./providers";
import { modelsFor, useProviders } from "./providers-client";
import { useHealth } from "./use-health";

export function useChatModels(): {
  /** Everything the picker may offer. */
  models: ChatModelChoice[];
  /** The model a request should use (falls back to the first available one). */
  model: string;
  /** The provider that owns `model` ("ollama" when none). */
  provider: string;
  /** Picks a model and remembers it across visits. */
  choose: (model: string, provider?: string) => void;
  /** Restores a saved conversation's choice without remembering it. */
  restore: (model?: string, provider?: string) => void;
} {
  const { health } = useHealth();
  const providerState = useProviders();
  const [model, setModel] = useState("");
  const [provider, setProvider] = useState("ollama");

  useEffect(() => {
    setModel(readStoredChatModel() ?? "");
    setProvider(readStoredChatProvider());
  }, []);

  const models = useMemo<ChatModelChoice[]>(() => {
    const fromHealth = health?.chatModels ?? [];
    const list =
      fromHealth.length > 0
        ? fromHealth
        : [
            { name: health?.text?.configured ?? "", badge: "text" as const, available: health?.text?.available ?? null },
            { name: health?.vision?.configured ?? "", badge: "vision" as const, available: health?.vision?.available ?? null },
          ];
    const ollama = list.filter((entry) => entry.name);
    const cloud: ChatModelChoice[] = [];
    for (const info of PROVIDERS) {
      if (!providerState.keys[info.id]) continue;
      const models = modelsFor(providerState, info.id).slice(0, 400);
      if (models.length === 0) continue;
      // Only a live /models listing proves the key can use them; the built-in
      // seeds stay badge-less until the first successful test.
      const listed = (providerState.models[info.id]?.length ?? 0) > 0;
      for (const name of models) {
        cloud.push({ name, badge: "text", available: listed ? true : null, provider: info.id });
      }
    }
    return [...ollama, ...cloud];
  }, [health, providerState]);

  const activeProvider = provider !== "ollama" && !providerState.keys[provider] ? "ollama" : provider;

  const activeModel = useMemo(() => {
    if (model && models.some((entry) => entry.name === model && (entry.provider ?? "ollama") === activeProvider)) {
      return model;
    }
    // The server puts the default first, but a different server may not have
    // it — fall back to the first model that is actually available.
    const firstAvailable = models.find(
      (entry) => (entry.provider ?? "ollama") === activeProvider && entry.available !== false,
    );
    return (
      firstAvailable ??
      models.find((entry) => (entry.provider ?? "ollama") === activeProvider) ??
      models[0]
    )?.name ?? model;
  }, [model, models, activeProvider]);

  const choose = useCallback((next: string, nextProvider?: string) => {
    const chosenProvider = nextProvider ?? "ollama";
    setProvider(chosenProvider);
    storeChatProvider(chosenProvider);
    setModel(next);
    storeChatModel(next);
  }, []);

  const restore = useCallback((nextModel?: string, nextProvider?: string) => {
    if (nextModel) setModel(nextModel);
    if (nextProvider) setProvider(nextProvider);
  }, []);

  return { models, model: activeModel, provider: activeProvider, choose, restore };
}