From 204020fd4328a5c6a988d3dacf2a665f53337acf Mon Sep 17 00:00:00 2001
From: Jorge Lapa <2780099+heyjorgedev@users.noreply.github.com>
Date: Tue, 4 Mar 2025 18:39:48 +0800
Subject: [PATCH] feat: allow fetching vectors using a prefix

---
 src/Contracts/IndexNamespaceInterface.php     |  3 +-
 src/Index.php                                 |  2 +-
 src/IndexNamespace.php                        |  2 +-
 src/Operations/FetchVectorsOperation.php      |  3 +-
 src/VectorFetchByPrefix.php                   | 33 +++++++++++++++++++
 .../Operations/FetchVectorsOperationTest.php  | 22 +++++++++++++
 6 files changed, 61 insertions(+), 4 deletions(-)
 create mode 100644 src/VectorFetchByPrefix.php

diff --git a/src/Contracts/IndexNamespaceInterface.php b/src/Contracts/IndexNamespaceInterface.php
index 78c508c..3df2ba6 100644
--- a/src/Contracts/IndexNamespaceInterface.php
+++ b/src/Contracts/IndexNamespaceInterface.php
@@ -9,6 +9,7 @@
 use Upstash\Vector\NamespaceInfo;
 use Upstash\Vector\VectorDeleteResult;
 use Upstash\Vector\VectorFetch;
+use Upstash\Vector\VectorFetchByPrefix;
 use Upstash\Vector\VectorFetchResult;
 use Upstash\Vector\VectorMatch;
 use Upstash\Vector\VectorQuery;
@@ -55,7 +56,7 @@ public function queryData(DataQuery $query): DataQueryResult;
      */
     public function delete(array $ids): VectorDeleteResult;
 
-    public function fetch(VectorFetch $vectorFetch): VectorFetchResult;
+    public function fetch(VectorFetch|VectorFetchByPrefix $vectorFetch): VectorFetchResult;
 
     public function random(): ?VectorMatch;
 
diff --git a/src/Index.php b/src/Index.php
index 11b9ec3..d8fdc98 100755
--- a/src/Index.php
+++ b/src/Index.php
@@ -118,7 +118,7 @@ public function delete(array $ids): VectorDeleteResult
         return $this->namespace('')->delete($ids);
     }
 
-    public function fetch(VectorFetch $vectorFetch): VectorFetchResult
+    public function fetch(VectorFetch|VectorFetchByPrefix $vectorFetch): VectorFetchResult
     {
         return $this->namespace('')->fetch($vectorFetch);
     }
diff --git a/src/IndexNamespace.php b/src/IndexNamespace.php
index ea379c5..787382e 100644
--- a/src/IndexNamespace.php
+++ b/src/IndexNamespace.php
@@ -79,7 +79,7 @@ public function delete(array $ids): VectorDeleteResult
             ->delete($ids);
     }
 
-    public function fetch(VectorFetch $vectorFetch): VectorFetchResult
+    public function fetch(VectorFetch|VectorFetchByPrefix $vectorFetch): VectorFetchResult
     {
         return (new FetchVectorsOperation($this->namespace, $this->transporter))
             ->fetch($vectorFetch);
diff --git a/src/Operations/FetchVectorsOperation.php b/src/Operations/FetchVectorsOperation.php
index 7cb513b..fbc757b 100644
--- a/src/Operations/FetchVectorsOperation.php
+++ b/src/Operations/FetchVectorsOperation.php
@@ -10,6 +10,7 @@
 use Upstash\Vector\Transporter\TransporterRequest;
 use Upstash\Vector\Transporter\TransporterResponse;
 use Upstash\Vector\VectorFetch;
+use Upstash\Vector\VectorFetchByPrefix;
 use Upstash\Vector\VectorFetchResult;
 use Upstash\Vector\VectorMatch;
 
@@ -23,7 +24,7 @@
 
     public function __construct(private string $namespace, private TransporterInterface $transporter) {}
 
-    public function fetch(VectorFetch $vectorFetch): VectorFetchResult
+    public function fetch(VectorFetch|VectorFetchByPrefix $vectorFetch): VectorFetchResult
     {
         $path = $this->getPath();
         $request = new TransporterRequest(
diff --git a/src/VectorFetchByPrefix.php b/src/VectorFetchByPrefix.php
new file mode 100644
index 0000000..ad662a5
--- /dev/null
+++ b/src/VectorFetchByPrefix.php
@@ -0,0 +1,33 @@
+<?php
+
+namespace Upstash\Vector;
+
+use Upstash\Vector\Contracts\Arrayable;
+
+final readonly class VectorFetchByPrefix implements Arrayable
+{
+    public function __construct(
+        public string $prefix,
+        public bool $includeMetadata = false,
+        public bool $includeVectors = false,
+        public bool $includeData = false,
+    ) {}
+
+    /**
+     * @return array{
+     *     prefix: string,
+     *     includeMetadata: bool,
+     *     includeVectors: bool,
+     *     includeData: bool,
+     * }
+     */
+    public function toArray(): array
+    {
+        return [
+            'prefix' => $this->prefix,
+            'includeMetadata' => $this->includeMetadata,
+            'includeVectors' => $this->includeVectors,
+            'includeData' => $this->includeData,
+        ];
+    }
+}
diff --git a/tests/Dense/Operations/FetchVectorsOperationTest.php b/tests/Dense/Operations/FetchVectorsOperationTest.php
index 5a963f7..2ffdac9 100644
--- a/tests/Dense/Operations/FetchVectorsOperationTest.php
+++ b/tests/Dense/Operations/FetchVectorsOperationTest.php
@@ -6,6 +6,7 @@
 use Upstash\Vector\Tests\Concerns\UsesDenseIndex;
 use Upstash\Vector\Tests\Concerns\WaitsForIndex;
 use Upstash\Vector\VectorFetch;
+use Upstash\Vector\VectorFetchByPrefix;
 use Upstash\Vector\VectorUpsert;
 
 use function Upstash\Vector\createRandomVector;
@@ -16,6 +17,27 @@ class FetchVectorsOperationTest extends TestCase
     use WaitsForIndex;
 
     public function test_can_fetch_vectors(): void
+    {
+        // Arrange
+        $this->namespace->upsertMany([
+            new VectorUpsert(id: 'user:1', vector: createRandomVector(2)),
+            new VectorUpsert(id: 'user:2', vector: createRandomVector(2)),
+            new VectorUpsert(id: 'test:3', vector: createRandomVector(2)),
+        ]);
+        $this->waitForIndex($this->namespace);
+
+        // Act
+        $results = $this->namespace->fetch(new VectorFetchByPrefix(
+            prefix: 'user:',
+            includeVectors: true,
+        ));
+
+        // Assert
+        $this->assertCount(2, $results);
+        $this->assertCount(2, $results[0]->vector);
+    }
+
+    public function test_can_fetch_vectors_using_a_filter(): void
     {
         // Arrange
         $this->namespace->upsertMany([