ai-agent/internal/ai/platform.go

package ai

import (
	"context"
	"errors"
	"fmt"
	"strings"
	"sync"

	"code.tczkiot.com/wlw/ai-agent/contract"
	"code.tczkiot.com/wlw/ai-agent/internal/models"
	"code.tczkiot.com/wlw/ai-agent/internal/pkg/enums"
	"code.tczkiot.com/wlw/ai-agent/internal/repositories"

	"github.com/mlogclub/simple/sqls"
)

// ErrPlatformModelUnsupported means platform mode intentionally does not
// expose this model capability. Callers may use errors.Is to apply a silent,
// deterministic fallback without treating it as an upstream outage.
var ErrPlatformModelUnsupported = errors.New("system built-in AI model type is unsupported")

var platformAIProviderRegistry struct {
	sync.RWMutex
	provider contract.PlatformAIProvider
}

// SetPlatformAIProvider registers the host-provided system AI gateway for
// chat, vision, and embedding calls. Passing nil keeps the standalone/custom
// model behavior unchanged.
func SetPlatformAIProvider(provider contract.PlatformAIProvider) {
	platformAIProviderRegistry.Lock()
	defer platformAIProviderRegistry.Unlock()
	platformAIProviderRegistry.provider = provider
}

func resolveDefaultAIConfig(ctx context.Context, modelType enums.AIModelType) (*models.AIConfig, error) {
	return ResolveAIConfig(ctx, modelType, 0)
}

// ResolveAIConfig resolves the effective model configuration for one logical
// AI call. Platform mode always uses the host gateway and never falls back to
// locally stored credentials. Custom mode uses customConfigID when provided,
// otherwise it selects the enabled configuration for modelType.
func ResolveAIConfig(ctx context.Context, modelType enums.AIModelType, customConfigID int64) (*models.AIConfig, error) {
	provider := currentPlatformAIProvider()
	if provider == nil {
		return resolveCustomAIConfig(modelType, customConfigID)
	}

	source, err := provider.ModelSource(ctx)
	if err != nil {
		return nil, fmt.Errorf("failed to resolve AI model source: %w", err)
	}
	if !strings.EqualFold(strings.TrimSpace(source), contract.ModelSourcePlatform) {
		return resolveCustomAIConfig(modelType, customConfigID)
	}
	if modelType != enums.AIModelTypeLLM && modelType != enums.AIModelTypeEmbedding {
		return nil, fmt.Errorf("%w: %s", ErrPlatformModelUnsupported, modelType)
	}

	platformConfig, err := provider.Config(ctx)
	if err != nil {
		return nil, fmt.Errorf("system built-in AI is unavailable: %w", err)
	}
	if platformConfig == nil {
		return nil, fmt.Errorf("system built-in AI is unavailable")
	}

	config := newPlatformRuntimeConfig(platformConfig, modelType)
	switch modelType {
	case enums.AIModelTypeLLM:
		config.ChatEnabled, config.ModelName = resolvePlatformChatCapability(platformConfig)
		if !config.ChatEnabled {
			return nil, fmt.Errorf("system built-in chat model is not enabled")
		}
	case enums.AIModelTypeEmbedding:
		config.EmbeddingEnabled = platformConfig.EmbeddingEnabled
		config.ModelName = strings.TrimSpace(platformConfig.EmbeddingModel)
		config.Dimension = platformConfig.EmbeddingDimension
		// PlatformAIConfig predates the explicit capability flags. Preserve
		// compatibility with hosts that still provide only a valid model and
		// dimension; new hosts clear these fields when embedding is disabled.
		if !config.EmbeddingEnabled && config.ModelName != "" && config.Dimension > 0 {
			config.EmbeddingEnabled = true
		}
		if !config.EmbeddingEnabled {
			return nil, fmt.Errorf("system built-in embedding model is not enabled")
		}
	}
	if config.BaseURL == "" || config.APIKey == "" || config.ModelName == "" {
		return nil, fmt.Errorf("system built-in %s model is not configured", modelType)
	}
	if modelType == enums.AIModelTypeEmbedding && config.Dimension <= 0 {
		return nil, fmt.Errorf("system built-in embedding dimension is invalid")
	}
	return config, nil
}

