diff --git a/api/src/helper/lib/base-nlp-helper.ts b/api/src/helper/lib/base-nlp-helper.ts
index c88ed6cf..5b0e8310 100644
--- a/api/src/helper/lib/base-nlp-helper.ts
+++ b/api/src/helper/lib/base-nlp-helper.ts
@@ -123,7 +123,7 @@ export default abstract class BaseNlpHelper<
   async format(
     samples: NlpSampleFull[],
     entities: NlpEntityFull[],
-  ): Promise<Record<string, any>[]> {
+  ): Promise<Record<string, any>[] | Record<string, any>> {
     const entityMap = NlpEntity.getEntityMap(entities);
     const valueMap = NlpValue.getValueMap(
       NlpValue.getValuesFromEntities(entities),