Skip to content

Commit

Permalink
test: add tests for hybrid indexes
Browse files Browse the repository at this point in the history
  • Loading branch information
heyjorgedev committed Jan 22, 2025
1 parent 9742ae4 commit 0a78694
Show file tree
Hide file tree
Showing 8 changed files with 321 additions and 1 deletion.
2 changes: 2 additions & 0 deletions .github/workflows/run-tests.yml
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,8 @@ env:
DENSE_EMBEDDING_UPSTASH_VECTOR_REST_TOKEN: ${{ secrets.DENSE_EMBEDDING_UPSTASH_VECTOR_REST_TOKEN }}
SPARSE_UPSTASH_VECTOR_REST_URL: ${{ secrets.SPARSE_UPSTASH_VECTOR_REST_URL }}
SPARSE_UPSTASH_VECTOR_REST_TOKEN: ${{ secrets.SPARSE_UPSTASH_VECTOR_REST_TOKEN }}
HYBRID_UPSTASH_VECTOR_REST_URL: ${{ secrets.HYBRID_UPSTASH_VECTOR_REST_URL }}
HYBRID_UPSTASH_VECTOR_REST_TOKEN: ${{ secrets.HYBRID_UPSTASH_VECTOR_REST_TOKEN }}

jobs:
test:
Expand Down
3 changes: 3 additions & 0 deletions phpunit.xml.dist
Original file line number Diff line number Diff line change
Expand Up @@ -52,5 +52,8 @@
<!-- For testing sparse index -->
<env name="SPARSE_UPSTASH_VECTOR_REST_URL" value="https://vector.upstash.com" />
<env name="SPARSE_UPSTASH_VECTOR_REST_TOKEN" value="test-token" />
<!-- For testing hybrid index -->
<env name="HYBRID_UPSTASH_VECTOR_REST_URL" value="https://vector.upstash.com" />
<env name="HYBRID_UPSTASH_VECTOR_REST_TOKEN" value="test-token" />
</php>
</phpunit>
14 changes: 13 additions & 1 deletion tests/Concerns/GeneratesVectors.php
Original file line number Diff line number Diff line change
Expand Up @@ -2,4 +2,16 @@

namespace Upstash\Vector\Tests\Concerns;

trait GeneratesVectors {}
trait GeneratesVectors
{
protected function generateVector(int $dimensions): array
{
$vector = [];

for ($i = 0; $i < $dimensions; $i++) {
$vector[] = random_int(0, 100) / 100;
}

return $vector;
}
}
33 changes: 33 additions & 0 deletions tests/Concerns/UsesHybridIndex.php
Original file line number Diff line number Diff line change
@@ -0,0 +1,33 @@
<?php

namespace Upstash\Vector\Tests\Concerns;

use Upstash\Vector\Contracts\IndexInterface;
use Upstash\Vector\Contracts\IndexNamespaceInterface;
use Upstash\Vector\Index;

trait UsesHybridIndex
{
protected IndexInterface $index;

protected IndexNamespaceInterface $namespace;

public function setUp(): void
{
parent::setUp();

$this->index = new Index(
url: getenv('HYBRID_UPSTASH_VECTOR_REST_URL'),
token: getenv('HYBRID_UPSTASH_VECTOR_REST_TOKEN'),
);

$this->namespace = $this->index->namespace(bin2hex(random_bytes(32)));
}

public function tearDown(): void
{
$this->namespace->delete();

parent::tearDown();
}
}
43 changes: 43 additions & 0 deletions tests/Hybrid/Operations/QueryDataTest.php
Original file line number Diff line number Diff line change
@@ -0,0 +1,43 @@
<?php

namespace Upstash\Vector\Tests\Hybrid\Operations;

use PHPUnit\Framework\TestCase;
use Upstash\Vector\DataQuery;
use Upstash\Vector\DataUpsert;
use Upstash\Vector\Tests\Concerns\UsesHybridIndex;
use Upstash\Vector\Tests\Concerns\WaitsForIndex;

class QueryDataTest extends TestCase
{
use UsesHybridIndex;
use WaitsForIndex;

public function test_query_data(): void
{
$this->namespace->upsertDataMany([
new DataUpsert(
id: '1',
data: 'The capital of Japan is Tokyo',
),
new DataUpsert(
id: '2',
data: 'The capital of France is Paris',
),
new DataUpsert(
id: '3',
data: 'The capital of Germany is Berlin',
),
]);

$this->waitForIndex($this->namespace);

$results = $this->namespace->queryData(new DataQuery(
data: 'capital of France',
topK: 1,
));

$this->assertCount(1, $results);
$this->assertEquals('2', $results[0]->id);
}
}
95 changes: 95 additions & 0 deletions tests/Hybrid/Operations/QueryVectorsTest.php
Original file line number Diff line number Diff line change
@@ -0,0 +1,95 @@
<?php

namespace Upstash\Vector\Tests\Hybrid\Operations;

use PHPUnit\Framework\TestCase;
use Upstash\Vector\SparseVector;
use Upstash\Vector\Tests\Concerns\GeneratesVectors;
use Upstash\Vector\Tests\Concerns\UsesHybridIndex;
use Upstash\Vector\Tests\Concerns\WaitsForIndex;
use Upstash\Vector\VectorQuery;
use Upstash\Vector\VectorUpsert;

