added float data type for cognition schema (#1362)

digiteinfotech · Jul 19, 2024 · 32ecad2 · 32ecad2
1 parent 2a1b4dd
commit 32ecad2
Show file tree

Hide file tree

Showing 6 changed files with 15 additions and 9 deletions.
diff --git a/kairon/api/models.py b/kairon/api/models.py
@@ -1116,8 +1116,8 @@ class ColumnMetadata(BaseModel):
     def check(cls, values):
         from kairon.shared.utils import Utility
 
-        if values.get('data_type') not in [CognitionMetadataType.str.value, CognitionMetadataType.int.value]:
-            raise ValueError("Only str and int data types are supported")
+        if values.get('data_type') not in [CognitionMetadataType.str.value, CognitionMetadataType.int.value, CognitionMetadataType.float.value]:
+            raise ValueError("Only str, int and float data types are supported")
         if Utility.check_empty_string(values.get('column_name')):
             raise ValueError("Column name cannot be empty")
         return values

diff --git a/kairon/shared/cognition/data_objects.py b/kairon/shared/cognition/data_objects.py
@@ -11,7 +11,7 @@
 class ColumnMetadata(EmbeddedDocument):
     column_name = StringField(required=True)
     data_type = StringField(required=True, default=CognitionMetadataType.str.value,
-                            choices=[CognitionMetadataType.str.value, CognitionMetadataType.int.value])
+                            choices=[CognitionMetadataType.str.value, CognitionMetadataType.int.value, CognitionMetadataType.float.value])
     enable_search = BooleanField(default=True)
     create_embeddings = BooleanField(default=True)
 
@@ -20,8 +20,8 @@ def validate(self, clean=True):
 
         if clean:
             self.clean()
-        if self.data_type not in [CognitionMetadataType.str.value, CognitionMetadataType.int.value]:
-            raise ValidationError("Only str and int data types are supported")
+        if self.data_type not in [CognitionMetadataType.str.value, CognitionMetadataType.int.value, CognitionMetadataType.float.value]:
+            raise ValidationError("Only str,int and float data types are supported")
         if Utility.check_empty_string(self.column_name):
             raise ValidationError("Column name cannot be empty")
 

diff --git a/kairon/shared/cognition/processor.py b/kairon/shared/cognition/processor.py
@@ -229,6 +229,11 @@ def validate_column_values(data: Any, schema: Dict):
                     return int(data[column_name])
                 except ValueError:
                     raise AppException("Invalid data type!")
+            elif column_name in data and data[column_name] and data_type == CognitionMetadataType.float.value:
+                try:
+                    return float(data[column_name])
+                except ValueError:
+                    raise AppException("Invalid data type!")
             else:
                 return data[column_name]
 

diff --git a/kairon/shared/models.py b/kairon/shared/models.py
@@ -113,3 +113,4 @@ class CognitionDataType(str, Enum):
 class CognitionMetadataType(str, Enum):
     str = "str"
     int = "int"
+    float = "float"
diff --git a/tests/integration_test/services_test.py b/tests/integration_test/services_test.py
@@ -2714,10 +2714,10 @@ def test_metadata_upload_invalid_data_type():
     )
     actual = response.json()
     assert actual["message"] == [{'loc': ['body', 'metadata', 0, 'data_type'],
-                                  'msg': "value is not a valid enumeration member; permitted: 'str', 'int'",
-                                  'type': 'type_error.enum', 'ctx': {'enum_values': ['str', 'int']}},
+                                  'msg': "value is not a valid enumeration member; permitted: 'str', 'int', 'float'",
+                                  'type': 'type_error.enum', 'ctx': {'enum_values': ['str', 'int', 'float']}},
                                  {'loc': ['body', 'metadata', 0, '__root__'],
-                                  'msg': 'Only str and int data types are supported', 'type': 'value_error'}]
+                                  'msg': 'Only str, int and float data types are supported', 'type': 'value_error'}]
     assert not actual["data"]
     assert actual["error_code"] == 422
 

diff --git a/tests/unit_test/data_processor/data_processor_test.py b/tests/unit_test/data_processor/data_processor_test.py
@@ -15025,7 +15025,7 @@ def test_save_payload_metadata_data_type_invalid(self):
             "bot": bot,
             "user": user
         }
-        with pytest.raises(ValidationError, match="Only str and int data types are supported"):
+        with pytest.raises(ValidationError, match="Only str,int and float data types are supported"):
             CognitionSchema(**schema).save()
 
     def test_get_payload_metadata(self):