1 files changed, 5 insertions, 11 deletions
diff --git a/src/main/scala/xyz/driver/pdsuidomain/formats/json/sprayformats/export.scala b/src/main/scala/xyz/driver/pdsuidomain/formats/json/sprayformats/export.scala
index 39d5c59..1c4b6b4 100644
--- a/src/main/scala/xyz/driver/pdsuidomain/formats/json/sprayformats/export.scala
+++ b/src/main/scala/xyz/driver/pdsuidomain/formats/json/sprayformats/export.scala
@@ -2,15 +2,16 @@ package xyz.driver.pdsuidomain.formats.json.sprayformats
 
 import spray.json._
 import xyz.driver.entities.labels.Label
-import xyz.driver.pdsuidomain.entities.{Arm, Criterion}
+import xyz.driver.formats.json.labels._
 import xyz.driver.pdsuidomain.entities.export.patient._
 import xyz.driver.pdsuidomain.entities.export.trial.{ExportTrialArm, ExportTrialLabelCriterion, ExportTrialWithLabels}
+import xyz.driver.pdsuidomain.entities.{Criterion, EligibilityArm}
 
 object export {
   import DefaultJsonProtocol._
   import common._
-  import record._
   import document._
+  import record._
 
   implicit val patientLabelEvidenceDocumentFormat: RootJsonFormat[ExportPatientLabelEvidenceDocument] =
     jsonFormat5(ExportPatientLabelEvidenceDocument.apply)
@@ -67,7 +68,7 @@ object export {
               longIdFormat[Criterion].read(criterionId),
               value,
               longIdFormat[Label].read(labelId),
-              armIdsVector.map(longIdFormat[Arm].read).toSet,
+              armIdsVector.map(longIdFormat[EligibilityArm].read).toSet,
               criterionText,
               isCompound,
               isDefining
@@ -81,12 +82,5 @@ object export {
     }
 
   implicit val trialWithLabelsFormat: RootJsonFormat[ExportTrialWithLabels] =
-    jsonFormat(ExportTrialWithLabels.apply,
-               "nctId",
-               "trialId",
-               "disease",
-               "lastReviewed",
-               "labelVersion",
-               "arms",
-               "criteria")
+    jsonFormat(ExportTrialWithLabels.apply, "nctId", "trialId", "lastReviewed", "labelVersion", "arms", "criteria")
 }