Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
37 changes: 37 additions & 0 deletions client/python/USAGE.md
Original file line number Diff line number Diff line change
Expand Up @@ -39,6 +39,24 @@ The client supports usage as a context manager, which automatically closes the u
Example available in:
```examples/context_manager_usage.py```

### Async Client Support

For async applications, use `AsyncVortexDB`. It mirrors the synchronous client API and uses `grpc.aio` under the hood.

Example available in:
```examples/async_usage.py```

```python
async with AsyncVortexDB(
grpc_url="localhost:50051",
api_key="your-api-key",
) as db:
point_id = await db.insert(
vector=DenseVector([0.1, 0.2, 0.3]),
payload=Payload.text("hello async vortex"),
)
```

---

## Client API
Expand All @@ -47,6 +65,25 @@ Example available in:

Main client class for interacting with the VortexDB gRPC server.

### `AsyncVortexDB`

Async client class for I/O-heavy applications. It has the same constructor and method names as `VortexDB`, but methods are awaitable:

```
await db.insert(...)
await db.get(...)
await db.search(...)
await db.delete(...)
await db.close()
```

It also supports async context manager usage:

```
async with AsyncVortexDB(...) as db:
...
```

#### **Constructor**

```
Expand Down
31 changes: 31 additions & 0 deletions client/python/examples/async_usage.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,31 @@
import asyncio

from vortexdb import AsyncVortexDB, DenseVector, Payload, Similarity


async def main():
async with AsyncVortexDB(
grpc_url="localhost:50051",
api_key="your-api-key",
) as db:
point_id = await db.insert(
vector=DenseVector([0.1, 0.2, 0.3]),
payload=Payload.text("hello async vortex"),
)

point = await db.get(point_id=point_id)
if point is not None:
print(point.pretty())

results = await db.search(
vector=DenseVector([0.1, 0.2, 0.3]),
similarity=Similarity.COSINE,
limit=5,
)
print(results)

await db.delete(point_id=point_id)


if __name__ == "__main__":
asyncio.run(main())
158 changes: 158 additions & 0 deletions client/python/tests/test_async_client.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,158 @@
import asyncio
from unittest.mock import AsyncMock, Mock

import pytest

from vortexdb.async_client import AsyncVortexDB
from vortexdb.async_connection import AsyncGRPCConnection
from vortexdb.models import ContentType, DenseVector, Payload, Point, Similarity


@pytest.fixture
def mock_connection(monkeypatch):
"""
Replace AsyncGRPCConnection with a mock instance.
"""
conn = Mock(spec=AsyncGRPCConnection)
conn.stub = Mock()
conn.call = AsyncMock()
conn.close = AsyncMock()
monkeypatch.setattr("vortexdb.async_client.AsyncGRPCConnection", lambda _: conn)
return conn


@pytest.fixture
def client(mock_connection):
return AsyncVortexDB(
grpc_url="localhost:50051",
api_key="secret",
)


def test_async_insert_success(client, mock_connection):
async def run():
response = Mock()
response.id = Mock()
response.id.value = "point-123"

mock_connection.call.return_value = response

point_id = await client.insert(
vector=DenseVector([1, 2, 3]),
payload=Payload.text("hello"),
)

assert point_id == "point-123"

asyncio.run(run())


def test_async_insert_rejects_invalid_vector(client):
async def run():
with pytest.raises(TypeError):
await client.insert(
vector=[1, 2, 3],
payload=Payload.text("hello"),
)

asyncio.run(run())


def test_async_get_point_success(client, mock_connection):
async def run():
proto_point = Mock()
proto_point.id.id.value = "point-123"
proto_point.vector.values = [1, 2, 3]
proto_point.payload.content_type = ContentType.TEXT.to_proto()
proto_point.payload.content = "hello"

mock_connection.call.return_value = proto_point

point = await client.get(point_id="point-123")

assert isinstance(point, Point)
assert point.id == "point-123"
assert point.payload.content == "hello"

asyncio.run(run())


def test_async_get_point_not_found(client, mock_connection):
async def run():
mock_connection.call.return_value = None

result = await client.get(point_id="missing")

assert result is None

asyncio.run(run())


def test_async_delete_success(client, mock_connection):
async def run():
mock_connection.call.return_value = None

await client.delete(point_id="point-123")

mock_connection.call.assert_awaited_once()

asyncio.run(run())


def test_async_search_success(client, mock_connection):
async def run():
mock_connection.call.return_value = Mock(
result_point_ids=[
Mock(id=Mock(value="p1")),
Mock(id=Mock(value="p2")),
]
)

results = await client.search(
vector=DenseVector([1, 2, 3]),
similarity=Similarity.COSINE,
limit=2,
)

assert results == ["p1", "p2"]

asyncio.run(run())


def test_async_search_invalid_vector(client):
async def run():
with pytest.raises(TypeError):
await client.search(
vector=[1, 2, 3],
similarity=Similarity.COSINE,
limit=2,
)

asyncio.run(run())


def test_async_close_closes_connection(client, mock_connection):
async def run():
await client.close()
mock_connection.close.assert_awaited_once()

asyncio.run(run())


def test_async_context_manager_closes_connection(monkeypatch):
async def run():
conn = Mock(spec=AsyncGRPCConnection)
conn.stub = Mock()
conn.call = AsyncMock()
conn.close = AsyncMock()
monkeypatch.setattr("vortexdb.async_client.AsyncGRPCConnection", lambda _: conn)

async with AsyncVortexDB(
grpc_url="localhost:50051",
api_key="secret",
) as db:
assert db is not None

conn.close.assert_awaited_once()

asyncio.run(run())
106 changes: 106 additions & 0 deletions client/python/tests/test_async_connection.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,106 @@
import asyncio
from unittest.mock import AsyncMock, Mock, patch

import grpc

from vortexdb._grpc_common import map_grpc_error
from vortexdb.async_connection import AsyncGRPCConnection
from vortexdb.config import VortexDBConfig
from vortexdb.exceptions import (
AuthenticationError,
InternalServerError,
InvalidArgumentError,
NotFoundError,
ServiceUnavailableError,
TimeoutError,
)


class FakeAioRpcError:
"""
Minimal AioRpcError-compatible object for mapping tests.
"""

def __init__(self, status_code: grpc.StatusCode, details: str):
self._status_code = status_code
self._details = details

def code(self):
return self._status_code

def details(self):
return self._details


def make_config() -> VortexDBConfig:
return VortexDBConfig(
grpc_url="localhost:50051",
api_key="secret",
timeout=3.0,
)


def test_async_channel_created_with_correct_url():
with patch("grpc.aio.insecure_channel") as mock_channel:
mock_channel.return_value = Mock()
AsyncGRPCConnection(make_config())
mock_channel.assert_called_once_with("localhost:50051")


def test_async_metadata_is_attached():
with patch("grpc.aio.insecure_channel") as mock_channel:
mock_channel.return_value = Mock()
connection = AsyncGRPCConnection(make_config())

assert ("authorization", "Bearer secret") in connection._metadata


def test_successful_async_rpc_call():
async def run():
with patch("grpc.aio.insecure_channel") as mock_channel:
mock_channel.return_value = Mock()
connection = AsyncGRPCConnection(make_config())

fake_rpc = AsyncMock(return_value="ok")

result = await connection.call(fake_rpc, request="req")

fake_rpc.assert_awaited_once_with(
"req",
timeout=3.0,
metadata=connection._metadata,
)
assert result == "ok"

asyncio.run(run())


def test_async_grpc_error_mapping():
cases = [
(grpc.StatusCode.UNAUTHENTICATED, AuthenticationError),
(grpc.StatusCode.NOT_FOUND, NotFoundError),
(grpc.StatusCode.INVALID_ARGUMENT, InvalidArgumentError),
(grpc.StatusCode.DEADLINE_EXCEEDED, TimeoutError),
(grpc.StatusCode.UNAVAILABLE, ServiceUnavailableError),
(grpc.StatusCode.UNKNOWN, InternalServerError),
]

for status_code, expected_exception in cases:
error = FakeAioRpcError(status_code, "boom")
mapped = map_grpc_error(error)
assert isinstance(mapped, expected_exception)


def test_async_close_closes_channel():
async def run():
with patch("grpc.aio.insecure_channel") as mock_channel:
channel = Mock()
channel.close = AsyncMock()
mock_channel.return_value = channel
connection = AsyncGRPCConnection(make_config())

await connection.close()

channel.close.assert_awaited_once()

asyncio.run(run())
12 changes: 6 additions & 6 deletions client/python/tests/test_connection.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@
import pytest
from unittest.mock import Mock, patch

from vortexdb._grpc_common import map_grpc_error
from vortexdb.connection import GRPCConnection
from vortexdb.config import VortexDBConfig
from vortexdb.exceptions import (
Expand Down Expand Up @@ -89,18 +90,17 @@ def test_successful_rpc_call(connection):
)
def test_grpc_error_mapping(status_code, expected_exception, connection):
error = FakeRpcError(status_code, "boom")
fake_rpc = Mock(side_effect=error)

with pytest.raises(expected_exception):
connection.call(fake_rpc, request="req")
mapped = map_grpc_error(error)

assert isinstance(mapped, expected_exception)


def test_unknown_grpc_error_maps_to_internal_error(connection):
error = FakeRpcError(grpc.StatusCode.UNKNOWN, "unknown")
fake_rpc = Mock(side_effect=error)
mapped = map_grpc_error(error)

with pytest.raises(InternalServerError):
connection.call(fake_rpc, request="req")
assert isinstance(mapped, InternalServerError)


# Clean connection closure test
Expand Down
2 changes: 2 additions & 0 deletions client/python/vortexdb/__init__.py
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
# vortexdb/__init__.py

from vortexdb.client import VortexDB
from vortexdb.async_client import AsyncVortexDB
from vortexdb.models import (
DenseVector,
Payload,
Expand All @@ -19,6 +20,7 @@

__all__ = [
"VortexDB",
"AsyncVortexDB",
"DenseVector",
"Payload",
"Point",
Expand Down
Loading
Loading