fix: model name unify

StartUpNationLabs · Oct 20, 2024 · a6f3c10 · a6f3c10
1 parent 9124d73
commit a6f3c10
Show file tree

Hide file tree

Showing 3 changed files with 37 additions and 25 deletions.
diff --git a/internal/controller/custommodel_controller.go b/internal/controller/custommodel_controller.go
@@ -75,6 +75,7 @@ func (r *CustomModelReconciler) Reconcile(ctx context.Context, req ctrl.Request)
 	// If the CustomModel is being deleted, delete it from the Ollama Client
 	CustomModelFinalizer := "CustomModel.finalizer.ollama.ollama.startupnation"
 
+	modelName := unifyModelName(CustomModel.Spec.ModelName)
 	if CustomModel.ObjectMeta.DeletionTimestamp.IsZero() {
 		// The object is not being deleted, so if it does not have our finalizer,
 		// then lets add the finalizer and update the object.
@@ -89,9 +90,9 @@ func (r *CustomModelReconciler) Reconcile(ctx context.Context, req ctrl.Request)
 		if containsString(CustomModel.ObjectMeta.Finalizers, CustomModelFinalizer) {
 			// our finalizer is present, so lets handle our external dependency
 			// first, we delete the external dependency
-			logger.Info("Deleting CustomModel", "CustomModel Name", CustomModel.Spec.ModelName, "Ollama URL", ollamaUrl)
+			logger.Info("Deleting CustomModel", "CustomModel Name", modelName, "Ollama URL", ollamaUrl)
 			_, err := ollamaClient.DeleteApiDelete(ctx, &ollama_client.DeleteApiDeleteParams{
-				Model: CustomModel.Spec.ModelName,
+				Model: modelName,
 			})
 			if err != nil {
 				logger.Error(err, "unable to delete CustomModel")
@@ -110,33 +111,33 @@ func (r *CustomModelReconciler) Reconcile(ctx context.Context, req ctrl.Request)
 	// if the CustomModel is not being deleted, start reconciliation
 
 	// get the CustomModel from the Ollama Client
-	logger.Info("Checking if CustomModel exists", "CustomModel Name", CustomModel.Spec.ModelName, "Ollama URL", ollamaUrl)
+	logger.Info("Checking if CustomModel exists", "CustomModel Name", modelName, "Ollama URL", ollamaUrl)
 	res, err := ollamaClient.PostApiShowWithResponse(ctx, ollama_client.PostApiShowJSONRequestBody{
-		Name: &CustomModel.Spec.ModelName,
+		Name: &modelName,
 	})
-	logger.Info("Checking if CustomModel exists", "CustomModel Name", CustomModel.Spec.ModelName, "Ollama URL", ollamaUrl, "status", res.Status())
+	logger.Info("Checking if CustomModel exists", "CustomModel Name", modelName, "Ollama URL", ollamaUrl, "status", res.Status())
 	if err == nil && res.StatusCode() == 200 {
-		logger.Info("CustomModel exists", "CustomModel Name", CustomModel.Spec.ModelName, "Ollama URL", ollamaUrl)
+		logger.Info("CustomModel exists", "CustomModel Name", modelName, "Ollama URL", ollamaUrl)
 		if res.JSON200 != nil {
-			logger.Info("Checking if the ModelFile is the same", "CustomModel Name", CustomModel.Spec.ModelName, "Ollama URL", ollamaUrl)
+			logger.Info("Checking if the ModelFile is the same", "CustomModel Name", modelName, "Ollama URL", ollamaUrl)
 			if *res.JSON200.Parameters == CustomModel.Spec.ModelFile {
-				logger.Info("ModelFile is the same", "CustomModel Name", CustomModel.Spec.ModelName, "Ollama URL", ollamaUrl)
+				logger.Info("ModelFile is the same", "CustomModel Name", modelName, "Ollama URL", ollamaUrl)
 				return ctrl.Result{}, nil
 			}
-			logger.Info("ModelFile is not the same", "CustomModel Name", CustomModel.Spec.ModelName, "Ollama URL", ollamaUrl)
+			logger.Info("ModelFile is not the same", "CustomModel Name", modelName, "Ollama URL", ollamaUrl)
 			// delete the CustomModel and create a new one
-			logger.Info("Deleting CustomModel", "CustomModel Name", CustomModel.Spec.ModelName, "Ollama URL", ollamaUrl)
+			logger.Info("Deleting CustomModel", "CustomModel Name", modelName, "Ollama URL", ollamaUrl)
 			_, err := ollamaClient.DeleteApiDelete(ctx, &ollama_client.DeleteApiDeleteParams{
-				Model: CustomModel.Spec.ModelName,
+				Model: modelName,
 			})
 			if err != nil {
 				logger.Error(err, "unable to delete CustomModel")
 			}
 			// Create the CustomModel
-			logger.Info("Creating CustomModel", "CustomModel Name", CustomModel.Spec.ModelName, "Ollama URL", ollamaUrl)
+			logger.Info("Creating CustomModel", "CustomModel Name", modelName, "Ollama URL", ollamaUrl)
 			stream := false
 			_, err = ollamaClient.PostApiCreate(ctx, ollama_client.PostApiCreateJSONRequestBody{
-				Name:      &CustomModel.Spec.ModelName,
+				Name:      &modelName,
 				Modelfile: &CustomModel.Spec.ModelFile,
 				Stream:    &stream,
 			})
@@ -148,18 +149,18 @@ func (r *CustomModelReconciler) Reconcile(ctx context.Context, req ctrl.Request)
 		return ctrl.Result{}, err
 	}
 	// if the CustomModel does not exist, create it
-	logger.Info("CustomModel does not exist, creating CustomModel", "CustomModel Name", CustomModel.Spec.ModelName, "Ollama URL", ollamaUrl, "ModelFile", CustomModel.Spec.ModelFile)
+	logger.Info("CustomModel does not exist, creating CustomModel", "CustomModel Name", modelName, "Ollama URL", ollamaUrl, "ModelFile", CustomModel.Spec.ModelFile)
 	stream := false
 	_, err = ollamaClient.PostApiCreate(ctx, ollama_client.PostApiCreateJSONRequestBody{
-		Name:      &CustomModel.Spec.ModelName,
+		Name:      &modelName,
 		Modelfile: &CustomModel.Spec.ModelFile,
 		Stream:    &stream,
 	})
 	if err != nil || res.StatusCode() != 200 {
 		logger.Error(err, "unable to create CustomModel")
 		return ctrl.Result{}, err
 	}
-	logger.Info("CustomModel created", "CustomModel Name", CustomModel.Spec.ModelName, "Ollama URL", ollamaUrl)
+	logger.Info("CustomModel created", "CustomModel Name", modelName, "Ollama URL", ollamaUrl)
 
 	return ctrl.Result{}, nil
 }

diff --git a/internal/controller/helper.go b/internal/controller/helper.go
@@ -1,5 +1,7 @@
 package controller
 
+import "strings"
+
 // Helper functions to check and remove string from a slice of strings.
 func containsString(slice []string, s string) bool {
 	for _, item := range slice {
@@ -19,3 +21,11 @@ func removeString(slice []string, s string) (result []string) {
 	}
 	return
 }
+
+func unifyModelName(modelName string) string {
+	// check if : is present in the modelName
+	if strings.Contains(modelName, ":") {
+		return modelName
+	}
+	return modelName + ":latest"
+}
diff --git a/internal/controller/model_controller.go b/internal/controller/model_controller.go
@@ -74,6 +74,7 @@ func (r *ModelReconciler) Reconcile(ctx context.Context, req ctrl.Request) (ctrl
 	// If the Model is being deleted, delete it from the Ollama Client
 	modelFinalizer := "model.finalizer.ollama.ollama.startupnation"
 
+	modelName := unifyModelName(model.Spec.ModelName)
 	if model.ObjectMeta.DeletionTimestamp.IsZero() {
 		// The object is not being deleted, so if it does not have our finalizer,
 		// then lets add the finalizer and update the object.
@@ -88,9 +89,9 @@ func (r *ModelReconciler) Reconcile(ctx context.Context, req ctrl.Request) (ctrl
 		if containsString(model.ObjectMeta.Finalizers, modelFinalizer) {
 			// our finalizer is present, so lets handle our external dependency
 			// first, we delete the external dependency
-			logger.Info("Deleting Model", "Model Name", model.Spec.ModelName, "Ollama URL", ollamaUrl)
+			logger.Info("Deleting Model", "Model Name", modelName, "Ollama URL", ollamaUrl)
 			_, err := ollamaClient.DeleteApiDelete(ctx, &ollama_client.DeleteApiDeleteParams{
-				Model: model.Spec.ModelName,
+				Model: modelName,
 			})
 			if err != nil {
 				logger.Error(err, "unable to delete Model")
@@ -109,31 +110,31 @@ func (r *ModelReconciler) Reconcile(ctx context.Context, req ctrl.Request) (ctrl
 	// if the Model is not being deleted, start reconciliation
 
 	// get the model from the Ollama Client
-	logger.Info("Checking if Model exists", "Model Name", model.Spec.ModelName, "Ollama URL", ollamaUrl)
+	logger.Info("Checking if Model exists", "Model Name", modelName, "Ollama URL", ollamaUrl)
 	res, err := ollamaClient.PostApiShowWithResponse(ctx, ollama_client.PostApiShowJSONRequestBody{
-		Name: &model.Spec.ModelName,
+		Name: &modelName,
 	})
-	logger.Info("Checking if Model exists", "Model Name", model.Spec.ModelName, "Ollama URL", ollamaUrl, "status", res.Status())
+	logger.Info("Checking if Model exists", "Model Name", modelName, "Ollama URL", ollamaUrl, "status", res.Status())
 	if err == nil && res.StatusCode() == 200 {
-		logger.Info("Model exists", "Model Name", model.Spec.ModelName, "Ollama URL", ollamaUrl)
+		logger.Info("Model exists", "Model Name", modelName, "Ollama URL", ollamaUrl)
 		if res.JSON200 != nil {
 			logger.Info("Model exists", "Model Params", res.JSON200.Parameters, "Ollama URL", ollamaUrl)
 			return ctrl.Result{}, nil
 		}
 		return ctrl.Result{}, err
 	}
 	// if the model does not exist, create it
-	logger.Info("Model does not exist, creating Model", "Model Name", model.Spec.ModelName, "Ollama URL", ollamaUrl)
+	logger.Info("Model does not exist, creating Model", "Model Name", modelName, "Ollama URL", ollamaUrl)
 	stream := false
 	_, err = ollamaClient.PostApiPull(ctx, ollama_client.PostApiPullJSONRequestBody{
-		Name:   &model.Spec.ModelName,
+		Name:   &modelName,
 		Stream: &stream,
 	})
 	if err != nil {
 		logger.Error(err, "unable to create Model")
 		return ctrl.Result{}, err
 	}
-	logger.Info("Model created", "Model Name", model.Spec.ModelName, "Ollama URL", ollamaUrl)
+	logger.Info("Model created", "Model Name", modelName, "Ollama URL", ollamaUrl)
 
 	return ctrl.Result{}, nil
 }