Send extra metadata to the Metadata Extractor service (#7088)

* Send extra metadata to the Metadata Extractor service
huridocs · Oct 18, 2024 · bfad3e3 · bfad3e3
1 parent 62cd8d1
commit bfad3e3
Show file tree

Hide file tree

Showing 2 changed files with 46 additions and 3 deletions.
diff --git a/app/api/services/informationextraction/InformationExtraction.ts b/app/api/services/informationextraction/InformationExtraction.ts
@@ -55,6 +55,7 @@ interface TaskParameters {
   id: string;
   multi_value?: boolean;
   options?: { label: string; id: string }[];
+  metadata?: { [key: string]: string };
 }
 
 type ResultParameters = TaskParameters;
@@ -354,7 +355,14 @@ class InformationExtraction {
     await this.taskManager.startTask({
       task: 'suggestions',
       tenant: tenants.current().name,
-      params: { id: extractorId.toString() },
+      params: {
+        id: extractorId.toString(),
+        metadata: {
+          extractor_name: extractor.name || '',
+          property: extractor.property || '',
+          templates: Array.isArray(extractor.templates) ? extractor.templates.join(',') : '',
+        },
+      },
     });
   };
 
@@ -395,6 +403,11 @@ class InformationExtraction {
     const params: TaskParameters = {
       id: extractorId.toString(),
       multi_value: property.type === 'multiselect' || property.type === 'relationship',
+      metadata: {
+        extractor_name: extractor.name || '',
+        property: extractor.property || '',
+        templates: extractor.templates ? extractor.templates.join(',') : '',
+      },
     };
 
     if (propertyTypeIsSelectOrMultiSelect(property.type)) {

diff --git a/app/api/services/informationextraction/specs/InformationExtraction.spec.ts b/app/api/services/informationextraction/specs/InformationExtraction.spec.ts
@@ -334,7 +334,15 @@ describe('InformationExtraction', () => {
       await informationExtraction.trainModel(factory.id('prop1extractor'));
 
       expect(informationExtraction.taskManager?.startTask).toHaveBeenCalledWith({
-        params: { id: factory.id('prop1extractor').toString(), multi_value: false },
+        params: {
+          id: factory.id('prop1extractor').toString(),
+          multi_value: false,
+          metadata: {
+            extractor_name: 'prop1extractor',
+            property: 'property1',
+            templates: expect.anything(),
+          },
+        },
         tenant: 'tenant1',
         task: 'create_model',
       });
@@ -369,6 +377,11 @@ describe('InformationExtraction', () => {
               label: '1B',
             },
           ],
+          metadata: {
+            extractor_name: 'extractorWithMultiselect',
+            property: 'property_multiselect',
+            templates: expect.anything(),
+          },
         },
         tenant: 'tenant1',
         task: 'create_model',
@@ -396,6 +409,11 @@ describe('InformationExtraction', () => {
               label: 'P3',
             },
           ],
+          metadata: {
+            extractor_name: 'extractorWithRelationship',
+            property: 'property_relationship',
+            templates: expect.anything(),
+          },
         },
         tenant: 'tenant1',
         task: 'create_model',
@@ -427,6 +445,11 @@ describe('InformationExtraction', () => {
               label: `A${i + 1}`,
             })),
           ],
+          metadata: {
+            extractor_name: 'extractorWithRelationshipToAny',
+            property: 'property_relationship_to_any',
+            templates: expect.anything(),
+          },
         },
         tenant: 'tenant1',
         task: 'create_model',
@@ -711,7 +734,14 @@ describe('InformationExtraction', () => {
       await informationExtraction.getSuggestions(factory.id('prop1extractor'));
 
       expect(informationExtraction.taskManager?.startTask).toHaveBeenCalledWith({
-        params: { id: factory.id('prop1extractor').toString() },
+        params: {
+          id: factory.id('prop1extractor').toString(),
+          metadata: {
+            extractor_name: 'prop1extractor',
+            property: 'property1',
+            templates: expect.anything(),
+          },
+        },
         tenant: 'tenant1',
         task: 'suggestions',
       });