GopherGate/internal/server/models_config.go

package server

import (
	"fmt"
	"net/http"

	"gophergate/internal/db"

	"github.com/gin-gonic/gin"
)

func (s *Server) handleGetModels(c *gin.Context) {
	usedOnly := c.Query("used_only") == "true"

	// Registry provider normalized name -> Proxy-internal provider ID
	allowedRegistryProviders := map[string]string{
		"openai":   "openai",
		"google":   "gemini",
		"deepseek": "deepseek",
		"xai":      "grok",
		"ollama":   "ollama",
		"xiaomi":   "xiaomi",
	}

	// Merge registry models with DB overrides
	var dbModels []db.ModelConfig
	_ = s.database.Select(&dbModels, "SELECT * FROM model_configs")

	dbMap := make(map[string]db.ModelConfig)
	for _, m := range dbModels {
		dbMap[m.ID] = m
	}

	// Fetch specific (model, provider) combinations that have been used
	type modelProvider struct {
		Model    string `db:"model"`
		Provider string `db:"provider"`
	}
	usedPairs := make(map[string]bool)
	if usedOnly {
		var pairs []modelProvider
		err := s.database.Select(&pairs, "SELECT DISTINCT model, provider FROM llm_requests WHERE status = 'success'")
		if err == nil {
			for _, p := range pairs {
				usedPairs[fmt.Sprintf("%s:%s", p.Model, p.Provider)] = true
			}
		}
	}

	var result []gin.H
	s.registryMu.RLock()
	if s.registry != nil {
		for pID, pInfo := range s.registry.Providers {
			proxyProvider, allowed := allowedRegistryProviders[pID]
			if !allowed {
				continue
			}

			for mID, mMeta := range pInfo.Models {
				if usedOnly && !usedPairs[fmt.Sprintf("%s:%s", mID, proxyProvider)] {
					continue
				}

				enabled := true
				promptCost := 0.0
				completionCost := 0.0
				var cacheReadCost *float64
				var cacheWriteCost *float64
				var mapping *string
				contextLimit := uint32(0)

				if mMeta.Cost != nil {
					promptCost = mMeta.Cost.Input
					completionCost = mMeta.Cost.Output
					cacheReadCost = mMeta.Cost.CacheRead
					cacheWriteCost = mMeta.Cost.CacheWrite
				}
				if mMeta.Limit != nil {
					contextLimit = mMeta.Limit.Context
				}

				// Override from DB
				if dbCfg, ok := dbMap[mID]; ok {
					enabled = dbCfg.Enabled
					if dbCfg.PromptCostPerM != nil {
						promptCost = *dbCfg.PromptCostPerM
					}
					if dbCfg.CompletionCostPerM != nil {
						completionCost = *dbCfg.CompletionCostPerM
					}
					if dbCfg.CacheReadCostPerM != nil {
						cacheReadCost = dbCfg.CacheReadCostPerM
					}
					if dbCfg.CacheWriteCostPerM != nil {
						cacheWriteCost = dbCfg.CacheWriteCostPerM
					}
					mapping = dbCfg.Mapping
				}

				result = append(result, gin.H{
					"id":               mID,
					"name":             mMeta.Name,
					"provider":         proxyProvider,
					"enabled":          enabled,
					"prompt_cost":      promptCost,
					"completion_cost":  completionCost,
					"cache_read_cost":  cacheReadCost,
					"cache_write_cost": cacheWriteCost,
					"context_limit":    contextLimit,
					"mapping":          mapping,
					"tool_call":        mMeta.ToolCall != nil && *mMeta.ToolCall,
					"reasoning":        mMeta.Reasoning != nil && *mMeta.Reasoning,
					"modalities":       mMeta.Modalities,
				})
			}
		}
	}

	// Add configured Ollama models if they aren't in registry
	if s.cfg.Providers.Ollama.Enabled {
		for _, mID := range s.cfg.Providers.Ollama.Models {
			// Check if already added from registry
			exists := false
			for _, r := range result {
				if r["id"] == mID {
					exists = true
					break
				}
			}
			if exists {
				continue
			}

			if usedOnly && !usedPairs[fmt.Sprintf("%s:ollama", mID)] {
				continue
			}

			enabled := true
			promptCost := 0.0
			completionCost := 0.0
			var cacheReadCost *float64
			var cacheWriteCost *float64
			var mapping *string
			contextLimit := uint32(0)

			// Override from DB
			if dbCfg, ok := dbMap[mID]; ok {
				enabled = dbCfg.Enabled
				if dbCfg.PromptCostPerM != nil {
					promptCost = *dbCfg.PromptCostPerM
				}
				if dbCfg.CompletionCostPerM != nil {
					completionCost = *dbCfg.CompletionCostPerM
				}
				if dbCfg.CacheReadCostPerM != nil {
					cacheReadCost = dbCfg.CacheReadCostPerM
				}
				if dbCfg.CacheWriteCostPerM != nil {
					cacheWriteCost = dbCfg.CacheWriteCostPerM
				}
				mapping = dbCfg.Mapping
			}

			result = append(result, gin.H{
				"id":               mID,
				"name":             mID,
				"provider":         "ollama",
				"enabled":          enabled,
				"prompt_cost":      promptCost,
				"completion_cost":  completionCost,
				"cache_read_cost":  cacheReadCost,
				"cache_write_cost": cacheWriteCost,
				"context_limit":    contextLimit,
				"modalities":       gin.H{"input": []string{"text"}, "output": []string{"text"}},
				"tool_call":        false,
				"reasoning":        false,
				"mapping":          mapping,
			})
		}
	}

	c.JSON(http.StatusOK, SuccessResponse(result))
}

