feat: add LanceDB vector database provider
- 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.
This commit is contained in:
@@ -9,56 +9,11 @@ import (
|
||||
"agent-desk/internal/pkg/config"
|
||||
)
|
||||
|
||||
type Vector struct {
|
||||
ID string `json:"id"`
|
||||
Vector []float32 `json:"vector"`
|
||||
Payload ChunkPayload `json:"payload"`
|
||||
}
|
||||
|
||||
type SearchRequest struct {
|
||||
CollectionName string `json:"collectionName"`
|
||||
Vector []float32 `json:"vector"`
|
||||
TopK int `json:"topK"`
|
||||
ScoreThreshold float32 `json:"scoreThreshold"`
|
||||
Filter *SearchFilter `json:"filter,omitempty"`
|
||||
}
|
||||
|
||||
type SearchFilter struct {
|
||||
KnowledgeBaseIDs []int64 `json:"knowledgeBaseIds,omitempty"`
|
||||
DocumentIDs []int64 `json:"documentIds,omitempty"`
|
||||
}
|
||||
|
||||
type SearchResult struct {
|
||||
ID string `json:"id"`
|
||||
Score float32 `json:"score"`
|
||||
Payload ChunkPayload `json:"payload"`
|
||||
}
|
||||
|
||||
type CollectionInfo struct {
|
||||
Name string `json:"name"`
|
||||
Dimension int `json:"dimension"`
|
||||
PointCount int `json:"pointCount"`
|
||||
Status string `json:"status"`
|
||||
}
|
||||
|
||||
type Provider interface {
|
||||
CreateCollection(ctx context.Context, name string, dimension int) error
|
||||
DeleteCollection(ctx context.Context, name string) error
|
||||
GetCollection(ctx context.Context, name string) (*CollectionInfo, error)
|
||||
ListCollections(ctx context.Context) ([]string, error)
|
||||
|
||||
UpsertVectors(ctx context.Context, collectionName string, vectors []Vector) error
|
||||
DeleteVectors(ctx context.Context, collectionName string, ids []string) error
|
||||
|
||||
Search(ctx context.Context, req *SearchRequest) ([]SearchResult, error)
|
||||
Close() error
|
||||
}
|
||||
|
||||
type QdrantProvider struct {
|
||||
client *qdrant.Client
|
||||
}
|
||||
|
||||
func NewQdrantProvider(cfg *config.VectorDBConfig) (*QdrantProvider, error) {
|
||||
func NewQdrantProvider(cfg *config.QdrantVectorDBConfig) (*QdrantProvider, error) {
|
||||
if cfg == nil {
|
||||
return nil, fmt.Errorf("vectordb config is nil")
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user