ai-agent/internal/ai/rag/vectordb/provider.go

package vectordb

import (
	"context"
	"fmt"

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

var defaultProvider Provider

func Init(cfg *config.VectorDBConfig) error {
	if cfg == nil {
		return fmt.Errorf("libsql vector database config is required")
	}
	provider, err := NewLibSQLProvider(cfg)
	if err != nil {
		return err
	}
	if defaultProvider != nil {
		_ = defaultProvider.Close()
	}
	defaultProvider = provider
	return nil
}

func GetProvider() Provider {
	return defaultProvider
}

func Close() error {
	if defaultProvider != nil {
		err := defaultProvider.Close()
		defaultProvider = nil
		return err
	}
	return nil
}

func CreateCollection(ctx context.Context, name string, dimension int) error {
	if defaultProvider == nil {
		return fmt.Errorf("vectordb provider not initialized")
	}
	return defaultProvider.CreateCollection(ctx, name, dimension)
}

func DeleteCollection(ctx context.Context, name string) error {
	if defaultProvider == nil {
		return fmt.Errorf("vectordb provider not initialized")
	}
	return defaultProvider.DeleteCollection(ctx, name)
}

func GetCollection(ctx context.Context, name string) (*CollectionInfo, error) {
	if defaultProvider == nil {
		return nil, fmt.Errorf("vectordb provider not initialized")
	}
	return defaultProvider.GetCollection(ctx, name)
}

func ListCollections(ctx context.Context) ([]string, error) {
	if defaultProvider == nil {
		return nil, fmt.Errorf("vectordb provider not initialized")
	}
	return defaultProvider.ListCollections(ctx)
}

func UpsertVectors(ctx context.Context, collectionName string, vectors []Vector) error {
	if defaultProvider == nil {
		return fmt.Errorf("vectordb provider not initialized")
	}
	return defaultProvider.UpsertVectors(ctx, collectionName, vectors)
}

func DeleteVectors(ctx context.Context, collectionName string, ids []string) error {
	if defaultProvider == nil {
		return fmt.Errorf("vectordb provider not initialized")
	}
	return defaultProvider.DeleteVectors(ctx, collectionName, ids)
}

func Search(ctx context.Context, req *SearchRequest) ([]SearchResult, error) {
	if defaultProvider == nil {
		return nil, fmt.Errorf("vectordb provider not initialized")
	}
	return defaultProvider.Search(ctx, req)
}