pull-model-form.tsx

"use client";

import React, { useState } from "react";
import {
  Form,
  FormField,
  FormItem,
  FormLabel,
  FormMessage,
} from "@/components/ui/form";
import { Button } from "./ui/button";
import { z } from "zod";
import { useForm } from "react-hook-form";
import { zodResolver } from "@hookform/resolvers/zod";
import { toast } from "sonner";
import { Loader2Icon } from "lucide-react";
import { Input } from "./ui/input";
import { useRouter } from "next/navigation";

const formSchema = z.object({
  name: z.string().min(1, {
    message: "Please select a model to pull",
  }),
});

export default function PullModelForm() {
  const [isDownloading, setIsDownloading] = useState(false);
  const [name, setName] = useState("");
  const router = useRouter();
  const env = process.env.NODE_ENV;

  const form = useForm<z.infer<typeof formSchema>>({
    resolver: zodResolver(formSchema),
  });

  function onSubmit(data: z.infer<typeof formSchema>) {
    // Trim whitespace
    data.name = data.name.trim();

    setIsDownloading(true);
    // Send the model name to the server
    if (env === "production") {
      // Make a post request to localhost
      const pullModel = async () => {
        const response = await fetch(process.env.NEXT_PUBLIC_OLLAMA_URL + "/api/pull", {
          method: "POST",
          headers: {
            "Content-Type": "application/json",
          },
          body: JSON.stringify(data),
        });
        const json = await response.json();
        if (json.error) {
          toast.error("Error: " + json.error);
          setIsDownloading(false);
          return;
        } else if (json.status === "success") {
          toast.success("Model pulled successfully");
          setIsDownloading(false);
          return;
        }
      }
      pullModel();
    } else {
      fetch("/api/model", {
        method: "POST",
        headers: {
          "Content-Type": "application/json",
        },
        body: JSON.stringify(data),
      })
        .then((response) => {
          // Check if response is successful
          if (!response.ok) {
            throw new Error("Network response was not ok");
          }
          if (!response.body) {
            throw new Error("Something went wrong");
          }
          // Create a new ReadableStream from the response body
          const reader = response.body.getReader();
  
          // Read the data in chunks
          reader.read().then(function processText({ done, value }) {
            if (done) {
              setIsDownloading(false);
              return;
            }
  
            // Convert the chunk of data to a string
            const text = new TextDecoder().decode(value);
  
            // Split the text into individual JSON objects
            const jsonObjects = text.trim().split("\n");
  
            jsonObjects.forEach((jsonObject) => {
              try {
                const responseJson = JSON.parse(jsonObject);
                if (responseJson.error) {
                  // Display an error toast if the response contains an error
                  toast.error("Error: " + responseJson.error);
                  setIsDownloading(false);
                  return;
                } else if (responseJson.status === "success") {
                  // Display a success toast if the response status is success
                  toast.success("Model pulled successfully");
                  setIsDownloading(false);
                  return;
                }
              } catch (error) {
                toast.error("Error parsing JSON");
                setIsDownloading(false);
                return;
              }
            });
  
            // Continue reading the next chunk
            reader.read().then(processText);
          });
        })
        .catch((error) => {
          setIsDownloading(false);
          console.error("Error pulling model:", error);
          toast.error("Error pulling model");
        });
    }
  }

  const handleChange = (e: React.ChangeEvent<HTMLInputElement>) => {
    e.preventDefault();
    form.setValue("name", e.currentTarget.value);
    setName(e.currentTarget.value);
  };

  return (
    <Form {...form}>
      <form onSubmit={form.handleSubmit(onSubmit)} className="w-full space-y-6">
        <FormField
          control={form.control}
          name="name"
          render={({ field }) => (
            <FormItem>
              <FormLabel>Model name</FormLabel>
              <Input
                {...field}
                type="text"
                placeholder="llama2"
                value={name}
                onChange={(e) => handleChange(e)}
              />
              <p className="text-xs pt-1">
                Check the{" "}
                <a
                  href="https://ollama.com/library"
                  target="_blank"
                  className="text-blue-500 underline"
                >
                  library
                </a>{" "}
                for a list of available models.
              </p>
              <FormMessage />
            </FormItem>
          )}
        />
        <div className="space-y-2 w-full">
          <Button type="submit" className="w-full " disabled={isDownloading}>
            {isDownloading ? (
              <div className="flex items-center gap-2">
                <Loader2Icon className="animate-spin w-4 h-4" />
                <span>Pulling model...</span>
              </div>
            ) : (
              "Pull model"
            )}
          </Button>
          <p className="text-xs text-center">
            {isDownloading
              ? "This may take a while. You can safely close this modal and continue using the app"
              : "Pressing the button will download the specified model to your device."}
          </p>
        </div>
      </form>
    </Form>
  );
}