ai-agent/internal/services/storage/provider.go

package storage

import (
	"context"
	"fmt"
	"io"

	"code.tczkiot.com/wlw/ai-agent/contract"
	"code.tczkiot.com/wlw/ai-agent/internal/pkg/config"
	"code.tczkiot.com/wlw/ai-agent/internal/pkg/enums"
	"code.tczkiot.com/wlw/ai-agent/internal/pkg/errorsx"
)

var providers = make(map[enums.AssetProvider]FileStorageProvider)
var hostStorage contract.FileStorage

type FileStorageProvider interface {
	ProviderType() enums.AssetProvider
	Upload(reader io.Reader, key string, info UploadInfo) (*StoredFile, error)
	GetURL(key string) string
	GetSignedURL(key string) string
	Delete(key string) error
	Read(key string) (io.ReadCloser, error)
}

// SetHostStorage lets the embedding system own file persistence. Passing nil
// preserves the standalone module's built-in local/OSS providers.
func SetHostStorage(value contract.FileStorage) {
	hostStorage = value
	providers = make(map[enums.AssetProvider]FileStorageProvider)
}

func GetDefault() (FileStorageProvider, error) {
	if hostStorage != nil {
		provider, err := hostStorage.DefaultProvider(context.Background())
		if err != nil {
			return nil, err
		}
		if provider == "" {
			return nil, fmt.Errorf("host file storage returned an empty default provider")
		}
		return GetProvider(enums.AssetProvider(provider))
	}
	return NewProvider(config.Current().Storage.Default)
}

func GetProvider(providerType enums.AssetProvider) (FileStorageProvider, error) {
	if provider, exists := providers[providerType]; exists {
		return provider, nil
	}
	provider, err := NewProvider(providerType)
	if err != nil {
		return nil, err
	}

	providers[providerType] = provider

	return provider, nil
}

func NewProvider(provider enums.AssetProvider) (FileStorageProvider, error) {
	if hostStorage != nil {
		if provider == "" {
			return GetDefault()
		}
		return &hostFileStorageProvider{provider: provider}, nil
	}
	cfg := config.Current().Storage

	switch provider {
	case "", enums.AssetProviderLocal:
		return NewLocalStorage(cfg.Local), nil
	case enums.AssetProviderOSS:
		return NewOSSStorage(cfg.OSS), nil
	default:
		return nil, errorsx.InvalidParamI18n("error.e0082")
	}
}

type hostFileStorageProvider struct {
	provider enums.AssetProvider
}

func (p *hostFileStorageProvider) ProviderType() enums.AssetProvider {
	return p.provider
}

func (p *hostFileStorageProvider) Upload(reader io.Reader, key string, info UploadInfo) (*StoredFile, error) {
	accessURL, err := hostStorage.Upload(
		context.Background(), string(p.provider), key, info.Filename, info.MimeType, info.FileSize, reader,
	)
	if err != nil {
		return nil, err
	}
	return &StoredFile{
		Provider:   p.provider,
		StorageKey: key,
		URL:        accessURL,
		Filename:   info.Filename,
		FileSize:   info.FileSize,
		MimeType:   info.MimeType,
	}, nil
}

func (p *hostFileStorageProvider) GetURL(key string) string {
	value, _ := hostStorage.URL(context.Background(), string(p.provider), key)
	return value
}

func (p *hostFileStorageProvider) GetSignedURL(key string) string {
	return p.GetURL(key)
}

func (p *hostFileStorageProvider) Delete(key string) error {
	return hostStorage.Delete(context.Background(), string(p.provider), key)
}

func (p *hostFileStorageProvider) Read(key string) (io.ReadCloser, error) {
	return hostStorage.Open(context.Background(), string(p.provider), key)
}