llmgw/internal/cli/mirror_import.go

package cli

import (
	"context"
	"fmt"
	"io"
	"path/filepath"
	"strconv"
	"strings"
	"time"

	"github.com/spf13/cobra"

	"gitea.mynodes.xyz/ilker/llmgw/internal/mirror"
)

// placeholderBase is how the mirror's address is written when the operator
// has not said it, as `mirror add` writes it.
const placeholderBase = "http://<this-host>:8081"

func newMirrorImportCmd(dir *string) *cobra.Command {
	var model, version, as, sum, base string

	c := &cobra.Command{
		Use:   "import <file>... --model NAME",
		Short: "Copy model files on this machine into the mirror",
		Long: "For weights already here: a GGUF you converted or downloaded by hand,\n" +
			"or the files of a safetensors model. The same discipline as add — copied\n" +
			"under a temporary name and renamed when whole, the checksum kept beside\n" +
			"it for nodes to check against, an existing version never replaced.\n\n" +
			"A file named .gguf is checked to be one. Several files go into one\n" +
			"version, as the files of one model; its source is then the version,\n" +
			"ending in a slash. Run it on the mirror; for a model in Ollama, use\n" +
			"import-ollama.",
		Example: "  llmgw mirror import ./qwen2.5-7b-instruct-q4_k_m.gguf --model qwen2.5-7b-instruct",
		Args:    cobra.MinimumNArgs(1),
		RunE: func(cmd *cobra.Command, args []string) error {
			if model == "" {
				return fmt.Errorf("--model is required: it is the mirror's directory for the model")
			}
			if as != "" && len(args) > 1 {
				return fmt.Errorf("--as names one file; with %d files each keeps its own name", len(args))
			}
			if sum != "" && len(args) > 1 {
				return fmt.Errorf("--sha256 is one file's checksum; with %d files, import them one at a time", len(args))
			}
			if version == "" {
				version = time.Now().Format("2006-01-02")
			}
			out := cmd.OutOrStdout()
			var added []mirror.Added
			for _, src := range args {
				abs, err := filepath.Abs(src)
				if err != nil {
					return err
				}
				file := as
				if file == "" {
					file = filepath.Base(src)
				}
				e := mirror.Entry{Model: model, Version: version, File: file}
				fmt.Fprintf(out, "copying %s\n", src)
				a, err := mirror.Import(cmd.Context(), *dir, abs, e, sum, "file://"+abs, progressTo(out))
				if err != nil {
					return err
				}
				fmt.Fprintf(out, "%s/%s/%s  %s\n  sha256 %s\n", model, version, file, fileSize(a.Entry.Bytes), a.SHA256)
				added = append(added, a)
			}
			fmt.Fprintf(out, "\nPoint a model at it with:\n  source: %s\n", importedSource(base, added))
			return nil
		},
	}
	c.Flags().StringVar(&model, "model", "", "the model's name, one directory level")
	c.Flags().StringVar(&version, "version", "", "the version, one directory level (default today's date)")
	c.Flags().StringVar(&as, "as", "", "the filename to keep it as, for one file; its own name when omitted")
	c.Flags().StringVar(&sum, "sha256", "", "the checksum its publisher gave, to refuse a file that does not match")
	c.Flags().StringVar(&base, "base", placeholderBase, "the address nodes reach this mirror at, for the source printed")
	return c
}

func newMirrorImportOllamaCmd(dir *string) *cobra.Command {
	var o ollamaImport

	c := &cobra.Command{
		Use:   "import-ollama [NAME[:TAG]]",
		Short: "Put a model you already have in Ollama into the mirror, for llama.cpp",
		Long: "Ollama keeps each model as a GGUF under its models directory, named by\n" +
			"its sha256. This copies that file into the mirror — checked against the\n" +
			"digest it is named by, and with a vision model's projector beside it —\n" +
			"so llama.cpp can serve it on any node, from your own network.\n\n" +
			"With no name it lists the models Ollama has. The mirror version is the\n" +
			"tag and the start of the digest, since a tag such as latest moves in\n" +
			"Ollama and a mirror version must not; --model and --version choose\n" +
			"otherwise.\n\n" +
			"Ollama's prompt template is not carried: llama.cpp renders the chat\n" +
			"template inside the GGUF. Run it on the mirror, where Ollama's directory\n" +
			"is readable — ~/.ollama/models, /usr/share/ollama/.ollama/models for\n" +
			"Ollama's Linux service, or OLLAMA_MODELS.",
		Example: "  llmgw mirror import-ollama\n" +
			"  sudo llmgw mirror import-ollama llama3.2:3b --ollama-dir /usr/share/ollama/.ollama/models",
		Args: cobra.MaximumNArgs(1),
		RunE: func(cmd *cobra.Command, args []string) error {
			o.mirrorDir = *dir
			if len(args) == 1 {
				o.name = args[0]
			}
			return o.run(cmd.Context(), cmd.OutOrStdout())
		},
	}
	c.Flags().StringVar(&o.ollamaDir, "ollama-dir", "", "Ollama's models directory (default: found where Ollama keeps it)")
	c.Flags().StringVar(&o.model, "model", "", "the mirror's directory for it (default the Ollama model's name)")
	c.Flags().StringVar(&o.version, "version", "", "the mirror version (default its tag and the start of its digest)")
	c.Flags().StringVar(&o.base, "base", placeholderBase, "the address nodes reach this mirror at, for the source printed")
	return c
}

