fix: edit validation json file

radicalbit · Dec 13, 2024 · e8e0a52 · e8e0a52
1 parent 8cf0812
commit e8e0a52
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 6 deletions.
diff --git a/api/app/services/file_service.py b/api/app/services/file_service.py
@@ -3,6 +3,7 @@
 import json
 import logging
 import pathlib
+from io import BytesIO
 from typing import List, Optional
 from uuid import UUID, uuid4
 
@@ -344,13 +345,13 @@ def upload_completion_file(
             logger.error('Model %s not found', model_uuid)
             raise ModelNotFoundError(f'Model {model_uuid} not found')
 
-        self.validate_json_file(json_file)
-        _f_name = json_file.filename
+        validated_json_file = self.validate_json_file(json_file)
+        _f_name = validated_json_file.filename
         _f_uuid = uuid4()
         try:
             object_name = f'{str(model_out.uuid)}/completion/{_f_uuid}/{_f_name}'
             self.s3_client.upload_fileobj(
-                json_file.file,
+                validated_json_file.file,
                 self.bucket_name,
                 object_name,
                 ExtraArgs={
@@ -560,11 +561,15 @@ def validate_file(
         csv_file.file.seek(0)
 
     @staticmethod
-    def validate_json_file(json_file: UploadFile) -> None:
+    def validate_json_file(json_file: UploadFile):
         try:
             content = json_file.file.read().decode('utf-8')
             json_data = json.loads(content)
-            CompletionResponses.model_validate(json_data)
+            validated_data = CompletionResponses.model_validate(json_data)
+            return UploadFile(
+                filename=json_file.filename,
+                file=BytesIO(validated_data.model_dump_json().encode()),
+            )
         except ValidationError as e:
             logger.error('Invalid json file: %s', str(e))
             raise InvalidFileException(f'Invalid json file: {str(e)}') from e

diff --git a/api/tests/services/file_service_test.py b/api/tests/services/file_service_test.py
@@ -83,7 +83,7 @@ def test_infer_schema_separator(self):
 
     def test_validate_completion_json_file_ok(self):
         json_file = json.get_completion_sample_json_file()
-        self.files_service.validate_json_file(json_file)
+        assert self.files_service.validate_json_file(json_file) is not None
 
     def test_validate_completion_json_file_without_logprobs_field(self):
         json_file = json.get_completion_sample_json_file_without_logprobs_field()