func (s *Server) handleUpdateModel(c *gin.Context) {
	id := c.Param("id")
	var req struct {
		Enabled        bool     `json:"enabled"`
		PromptCost     float64  `json:"prompt_cost"`
		CompletionCost float64  `json:"completion_cost"`
		CacheReadCost  *float64 `json:"cache_read_cost"`
		CacheWriteCost *float64 `json:"cache_write_cost"`
		Mapping        *string  `json:"mapping"`
	}
	if err := c.ShouldBindJSON(&req); err != nil {
		c.JSON(http.StatusBadRequest, ErrorResponse("Invalid request"))
		return
	}

	// Find provider for this model
	providerID := "unknown"
	s.registryMu.RLock()
	if s.registry != nil {
		for pID, pInfo := range s.registry.Providers {
			if _, ok := pInfo.Models[id]; ok {
				providerID = pID
				break
			}
		}
	}

	_, err := s.database.Exec(`
		INSERT INTO model_configs (id, provider_id, enabled, prompt_cost_per_m, completion_cost_per_m, cache_read_cost_per_m, cache_write_cost_per_m, mapping)
		VALUES (?, ?, ?, ?, ?, ?, ?, ?)
		ON CONFLICT(id) DO UPDATE SET
			enabled = excluded.enabled,
			prompt_cost_per_m = excluded.prompt_cost_per_m,
			completion_cost_per_m = excluded.completion_cost_per_m,
			cache_read_cost_per_m = excluded.cache_read_cost_per_m,
			cache_write_cost_per_m = excluded.cache_write_cost_per_m,
			mapping = excluded.mapping,
			updated_at = CURRENT_TIMESTAMP
	`, id, providerID, req.Enabled, req.PromptCost, req.CompletionCost, req.CacheReadCost, req.CacheWriteCost, req.Mapping)

	if err != nil {
		c.JSON(http.StatusInternalServerError, ErrorResponse(err.Error()))
		return
	}

	c.JSON(http.StatusOK, SuccessResponse(gin.H{"message": "Model updated"}))
}