From 666248ee40410d6cd585e5a3423922c046affb7f Mon Sep 17 00:00:00 2001
From: Gabriel Luiz Freitas Almeida <gabriel@logspace.ai>
Date: Sun, 31 Mar 2024 22:05:00 -0300
Subject: [PATCH] Fix range_spec step type validation

---
 .../base/langflow/components/helpers/CreateRecord.py  |  2 +-
 src/backend/base/langflow/field_typing/range_spec.py  | 11 ++++++++++-
 src/backend/base/langflow/template/field/base.py      |  9 ++++++++-
 3 files changed, 19 insertions(+), 3 deletions(-)

diff --git a/src/backend/base/langflow/components/helpers/CreateRecord.py b/src/backend/base/langflow/components/helpers/CreateRecord.py
index c85f0e583..744282f7d 100644
--- a/src/backend/base/langflow/components/helpers/CreateRecord.py
+++ b/src/backend/base/langflow/components/helpers/CreateRecord.py
@@ -52,7 +52,7 @@ class CreateRecordComponent(CustomComponent):
                 "display_name": "Number of Fields",
                 "info": "Number of fields to be added to the record.",
                 "real_time_refresh": True,
-                "rangeSpec": RangeSpec(min=1, max=15, step=1),
+                "rangeSpec": RangeSpec(min=1, max=15, step=1, step_type="int"),
             },
             "text_key": {
                 "display_name": "Text Key",
diff --git a/src/backend/base/langflow/field_typing/range_spec.py b/src/backend/base/langflow/field_typing/range_spec.py
index b81cf7a22..7147f9c20 100644
--- a/src/backend/base/langflow/field_typing/range_spec.py
+++ b/src/backend/base/langflow/field_typing/range_spec.py
@@ -1,3 +1,5 @@
+from typing import Literal
+
 from pydantic import BaseModel, field_validator
 
 
@@ -5,6 +7,7 @@ class RangeSpec(BaseModel):
     min: float = -1.0
     max: float = 1.0
     step: float = 0.1
+    step_type: Literal["int", "float"] = "float"
 
     @field_validator("max")
     @classmethod
@@ -15,7 +18,13 @@ class RangeSpec(BaseModel):
 
     @field_validator("step")
     @classmethod
-    def step_must_be_positive(cls, v):
+    def step_must_be_positive(cls, v, values, **kwargs):
         if v <= 0:
             raise ValueError("Step must be positive")
+        if values.data["step_type"] == "int" and isinstance(v, float) and not v.is_integer():
+            raise ValueError("When step_type is int, step must be an integer")
         return v
+
+    @classmethod
+    def set_step_type(cls, step_type: Literal["int", "float"], range_spec: "RangeSpec") -> "RangeSpec":
+        return cls(min=range_spec.min, max=range_spec.max, step=range_spec.step, step_type=step_type)
diff --git a/src/backend/base/langflow/template/field/base.py b/src/backend/base/langflow/template/field/base.py
index 85c9affdd..c68a5c476 100644
--- a/src/backend/base/langflow/template/field/base.py
+++ b/src/backend/base/langflow/template/field/base.py
@@ -1,6 +1,6 @@
 from typing import Any, Callable, Optional, Union
 
-from pydantic import BaseModel, ConfigDict, Field, field_serializer, field_validator, model_serializer
+from pydantic import BaseModel, ConfigDict, Field, field_serializer, field_validator, model_serializer, model_validator
 
 from langflow.field_typing.range_spec import RangeSpec
 
@@ -91,6 +91,13 @@ class TemplateField(BaseModel):
             result["type"] = self.field_type
         return result
 
+    @model_validator(mode="after")
+    def validate_model(self):
+        # if field_type is int, we need to set the range_spec
+        if self.field_type == "int" and self.range_spec is not None:
+            self.range_spec = RangeSpec.set_step_type("int", self.range_spec)
+        return self
+
     @field_serializer("file_path")
     def serialize_file_path(self, value):
         return value if self.field_type == "file" else ""