// ollamaImport is one import-ollama, from the command line or the console.
type ollamaImport struct {
	mirrorDir, ollamaDir, name, model, version, base string
	// quiet leaves out the copy's progress, for the console, which shows
	// the output once the import is over.
	quiet bool
}

func (o ollamaImport) run(ctx context.Context, out io.Writer) error {
	dir, err := mirror.OllamaDir(o.ollamaDir)
	if err != nil {
		return err
	}
	if o.name == "" {
		names, err := mirror.ListOllama(dir)
		if err != nil {
			return err
		}
		if len(names) == 0 {
			fmt.Fprintf(out, "Ollama in %s has no models.\n", dir)
			return nil
		}
		fmt.Fprintf(out, "Ollama's models in %s:\n", dir)
		for _, n := range names {
			fmt.Fprintf(out, "  %s\n", n)
		}
		fmt.Fprintln(out, "\nImport one with: llmgw mirror import-ollama NAME:TAG")
		return nil
	}

	n, err := mirror.ParseOllamaName(o.name)
	if err != nil {
		return err
	}
	m, err := mirror.ReadOllama(dir, n)
	if err != nil {
		return err
	}
	fmt.Fprintf(out, "importing %s from %s\n", n, dir)
	last := ""
	added, err := mirror.ImportOllama(ctx, o.mirrorDir, dir, m, o.model, o.version, func(file string, done, total int64) {
		if file != last {
			fmt.Fprintf(out, "copying %s\n", file)
			last = file
		}
		if !o.quiet {
			progressTo(out)(done, total)
		}
	})
	if err != nil {
		return err
	}
	for _, a := range added {
		fmt.Fprintf(out, "%s/%s/%s  %s\n  sha256 %s, as Ollama named it\n",
			a.Entry.Model, a.Entry.Version, a.Entry.File, fileSize(a.Entry.Bytes), a.SHA256)
	}
	for _, l := range m.Left {
		fmt.Fprintf(out, "Not carried: %s.\n", l)
	}

	name := n.Model
	if n.Tag != "latest" {
		name += "-" + n.Tag
	}
	fmt.Fprintf(out, "\nA model definition for it, for llmgw ctl model add -f:\n\n"+
		"name: %s\nsource: %s\nprofiles:\n  - name: cpu\n    engine: llamacpp\n",
		name, importedSource(o.base, added))
	if m.ContextLength > 0 {
		// Ollama's context, which llama.cpp's own default would not match;
		// llama.cpp divides it between its parallel slots.
		fmt.Fprintf(out, "    args: [\"--ctx-size\", %q]   # Ollama ran it with num_ctx %d\n",
			strconv.Itoa(m.ContextLength), m.ContextLength)
	}
	return nil
}

// importedSource is what a model's source says for what was just imported:
// the file, or the version's directory when it is several.
func importedSource(base string, added []mirror.Added) string {
	if base == "" {
		base = placeholderBase
	}
	if len(added) == 1 {
		return added[0].Entry.URL(base)
	}
	e := added[0].Entry
	return strings.TrimSuffix(mirror.Entry{Model: e.Model, Version: e.Version}.URL(base), "/") + "/"
}

// progressTo reports a copy's progress the way add reports a download's.
func progressTo(out io.Writer) func(done, total int64) {
	return func(done, total int64) {
		if total > 0 {
			fmt.Fprintf(out, "  %s of %s (%d%%)\n", fileSize(done), fileSize(total), done*100/total)
			return
		}
		fmt.Fprintf(out, "  %s so far\n", fileSize(done))
	}
}