Add withPayload and withVector properties on RecommendRequest c…

…lass.
hkulekci · Sep 25, 2024 · 70ed69d · 70ed69d
1 parent 81e01e9
commit 70ed69d
Show file tree

Hide file tree

Showing 2 changed files with 68 additions and 1 deletion.
diff --git a/src/Models/Request/Points/RecommendRequest.php b/src/Models/Request/Points/RecommendRequest.php
@@ -35,6 +35,8 @@ class RecommendRequest
     protected ?int $limit = null;
     protected ?int $offset = null;
     protected ?float $scoreThreshold = null;
+    protected bool|array|null $withVector = null;
+    protected bool|array|null $withPayload = null;
 
     public function __construct(protected array $positive, protected array $negative = [])
     {
@@ -96,6 +98,20 @@ public function setOffset(int $offset): static
         return $this;
     }
 
+    public function setWithPayload($withPayload): static
+    {
+        $this->withPayload = $withPayload;
+
+        return $this;
+    }
+
+    public function setWithVector($withVector): static
+    {
+        $this->withVector = $withVector;
+
+        return $this;
+    }
+
     public function toArray(): array
     {
         $body = [
@@ -124,7 +140,13 @@ public function toArray(): array
         if ($this->offset !== null) {
             $body['offset'] = $this->offset;
         }
+        if ($this->withVector !== null) {
+            $body['with_vector'] = $this->withVector;
+        }
+        if ($this->withPayload !== null) {
+            $body['with_payload'] = $this->withPayload;
+        }
 
         return $body;
     }
-}
+}
diff --git a/tests/Integration/Endpoints/Collections/Points/RecommendTest.php b/tests/Integration/Endpoints/Collections/Points/RecommendTest.php
@@ -15,6 +15,7 @@
 use Qdrant\Models\Request\Points\RecommendRequest;
 use Qdrant\Models\VectorStruct;
 use Qdrant\Tests\Integration\AbstractIntegration;
+use function var_dump;
 
 class RecommendTest extends AbstractIntegration
 {
@@ -94,6 +95,50 @@ public function testRecommendPoint(array $positive, array $negative): void
         $this->assertEquals('ok', $response['status']);
     }
 
+    /**
+     * @dataProvider recommendQueryProvider
+     */
+    public function testRecommendPointWithPayload(array $positive, array $negative): void
+    {
+        $recommendRequest = (new RecommendRequest($positive, $negative))
+            ->setLimit(3)
+            ->setUsing('image')
+            ->setFilter(
+                (new Filter())->addMust(
+                    new MatchString('image', 'sample image')
+                )
+            )
+            ->setWithPayload(true);
+
+        $response = $this->getCollections('sample-collection')->points()->recommend()->recommend($recommendRequest);
+
+        $this->assertEquals('ok', $response['status']);
+        $this->assertCount(2, $response['result']);
+        $this->assertArrayHasKey('payload', $response['result'][0]);
+    }
+
+    /**
+     * @dataProvider recommendQueryProvider
+     */
+    public function testRecommendPointWithVector(array $positive, array $negative): void
+    {
+        $recommendRequest = (new RecommendRequest($positive, $negative))
+            ->setLimit(3)
+            ->setUsing('image')
+            ->setFilter(
+                (new Filter())->addMust(
+                    new MatchString('image', 'sample image')
+                )
+            )
+            ->setWithVector(true);
+
+        $response = $this->getCollections('sample-collection')->points()->recommend()->recommend($recommendRequest);
+
+        $this->assertEquals('ok', $response['status']);
+        $this->assertCount(2, $response['result']);
+        $this->assertArrayHasKey('vector', $response['result'][0]);
+    }
+
     public function testRecommendWithThreshold(): void
     {
         // Upsert points