// ResolveVisionAIConfig resolves the image-message model independently from
// the ordinary chat capability. Custom mode keeps using the configured LLM;
// the runtime's conservative model-name check decides whether it can receive
// image parts. Platform mode requires the explicit vision task route.
func ResolveVisionAIConfig(ctx context.Context, customConfigID int64) (*models.AIConfig, error) {
	provider := currentPlatformAIProvider()
	if provider == nil {
		return resolveCustomAIConfig(enums.AIModelTypeLLM, customConfigID)
	}
	source, err := provider.ModelSource(ctx)
	if err != nil {
		return nil, fmt.Errorf("failed to resolve AI model source: %w", err)
	}
	if !strings.EqualFold(strings.TrimSpace(source), contract.ModelSourcePlatform) {
		return resolveCustomAIConfig(enums.AIModelTypeLLM, customConfigID)
	}
	platformConfig, err := provider.Config(ctx)
	if err != nil {
		return nil, fmt.Errorf("system built-in AI is unavailable: %w", err)
	}
	if platformConfig == nil {
		return nil, fmt.Errorf("system built-in AI is unavailable")
	}
	config := newPlatformRuntimeConfig(platformConfig, enums.AIModelTypeLLM)
	config.ChatEnabled, _ = resolvePlatformChatCapability(platformConfig)
	if !config.VisionEnabled {
		return nil, fmt.Errorf("system built-in vision model is not enabled")
	}
	config.ModelName = strings.TrimSpace(config.VisionModel)
	if config.ModelName == "" {
		return nil, fmt.Errorf("system built-in vision model is not configured")
	}
	if config.BaseURL == "" || config.APIKey == "" {
		return nil, fmt.Errorf("system built-in vision model is not configured")
	}
	return config, nil
}

func newPlatformRuntimeConfig(platformConfig *contract.PlatformAIConfig, modelType enums.AIModelType) *models.AIConfig {
	return &models.AIConfig{
		Provider:        enums.AIProviderOpenAI,
		BaseURL:         strings.TrimRight(strings.TrimSpace(platformConfig.BaseURL), "/"),
		APIKey:          platformConfig.APIKey,
		ModelType:       modelType,
		MaxOutputTokens: platformConfig.MaxOutputTokens,
		TimeoutMS:       platformConfig.TimeoutMS,
		MaxRetryCount:   platformConfig.MaxRetryCount,
		Status:          enums.StatusOk,
		Platform:        true,
		HTTPClient:      platformConfig.HTTPClient,
		VisionEnabled:   platformConfig.VisionEnabled,
		VisionModel:     strings.TrimSpace(platformConfig.VisionModel),
	}
}

func resolvePlatformChatCapability(platformConfig *contract.PlatformAIConfig) (bool, string) {
	if platformConfig == nil {
		return false, ""
	}
	chatModel := strings.TrimSpace(platformConfig.ChatModel)
	if chatModel != "" {
		return platformConfig.ChatEnabled, chatModel
	}
	// ModelName is the legacy chat field. A non-empty legacy value remains an
	// enabled chat capability so existing host implementations keep working.
	legacyModel := strings.TrimSpace(platformConfig.ModelName)
	if legacyModel != "" {
		return true, legacyModel
	}
	return platformConfig.ChatEnabled, ""
}

func resolveCustomAIConfig(modelType enums.AIModelType, customConfigID int64) (*models.AIConfig, error) {
	if customConfigID <= 0 {
		return GetEnabledAIConfig(modelType)
	}
	config := repositories.AIConfigRepository.Get(sqls.DB(), customConfigID)
	if config == nil || config.Status != enums.StatusOk {
		return nil, fmt.Errorf("ai config is unavailable")
	}
	return config, nil
}

func currentPlatformAIProvider() contract.PlatformAIProvider {
	platformAIProviderRegistry.RLock()
	defer platformAIProviderRegistry.RUnlock()
	return platformAIProviderRegistry.provider
}