bolt.diy/app/components/settings/providers/OllamaModelUpdater.tsx

import React, { useEffect, useState } from 'react';
import { motion } from 'framer-motion';
import { toast } from 'react-toastify';
import { classNames } from '~/utils/classNames';

interface OllamaModel {
  name: string;
  digest: string;
  size: number;
  modified_at: string;
  details?: {
    family: string;
    parameter_size: string;
    quantization_level: string;
  };
  status?: 'idle' | 'updating' | 'updated' | 'error' | 'checking';
  error?: string;
  newDigest?: string;
  progress?: {
    current: number;
    total: number;
    status: string;
  };
}

interface OllamaTagResponse {
  models: Array<{
    name: string;
    digest: string;
    size: number;
    modified_at: string;
    details?: {
      family: string;
      parameter_size: string;
      quantization_level: string;
    };
  }>;
}

interface OllamaPullResponse {
  status: string;
  digest?: string;
  total?: number;
  completed?: number;
}

export default function OllamaModelUpdater() {
  const [models, setModels] = useState<OllamaModel[]>([]);
  const [isLoading, setIsLoading] = useState(true);
  const [isBulkUpdating, setIsBulkUpdating] = useState(false);

  useEffect(() => {
    fetchModels();
  }, []);

  const fetchModels = async () => {
    try {
      setIsLoading(true);

      const response = await fetch('http://localhost:11434/api/tags');
      const data = (await response.json()) as OllamaTagResponse;
      setModels(
        data.models.map((model) => ({
          name: model.name,
          digest: model.digest,
          size: model.size,
          modified_at: model.modified_at,
          details: model.details,
          status: 'idle' as const,
        })),
      );
    } catch (error) {
      toast.error('Failed to fetch Ollama models');
      console.error('Error fetching models:', error);
    } finally {
      setIsLoading(false);
    }
  };

  const updateModel = async (modelName: string): Promise<{ success: boolean; newDigest?: string }> => {
    try {
      const response = await fetch('http://localhost:11434/api/pull', {
        method: 'POST',
        headers: {
          'Content-Type': 'application/json',
        },
        body: JSON.stringify({ name: modelName }),
      });

      if (!response.ok) {
        throw new Error(`Failed to update ${modelName}`);
      }

      const reader = response.body?.getReader();

      if (!reader) {
        throw new Error('No response reader available');
      }

      while (true) {
        const { done, value } = await reader.read();

        if (done) {
          break;
        }

        const text = new TextDecoder().decode(value);
        const lines = text.split('\n').filter(Boolean);

        for (const line of lines) {
          const data = JSON.parse(line) as OllamaPullResponse;

          setModels((current) =>
            current.map((m) =>
              m.name === modelName
                ? {
                    ...m,
                    progress: {
                      current: data.completed || 0,
                      total: data.total || 0,
                      status: data.status,
                    },
                    newDigest: data.digest,
                  }
                : m,
            ),
          );
        }
      }

      setModels((current) => current.map((m) => (m.name === modelName ? { ...m, status: 'checking' } : m)));

      const updatedResponse = await fetch('http://localhost:11434/api/tags');
      const data = (await updatedResponse.json()) as OllamaTagResponse;
      const updatedModel = data.models.find((m) => m.name === modelName);

      return { success: true, newDigest: updatedModel?.digest };
    } catch (error) {
      console.error(`Error updating ${modelName}:`, error);
      return { success: false };
    }
  };

  const handleBulkUpdate = async () => {
    setIsBulkUpdating(true);

    for (const model of models) {
      setModels((current) => current.map((m) => (m.name === model.name ? { ...m, status: 'updating' } : m)));

      const { success, newDigest } = await updateModel(model.name);

      setModels((current) =>
        current.map((m) =>
          m.name === model.name
            ? {
                ...m,
                status: success ? 'updated' : 'error',
                error: success ? undefined : 'Update failed',
                newDigest,
              }
            : m,
        ),
      );
    }

    setIsBulkUpdating(false);
    toast.success('Bulk update completed');
  };

  const handleSingleUpdate = async (modelName: string) => {
    setModels((current) => current.map((m) => (m.name === modelName ? { ...m, status: 'updating' } : m)));

    const { success, newDigest } = await updateModel(modelName);

    setModels((current) =>
      current.map((m) =>
        m.name === modelName
          ? {
              ...m,
              status: success ? 'updated' : 'error',
              error: success ? undefined : 'Update failed',
              newDigest,
            }
          : m,
      ),
    );

    if (success) {
      toast.success(`Updated ${modelName}`);
    } else {
      toast.error(`Failed to update ${modelName}`);
    }
  };

  if (isLoading) {
    return (
      <div className="flex items-center justify-center p-4">
        <div
          className={classNames(
            'rounded-full border-4 border-t-4 border-b-4 border-purple-500',
            'h-16 w-16 animate-spin',
          )}
        />
        <span className="ml-2 text-bolt-elements-textSecondary">Loading models...</span>
      </div>
    );
  }

  return (
    <div
      className={classNames(
        'rounded-lg border bg-bolt-elements-background text-bolt-elements-textPrimary shadow-sm p-4',
        'hover:bg-bolt-elements-background-depth-2',
        'transition-all duration-200',
      )}
    >
      <div className="space-y-4">
        <div className="space-y-2">
          <h2 className="text-2xl font-bold">Ollama Model Manager</h2>
          <p>Update your local Ollama models to their latest versions</p>
        </div>

        <div className="flex items-center justify-between">
          <div className="flex items-center gap-2">
            <div className="i-ph:arrows-clockwise text-purple-500" />
            <span className="text-sm text-bolt-elements-textPrimary">{models.length} models available</span>
          </div>
          <motion.button
            onClick={handleBulkUpdate}
            disabled={isBulkUpdating}
            className={classNames(
              'rounded-md px-4 py-2 text-sm',
              'bg-purple-500 text-white',
              'hover:bg-purple-600',
              'dark:bg-purple-500 dark:hover:bg-purple-600',
              'transition-all duration-200',
            )}
            whileHover={{ scale: 1.02 }}
            whileTap={{ scale: 0.98 }}
          >
            {isBulkUpdating ? (
              <>
                <div
                  className={classNames(
                    'rounded-full border-4 border-t-4 border-b-4 border-purple-500',
                    'h-4 w-4 animate-spin mr-2',
                  )}
                />
                Updating All...
              </>
            ) : (
              <>
                <div className="i-ph:arrows-clockwise" />
                Update All Models
              </>
            )}
          </motion.button>
        </div>

        <div className="space-y-2">
          {models.map((model) => (
            <div
              key={model.name}
              className={classNames(
                'flex items-center justify-between p-3 rounded-lg',
                'bg-[#F8F8F8] dark:bg-[#1A1A1A]',
                'border border-[#E5E5E5] dark:border-[#333333]',
              )}
            >
              <div className="flex flex-col gap-1">
                <div className="flex items-center gap-2">
                  <div className="i-ph:cube text-purple-500" />
                  <span className="text-sm text-bolt-elements-textPrimary">{model.name}</span>
                  {model.status === 'updating' && (
                    <div
                      className={classNames(
                        'rounded-full border-4 border-t-4 border-b-4 border-purple-500',
                        'h-4 w-4 animate-spin',
                      )}
                    />
                  )}
                  {model.status === 'updated' && <div className="i-ph:check-circle text-green-500" />}
                  {model.status === 'error' && <div className="i-ph:x-circle text-red-500" />}
                </div>
                <div className="flex items-center gap-2 text-xs text-bolt-elements-textSecondary">
                  <span>Version: {model.digest.substring(0, 7)}</span>
                  {model.status === 'updated' && model.newDigest && (
                    <>
                      <div className="i-ph:arrow-right w-3 h-3" />
                      <span className="text-green-500">{model.newDigest.substring(0, 7)}</span>
                    </>
                  )}
                  {model.progress && (
                    <span className="ml-2">
                      {model.progress.status}{' '}
                      {model.progress.total > 0 && (
                        <>({Math.round((model.progress.current / model.progress.total) * 100)}%)</>
                      )}
                    </span>
                  )}
                  {model.details && (
                    <span className="ml-2">
                      ({model.details.parameter_size}, {model.details.quantization_level})
                    </span>
                  )}
                </div>
              </div>
              <motion.button
                onClick={() => handleSingleUpdate(model.name)}
                disabled={model.status === 'updating'}
                className={classNames(
                  'rounded-md px-4 py-2 text-sm',
                  'bg-purple-500 text-white',
                  'hover:bg-purple-600',
                  'dark:bg-purple-500 dark:hover:bg-purple-600',
                  'transition-all duration-200',
                )}
                whileHover={{ scale: 1.02 }}
                whileTap={{ scale: 0.98 }}
              >
                <div className="i-ph:arrows-clockwise" />
                Update
              </motion.button>
            </div>
          ))}
        </div>
      </div>
    </div>
  );
}