Merge pull request #1098 from trheyi/main

Add provider settings retrieval and update CreateCollection request structure
This commit is contained in:
Max 2025-08-12 11:11:29 +08:00 committed by GitHub
commit 4d9d4f9ae4
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 106 additions and 23 deletions

View file

@ -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

View file

@ -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
},
}