diff --git a/src/store/src/Bridge/Pinecone/Store.php b/src/store/src/Bridge/Pinecone/Store.php index deabd8bbb..e5003ca82 100644 --- a/src/store/src/Bridge/Pinecone/Store.php +++ b/src/store/src/Bridge/Pinecone/Store.php @@ -16,13 +16,15 @@ use Symfony\AI\Platform\Vector\Vector; use Symfony\AI\Store\Document\Metadata; use Symfony\AI\Store\Document\VectorDocument; +use Symfony\AI\Store\Exception\InvalidArgumentException; +use Symfony\AI\Store\ManagedStoreInterface; use Symfony\AI\Store\StoreInterface; use Symfony\Component\Uid\Uuid; /** * @author Christopher Hertel */ -final class Store implements StoreInterface +final class Store implements StoreInterface, ManagedStoreInterface { /** * @param array $filter @@ -35,6 +37,23 @@ public function __construct( ) { } + public function setup(array $options = []): void + { + if (false === isset($options['indexName']) || false === isset($options['dimension'])) { + throw new InvalidArgumentException('No supported options.'); + } + + $this->pinecone + ->control() + ->index($options['indexName']) + ->createServerless( + $options['dimension'], + $options['metric'] ?? null, + $options['cloud'] ?? null, + $options['region'] ?? null, + ); + } + public function add(VectorDocument ...$documents): void { $vectors = []; @@ -73,6 +92,18 @@ public function query(Vector $vector, array $options = []): iterable } } + public function drop(array $options = []): void + { + if (false === isset($options['indexName'])) { + throw new InvalidArgumentException('No supported options.'); + } + + $this->pinecone + ->control() + ->index($options['indexName']) + ->delete(); + } + private function getVectors(): VectorResource { return $this->pinecone->data()->vectors(); diff --git a/src/store/src/Bridge/Pinecone/Tests/StoreTest.php b/src/store/src/Bridge/Pinecone/Tests/StoreTest.php index a3fb23bde..6674a6902 100644 --- a/src/store/src/Bridge/Pinecone/Tests/StoreTest.php +++ b/src/store/src/Bridge/Pinecone/Tests/StoreTest.php @@ -13,6 +13,8 @@ use PHPUnit\Framework\TestCase; use Probots\Pinecone\Client; +use Probots\Pinecone\Resources\Control\IndexResource; +use Probots\Pinecone\Resources\ControlResource; use Probots\Pinecone\Resources\Data\VectorResource; use Probots\Pinecone\Resources\DataResource; use Saloon\Http\Response; @@ -24,6 +26,76 @@ final class StoreTest extends TestCase { + public function testStoreCantSetupWithInvalidOptions() + { + $pinecone = $this->createMock(Client::class); + $store = new Store($pinecone); + + $this->expectException(\InvalidArgumentException::class); + + $store->setup(); + } + + public function testStoreCanSetup() + { + $pinecone = $this->createMock(Client::class); + $indexResource = $this->createMock(IndexResource::class); + $controlResource = $this->createMock(ControlResource::class); + + $pinecone->expects($this->once()) + ->method('control') + ->willReturn($controlResource); + + $controlResource->expects($this->once()) + ->method('index') + ->with('test-index') + ->willReturn($indexResource); + + $indexResource->expects($this->once()) + ->method('createServerless') + ->with(1536, null, null, null); + + $store = new Store($pinecone); + $store->setup([ + 'indexName' => 'test-index', + 'dimension' => 1536, + ]); + } + + public function testStoreCantDropWithInvalidOptions() + { + $pinecone = $this->createMock(Client::class); + $store = new Store($pinecone); + + $this->expectException(\InvalidArgumentException::class); + + $store->drop(); + } + + public function testStoreCanDrop() + { + $pinecone = $this->createMock(Client::class); + $indexResource = $this->createMock(IndexResource::class); + $controlResource = $this->createMock(ControlResource::class); + + $pinecone->expects($this->once()) + ->method('control') + ->willReturn($controlResource); + + $controlResource->expects($this->once()) + ->method('index') + ->with('test-index') + ->willReturn($indexResource); + + $indexResource->expects($this->once()) + ->method('delete'); + + $store = new Store($pinecone); + $store->drop([ + 'indexName' => 'test-index', + ]); + } + public function testAddSingleDocument() { $vectorResource = $this->createMock(VectorResource::class);