PCBC-994 Support for base64 encoded vector types (#169)

* Add base64 vector search * fix test
couchbase · Jun 4, 2024 · 998b6ac · 998b6ac
1 parent 450d1b3
commit 998b6ac
Show file tree

Hide file tree

Showing 3 changed files with 44 additions and 12 deletions.
diff --git a/Couchbase/VectorQuery.php b/Couchbase/VectorQuery.php
@@ -25,35 +25,44 @@
 class VectorQuery
 {
     private string $vectorFieldName;
-    private array $vectorQuery;
     private int $numCandidates;
+    private ?array $vectorQuery = null;
+    private ?string $base64VectorQuery = null;
     private ?float $boost = null;
 
     /**
      * @param string $vectorFieldName the document field that contains the vector
-     * @param array $vectorQuery the vector query to run. Cannot be empty.
+     * @param array<float>|string $vectorQuery the vector query to run. Cannot be empty. Either a vector array,
+     * or the vector query encoded into a base64 string.
      *
      * @since 4.1.7
      *
      * @throws InvalidArgumentException
      *
      * @UNCOMMITTED: This API may change in the future.
      */
-    public function __construct(string $vectorFieldName, array $vectorQuery)
+    public function __construct(string $vectorFieldName, array|string $vectorQuery)
     {
         if (empty($vectorQuery)) {
             throw new InvalidArgumentException("The vectorQuery cannot be empty");
         }
+
+        if (is_array($vectorQuery)) {
+            $this->vectorQuery = $vectorQuery;
+        } else {
+            $this->base64VectorQuery = $vectorQuery;
+        }
+
         $this->vectorFieldName = $vectorFieldName;
-        $this->vectorQuery = $vectorQuery;
         $this->numCandidates = 3;
     }
 
     /**
      * Static helper to keep code more readable
      *
      * @param string $vectorFieldName the document field that contains the vector
-     * @param array $vectorQuery the vector query to run. Cannot be empty.
+     * @param array<float>|string $vectorQuery the vector query to run. Cannot be empty. Either a vector array,
+     * or the vector query encoded into a base64 string.
      *
      * @since 4.1.7
      * @return VectorQuery
@@ -62,15 +71,15 @@ public function __construct(string $vectorFieldName, array $vectorQuery)
      *
      * @UNCOMMITTED: This API may change in the future.
      */
-    static function build(string $vectorFieldName, array $vectorQuery): VectorQuery
+    static function build(string $vectorFieldName, array|string $vectorQuery): VectorQuery
     {
         return new VectorQuery($vectorFieldName, $vectorQuery);
     }
 
     /**
      * Sets the number of results that will be returned from this vector query. Defaults to 3.
      *
-     * @param int|null $numCandidates the number of results returned.
+     * @param int $numCandidates the number of results returned.
      *
      * @since 4.1.7
      * @return VectorQuery
@@ -122,11 +131,16 @@ public static function export(VectorQuery $query): array
             $json['boost'] = $query->boost;
         }
 
-        $vectorQueries = [];
-        foreach ($query->vectorQuery as $value) {
-            $vectorQueries[] = $value;
+        if ($query->vectorQuery != null) {
+            $vectorQueries = [];
+            foreach ($query->vectorQuery as $value) {
+                $vectorQueries[] = $value;
+            }
+            $json['vector'] = $vectorQueries;
+        } else {
+            $json['vector_base64'] = $query->base64VectorQuery;
         }
-        $json['vector'] = $vectorQueries;
+
         $json['k'] = $query->numCandidates;
         return $json;
     }

diff --git a/tests/Helpers/CouchbaseTestCase.php b/tests/Helpers/CouchbaseTestCase.php
@@ -306,7 +306,8 @@ protected function assertErrorCode($code, $ex)
         );
     }
 
-    protected function fixCavesTimeResolutionOnWindows() {
+    protected function fixCavesTimeResolutionOnWindows()
+    {
         if (PHP_OS_FAMILY === 'Windows' && self::env()->useCaves()) {
             usleep(1);
         }

diff --git a/tests/SearchTest.php b/tests/SearchTest.php
@@ -597,6 +597,23 @@ public function testVectorSearchEncoding()
         $this->assertEquals('{"match_none":"null"}', $encodedSearchQuery);
     }
 
+    public function testVectorSearchEmptyStringThrowsInvalidArgument()
+    {
+        $this->expectException(\Couchbase\Exception\InvalidArgumentException::class);
+        VectorQuery::build("vectorField", "");
+    }
+
+    public function testVectorSearchEncodingWithBase64()
+    {
+        $base64EncodedVector = "aOeYBEXJ4kI=";
+        $vectorQueryOne = VectorQuery::build("foo", $base64EncodedVector)->boost(0.5)->numCandidates(4);
+        $vectorQueryTwo = VectorQuery::build("bar", [-0.00810353, 0.6433, 0.52364]);
+        $searchRequest = SearchRequest::export(SearchRequest::build(VectorSearch::build([$vectorQueryOne, $vectorQueryTwo])));
+        $encodedVectorQuery = json_encode($searchRequest['vectorSearch']);
+        $this->assertEquals(JSON_ERROR_NONE, json_last_error());
+        $this->assertEquals(sprintf('[{"field":"foo","boost":0.5,"vector_base64":"%s","k":4},{"field":"bar","vector":[-0.00810353,0.6433,0.52364],"k":3}]', $base64EncodedVector), $encodedVectorQuery);
+    }
+
     public function testScopeSearch()
     {
         $this->skipIfCaves();