diff --git a/openapi/kb/collection.go b/openapi/kb/collection.go index 3f0a680b..a8f0e250 100644 --- a/openapi/kb/collection.go +++ b/openapi/kb/collection.go @@ -11,6 +11,63 @@ import ( // Collection Management Handlers +// ProviderSettings represents the resolved provider configuration +type ProviderSettings struct { + Dimension int `json:"dimension"` + Connector string `json:"connector"` + Properties map[string]interface{} `json:"properties"` +} + +// getProviderSettings reads and resolves provider settings by provider ID and option value +func getProviderSettings(providerID, optionValue, locale string) (*ProviderSettings, error) { + // Default locale to "en" if empty + if locale == "" { + locale = "en" + } + + // Get the specific provider using KB API + provider, err := kb.GetProviderWithLanguage("embedding", providerID, locale) + if err != nil { + return nil, fmt.Errorf("failed to get provider %s: %v", providerID, err) + } + + // Find the target option + targetOption, found := provider.GetOption(optionValue) + if !found { + return nil, fmt.Errorf("option not found: %s for provider %s", optionValue, providerID) + } + + // Extract settings from option properties + settings := &ProviderSettings{ + Properties: make(map[string]interface{}), + } + + // Copy all properties + if targetOption.Properties != nil { + for key, value := range targetOption.Properties { + settings.Properties[key] = value + } + } + + // Extract dimension + if dim, ok := targetOption.Properties["dimensions"]; ok { + if dimInt, ok := dim.(int); ok { + settings.Dimension = dimInt + } else if dimFloat, ok := dim.(float64); ok { + settings.Dimension = int(dimFloat) + } + } + + // Extract connector + if connector, ok := targetOption.Properties["connector"]; ok { + if connStr, ok := connector.(string); ok { + settings.Connector = connStr + } + } + + return settings, nil +} + // CreateCollection creates a new collection func CreateCollection(c *gin.Context) { var req CreateCollectionRequest @@ -26,6 +83,28 @@ func CreateCollection(c *gin.Context) { return } + // Get provider settings by provider id and option value + providerSettings, err := getProviderSettings(req.Config.EmbeddingProvider, req.Config.EmbeddingOption, req.Config.Locale) + if err != nil { + errorResp := &response.ErrorResponse{ + Code: response.ErrInvalidRequest.Code, + ErrorDescription: fmt.Sprintf("Failed to resolve provider settings: %v", err), + } + response.RespondWithError(c, response.StatusBadRequest, errorResp) + return + } + + // Set dimension by provider settings and add original provider id and option value to metadata with prefix __ + req.Config.Dimension = providerSettings.Dimension + if req.Metadata == nil { + req.Metadata = make(map[string]interface{}) + } + req.Metadata["__embedding_provider"] = req.Config.EmbeddingProvider + req.Metadata["__embedding_option"] = req.Config.EmbeddingOption + if req.Config.Locale != "" { + req.Metadata["__locale"] = req.Config.Locale + } + // Validate request parameters if err := validateCreateCollectionRequest(&req); err != nil { // Create a custom error with the same structure but specific message @@ -52,7 +131,7 @@ func CreateCollection(c *gin.Context) { collectionConfig := types.CollectionConfig{ ID: req.ID, Metadata: req.Metadata, - Config: req.Config, + Config: req.Config.CreateCollectionOptions, } // Call the actual CreateCollection method @@ -269,9 +348,17 @@ func UpdateCollectionMetadata(c *gin.Context) { // CreateCollectionRequest represents the request structure for creating a collection type CreateCollectionRequest struct { - ID string `json:"id" binding:"required"` - Metadata map[string]interface{} `json:"metadata"` - Config *types.CreateCollectionOptions `json:"config" binding:"required"` + ID string `json:"id" binding:"required"` + Metadata map[string]interface{} `json:"metadata"` + Config *CreateCollectionConfig `json:"config" binding:"required"` +} + +// CreateCollectionConfig represents the request structure for creating a collection +type CreateCollectionConfig struct { + EmbeddingProvider string `json:"embedding_provider" binding:"required"` // embedding provider id + EmbeddingOption string `json:"embedding_option" binding:"required"` // embedding option value + Locale string `json:"locale,omitempty"` // locale for provider reading + *types.CreateCollectionOptions } // UpdateCollectionMetadataRequest represents the request structure for updating collection metadata diff --git a/openapi/tests/kb/collection_test.go b/openapi/tests/kb/collection_test.go index a8980239..5d0055b2 100644 --- a/openapi/tests/kb/collection_test.go +++ b/openapi/tests/kb/collection_test.go @@ -45,13 +45,11 @@ func TestCreateCollection(t *testing.T) { "created_by": "test_user", }, "config": map[string]interface{}{ - "collection_name": testCollectionID + "_vector", // Required: collection name - "embedding_model": "text-embedding-ada-002", - "chunk_size": 1000, - "chunk_overlap": 200, - "index_type": "hnsw", // Required: valid index type - "distance": "cosine", // Required: distance metric - "dimension": 1536, // Required: embedding dimensions (note: singular not plural) + "embedding_provider": "__yao.openai", // Required: embedding provider ID + "embedding_option": "text-embedding-3-small", // Required: embedding option value + "locale": "en", // Optional: locale for provider reading + "index_type": "hnsw", // Required: valid index type + "distance": "cosine", // Required: distance metric }, } @@ -171,11 +169,11 @@ func TestRemoveCollection(t *testing.T) { "category": "test_remove", }, "config": map[string]interface{}{ - "collection_name": testCollectionID + "_vector", - "embedding_model": "text-embedding-ada-002", - "index_type": "hnsw", - "distance": "cosine", - "dimension": 1536, + "embedding_provider": "__yao.openai", // Required: embedding provider ID + "embedding_option": "text-embedding-3-small", // Required: embedding option value + "locale": "en", // Optional: locale for provider reading + "index_type": "hnsw", // Required: valid index type + "distance": "cosine", // Required: distance metric }, } @@ -503,13 +501,11 @@ func TestCollectionIntegration(t *testing.T) { "purpose": "full_lifecycle_test", }, "config": map[string]interface{}{ - "collection_name": testCollectionID + "_vector", - "embedding_model": "text-embedding-ada-002", - "chunk_size": 1000, - "chunk_overlap": 200, - "index_type": "hnsw", - "distance": "cosine", - "dimension": 1536, + "embedding_provider": "__yao.openai", // Required: embedding provider ID + "embedding_option": "text-embedding-3-small", // Required: embedding option value + "locale": "en", // Optional: locale for provider reading + "index_type": "hnsw", // Required: valid index type + "distance": "cosine", // Required: distance metric }, }