mlogclub
c3387ec4a9
- Implemented LanceDBProvider for managing vector collections, including methods for creating, deleting, and searching collections. - Added support for LanceDB configuration in VectorDBConfig. - Introduced stub for LanceDB provider when not built with the appropriate tag. - Updated provider initialization to include LanceDB as a supported type. - Created types and interfaces for vector operations, including Vector, SearchRequest, and SearchResult. - Added tests for VectorDBConfig unmarshalling and for ensuring proper error handling when LanceDB is not built. - Updated enums to include LanceDB type and labels. - Modified frontend enums to reflect the addition of LanceDB.
89 lines
2.3 KiB
Go
89 lines
2.3 KiB
Go
package vectordb
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
|
|
"agent-desk/internal/pkg/config"
|
|
"agent-desk/internal/pkg/enums"
|
|
)
|
|
|
|
var defaultProvider Provider
|
|
|
|
func Init(cfg *config.VectorDBConfig) error {
|
|
if cfg == nil || cfg.Type == "" {
|
|
return nil
|
|
}
|
|
|
|
var err error
|
|
switch enums.VectorDBType(cfg.Type) {
|
|
case enums.VectorDBTypeQdrant:
|
|
defaultProvider, err = NewQdrantProvider(&cfg.Qdrant)
|
|
case enums.VectorDBTypeLanceDB:
|
|
defaultProvider, err = NewLanceDBProvider(&cfg.LanceDB)
|
|
default:
|
|
return fmt.Errorf("unsupported vectordb type: %s", cfg.Type)
|
|
}
|
|
return err
|
|
}
|
|
|
|
func GetProvider() Provider {
|
|
return defaultProvider
|
|
}
|
|
|
|
func Close() error {
|
|
if defaultProvider != nil {
|
|
return defaultProvider.Close()
|
|
}
|
|
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)
|
|
}
|