llama-link/internal/api/handlers/models.go

package handlers

import (
	"encoding/json"
	"net/http"

	"github.com/gin-gonic/gin"
	"github.com/google/uuid"

	"github.com/llamalink/llamalink/internal/db"
	"github.com/llamalink/llamalink/internal/llama"
)

type ModelsHandler struct {
	manager     *llama.Manager
}

func NewModelsHandler(manager *llama.Manager) *ModelsHandler {
	return &ModelsHandler{manager: manager}
}

type CreateModelRequest struct {
	Name       string                 `json:"name" binding:"required"`
	ModelPath string                 `json:"model_path" binding:"required"`
	Alias      string                 `json:"alias" binding:"required"`
	CtxSize    int                    `json:"ctx_size"`
	NGPULayers int                   `json:"n_gpu_layers"`
	ExtraArgs  map[string]interface{} `json:"extra_args"`
	IsDefault  bool                   `json:"is_default"`
}

func (h *ModelsHandler) ListModels(c *gin.Context) {
	var models []db.Model
	db := h.manager.GetDB()
	if err := db.Where("is_enabled = ?", true).Order("name").Find(&models).Error; err != nil {
		c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
		return
	}

	status := h.manager.GetStatus()
	currentModel := status.CurrentModel

	result := make([]gin.H, len(models))
	for i, m := range models {
		result[i] = gin.H{
			"id":          m.ID,
			"name":        m.Name,
			"model_path":  m.ModelPath,
			"alias":       m.Alias,
			"ctx_size":    m.CtxSize,
			"n_gpu_layers": m.NGPULayers,
			"is_default":  m.IsDefault,
			"is_active":   m.Name == currentModel && status.Status == llama.StatusReady,
			"loaded_at":   m.LoadedAt,
		}
	}

	c.JSON(http.StatusOK, gin.H{"data": result})
}

func (h *ModelsHandler) GetActiveModel(c *gin.Context) {
	status := h.manager.GetStatus()

	if status.CurrentModel == "" {
		c.JSON(http.StatusOK, gin.H{
			"data": gin.H{
				"status":        status.Status,
				"current_model": nil,
			},
		})
		return
	}

	c.JSON(http.StatusOK, gin.H{
		"data": gin.H{
			"status":        status.Status,
			"current_model": status.CurrentModel,
			"loaded_at":      status.LoadedAt,
		},
	})
}

func (h *ModelsHandler) CreateModel(c *gin.Context) {
	var req CreateModelRequest
	if err := c.ShouldBindJSON(&req); err != nil {
		c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
		return
	}

	if req.CtxSize == 0 {
		req.CtxSize = 8192
	}
	if req.NGPULayers == 0 {
		req.NGPULayers = -1
	}

	extraArgsJSON, _ := json.Marshal(req.ExtraArgs)

	model := &db.Model{
		ID:          uuid.New(),
		Name:        req.Name,
		ModelPath:   req.ModelPath,
		Alias:       req.Alias,
		CtxSize:     req.CtxSize,
		NGPULayers:  req.NGPULayers,
		ExtraArgs:   db.StringArray{string(extraArgsJSON)},
		IsDefault:   req.IsDefault,
		IsEnabled:   true,
		IsActive:    false,
	}

	db := h.manager.GetDB()
	if err := db.Create(model).Error; err != nil {
		c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
		return
	}

	c.JSON(http.StatusCreated, gin.H{
		"id":         model.ID,
		"name":       model.Name,
		"model_path": model.ModelPath,
		"alias":      model.Alias,
		"ctx_size":   model.CtxSize,
		"is_default": model.IsDefault,
	})
}

func (h *ModelsHandler) LoadModel(c *gin.Context) {
	name := c.Param("name")

	if err := h.manager.LoadModel(name); err != nil {
		if err == llama.ErrModelNotFound {
			c.JSON(http.StatusNotFound, gin.H{"error": "model not found"})
			return
		}
		if err == llama.ErrSwapInProgress {
			c.JSON(http.StatusConflict, gin.H{"error": err.Error()})
			return
		}
		c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
		return
	}

	c.JSON(http.StatusAccepted, gin.H{
		"status":  "loading",
		"model":   name,
		"message": "Model loading initiated",
	})
}