core/server/handlers/capabilities.go

package handlers

import (
	"encoding/json"
	"log"

	"github.com/gin-gonic/gin"

	"git.gobha.me/xcaliber/chat-switchboard/database"
	"git.gobha.me/xcaliber/chat-switchboard/providers"
)

// ResolveModelCaps is the canonical capability resolver for any model.
// It walks a priority chain and returns the best capabilities available:
//
//  1. model_configs DB — exact match (model_id + api_config_id)
//  2. model_configs DB — any provider (same model, different config)
//  3. Known model table (static, curated)
//  4. Heuristic inference (name-based fallback)
//
// configID is optional — pass "" to skip the exact-match step.
func ResolveModelCaps(c *gin.Context, modelID, configID string) providers.ModelCapabilities {
	// ── 1. Exact match: model_id + api_config_id ──
	if configID != "" {
		caps, ok := capsFromModelConfigs(modelID, configID)
		if ok {
			caps.MaxOutputTokens = providers.ResolveMaxOutput(modelID, caps)
			return caps
		}
	}

	// ── 2. Any provider: same model_id, any config ──
	caps, ok := capsFromModelConfigs(modelID, "")
	if ok {
		caps.MaxOutputTokens = providers.ResolveMaxOutput(modelID, caps)
		return caps
	}

	// ── 3. Known model table (static, curated) ──
	caps, found := providers.LookupKnownModel(modelID)
	if found {
		caps.MaxOutputTokens = providers.ResolveMaxOutput(modelID, caps)
		return caps
	}

	// ── 4. Heuristic inference ──
	caps = providers.InferCapabilities(modelID)
	caps.MaxOutputTokens = providers.ResolveMaxOutput(modelID, caps)
	return caps
}

// capsFromModelConfigs looks up capabilities from the model_configs table.
// If configID is non-empty, it matches exactly; otherwise it finds any entry
// for the model_id (capabilities for the same model are provider-agnostic).
func capsFromModelConfigs(modelID, configID string) (providers.ModelCapabilities, bool) {
	var capsJSON []byte
	var err error
	if configID != "" {
		err = database.DB.QueryRow(`
			SELECT capabilities FROM model_configs
			WHERE model_id = $1 AND api_config_id = $2
		`, modelID, configID).Scan(&capsJSON)
	} else {
		err = database.DB.QueryRow(`
			SELECT capabilities FROM model_configs
			WHERE model_id = $1 ORDER BY updated_at DESC LIMIT 1
		`, modelID).Scan(&capsJSON)
	}
	if err != nil || len(capsJSON) == 0 {
		return providers.ModelCapabilities{}, false
	}
	var caps providers.ModelCapabilities
	if json.Unmarshal(capsJSON, &caps) != nil || !caps.HasProviderData() {
		return providers.ModelCapabilities{}, false
	}
	return providers.MergeCapabilities(caps, modelID), true
}

// ResolveModelCapsFromLoaded checks an existing slice of models first (avoids
// redundant DB/network calls when we already have models in memory).
func ResolveModelCapsFromLoaded(c *gin.Context, modelID, configID string, loaded []enabledModel) providers.ModelCapabilities {
	// Check already-loaded models first
	for _, m := range loaded {
		if m.ModelID == modelID {
			return m.Capabilities
		}
	}
	// Fall through to canonical resolver
	return ResolveModelCaps(c, modelID, configID)
}

// liveQueryModelCaps queries a provider API to get capabilities for a specific model.
// Used for team provider presets whose base model isn't in model_configs.
func liveQueryModelCaps(c *gin.Context, configID, modelID string) (providers.ModelCapabilities, bool) {
	var providerID, endpoint string
	var apiKey *string
	var headersJSON []byte
	err := database.DB.QueryRow(`
		SELECT provider, endpoint, api_key_encrypted, custom_headers
		FROM api_configs WHERE id = $1 AND is_active = true
	`, configID).Scan(&providerID, &endpoint, &apiKey, &headersJSON)
	if err != nil {
		return providers.ModelCapabilities{}, false
	}

	provider, err := providers.Get(providerID)
	if err != nil {
		return providers.ModelCapabilities{}, false
	}

	key := ""
	if apiKey != nil {
		key = *apiKey
	}

	var customHeaders map[string]string
	_ = json.Unmarshal(headersJSON, &customHeaders)

	modelList, err := provider.ListModels(c.Request.Context(), providers.ProviderConfig{
		Endpoint:      endpoint,
		APIKey:        key,
		CustomHeaders: customHeaders,
	})
	if err != nil {
		log.Printf("[caps] live query for %s via config %s failed: %v", modelID, configID, err)
		return providers.ModelCapabilities{}, false
	}

	for _, m := range modelList {
		if m.ID == modelID {
			caps := providers.MergeCapabilities(m.Capabilities, modelID)
			return caps, true
		}
	}

	return providers.ModelCapabilities{}, false
}