go-ollama/models.go

package gollama

import (
	"context"
	"sort"
	"time"

	"github.com/ollama/ollama/api"
)

// Model describes a locally available Ollama model.
type Model struct {
	Name         string
	Size         int64
	ModifiedAt   time.Time
	Digest       string
	Family       string
	ParameterSize string
	Quantization string
}

// ModelInfo contains detailed information from Ollama's show endpoint.
type ModelInfo struct {
	Name       string
	License    string
	Modelfile  string
	Parameters string
	Template   string
	System     string
	ModifiedAt time.Time
}

// ModelNames returns sorted names of locally available models.
func (c *Client) ModelNames(ctx context.Context) ([]string, error) {
	models, err := c.Models(ctx)
	if err != nil {
		return nil, err
	}
	names := make([]string, len(models))
	for i, m := range models {
		names[i] = m.Name
	}
	sort.Strings(names)
	return names, nil
}

// Models returns locally available models with metadata.
func (c *Client) Models(ctx context.Context) ([]Model, error) {
	resp, err := c.api.List(ctx)
	if err != nil {
		return nil, wrapError("list", "", err)
	}
	out := make([]Model, 0, len(resp.Models))
	for _, m := range resp.Models {
		out = append(out, listModelToModel(m))
	}
	return out, nil
}

func listModelToModel(m api.ListModelResponse) Model {
	name := m.Name
	if name == "" {
		name = m.Model
	}
	return Model{
		Name:          name,
		Size:          m.Size,
		ModifiedAt:    m.ModifiedAt,
		Digest:        m.Digest,
		Family:        m.Details.Family,
		ParameterSize: m.Details.ParameterSize,
		Quantization:  m.Details.QuantizationLevel,
	}
}

// Pull downloads a model and blocks until complete.
func (c *Client) Pull(ctx context.Context, name string) error {
	return c.PullWithProgress(ctx, name, nil)
}

// PullWithProgress downloads a model, optionally reporting progress.
func (c *Client) PullWithProgress(ctx context.Context, name string, fn func(Progress) error) error {
	stream := false
	req := &api.PullRequest{Model: name, Stream: &stream}
	err := c.api.Pull(ctx, req, func(p api.ProgressResponse) error {
		if fn == nil {
			return nil
		}
		return fn(Progress{
			Status:    p.Status,
			Digest:    p.Digest,
			Total:     p.Total,
			Completed: p.Completed,
		})
	})
	return wrapError("pull", name, err)
}

// Show returns detailed information about a model.
func (c *Client) Show(ctx context.Context, name string) (*ModelInfo, error) {
	resp, err := c.api.Show(ctx, &api.ShowRequest{Model: name})
	if err != nil {
		return nil, wrapError("show", name, err)
	}
	return &ModelInfo{
		Name:       name,
		License:    resp.License,
		Modelfile:  resp.Modelfile,
		Parameters: resp.Parameters,
		Template:   resp.Template,
		System:     resp.System,
		ModifiedAt: resp.ModifiedAt,
	}, nil
}

// ModelDefaults parses a model's modelfile parameters into an OptionSet.
func (c *Client) ModelDefaults(ctx context.Context, name string) (OptionSet, error) {
	info, err := c.Show(ctx, name)
	if err != nil {
		return OptionSet{}, err
	}
	return parseModelfileParameters(info.Parameters)
}

// Delete removes a model from the local Ollama instance.
func (c *Client) Delete(ctx context.Context, name string) error {
	return wrapError("delete", name, c.api.Delete(ctx, &api.DeleteRequest{Model: name}))
}

// RunningModel describes a model currently loaded in memory.
type RunningModel struct {
	Name    string
	Size    int64
	Expires time.Time
}

// RunningModels returns models currently loaded in Ollama memory.
func (c *Client) RunningModels(ctx context.Context) ([]RunningModel, error) {
	resp, err := c.api.ListRunning(ctx)
	if err != nil {
		return nil, wrapError("list_running", "", err)
	}
	out := make([]RunningModel, 0, len(resp.Models))
	for _, m := range resp.Models {
		name := m.Name
		if name == "" {
			name = m.Model
		}
		out = append(out, RunningModel{
			Name:    name,
			Size:    m.Size,
			Expires: m.ExpiresAt,
		})
	}
	return out, nil
}