Bugfix - send in output schema for refuel models (#910)

refuel-ai · Oct 1, 2024 · d1417b2 · d1417b2
1 parent 5cedaff
commit d1417b2
Showing 1 changed file with 1 addition and 2 deletions.
diff --git a/src/autolabel/models/refuelV2.py b/src/autolabel/models/refuelV2.py
@@ -277,8 +277,7 @@ def _prepare_output_schema(self, schema: Dict) -> Dict:
                     del curr_schema[key]
                 else:
                     curr_schema[key] = self._prepare_output_schema(curr_schema[key])
-        else:
-            return curr_schema
+        return curr_schema
 
     def get_cost(self, prompt: str, label: Optional[str] = "") -> float:
         return 0