Fix build pipeline (#132)

Signed-off-by: Software Developer <[email protected]>
Signed-off-by: dsuhinin <[email protected]>
Co-authored-by: DSuhinin <[email protected]>

Author: dsuhinin

pkg/model_registry/store/sql/helpers.go

@@ -0,0 +1,40 @@
+package sql
+
+import (
+	"fmt"
+	"strings"
+
+	"github.com/mlflow/mlflow-go-backend/pkg/contract"
+	"github.com/mlflow/mlflow-go-backend/pkg/protos"
+)
+
+// HandleResourceAlreadyExistError handles name conflicts in the Model Registry.
+func HandleResourceAlreadyExistError(name string, isExistingEntityPrompt, isNewEntityPrompt bool) *contract.Error {
+	// Determine the entity types
+	oldEntity := "Registered Model"
+	if isExistingEntityPrompt {
+		oldEntity = "Prompt"
+	}
+
+	newEntity := "Registered Model"
+	if isNewEntityPrompt {
+		newEntity = "Prompt"
+	}
+
+	// Check if there is a conflict between different entity types
+	if oldEntity != newEntity {
+		return contract.NewError(
+			protos.ErrorCode_RESOURCE_ALREADY_EXISTS,
+			fmt.Sprintf(
+				"Tried to create a %s with name '%s', but the name is already taken by a %s. "+
+					"MLflow does not allow creating a model and a prompt with the same name.",
+				strings.ToLower(newEntity), name, strings.ToLower(oldEntity)),
+		)
+	}
+
+	// Raise an error if the entity already exists
+	return contract.NewError(
+		protos.ErrorCode_RESOURCE_ALREADY_EXISTS,
+		fmt.Sprintf("%s (name=%s) already exists.", newEntity, name),
+	)
+}

pkg/model_registry/store/sql/registered_models.go

@@ -15,6 +15,11 @@ import (
 	"github.com/mlflow/mlflow-go-backend/pkg/entities"
 	"github.com/mlflow/mlflow-go-backend/pkg/model_registry/store/sql/models"
 	"github.com/mlflow/mlflow-go-backend/pkg/protos"
+	"github.com/mlflow/mlflow-go-backend/pkg/utils"
+)
+
+const (
+	IsPromptTagKey = "mlflow.prompt.is_prompt"
 )
 
 func (m *ModelRegistrySQLStore) GetRegisteredModel(
@@ -251,9 +256,19 @@ func (m *ModelRegistrySQLStore) CreateRegisteredModel(
 
 	if err := m.db.WithContext(ctx).Create(&registeredModel).Error; err != nil {
 		if errors.Is(err, gorm.ErrDuplicatedKey) {
-			return nil, contract.NewError(
-				protos.ErrorCode_RESOURCE_ALREADY_EXISTS,
-				fmt.Sprintf("Registered Model (name=%s) already exists.", registeredModel.Name),
+			existingRegisteredModel, err := m.GetRegisteredModel(ctx, name)
+			if err != nil {
+				return nil, err
+			}
+
+			return nil, HandleResourceAlreadyExistError(
+				name,
+				utils.FindElementByProperty(existingRegisteredModel.Tags, func(tag *entities.RegisteredModelTag) bool {
+					return tag.Key == IsPromptTagKey
+				}),
+				utils.FindElementByProperty(tags, func(tag *entities.RegisteredModelTag) bool {
+					return tag.Key == IsPromptTagKey
+				}),
 			)
 		}
 

pkg/utils/slice.go

@@ -0,0 +1,12 @@
+package utils
+
+// FindElementByProperty searches for an element in a slice based on a given match function.
+func FindElementByProperty[T any](items []T, matchFunc func(T) bool) bool {
+	for _, item := range items {
+		if matchFunc(item) {
+			return true
+		}
+	}
+
+	return false
+}