Merge pull request #60 from weaviate/properties-field-for-groupedresult

Support properties field in generative groupedResult field

Author: antas-marcin

src/graphql/generate.ts

@@ -1,5 +1,6 @@
 export interface GenerateArgs {
   groupedTask?: string;
+  groupedProperties?: string[];
   singlePrompt?: string;
 }
 
@@ -11,10 +12,12 @@ export interface GenerateParts {
 
 export class GraphQLGenerate {
   private groupedTask?: string;
+  private groupedProperties?: string[];
   private singlePrompt?: string;
 
   constructor(args: GenerateArgs) {
     this.groupedTask = args.groupedTask;
+    this.groupedProperties = args.groupedProperties;
     this.singlePrompt = args.singlePrompt;
   }
 
@@ -27,8 +30,15 @@ export class GraphQLGenerate {
       str += `singleResult:{prompt:"${this.singlePrompt.replace(/[\n\r]+/g, '')}"}`;
       results.push('singleResult');
     }
-    if (this.groupedTask) {
-      str += `groupedResult:{task:"${this.groupedTask.replace(/[\n\r]+/g, '')}"}`;
+    if (this.groupedTask || (this.groupedProperties !== undefined && this.groupedProperties.length > 0)) {
+      const args: string[] = [];
+      if (this.groupedTask) {
+        args.push(`task:"${this.groupedTask.replace(/[\n\r]+/g, '')}"`);
+      }
+      if (this.groupedProperties !== undefined && this.groupedProperties.length > 0) {
+        args.push(`properties:${JSON.stringify(this.groupedProperties)}`);
+      }
+      str += `groupedResult:{${args.join(',')}}`;
       results.push('groupedResult');
     }
     str += `){${results.join(' ')}}`;

src/graphql/getter.test.ts

@@ -1339,6 +1339,43 @@ describe('generative search', () => {
 
     expect(mockClient.query).toHaveBeenCalledWith(expectedQuery);
   });
+
+  test('groupedTask with properties', () => {
+    const expectedQuery =
+      '{Get{Mammal{name taxonomy _additional{generate(groupedResult:' +
+      '{task:"Explain why platypi can lay eggs",properties:["title","description"]}){error groupedResult}}}}}';
+
+    new Getter(mockClient)
+      .withClassName('Mammal')
+      .withGenerate({
+        groupedTask: 'Explain why platypi can lay eggs',
+        groupedProperties: ['title', 'description'],
+      })
+      .withFields('name taxonomy')
+      .do();
+
+    expect(mockClient.query).toHaveBeenCalledWith(expectedQuery);
+  });
+
+  test('single prompt and grouped task with properties', () => {
+    const expectedQuery =
+      '{Get{Mammal{name taxonomy _additional{generate(singleResult:' +
+      '{prompt:"How tall is a baby giraffe?"}groupedResult:{task:' +
+      '"Explain how the heights of mammals relate to their prefferred food sources",properties:["property"]})' +
+      '{error singleResult groupedResult}}}}}';
+
+    new Getter(mockClient)
+      .withClassName('Mammal')
+      .withFields('name taxonomy')
+      .withGenerate({
+        singlePrompt: 'How tall is a baby giraffe?',
+        groupedTask: 'Explain how the heights of mammals relate to their prefferred food sources',
+        groupedProperties: ['property'],
+      })
+      .do();
+
+    expect(mockClient.query).toHaveBeenCalledWith(expectedQuery);
+  });
 });
 
 describe('groupBy valid searchers', () => {

src/graphql/journey.test.ts

@@ -1343,6 +1343,22 @@ Tastes like a fresh ocean breeze: {review}`,
       });
   });
 
+  test('groupedTask with groupedProperties', async () => {
+    await client.graphql
+      .get()
+      .withClassName('Wine')
+      .withFields('name review')
+      .withGenerate({
+        groupedTask: 'Describe the following as a LinkedIn Ad:',
+        groupedProperties: ['name', 'review'],
+      })
+      .do()
+      .then((res: any) => {
+        expect(res.data.Get.Wine[0]._additional.generate.groupedResult).toBeDefined();
+        expect(res.data.Get.Wine[0]._additional.generate.error).toBeNull();
+      });
+  });
+
   test('singlePrompt and groupedTask', async () => {
     await client.graphql
       .get()