Include service connector requirements in custom flavor registration

zenml-io · Dec 16, 2024 · c052c59 · c052c59
1 parent 384cb8b
commit c052c59
Showing 1 changed file with 2 additions and 12 deletions.
diff --git a/src/zenml/client.py b/src/zenml/client.py
@@ -109,7 +109,6 @@
     EventSourceResponse,
     EventSourceUpdate,
     FlavorFilter,
-    FlavorRequest,
     FlavorResponse,
     ModelFilter,
     ModelRequest,
@@ -2201,17 +2200,8 @@ def create_flavor(
                 "configuration class' docstring."
             )
 
-        create_flavor_request = FlavorRequest(
-            source=source,
-            type=flavor.type,
-            name=flavor.name,
-            config_schema=flavor.config_schema,
-            integration="custom",
-            user=self.active_user.id,
-            workspace=self.active_workspace.id,
-        )
-
-        return self.zen_store.create_flavor(flavor=create_flavor_request)
+        flavor_request = flavor.to_model(integration="custom", is_custom=True)
+        return self.zen_store.create_flavor(flavor=flavor_request)
 
     def get_flavor(
         self,