From 200be59c514c6a12f6340b7599f148e3b41bd133 Mon Sep 17 00:00:00 2001
From: Rhys <Anemy@users.noreply.github.com>
Date: Fri, 10 Nov 2023 09:57:31 -0500
Subject: [PATCH] chore(generative-ai): add additional accuracy tests
 COMPASS-7299 (#5091)

---
 .../scripts/ai-accuracy-tests.js              | 80 ++++++++++++++++++-
 1 file changed, 79 insertions(+), 1 deletion(-)

diff --git a/packages/compass-generative-ai/scripts/ai-accuracy-tests.js b/packages/compass-generative-ai/scripts/ai-accuracy-tests.js
index 08c059749c9..8435aa32290 100644
--- a/packages/compass-generative-ai/scripts/ai-accuracy-tests.js
+++ b/packages/compass-generative-ai/scripts/ai-accuracy-tests.js
@@ -44,7 +44,7 @@ const TESTS_TO_RUN_CONCURRENTLY = 3;
 // To avoid rate limit we also reduce the time between tests running
 // when the test returns a result quickly.
 const ADD_TIMEOUT_BETWEEN_TESTS_THRESHOLD_MS = 5000;
-const TIMEOUT_BETWEEN_TESTS_MS = 2000;
+const TIMEOUT_BETWEEN_TESTS_MS = 3000;
 
 let PQueue;
 
@@ -500,6 +500,84 @@ const tests = [
       },
     ]),
   },
+  {
+    type: 'query',
+    databaseName: 'sample_airbnb',
+    collectionName: 'listingsAndReviews',
+    userInput:
+      'give me just the price and the first 3 amenities of the listing has "Step-free access" in its amenities.',
+    assertResult: anyOf([
+      isDeepStrictEqualTo([
+        {
+          _id: '10108388',
+          price: {
+            $numberDecimal: '185.00',
+          },
+          amenities: ['TV', 'Wifi', 'Air conditioning'],
+        },
+      ]),
+      isDeepStrictEqualTo([
+        {
+          price: {
+            $numberDecimal: '185.00',
+          },
+          amenities: ['TV', 'Wifi', 'Air conditioning'],
+        },
+      ]),
+    ]),
+  },
+  {
+    type: 'aggregation',
+    databaseName: 'sample_airbnb',
+    collectionName: 'listingsAndReviews',
+    userInput:
+      '¿Qué alojamiento tiene el precio más bajo? devolver el número en un campo llamado "precio"',
+    assertResult: anyOf([
+      isDeepStrictEqualTo([
+        {
+          precio: {
+            $numberDecimal: '40.00',
+          },
+        },
+      ]),
+      isDeepStrictEqualTo([
+        {
+          _id: '10117617',
+          precio: {
+            $numberDecimal: '40.00',
+          },
+        },
+      ]),
+      isDeepStrictEqualTo([
+        {
+          precio: 40,
+        },
+      ]),
+    ]),
+  },
+  {
+    type: 'aggregation',
+    databaseName: 'sample_airbnb',
+    collectionName: 'listingsAndReviews',
+    userInput:
+      'give only me the cancellation policy and host url of the most expensive listing',
+    assertResult: anyOf([
+      isDeepStrictEqualTo([
+        {
+          cancellation_policy: 'moderate',
+          host: {
+            host_url: 'https://www.airbnb.com/users/show/51471538',
+          },
+        },
+      ]),
+      isDeepStrictEqualTo([
+        {
+          cancellation_policy: 'moderate',
+          host_url: 'https://www.airbnb.com/users/show/51471538',
+        },
+      ]),
+    ]),
+  },
   {
     type: 'aggregation',
     databaseName: 'netflix',