class QueryVectorsTest extends TestCase
{
use GeneratesVectors;
use UsesHybridIndex;
use WaitsForIndex;

public function test_query_vectors(): void
{
$vector = $this->generateVector(384);

$this->namespace->upsert(new VectorUpsert(
id: '1',
vector: $vector,
sparseVector: new SparseVector(
indices: [1, 2, 3],
values: [5, 6, 7],
),
));

$this->waitForIndex($this->namespace);

$vectors = $this->namespace->query(new VectorQuery(
vector: $vector,
sparseVector: new SparseVector(
indices: [1, 2, 3],
values: [5, 6, 7],
),
topK: 1,
));

$this->assertCount(1, $vectors);
$this->assertSame('1', $vectors[0]->id);
}

public function test_query_with_no_vectors(): void
{
$this->namespace->upsert(new VectorUpsert(
id: '1',
vector: $this->generateVector(384),
sparseVector: new SparseVector(
indices: [1, 2, 3],
values: [5, 6, 7],
),
));

$this->waitForIndex($this->namespace);

$vectors = $this->namespace->query(new VectorQuery(
sparseVector: new SparseVector(
indices: [1, 2, 3],
values: [5, 6, 7],
),
topK: 1,
));

$this->assertCount(1, $vectors);
$this->assertSame('1', $vectors[0]->id);
}

public function test_query_with_no_sparse_vectors(): void
{
$vector = $this->generateVector(384);

$this->namespace->upsert(new VectorUpsert(
id: '1',
vector: $vector,
sparseVector: new SparseVector(
indices: [1, 2, 3],
values: [5, 6, 7],
),
));

$this->waitForIndex($this->namespace);

$vectors = $this->namespace->query(new VectorQuery(
vector: $vector,
topK: 1,
));

$this->assertCount(1, $vectors);
$this->assertSame('1', $vectors[0]->id);
}
}
43 changes: 43 additions & 0 deletions tests/Hybrid/Operations/UpsertDataTest.php
Original file line number Diff line number Diff line change
@@ -0,0 +1,43 @@
<?php

namespace Upstash\Vector\Tests\Hybrid\Operations;

use PHPUnit\Framework\TestCase;
use Upstash\Vector\DataUpsert;
use Upstash\Vector\Tests\Concerns\UsesHybridIndex;
use Upstash\Vector\Tests\Concerns\WaitsForIndex;

class UpsertDataTest extends TestCase
{
use UsesHybridIndex;
use WaitsForIndex;

public function test_upsert_data(): void
{
$this->namespace->upsertData(new DataUpsert(
id: '1',
data: 'The capital of Japan is Tokyo',
));

$this->waitForIndex($this->namespace);

$info = $this->namespace->getNamespaceInfo();

$this->assertSame(1, $info->vectorCount);
}

public function test_upsert_many_data(): void
{
$this->namespace->upsertDataMany([
new DataUpsert(id: '1', data: 'The capital of Japan is Tokyo'),
new DataUpsert(id: '2', data: 'The capital of France is Paris'),
new DataUpsert(id: '3', data: 'The capital of Germany is Berlin'),
]);

$this->waitForIndex($this->namespace);

$info = $this->namespace->getNamespaceInfo();

$this->assertSame(3, $info->vectorCount);
}
}
89 changes: 89 additions & 0 deletions tests/Hybrid/Operations/UpsertVectorTest.php
Original file line number Diff line number Diff line change
@@ -0,0 +1,89 @@
<?php

namespace Upstash\Vector\Tests\Hybrid\Operations;

use PHPUnit\Framework\TestCase;
use Upstash\Vector\Exceptions\OperationFailedException;
use Upstash\Vector\SparseVector;
use Upstash\Vector\Tests\Concerns\GeneratesVectors;
use Upstash\Vector\Tests\Concerns\UsesHybridIndex;
use Upstash\Vector\Tests\Concerns\WaitsForIndex;
use Upstash\Vector\VectorUpsert;

class UpsertVectorTest extends TestCase
{
use GeneratesVectors;
use UsesHybridIndex;
use WaitsForIndex;

public function test_upsert_vector(): void
{
$this->namespace->upsert(new VectorUpsert(
id: '1',
vector: $this->generateVector(384),
sparseVector: new SparseVector(
indices: [1, 2, 3],
values: [5, 6, 7],
),
));

$this->waitForIndex($this->namespace);

$info = $this->namespace->getNamespaceInfo();

$this->assertSame(1, $info->vectorCount);
}

public function test_upsert_without_sparse_vector_throws_exception(): void
{
$this->expectException(OperationFailedException::class);
$this->expectExceptionMessage('This index requires sparse vectors');

$this->namespace->upsert(new VectorUpsert(
id: '1',
vector: $this->generateVector(384),
));
}

public function test_upsert_without_vector_throws_exception(): void
{
$this->expectException(OperationFailedException::class);
$this->expectExceptionMessage('This index requires dense vectors');

$this->namespace->upsert(new VectorUpsert(
id: '1',
sparseVector: new SparseVector(
indices: [1, 2, 3],
values: [5, 6, 7],
),
));
}

public function test_upsert_many_vectors(): void
{
$this->namespace->upsertMany([
new VectorUpsert(
id: '1',
vector: $this->generateVector(384),
sparseVector: new SparseVector(
indices: [1, 2, 3],
values: [5, 6, 7],
),
),
new VectorUpsert(
id: '2',
vector: $this->generateVector(384),
sparseVector: new SparseVector(
indices: [4, 5, 6],
values: [8, 9, 10],
),
),
]);

$this->waitForIndex($this->namespace);

$info = $this->namespace->getNamespaceInfo();

$this->assertSame(2, $info->vectorCount);
}
}

0 comments on commit 0a78694

Please sign in to comment.