diff --git a/Cargo.lock b/Cargo.lock index 3471623..f26b2e8 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -488,7 +488,6 @@ checksum = "d7a1e2f27636f116493b8b860f5546edb47c8d8f8ea73e1d2a20be88e28d1fea" name = "defs" version = "0.1.0" dependencies = [ - "axum", "serde", "snafu", "uuid", diff --git a/client/python/USAGE.md b/client/python/USAGE.md index 18fd497..57b5223 100644 --- a/client/python/USAGE.md +++ b/client/python/USAGE.md @@ -52,7 +52,7 @@ async with AsyncVortexDB( api_key="your-api-key", ) as db: point_id = await db.insert( - vector=DenseVector([0.1, 0.2, 0.3]), + dense_vector=DenseVector([0.1, 0.2, 0.3]), payload=Payload.text("hello async vortex"), ) ``` @@ -109,32 +109,40 @@ VortexDB( #### **Insert** -Insert a vector with an associated payload +Insert a point with a dense vector, a sparse vector, or both, and an optional payload ``` -insert(*, vector: DenseVector, payload: Payload) -> str +insert( + *, + dense_vector: DenseVector | None = None, + sparse_vector: SparseVector | None = None, + payload: Payload | None = None, + point_id: str | None = None, +) -> str ``` Returns -- `point_id` (UUID string) +- `point_id` (UUID string, generated by the server when not given) Raises -- `TypeError` if `vector` is not a `DenseVector` +- `ValueError` if neither vector is given +- `TypeError` if an argument has the wrong type - gRPC-mapped errors (see Error Handling) --- #### **Batch Insert** -Insert multiple vectors with payloads in a single request +Insert multiple points in a single request. The server validates the whole batch before writing anything ``` -batch_insert(*, items: list[tuple[DenseVector, Payload]]) -> list[str] +batch_insert(*, points: list[PointInput]) -> list[str] ``` Returns -- List of `point_id` (UUID string) +- List of `point_id` (UUID string), in input order Raises -- `TypeError` if input structure is invalid +- `TypeError` if `points` is not a list of `PointInput` +- `ValueError` if `points` is empty - gRPC-mapped errors (see Error Handling) --- @@ -154,48 +162,33 @@ Returns #### **Search** -Search for nearest neighbours to a query vector +Search for the nearest points to a dense or a sparse query ``` -search( - *, - vector: DenseVector, - similarity: Similarity, - limit: int, -) -> list[tuple[str, float]] +search(*, query: SearchQuery) -> list[SearchHit] ``` Returns -- List of `(point_id, score)` tuples, where score is the distance (lower is closer) +- List of `SearchHit(id, score)`, best first. For dense queries the score is a distance (lower is closer). For sparse queries it is a dot product (higher is closer) Raises -- `TypeError` if `vector` is not a `DenseVector` +- `TypeError` if `query` is not a `SearchQuery` - `InvalidArgumentError` for invalid parameters --- #### **Batch Search** -Search for nearest neighbours for multiple queries in a single request +Run multiple queries in a single request. Dense and sparse queries can be mixed ``` -batch_search( - *, - queries, - similarity: Similarity | None = None, - limit: int | None = None, -) -> list[list[tuple[str, float]]] +batch_search(*, queries: list[SearchQuery]) -> list[list[SearchHit]] ``` Returns -- `TypeError` for invalid query formats -- `ValueError` if required parameters are missing +- One list of `SearchHit` per query, in order -Supported Input Formats: -The `queries` parameter is flexible and supports multiple formats: -- List of `SearchQuery` objects -- List of `(DenseVector, Similarity, Limit)` tuples -- List of `(DenseVector, Similarity)` tuples with a global `Limit` -- List of `(DenseVector, Limit)` tuples with a global `Similarity` -- List of `DenseVector` with global `Similarity` and `Limit` +Raises +- `TypeError` if `queries` is not a list of `SearchQuery` +- `ValueError` if `queries` is empty --- @@ -247,20 +240,43 @@ Factory Helpers: --- +### `SparseVector` + +``` +SparseVector(token_ids: list[int], weights: list[float]) +``` +- `token_ids` must be non-empty, non-negative and strictly increasing +- one finite, non-negative weight per token id +- Immutable (`frozen=True`) + +--- + +### `PointInput` + +``` +PointInput( + dense_vector: DenseVector | None = None, + sparse_vector: SparseVector | None = None, + payload: Payload | None = None, + point_id: str | None = None, +) +``` +A point to insert with `batch_insert`. At least one vector is required + +--- + ### `Point` ``` Point( id: str, - vector: DenseVector, - payload: Payload, + dense_vector: DenseVector | None, + sparse_vector: SparseVector | None, + payload: Payload | None, ) ``` Additional `pretty()` method provided to properly format output -All fields are directly accessible: -- `point.id` -- `point.vector` -- `point.payload` +Fields that the point does not have are `None` --- @@ -268,12 +284,26 @@ All fields are directly accessible: ``` SearchQuery( + limit: int, + query: DenseQuery | SparseVector, +) + +DenseQuery( vector: DenseVector, similarity: Similarity, - limit: int, + ef: int | None = None, ) ``` -Structured representation of a search request +Structured representation of a search request. `ef` is the HNSW search breadth and uses the server default when not set + +--- + +### `SearchHit` + +``` +SearchHit(id: str, score: float) +``` +One search result --- @@ -307,6 +337,7 @@ All client exceptions inherit from `VortexDBError` | `UNAUTHENTICATED` | `AuthenticationError` | | `NOT_FOUND` | `NotFoundError` | | `INVALID_ARGUMENT` | `InvalidArgumentError` | +| `ALREADY_EXISTS` | `AlreadyExistsError` | | `DEADLINE_EXCEEDED` | `TimeoutError` | | `UNAVAILABLE` | `ServiceUnavailableError` | | Any other error | `InternalServerError` | diff --git a/client/python/examples/async_batch_usage.py b/client/python/examples/async_batch_usage.py index 72d2493..9f87436 100644 --- a/client/python/examples/async_batch_usage.py +++ b/client/python/examples/async_batch_usage.py @@ -1,7 +1,8 @@ import asyncio from vortexdb import AsyncVortexDB -from vortexdb import Payload, Similarity, SearchQuery, to_dense_vectors +from vortexdb import DenseQuery, PointInput, Similarity, SearchQuery, SparseVector +from vortexdb import to_dense_vectors async def main(): @@ -16,66 +17,35 @@ async def main(): ] vectors = to_dense_vectors(raw_vectors) - p1 = Payload.text("hello world") - p2 = Payload.image("/img/a.png") - p3 = Payload.text("foo bar") - - items = [ - (vectors[0], p1), - (vectors[1], p2), - (vectors[2], p3), - ] - - # Batch Insert - point_ids = await db.batch_insert(items=items) - print("Inserted ids:\n", point_ids) - - q = SearchQuery( - vector=vectors[0], - similarity=Similarity.COSINE, - limit=3, + point_ids = await db.batch_insert( + points=[ + PointInput( + dense_vector=vectors[0], + sparse_vector=SparseVector([1, 5], [1.0, 0.5]), + ), + PointInput( + dense_vector=vectors[1], + sparse_vector=SparseVector([5, 9], [2.0, 1.0]), + ), + PointInput(dense_vector=vectors[2]), + ] ) - res = await db.search(query=q) - print("\nSingle SearchQuery:\n", res) - - # List of SearchQuery - queries = [ - SearchQuery(vectors[0], Similarity.HAMMING, 3), - SearchQuery(vectors[1], Similarity.EUCLIDEAN, 2), - q, - ] - res = await db.batch_search(queries=queries) - print("\nBatch SearchQuery:\n", res) - # List of vectors with global Similarity and Limit - res = await db.batch_search( - queries=vectors, - similarity=Similarity.COSINE, - limit=3, - ) - print("\nList of DenseVectors:\n", res) + q = SearchQuery(limit=3, query=DenseQuery(vectors[0], Similarity.COSINE)) + res = await db.search(query=q) + print("Dense search:\n", res) - # List of tuple (DenseVector, Similarity) with global Limit - queries = [ - (vectors[0], Similarity.COSINE), - (vectors[1], Similarity.MANHATTAN), - ] - res = await db.batch_search( - queries=queries, - limit=3, - ) - print("\nList of (DenseVector, Similarity):\n", res) + sparse_q = SearchQuery(limit=3, query=SparseVector([5], [1.0])) + res = await db.search(query=sparse_q) + print("\nSparse search:\n", res) - # List of tuple (DenseVector, Limit) with global Similarity queries = [ - (vectors[0], 2), - (vectors[1], 4), + SearchQuery(3, DenseQuery(vectors[0], Similarity.HAMMING)), + SearchQuery(2, DenseQuery(vectors[1], Similarity.EUCLIDEAN)), + sparse_q, ] - res = await db.batch_search( - queries=queries, - similarity=Similarity.COSINE, - ) - print("\nList of (DenseVector, Limit):\n", res) + res = await db.batch_search(queries=queries) + print("\nBatch search:\n", res) for pid in point_ids: await db.delete(point_id=pid) diff --git a/client/python/examples/async_usage.py b/client/python/examples/async_usage.py index 882a926..d159287 100644 --- a/client/python/examples/async_usage.py +++ b/client/python/examples/async_usage.py @@ -1,6 +1,13 @@ import asyncio -from vortexdb import AsyncVortexDB, DenseVector, Payload, Similarity +from vortexdb import ( + AsyncVortexDB, + DenseQuery, + DenseVector, + Payload, + SearchQuery, + Similarity, +) async def main(): @@ -9,7 +16,7 @@ async def main(): api_key="my-secret-password", ) as db: point_id = await db.insert( - vector=DenseVector([0.1, 0.2, 0.3]), + dense_vector=DenseVector([0.1, 0.2, 0.3]), payload=Payload.text("hello async vortex"), ) @@ -18,9 +25,10 @@ async def main(): print(point.pretty()) results = await db.search( - vector=DenseVector([0.1, 0.2, 0.3]), - similarity=Similarity.COSINE, - limit=5, + query=SearchQuery( + limit=5, + query=DenseQuery(DenseVector([0.1, 0.2, 0.3]), Similarity.COSINE), + ) ) print(results) diff --git a/client/python/examples/basic_usage.py b/client/python/examples/basic_usage.py index f58d8c5..45631df 100644 --- a/client/python/examples/basic_usage.py +++ b/client/python/examples/basic_usage.py @@ -1,5 +1,11 @@ from vortexdb import VortexDB -from vortexdb import DenseVector, Payload, Similarity # from vortexdb.models +from vortexdb import ( + DenseQuery, + DenseVector, + Payload, + SearchQuery, + Similarity, +) # from vortexdb.models def main(): @@ -11,7 +17,7 @@ def main(): # Insert a vector point_id = db.insert( - vector=DenseVector([0.1, 0.2, 0.3]), + dense_vector=DenseVector([0.1, 0.2, 0.3]), payload=Payload.text("hello world"), ) @@ -21,9 +27,10 @@ def main(): # Search results = db.search( - vector=DenseVector([0.1, 0.2, 0.3]), - similarity=Similarity.COSINE, - limit=3, + query=SearchQuery( + limit=3, + query=DenseQuery(DenseVector([0.1, 0.2, 0.3]), Similarity.COSINE), + ) ) print("Search results:", results) diff --git a/client/python/examples/batch_insert_usage.py b/client/python/examples/batch_insert_usage.py index e021b1a..66f51bf 100644 --- a/client/python/examples/batch_insert_usage.py +++ b/client/python/examples/batch_insert_usage.py @@ -1,5 +1,5 @@ from vortexdb import VortexDB -from vortexdb import Payload, to_dense_vectors +from vortexdb import Payload, PointInput, SparseVector, to_dense_vectors def main(): @@ -15,18 +15,20 @@ def main(): ] vectors = to_dense_vectors(raw_vectors) - p1 = Payload.text("hello world") - p2 = Payload.image("/img/a.png") - p3 = Payload.text("foo bar") - - items = [ - (vectors[0], p1), - (vectors[1], p2), - (vectors[2], p3), + points = [ + PointInput(dense_vector=vectors[0], payload=Payload.text("hello world")), + PointInput( + sparse_vector=SparseVector([3, 17, 42], [0.5, 1.2, 0.8]), + payload=Payload.image("/img/a.png"), + ), + PointInput( + dense_vector=vectors[2], + sparse_vector=SparseVector([17], [2.0]), + payload=Payload.text("foo bar"), + ), ] - # Batch Insert - point_ids = db.batch_insert(items=items) + point_ids = db.batch_insert(points=points) print("Inserted ids:\n", point_ids) for pid in point_ids: diff --git a/client/python/examples/context_manager_usage.py b/client/python/examples/context_manager_usage.py index d2f1de2..3b4cd93 100644 --- a/client/python/examples/context_manager_usage.py +++ b/client/python/examples/context_manager_usage.py @@ -1,4 +1,4 @@ -from vortexdb import VortexDB, DenseVector, Payload, Similarity +from vortexdb import VortexDB, DenseQuery, DenseVector, Payload, SearchQuery, Similarity def main(): @@ -8,7 +8,7 @@ def main(): ) as db: # Insert a vector point_id = db.insert( - vector=DenseVector([0.1, 0.2, 0.3]), + dense_vector=DenseVector([0.1, 0.2, 0.3]), payload=Payload.text("hello world"), ) @@ -18,9 +18,10 @@ def main(): # Search results = db.search( - vector=DenseVector([0.1, 0.2, 0.3]), - similarity=Similarity.COSINE, - limit=3, + query=SearchQuery( + limit=3, + query=DenseQuery(DenseVector([0.1, 0.2, 0.3]), Similarity.COSINE), + ) ) print("Search results:", results) diff --git a/client/python/examples/search_query_usage.py b/client/python/examples/search_query_usage.py index 3044056..64500b6 100644 --- a/client/python/examples/search_query_usage.py +++ b/client/python/examples/search_query_usage.py @@ -1,5 +1,6 @@ from vortexdb import VortexDB -from vortexdb import Similarity, SearchQuery, to_dense_vectors +from vortexdb import DenseQuery, PointInput, Similarity, SearchQuery, SparseVector +from vortexdb import to_dense_vectors def main(): @@ -15,52 +16,36 @@ def main(): ] vectors = to_dense_vectors(raw_vectors) - q = SearchQuery( - vector=vectors[0], - similarity=Similarity.COSINE, - limit=3, + point_ids = db.batch_insert( + points=[ + PointInput( + dense_vector=vectors[0], sparse_vector=SparseVector([1, 5], [1.0, 0.5]) + ), + PointInput( + dense_vector=vectors[1], sparse_vector=SparseVector([5, 9], [2.0, 1.0]) + ), + PointInput(dense_vector=vectors[2]), + ] ) - res = db.search(query=q) - print("Single SearchQuery:\n", res) - # List of SearchQuery - queries = [ - SearchQuery(vectors[0], Similarity.HAMMING, 3), - SearchQuery(vectors[1], Similarity.EUCLIDEAN, 2), - q, - ] - res = db.batch_search(queries=queries) - print("\nBatch SearchQuery:\n", res) + q = SearchQuery(limit=3, query=DenseQuery(vectors[0], Similarity.COSINE)) + res = db.search(query=q) + print("Dense search:\n", res) - # List of vectors with global Similarity and Limit - res = db.batch_search( - queries=vectors, - similarity=Similarity.COSINE, - limit=3, - ) - print("\nList of DenseVectors:\n", res) + sparse_q = SearchQuery(limit=3, query=SparseVector([5], [1.0])) + res = db.search(query=sparse_q) + print("\nSparse search:\n", res) - # List of tuple (DenseVector, Similarity) with global Limit queries = [ - (vectors[0], Similarity.COSINE), - (vectors[1], Similarity.MANHATTAN), + SearchQuery(3, DenseQuery(vectors[0], Similarity.HAMMING)), + SearchQuery(2, DenseQuery(vectors[1], Similarity.EUCLIDEAN)), + sparse_q, ] - res = db.batch_search( - queries=queries, - limit=3, - ) - print("\nList of (DenseVector, Similarity):\n", res) + res = db.batch_search(queries=queries) + print("\nBatch search:\n", res) - # List of tuple (DenseVector, Limit) with global Similarity - queries = [ - (vectors[0], 2), - (vectors[1], 4), - ] - res = db.batch_search( - queries=queries, - similarity=Similarity.COSINE, - ) - print("\nList of (DenseVector, Limit):\n", res) + for pid in point_ids: + db.delete(point_id=pid) db.close() diff --git a/client/python/tests/test_async_client.py b/client/python/tests/test_async_client.py index dfbee27..6aff1c4 100644 --- a/client/python/tests/test_async_client.py +++ b/client/python/tests/test_async_client.py @@ -5,13 +5,17 @@ from vortexdb.async_client import AsyncVortexDB from vortexdb.async_connection import AsyncGRPCConnection +from vortexdb.grpc import vector_db_pb2 as pb from vortexdb.models import ( - ContentType, + DenseQuery, DenseVector, Payload, Point, - Similarity, + PointInput, + SearchHit, SearchQuery, + Similarity, + SparseVector, ) @@ -36,20 +40,37 @@ def client(mock_connection): ) -def test_async_insert_success(client, mock_connection): +def test_async_insert_dense_point(client, mock_connection): async def run(): - response = Mock() - response.id = Mock() - response.id.value = "point-123" - - mock_connection.call.return_value = response + mock_connection.call.return_value = pb.PointID(id=pb.UUID(value="point-123")) point_id = await client.insert( - vector=DenseVector([1, 2, 3]), + dense_vector=DenseVector([1, 2, 3]), payload=Payload.text("hello"), ) assert point_id == "point-123" + request = mock_connection.call.call_args.args[1] + assert list(request.dense_vector.values) == [1.0, 2.0, 3.0] + assert not request.HasField("sparse_vector") + assert not request.HasField("id") + + asyncio.run(run()) + + +def test_async_insert_sparse_point_with_id(client, mock_connection): + async def run(): + mock_connection.call.return_value = pb.PointID(id=pb.UUID(value="point-1")) + + await client.insert( + sparse_vector=SparseVector([1, 5], [0.5, 2.0]), + point_id="point-1", + ) + + request = mock_connection.call.call_args.args[1] + assert request.id.id.value == "point-1" + assert list(request.sparse_vector.token_ids) == [1, 5] + assert not request.HasField("dense_vector") asyncio.run(run()) @@ -58,61 +79,65 @@ def test_async_insert_rejects_invalid_vector(client): async def run(): with pytest.raises(TypeError): await client.insert( - vector=[1, 2, 3], + dense_vector=[1, 2, 3], # not DenseVector payload=Payload.text("hello"), ) asyncio.run(run()) +def test_async_insert_rejects_point_without_vectors(client): + async def run(): + with pytest.raises(ValueError): + await client.insert(payload=Payload.text("hello")) + + asyncio.run(run()) + + # Batch Insert def test_async_batch_insert_success(client, mock_connection): async def run(): - response = Mock() - response.ids = [ - Mock(id=Mock(value="p1")), - Mock(id=Mock(value="p2")), + mock_connection.call.return_value = pb.InsertPointsBatchResponse( + ids=[pb.PointID(id=pb.UUID(value="p1")), pb.PointID(id=pb.UUID(value="p2"))] + ) + points = [ + PointInput(dense_vector=DenseVector([1, 2, 3]), payload=Payload.text("a")), + PointInput(sparse_vector=SparseVector([2], [1.0])), ] - mock_connection.call.return_value = response - items = [ - (DenseVector([1, 2, 3]), Payload.text("a")), - (DenseVector([4, 5, 6]), Payload.text("b")), - ] - result = await client.batch_insert(items=items) + result = await client.batch_insert(points=points) + assert result == ["p1", "p2"] + request = mock_connection.call.call_args.args[1] + assert len(request.points) == 2 asyncio.run(run()) -def test_async_batch_insert_invalid_items_type(client): +def test_async_batch_insert_invalid_points_type(client): async def run(): with pytest.raises(TypeError): - await client.batch_insert(items="not-a-list") + await client.batch_insert(points="not-a-list") asyncio.run(run()) -def test_async_batch_insert_invalid_tuple_structure(client): +def test_async_batch_insert_rejects_non_point_input(client): async def run(): - items = [ - (DenseVector([1, 2, 3]),), # only one element - ] with pytest.raises(TypeError): - await client.batch_insert(items=items) + await client.batch_insert( + points=[(DenseVector([1, 2, 3]), Payload.text("a"))] + ) asyncio.run(run()) -def test_async_batch_insert_invalid_vector(client): +def test_async_batch_insert_rejects_empty(client): async def run(): - items = [ - ([1, 2, 3], Payload.text("a")), # not DenseVector - ] - with pytest.raises(TypeError): - await client.batch_insert(items=items) + with pytest.raises(ValueError): + await client.batch_insert(points=[]) asyncio.run(run()) @@ -122,18 +147,18 @@ async def 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 + mock_connection.call.return_value = pb.Point( + id=pb.PointID(id=pb.UUID(value="point-123")), + sparse_vector=pb.SparseVector(token_ids=[1, 4], weights=[0.5, 2.0]), + payload=pb.Payload(content_type=pb.Text, content="hello"), + ) point = await client.get(point_id="point-123") assert isinstance(point, Point) assert point.id == "point-123" + assert point.dense_vector is None + assert point.sparse_vector == SparseVector([1, 4], [0.5, 2.0]) assert point.payload.content == "hello" asyncio.run(run()) @@ -159,7 +184,7 @@ async def run(): await client.delete(point_id="point-123") - mock_connection.call.assert_awaited_once() + mock_connection.call.assert_called_once() asyncio.run(run()) @@ -167,69 +192,49 @@ async def run(): # Search -def test_async_search_success(client, mock_connection): +def test_async_dense_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")), - ], - scores=[0.1, 0.2], - ) - - results = await client.search( - vector=DenseVector([1, 2, 3]), - similarity=Similarity.COSINE, - limit=2, + mock_connection.call.return_value = pb.SearchResponse( + hits=[ + pb.SearchHit(id=pb.PointID(id=pb.UUID(value="p1")), score=0.25), + pb.SearchHit(id=pb.PointID(id=pb.UUID(value="p2")), score=0.5), + ] ) - - assert results == [("p1", 0.1), ("p2", 0.2)] - - asyncio.run(run()) - - -def test_async_search_with_query_object(client, mock_connection): - async def run(): - mock_connection.call.return_value = Mock( - result_point_ids=[ - Mock(id=Mock(value="p1")), - ], - scores=[0.1], + query = SearchQuery( + limit=2, query=DenseQuery(DenseVector([1, 2, 3]), Similarity.COSINE, ef=128) ) - q = SearchQuery(DenseVector([1, 2, 3]), Similarity.COSINE, 2) - results = await client.search(query=q) + results = await client.search(query=query) - assert results == [("p1", 0.1)] + assert results == [SearchHit("p1", 0.25), SearchHit("p2", 0.5)] + request = mock_connection.call.call_args.args[1] + assert request.limit == 2 + assert request.dense.ef == 128 asyncio.run(run()) -def test_async_search_accepts_ef(client, mock_connection): +def test_async_sparse_search_success(client, mock_connection): async def run(): - mock_connection.call.return_value = Mock(result_point_ids=[], scores=[]) + mock_connection.call.return_value = pb.SearchResponse( + hits=[pb.SearchHit(id=pb.PointID(id=pb.UUID(value="p1")), score=2.0)] + ) - await client.search( - vector=DenseVector([1, 2, 3]), - similarity=Similarity.COSINE, - limit=2, - ef=128, + results = await client.search( + query=SearchQuery(limit=5, query=SparseVector([4], [1.0])) ) + assert results == [SearchHit("p1", 2.0)] request = mock_connection.call.call_args.args[1] - assert request.ef == 128 + assert request.WhichOneof("query") == "sparse" asyncio.run(run()) -def test_async_search_invalid_vector(client): +def test_async_search_rejects_non_search_query(client): async def run(): with pytest.raises(TypeError): - await client.search( - vector=[1, 2, 3], - similarity=Similarity.COSINE, - limit=2, - ) + await client.search(query=DenseVector([1, 2, 3])) asyncio.run(run()) @@ -237,122 +242,44 @@ async def run(): # Batch Search -def test_async_batch_search_full_tuple(client, mock_connection): +def test_async_batch_search_mixes_dense_and_sparse(client, mock_connection): async def run(): - mock_connection.call.return_value = Mock( + mock_connection.call.return_value = pb.SearchBatchResponse( results=[ - Mock(result_point_ids=[Mock(id=Mock(value="p1"))], scores=[0.5]), - Mock(result_point_ids=[Mock(id=Mock(value="p2"))], scores=[0.5]), + pb.SearchResponse( + hits=[ + pb.SearchHit(id=pb.PointID(id=pb.UUID(value="p1")), score=0.5) + ] + ), + pb.SearchResponse(), ] ) queries = [ - (DenseVector([1, 2, 3]), Similarity.COSINE, 2), - (DenseVector([4, 5, 6]), Similarity.EUCLIDEAN, 1), + SearchQuery(2, DenseQuery(DenseVector([1, 2, 3]), Similarity.EUCLIDEAN)), + SearchQuery(1, SparseVector([7], [1.0])), ] - result = await client.batch_search(queries=queries) - assert result == [[("p1", 0.5)], [("p2", 0.5)]] - asyncio.run(run()) - - -def test_async_batch_search_searchquery_objects(client, mock_connection): - async def run(): - mock_connection.call.return_value = Mock( - results=[ - Mock(result_point_ids=[Mock(id=Mock(value="p1"))], scores=[0.5]), - ] - ) - queries = [ - SearchQuery(DenseVector([1, 2, 3]), Similarity.COSINE, 2), - ] result = await client.batch_search(queries=queries) - assert result == [[("p1", 0.5)]] - - asyncio.run(run()) - - -def test_async_batch_search_vectors_with_global_params(client, mock_connection): - async def run(): - mock_connection.call.return_value = Mock( - results=[ - Mock(result_point_ids=[Mock(id=Mock(value="p1"))], scores=[0.5]), - ] - ) - queries = [DenseVector([1, 2, 3])] - result = await client.batch_search( - queries=queries, - similarity=Similarity.MANHATTAN, - limit=2, - ) - assert result == [[("p1", 0.5)]] - - asyncio.run(run()) - - -def test_async_batch_search_vector_similarity_with_global_limit( - client, mock_connection -): - async def run(): - mock_connection.call.return_value = Mock( - results=[ - Mock(result_point_ids=[Mock(id=Mock(value="p1"))], scores=[0.5]), - ] - ) - queries = [ - (DenseVector([1, 2, 3]), Similarity.COSINE), - ] - result = await client.batch_search( - queries=queries, - limit=2, - ) - assert result == [[("p1", 0.5)]] - - asyncio.run(run()) - - -def test_async_batch_search_accepts_ef(client, mock_connection): - async def run(): - mock_connection.call.return_value = Mock(results=[]) - - await client.batch_search( - queries=[ - (DenseVector([1, 2, 3]), Similarity.COSINE, 2), - (DenseVector([4, 5, 6]), Similarity.COSINE, 1), - ], - ef=256, - ) + assert result == [[SearchHit("p1", 0.5)], []] request = mock_connection.call.call_args.args[1] - assert [query.ef for query in request.queries] == [256, 256] + assert [q.WhichOneof("query") for q in request.queries] == ["dense", "sparse"] asyncio.run(run()) -def test_async_batch_search_missing_globals_for_vector(client): +def test_async_batch_search_rejects_invalid_format(client): async def run(): - queries = [DenseVector([1, 2, 3])] - with pytest.raises(ValueError): - await client.batch_search(queries=queries) + with pytest.raises(TypeError): + await client.batch_search(queries=["invalid"]) asyncio.run(run()) -def test_async_batch_search_missing_limit(client): +def test_async_batch_search_rejects_empty(client): async def run(): - queries = [ - (DenseVector([1, 2, 3]), Similarity.COSINE), - ] with pytest.raises(ValueError): - await client.batch_search(queries=queries) - - asyncio.run(run()) - - -def test_async_batch_search_invalid_format(client): - async def run(): - queries = ["invalid"] - with pytest.raises(TypeError): - await client.batch_search(queries=queries) + await client.batch_search(queries=[]) asyncio.run(run()) diff --git a/client/python/tests/test_async_connection.py b/client/python/tests/test_async_connection.py index 62f1043..a0c798a 100644 --- a/client/python/tests/test_async_connection.py +++ b/client/python/tests/test_async_connection.py @@ -7,6 +7,7 @@ from vortexdb.async_connection import AsyncGRPCConnection from vortexdb.config import VortexDBConfig from vortexdb.exceptions import ( + AlreadyExistsError, AuthenticationError, InternalServerError, InvalidArgumentError, @@ -80,6 +81,7 @@ def test_async_grpc_error_mapping(): (grpc.StatusCode.UNAUTHENTICATED, AuthenticationError), (grpc.StatusCode.NOT_FOUND, NotFoundError), (grpc.StatusCode.INVALID_ARGUMENT, InvalidArgumentError), + (grpc.StatusCode.ALREADY_EXISTS, AlreadyExistsError), (grpc.StatusCode.DEADLINE_EXCEEDED, TimeoutError), (grpc.StatusCode.UNAVAILABLE, ServiceUnavailableError), (grpc.StatusCode.UNKNOWN, InternalServerError), diff --git a/client/python/tests/test_client.py b/client/python/tests/test_client.py index 5fe3518..3176b3f 100644 --- a/client/python/tests/test_client.py +++ b/client/python/tests/test_client.py @@ -3,8 +3,18 @@ from vortexdb.client import VortexDB from vortexdb.connection import GRPCConnection -from vortexdb.models import DenseVector, Payload, Similarity, ContentType, Point -from vortexdb.models import SearchQuery +from vortexdb.grpc import vector_db_pb2 as pb +from vortexdb.models import ( + DenseQuery, + DenseVector, + Payload, + Point, + PointInput, + SearchHit, + SearchQuery, + Similarity, + SparseVector, +) # Fixtures for a mock connection and client layer @@ -31,84 +41,98 @@ def client(mock_connection): # Insert -def test_insert_success(client, mock_connection): - response = Mock() - response.id = Mock() - response.id.value = "point-123" +def test_insert_dense_point(client, mock_connection): + mock_connection.call.return_value = pb.PointID(id=pb.UUID(value="point-123")) - mock_connection.call.return_value = response + point_id = client.insert( + dense_vector=DenseVector([1, 2, 3]), + payload=Payload.text("hello"), + ) - vector = DenseVector([1, 2, 3]) - payload = Payload.text("hello") + assert point_id == "point-123" + request = mock_connection.call.call_args.args[1] + assert list(request.dense_vector.values) == [1.0, 2.0, 3.0] + assert not request.HasField("sparse_vector") + assert not request.HasField("id") - point_id = client.insert(vector=vector, payload=payload) - assert point_id == "point-123" +def test_insert_sparse_point_with_id(client, mock_connection): + mock_connection.call.return_value = pb.PointID(id=pb.UUID(value="point-1")) + + client.insert( + sparse_vector=SparseVector([1, 5], [0.5, 2.0]), + point_id="point-1", + ) + + request = mock_connection.call.call_args.args[1] + assert request.id.id.value == "point-1" + assert list(request.sparse_vector.token_ids) == [1, 5] + assert not request.HasField("dense_vector") def test_insert_rejects_invalid_vector(client): with pytest.raises(TypeError): client.insert( - vector=[1, 2, 3], # not DenseVector + dense_vector=[1, 2, 3], # not DenseVector payload=Payload.text("hello"), ) +def test_insert_rejects_point_without_vectors(client): + with pytest.raises(ValueError): + client.insert(payload=Payload.text("hello")) + + # Batch Insert def test_batch_insert_success(client, mock_connection): - response = Mock() - response.ids = [ - Mock(id=Mock(value="p1")), - Mock(id=Mock(value="p2")), - ] - mock_connection.call.return_value = response - items = [ - (DenseVector([1, 2, 3]), Payload.text("a")), - (DenseVector([4, 5, 6]), Payload.text("b")), + mock_connection.call.return_value = pb.InsertPointsBatchResponse( + ids=[pb.PointID(id=pb.UUID(value="p1")), pb.PointID(id=pb.UUID(value="p2"))] + ) + points = [ + PointInput(dense_vector=DenseVector([1, 2, 3]), payload=Payload.text("a")), + PointInput(sparse_vector=SparseVector([2], [1.0])), ] - result = client.batch_insert(items=items) + + result = client.batch_insert(points=points) + assert result == ["p1", "p2"] + request = mock_connection.call.call_args.args[1] + assert len(request.points) == 2 -def test_batch_insert_invalid_items_type(client): +def test_batch_insert_invalid_points_type(client): with pytest.raises(TypeError): - client.batch_insert(items="not-a-list") + client.batch_insert(points="not-a-list") -def test_batch_insert_invalid_tuple_structure(client): - items = [ - (DenseVector([1, 2, 3]),), # only one element - ] +def test_batch_insert_rejects_non_point_input(client): with pytest.raises(TypeError): - client.batch_insert(items=items) + client.batch_insert(points=[(DenseVector([1, 2, 3]), Payload.text("a"))]) -def test_batch_insert_invalid_vector(client): - items = [ - ([1, 2, 3], Payload.text("a")), # not DenseVector - ] - with pytest.raises(TypeError): - client.batch_insert(items=items) +def test_batch_insert_rejects_empty(client): + with pytest.raises(ValueError): + client.batch_insert(points=[]) # Get def test_get_point_success(client, mock_connection): - 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 + mock_connection.call.return_value = pb.Point( + id=pb.PointID(id=pb.UUID(value="point-123")), + sparse_vector=pb.SparseVector(token_ids=[1, 4], weights=[0.5, 2.0]), + payload=pb.Payload(content_type=pb.Text, content="hello"), + ) point = client.get(point_id="point-123") assert isinstance(point, Point) assert point.id == "point-123" + assert point.dense_vector is None + assert point.sparse_vector == SparseVector([1, 4], [0.5, 2.0]) assert point.payload.content == "hello" @@ -134,142 +158,74 @@ def test_delete_success(client, mock_connection): # Search -def test_search_success(client, mock_connection): - mock_connection.call.return_value = Mock( - result_point_ids=[ - Mock(id=Mock(value="p1")), - Mock(id=Mock(value="p2")), - ], - scores=[0.1, 0.2], +def test_dense_search_success(client, mock_connection): + mock_connection.call.return_value = pb.SearchResponse( + hits=[ + pb.SearchHit(id=pb.PointID(id=pb.UUID(value="p1")), score=0.25), + pb.SearchHit(id=pb.PointID(id=pb.UUID(value="p2")), score=0.5), + ] ) - - results = client.search( - vector=DenseVector([1, 2, 3]), - similarity=Similarity.COSINE, - limit=2, + query = SearchQuery( + limit=2, query=DenseQuery(DenseVector([1, 2, 3]), Similarity.COSINE, ef=128) ) - assert results == [("p1", 0.1), ("p2", 0.2)] + results = client.search(query=query) + assert results == [SearchHit("p1", 0.25), SearchHit("p2", 0.5)] + request = mock_connection.call.call_args.args[1] + assert request.limit == 2 + assert request.dense.ef == 128 -def test_search_accepts_ef(client, mock_connection): - mock_connection.call.return_value = Mock(result_point_ids=[], scores=[]) - client.search( - vector=DenseVector([1, 2, 3]), - similarity=Similarity.COSINE, - limit=2, - ef=128, +def test_sparse_search_success(client, mock_connection): + mock_connection.call.return_value = pb.SearchResponse( + hits=[pb.SearchHit(id=pb.PointID(id=pb.UUID(value="p1")), score=2.0)] ) + results = client.search(query=SearchQuery(limit=5, query=SparseVector([4], [1.0]))) + + assert results == [SearchHit("p1", 2.0)] request = mock_connection.call.call_args.args[1] - assert request.ef == 128 + assert request.WhichOneof("query") == "sparse" -def test_search_invalid_vector(client): +def test_search_rejects_non_search_query(client): with pytest.raises(TypeError): - client.search( - vector=[1, 2, 3], - similarity=Similarity.COSINE, - limit=2, - ) - - -def test_batch_search_accepts_ef(client, mock_connection): - mock_connection.call.return_value = Mock(results=[]) - - client.batch_search( - queries=[ - (DenseVector([1, 2, 3]), Similarity.COSINE, 2), - (DenseVector([4, 5, 6]), Similarity.COSINE, 1), - ], - ef=256, - ) - - request = mock_connection.call.call_args.args[1] - assert [query.ef for query in request.queries] == [256, 256] + client.search(query=DenseVector([1, 2, 3])) # Batch Search -def test_batch_search_full_tuple(client, mock_connection): - mock_connection.call.return_value = Mock( - results=[ - Mock(result_point_ids=[Mock(id=Mock(value="p1"))], scores=[0.5]), - Mock(result_point_ids=[Mock(id=Mock(value="p2"))], scores=[0.5]), - ] - ) - queries = [ - (DenseVector([1, 2, 3]), Similarity.COSINE, 2), - (DenseVector([4, 5, 6]), Similarity.EUCLIDEAN, 1), - ] - result = client.batch_search(queries=queries) - assert result == [[("p1", 0.5)], [("p2", 0.5)]] - - -def test_batch_search_vectors_with_global_params(client, mock_connection): - mock_connection.call.return_value = Mock( - results=[ - Mock(result_point_ids=[Mock(id=Mock(value="p1"))], scores=[0.5]), - ] - ) - queries = [DenseVector([1, 2, 3])] - result = client.batch_search( - queries=queries, - similarity=Similarity.MANHATTAN, - limit=2, - ) - assert result == [[("p1", 0.5)]] - - -def test_batch_search_vector_similarity_with_global_limit(client, mock_connection): - mock_connection.call.return_value = Mock( +def test_batch_search_mixes_dense_and_sparse(client, mock_connection): + mock_connection.call.return_value = pb.SearchBatchResponse( results=[ - Mock(result_point_ids=[Mock(id=Mock(value="p1"))], scores=[0.5]), + pb.SearchResponse( + hits=[pb.SearchHit(id=pb.PointID(id=pb.UUID(value="p1")), score=0.5)] + ), + pb.SearchResponse(), ] ) queries = [ - (DenseVector([1, 2, 3]), Similarity.COSINE), + SearchQuery(2, DenseQuery(DenseVector([1, 2, 3]), Similarity.EUCLIDEAN)), + SearchQuery(1, SparseVector([7], [1.0])), ] - result = client.batch_search( - queries=queries, - limit=2, - ) - assert result == [[("p1", 0.5)]] - -def test_batch_search_searchquery_objects(client, mock_connection): - mock_connection.call.return_value = Mock( - results=[ - Mock(result_point_ids=[Mock(id=Mock(value="p1"))], scores=[0.5]), - ] - ) - queries = [ - SearchQuery(DenseVector([1, 2, 3]), Similarity.COSINE, 2), - ] result = client.batch_search(queries=queries) - assert result == [[("p1", 0.5)]] - -def test_batch_search_missing_globals_for_vector(client): - queries = [DenseVector([1, 2, 3])] - with pytest.raises(ValueError): - client.batch_search(queries=queries) + assert result == [[SearchHit("p1", 0.5)], []] + request = mock_connection.call.call_args.args[1] + assert [q.WhichOneof("query") for q in request.queries] == ["dense", "sparse"] -def test_batch_search_missing_limit(client): - queries = [ - (DenseVector([1, 2, 3]), Similarity.COSINE), - ] - with pytest.raises(ValueError): - client.batch_search(queries=queries) +def test_batch_search_rejects_invalid_format(client): + with pytest.raises(TypeError): + client.batch_search(queries=["invalid"]) -def test_batch_search_invalid_format(client): - queries = ["invalid"] - with pytest.raises(TypeError): - client.batch_search(queries=queries) +def test_batch_search_rejects_empty(client): + with pytest.raises(ValueError): + client.batch_search(queries=[]) # Close diff --git a/client/python/tests/test_connection.py b/client/python/tests/test_connection.py index 6d5c108..70067a4 100644 --- a/client/python/tests/test_connection.py +++ b/client/python/tests/test_connection.py @@ -9,6 +9,7 @@ AuthenticationError, NotFoundError, InvalidArgumentError, + AlreadyExistsError, TimeoutError, ServiceUnavailableError, InternalServerError, @@ -88,6 +89,7 @@ def test_successful_rpc_call(connection): (grpc.StatusCode.UNAUTHENTICATED, AuthenticationError), (grpc.StatusCode.NOT_FOUND, NotFoundError), (grpc.StatusCode.INVALID_ARGUMENT, InvalidArgumentError), + (grpc.StatusCode.ALREADY_EXISTS, AlreadyExistsError), (grpc.StatusCode.DEADLINE_EXCEEDED, TimeoutError), (grpc.StatusCode.UNAVAILABLE, ServiceUnavailableError), ], diff --git a/client/python/tests/test_models.py b/client/python/tests/test_models.py index 796631f..cfe7524 100644 --- a/client/python/tests/test_models.py +++ b/client/python/tests/test_models.py @@ -1,11 +1,15 @@ import pytest from vortexdb.models import ( + DenseQuery, DenseVector, Payload, Point, + PointInput, + SearchQuery, Similarity, ContentType, + SparseVector, ) from vortexdb.grpc import vector_db_pb2 @@ -101,30 +105,97 @@ def test_payload_rejects_invalid_content_type(): Payload("text", "hello") +def test_sparse_vector_valid(): + v = SparseVector((1, 4), (1, 0.5)) + assert v.token_ids == [1, 4] + assert v.weights == [1.0, 0.5] + assert SparseVector([1], [0]).weights == [0.0] + + +@pytest.mark.parametrize( + "token_ids, weights", + [ + ([], []), + ([1, 2], [1.0]), + ([2, 1], [1.0, 1.0]), + ([1, 1], [1.0, 1.0]), + ([1], [float("inf")]), + ([1, 2], [1.0, -0.5]), + ], +) +def test_sparse_vector_rejects_invalid_values(token_ids, weights): + with pytest.raises(ValueError): + SparseVector(token_ids, weights) + + +@pytest.mark.parametrize( + "token_ids, weights", + [([1.5], [1.0]), ([-1], [1.0]), ([1], ["a"])], +) +def test_sparse_vector_rejects_invalid_types(token_ids, weights): + with pytest.raises(TypeError): + SparseVector(token_ids, weights) + + +def test_point_input_requires_a_vector(): + with pytest.raises(ValueError): + PointInput(payload=Payload.text("hello")) + + +def test_point_input_to_proto_leaves_missing_fields_unset(): + proto = PointInput(sparse_vector=SparseVector([3], [1.0])).to_proto() + assert list(proto.sparse_vector.token_ids) == [3] + assert not proto.HasField("dense_vector") + assert not proto.HasField("payload") + assert not proto.HasField("id") + + +def test_dense_search_query_to_proto(): + query = SearchQuery(3, DenseQuery(DenseVector([1, 2]), Similarity.COSINE)) + proto = query.to_proto() + assert proto.limit == 3 + assert proto.WhichOneof("query") == "dense" + assert proto.dense.similarity == vector_db_pb2.Cosine + assert not proto.dense.HasField("ef") + + +def test_sparse_search_query_to_proto(): + proto = SearchQuery(3, SparseVector([1, 2], [0.5, 1.0])).to_proto() + assert proto.WhichOneof("query") == "sparse" + assert list(proto.sparse.weights) == [0.5, 1.0] + + +def test_search_query_rejects_raw_vector(): + with pytest.raises(TypeError): + SearchQuery(3, DenseVector([1, 2])) + + # Point Test def test_point_from_proto(): proto = vector_db_pb2.Point( id=vector_db_pb2.PointID(id=vector_db_pb2.UUID(value="point-123")), - vector=vector_db_pb2.DenseVector(values=[1, 2, 3]), + dense_vector=vector_db_pb2.DenseVector(values=[1, 2, 3]), + sparse_vector=vector_db_pb2.SparseVector(token_ids=[2], weights=[0.5]), payload=vector_db_pb2.Payload(content_type=vector_db_pb2.Text, content="hello"), ) point = Point.from_proto(proto) assert point.id == "point-123" - assert point.vector.values == [1.0, 2.0, 3.0] + assert point.dense_vector.values == [1.0, 2.0, 3.0] + assert point.sparse_vector == SparseVector([2], [0.5]) assert point.payload.content_type == ContentType.TEXT assert point.payload.content == "hello" -def test_point_from_proto_without_payload(): +def test_point_from_proto_with_missing_fields(): proto = vector_db_pb2.Point( id=vector_db_pb2.PointID(id=vector_db_pb2.UUID(value="p1")), - vector=vector_db_pb2.DenseVector(values=[1, 2, 3]), - payload=None, + dense_vector=vector_db_pb2.DenseVector(values=[1, 2, 3]), ) point = Point.from_proto(proto) - assert point.payload.content == "" + assert point.sparse_vector is None + assert point.payload is None diff --git a/client/python/vortexdb/__init__.py b/client/python/vortexdb/__init__.py index 9cc9dd4..123e06d 100644 --- a/client/python/vortexdb/__init__.py +++ b/client/python/vortexdb/__init__.py @@ -3,11 +3,15 @@ from vortexdb.client import VortexDB from vortexdb.async_client import AsyncVortexDB from vortexdb.models import ( + DenseQuery, DenseVector, Payload, Point, + PointInput, + SearchHit, Similarity, SearchQuery, + SparseVector, to_dense_vectors, ) from vortexdb.exceptions import ( @@ -15,6 +19,7 @@ AuthenticationError, NotFoundError, InvalidArgumentError, + AlreadyExistsError, TimeoutError, ServiceUnavailableError, InternalServerError, @@ -23,16 +28,21 @@ __all__ = [ "VortexDB", "AsyncVortexDB", + "DenseQuery", "DenseVector", "Payload", "Point", + "PointInput", + "SearchHit", "Similarity", "SearchQuery", + "SparseVector", "to_dense_vectors", "VortexDBError", "AuthenticationError", "NotFoundError", "InvalidArgumentError", + "AlreadyExistsError", "TimeoutError", "ServiceUnavailableError", "InternalServerError", diff --git a/client/python/vortexdb/_grpc_common.py b/client/python/vortexdb/_grpc_common.py index 2257f3d..6d4cc2b 100644 --- a/client/python/vortexdb/_grpc_common.py +++ b/client/python/vortexdb/_grpc_common.py @@ -3,6 +3,7 @@ import grpc from vortexdb.exceptions import ( + AlreadyExistsError, AuthenticationError, InternalServerError, InvalidArgumentError, @@ -29,6 +30,9 @@ def map_grpc_error(error: Any) -> VortexDBError: if code == grpc.StatusCode.INVALID_ARGUMENT: return InvalidArgumentError(error.details()) + if code == grpc.StatusCode.ALREADY_EXISTS: + return AlreadyExistsError(error.details()) + if code == grpc.StatusCode.DEADLINE_EXCEEDED: return TimeoutError(error.details()) diff --git a/client/python/vortexdb/async_client.py b/client/python/vortexdb/async_client.py index 3ce0aaf..7e70fb3 100644 --- a/client/python/vortexdb/async_client.py +++ b/client/python/vortexdb/async_client.py @@ -1,9 +1,15 @@ -from typing import List, Tuple - from vortexdb import protoutils as proto from vortexdb.async_connection import AsyncGRPCConnection from vortexdb.config import VortexDBConfig -from vortexdb.models import DenseVector, Payload, Point, Similarity, SearchQuery +from vortexdb.models import ( + DenseVector, + Payload, + Point, + PointInput, + SearchHit, + SearchQuery, + SparseVector, +) class AsyncVortexDB: @@ -25,42 +31,45 @@ def __init__( self._conn = AsyncGRPCConnection(self._config) - async def insert(self, *, vector: DenseVector, payload: Payload) -> str: + async def insert( + self, + *, + dense_vector: DenseVector | None = None, + sparse_vector: SparseVector | None = None, + payload: Payload | None = None, + point_id: str | None = None, + ) -> str: """ - Insert a vector with payload. + Insert a point with a dense vector, a sparse vector or both. Returns: point_id (str) """ - if not isinstance(vector, DenseVector): - raise TypeError( - "vector must be a DenseVector. Use: DenseVector([1.0, 2.0, 3.0])" - ) - request = proto.build_insert_request( - vector=vector, - payload=payload, + PointInput( + dense_vector=dense_vector, + sparse_vector=sparse_vector, + payload=payload, + point_id=point_id, + ) ) response = await self._conn.call( - self._conn.stub.InsertVector, + self._conn.stub.InsertPoint, request, ) return response.id.value - async def batch_insert( - self, *, items: list[tuple[DenseVector, Payload]] - ) -> list[str]: + async def batch_insert(self, *, points: list[PointInput]) -> list[str]: """ - Insert multiple vectors. + Insert multiple points. Returns: list of point_id (str) """ - request = proto.build_batch_insert_request(items=items) + request = proto.build_batch_insert_request(points=points) response = await self._conn.call( - self._conn.stub.InsertVectorsBatch, + self._conn.stub.InsertPointsBatch, request, ) - return [pid.id.value for pid in response.ids] async def get(self, *, point_id: str) -> Point | None: @@ -90,119 +99,34 @@ async def delete(self, *, point_id: str) -> None: request, ) - async def search( - self, - *, - vector: DenseVector | None = None, - similarity: Similarity | None = None, - limit: int | None = None, - query: SearchQuery | None = None, - ef: int | None = None, - ) -> List[Tuple[str, float]]: + async def search(self, *, query: SearchQuery) -> list[SearchHit]: """ - Search for nearest neighbors. - Returns: List of (point ID, score) tuples + Search with a dense or sparse query. + Returns: list of SearchHit (dense scores are distances, lower is better; + sparse scores are dot products, higher is better) """ - if query is not None: - if not isinstance(query, SearchQuery): - raise TypeError("query must be a SearchQuery") - vector = query.vector - similarity = query.similarity - limit = query.limit - else: - if not isinstance(vector, DenseVector): - raise TypeError( - "vector must be a DenseVector. Use: DenseVector([1.0, 2.0, 3.0])" - ) - if not isinstance(similarity, Similarity): - raise TypeError("similarity must be a Similarity enum") - if not isinstance(limit, int): - raise TypeError("limit must be an int") - - request = proto.build_search_request( - vector=vector, - similarity=similarity, - limit=limit, - ef=ef, - ) + request = proto.build_search_request(query) response = await self._conn.call( - self._conn.stub.SearchPoints, + self._conn.stub.Search, request, ) - - return [ - (pid.id.value, score) - for pid, score in zip(response.result_point_ids, response.scores) - ] + return proto.parse_search_response(response) async def batch_search( - self, - *, - queries, - similarity: Similarity | None = None, - limit: int | None = None, - ef: int | None = None, - ) -> List[List[Tuple[str, float]]]: + self, *, queries: list[SearchQuery] + ) -> list[list[SearchHit]]: """ - Flexible batch search. - - Accepts: - - List[SearchQuery] - - List[(DenseVector, Similarity, int)] - - List[(DenseVector, Similarity)] + global limit - - List[(DenseVector, int)] + global similarity - - List[DenseVector] + global similarity + limit + Run multiple dense and/or sparse queries in one call. + Returns: one list of SearchHit per query, in order """ - normalized = [] - - for i, q in enumerate(queries): - if ( - hasattr(q, "vector") - and hasattr(q, "similarity") - and hasattr(q, "limit") - ): - normalized.append((q.vector, q.similarity, q.limit)) - continue - - if isinstance(q, DenseVector): - if similarity is None or limit is None: - raise ValueError( - f"queries[{i}] requires global similarity and limit" - ) - normalized.append((q, similarity, limit)) - continue - - if isinstance(q, (list, tuple)): - if len(q) == 3: - normalized.append(q) - continue - if len(q) == 2: - a, b = q - - if isinstance(a, DenseVector) and isinstance(b, Similarity): - if limit is None: - raise ValueError(f"queries[{i}] missing global limit") - normalized.append((a, b, limit)) - continue - - if isinstance(a, DenseVector) and isinstance(b, int): - if similarity is None: - raise ValueError(f"queries[{i}] missing global similarity") - normalized.append((a, similarity, b)) - continue - - raise TypeError(f"Invalid query format at index {i}") - - request = proto.build_batch_search_request(queries=normalized, ef=ef) - response = await self._conn.call(self._conn.stub.SearchPointsBatch, request) - return [ - [ - (pid.id.value, score) - for pid, score in zip(result.result_point_ids, result.scores) - ] - for result in response.results - ] + request = proto.build_batch_search_request(queries=queries) + + response = await self._conn.call( + self._conn.stub.SearchBatch, + request, + ) + return [proto.parse_search_response(result) for result in response.results] async def close(self) -> None: """ diff --git a/client/python/vortexdb/client.py b/client/python/vortexdb/client.py index f38c934..7850aa0 100644 --- a/client/python/vortexdb/client.py +++ b/client/python/vortexdb/client.py @@ -1,13 +1,13 @@ -from typing import List, Tuple - from vortexdb.connection import GRPCConnection from vortexdb.config import VortexDBConfig from vortexdb.models import ( DenseVector, Payload, Point, - Similarity, + PointInput, + SearchHit, SearchQuery, + SparseVector, ) from vortexdb import protoutils as proto @@ -34,34 +34,43 @@ def __init__( # The basic operations - def insert(self, *, vector: DenseVector, payload: Payload) -> str: + def insert( + self, + *, + dense_vector: DenseVector | None = None, + sparse_vector: SparseVector | None = None, + payload: Payload | None = None, + point_id: str | None = None, + ) -> str: """ - Insert a vector with payload. + Insert a point with a dense vector, a sparse vector or both. Returns: point_id (str) """ - self._validate_dense_vector(vector) - request = proto.build_insert_request( - vector=vector, - payload=payload, + PointInput( + dense_vector=dense_vector, + sparse_vector=sparse_vector, + payload=payload, + point_id=point_id, + ) ) response = self._conn.call( - self._conn.stub.InsertVector, + self._conn.stub.InsertPoint, request, ) return response.id.value - def batch_insert(self, *, items: list[tuple[DenseVector, Payload]]) -> list[str]: + def batch_insert(self, *, points: list[PointInput]) -> list[str]: """ - Insert multiple vectors. + Insert multiple points. Returns: list of point_id (str) """ - request = proto.build_batch_insert_request(items=items) + request = proto.build_batch_insert_request(points=points) response = self._conn.call( - self._conn.stub.InsertVectorsBatch, + self._conn.stub.InsertPointsBatch, request, ) return [pid.id.value for pid in response.ids] @@ -93,121 +102,32 @@ def delete(self, *, point_id: str) -> None: request, ) - def search( - self, - *, - vector: DenseVector | None = None, - similarity: Similarity | None = None, - limit: int | None = None, - query: SearchQuery | None = None, - ef: int | None = None, - ) -> List[Tuple[str, float]]: + def search(self, *, query: SearchQuery) -> list[SearchHit]: """ - Search for nearest neighbors. - Returns: List of (point ID, score) tuples + Search with a dense or sparse query. + Returns: list of SearchHit (dense scores are distances, lower is better; + sparse scores are dot products, higher is better) """ - if query is not None: - if not isinstance(query, SearchQuery): - raise TypeError("query must be a SearchQuery") - vector = query.vector - similarity = query.similarity - limit = query.limit - else: - self._validate_dense_vector(vector) - if not isinstance(similarity, Similarity): - raise TypeError("similarity must be a Similarity enum") - if not isinstance(limit, int): - raise TypeError("limit must be an int") - - request = proto.build_search_request( - vector=vector, - similarity=similarity, - limit=limit, - ef=ef, - ) + request = proto.build_search_request(query) + response = self._conn.call( - self._conn.stub.SearchPoints, + self._conn.stub.Search, request, ) - return [ - (pid.id.value, score) - for pid, score in zip(response.result_point_ids, response.scores) - ] + return proto.parse_search_response(response) - def batch_search( - self, - *, - queries, - similarity: Similarity | None = None, - limit: int | None = None, - ef: int | None = None, - ) -> List[List[Tuple[str, float]]]: + def batch_search(self, *, queries: list[SearchQuery]) -> list[list[SearchHit]]: """ - Flexible batch search. - - Accepts: - - List[SearchQuery] - - List[(DenseVector, Similarity, int)] - - List[(DenseVector, Similarity)] + global limit - - List[(DenseVector, int)] + global similarity - - List[DenseVector] + global similarity + limit + Run multiple dense and/or sparse queries in one call. + Returns: one list of SearchHit per query, in order """ - normalized = [] - - for i, q in enumerate(queries): - if ( - hasattr(q, "vector") - and hasattr(q, "similarity") - and hasattr(q, "limit") - ): - normalized.append((q.vector, q.similarity, q.limit)) - continue - - if isinstance(q, DenseVector): - if similarity is None or limit is None: - raise ValueError( - f"queries[{i}] requires global similarity and limit" - ) - normalized.append((q, similarity, limit)) - continue - - if isinstance(q, (list, tuple)): - if len(q) == 3: - normalized.append(q) - continue - if len(q) == 2: - a, b = q - - if isinstance(a, DenseVector) and isinstance(b, Similarity): - if limit is None: - raise ValueError(f"queries[{i}] missing global limit") - normalized.append((a, b, limit)) - continue - - if isinstance(a, DenseVector) and isinstance(b, int): - if similarity is None: - raise ValueError(f"queries[{i}] missing global similarity") - normalized.append((a, similarity, b)) - continue - - raise TypeError(f"Invalid query format at index {i}") - - request = proto.build_batch_search_request(queries=normalized, ef=ef) - response = self._conn.call(self._conn.stub.SearchPointsBatch, request) - return [ - [ - (pid.id.value, score) - for pid, score in zip(result.result_point_ids, result.scores) - ] - for result in response.results - ] - - @staticmethod - def _validate_dense_vector(vector: DenseVector) -> None: - if not isinstance(vector, DenseVector): - raise TypeError( - "vector must be a DenseVector. Use: DenseVector([1.0, 2.0, 3.0])" - ) + request = proto.build_batch_search_request(queries=queries) + + response = self._conn.call( + self._conn.stub.SearchBatch, + request, + ) + return [proto.parse_search_response(result) for result in response.results] def close(self) -> None: """ diff --git a/client/python/vortexdb/exceptions.py b/client/python/vortexdb/exceptions.py index 07d914a..37ef093 100644 --- a/client/python/vortexdb/exceptions.py +++ b/client/python/vortexdb/exceptions.py @@ -14,6 +14,10 @@ class InvalidArgumentError(VortexDBError): """Invalid input was provided""" +class AlreadyExistsError(VortexDBError): + """A point with the given id already exists""" + + class TimeoutError(VortexDBError): """Request timed out while communicating with the server""" diff --git a/client/python/vortexdb/grpc/vector_db_pb2.py b/client/python/vortexdb/grpc/vector_db_pb2.py index 3cdf444..be3a175 100644 --- a/client/python/vortexdb/grpc/vector_db_pb2.py +++ b/client/python/vortexdb/grpc/vector_db_pb2.py @@ -25,41 +25,47 @@ from google.protobuf import empty_pb2 as google_dot_protobuf_dot_empty__pb2 -DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n\x0fvector-db.proto\x12\x08vectordb\x1a\x1bgoogle/protobuf/empty.proto\"\x15\n\x04UUID\x12\r\n\x05value\x18\x01 \x01(\t\"`\n\x13InsertVectorRequest\x12%\n\x06vector\x18\x01 \x01(\x0b\x32\x15.vectordb.DenseVector\x12\"\n\x07payload\x18\x02 \x01(\x0b\x32\x11.vectordb.Payload\"\x81\x01\n\rSearchRequest\x12+\n\x0cquery_vector\x18\x01 \x01(\x0b\x32\x15.vectordb.DenseVector\x12(\n\nsimilarity\x18\x02 \x01(\x0e\x32\x14.vectordb.Similarity\x12\r\n\x05limit\x18\x03 \x01(\x04\x12\n\n\x02\x65\x66\x18\x04 \x01(\x04\"M\n\x0eSearchResponse\x12+\n\x10result_point_ids\x18\x01 \x03(\x0b\x32\x11.vectordb.PointID\x12\x0e\n\x06scores\x18\x02 \x03(\x02\"\x1d\n\x0b\x44\x65nseVector\x12\x0e\n\x06values\x18\x01 \x03(\x02\"q\n\x05Point\x12\x1d\n\x02id\x18\x01 \x01(\x0b\x32\x11.vectordb.PointID\x12\"\n\x07payload\x18\x02 \x01(\x0b\x32\x11.vectordb.Payload\x12%\n\x06vector\x18\x03 \x01(\x0b\x32\x15.vectordb.DenseVector\"%\n\x07PointID\x12\x1a\n\x02id\x18\x01 \x01(\x0b\x32\x0e.vectordb.UUID\"G\n\x07Payload\x12+\n\x0c\x63ontent_type\x18\x01 \x01(\x0e\x32\x15.vectordb.ContentType\x12\x0f\n\x07\x63ontent\x18\x02 \x01(\t\"K\n\x19InsertVectorsBatchRequest\x12.\n\x07vectors\x18\x01 \x03(\x0b\x32\x1d.vectordb.InsertVectorRequest\"<\n\x1aInsertVectorsBatchResponse\x12\x1e\n\x03ids\x18\x01 \x03(\x0b\x32\x11.vectordb.PointID\"D\n\x18SearchPointsBatchRequest\x12(\n\x07queries\x18\x01 \x03(\x0b\x32\x17.vectordb.SearchRequest\"F\n\x19SearchPointsBatchResponse\x12)\n\x07results\x18\x01 \x03(\x0b\x32\x18.vectordb.SearchResponse*C\n\nSimilarity\x12\r\n\tEuclidean\x10\x00\x12\r\n\tManhattan\x10\x01\x12\x0b\n\x07Hamming\x10\x02\x12\n\n\x06\x43osine\x10\x03*\"\n\x0b\x43ontentType\x12\t\n\x05Image\x10\x00\x12\x08\n\x04Text\x10\x01\x32\xc4\x03\n\x08VectorDB\x12\x42\n\x0cInsertVector\x12\x1d.vectordb.InsertVectorRequest\x1a\x11.vectordb.PointID\"\x00\x12:\n\x0b\x44\x65letePoint\x12\x11.vectordb.PointID\x1a\x16.google.protobuf.Empty\"\x00\x12\x30\n\x08GetPoint\x12\x11.vectordb.PointID\x1a\x0f.vectordb.Point\"\x00\x12\x43\n\x0cSearchPoints\x12\x17.vectordb.SearchRequest\x1a\x18.vectordb.SearchResponse\"\x00\x12\x61\n\x12InsertVectorsBatch\x12#.vectordb.InsertVectorsBatchRequest\x1a$.vectordb.InsertVectorsBatchResponse\"\x00\x12^\n\x11SearchPointsBatch\x12\".vectordb.SearchPointsBatchRequest\x1a#.vectordb.SearchPointsBatchResponse\"\x00\x62\x06proto3') +DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n\x0fvector-db.proto\x12\x08vectordb\x1a\x1bgoogle/protobuf/empty.proto\"\x15\n\x04UUID\x12\r\n\x05value\x18\x01 \x01(\t\"\xab\x01\n\nPointInput\x12\x1d\n\x02id\x18\x01 \x01(\x0b\x32\x11.vectordb.PointID\x12+\n\x0c\x64\x65nse_vector\x18\x02 \x01(\x0b\x32\x15.vectordb.DenseVector\x12-\n\rsparse_vector\x18\x03 \x01(\x0b\x32\x16.vectordb.SparseVector\x12\"\n\x07payload\x18\x04 \x01(\x0b\x32\x11.vectordb.Payload\"u\n\nDenseQuery\x12%\n\x06vector\x18\x01 \x01(\x0b\x32\x15.vectordb.DenseVector\x12(\n\nsimilarity\x18\x02 \x01(\x0e\x32\x14.vectordb.Similarity\x12\x0f\n\x02\x65\x66\x18\x03 \x01(\x04H\x00\x88\x01\x01\x42\x05\n\x03_ef\"x\n\rSearchRequest\x12\r\n\x05limit\x18\x01 \x01(\x04\x12%\n\x05\x64\x65nse\x18\x02 \x01(\x0b\x32\x14.vectordb.DenseQueryH\x00\x12(\n\x06sparse\x18\x03 \x01(\x0b\x32\x16.vectordb.SparseVectorH\x00\x42\x07\n\x05query\"9\n\tSearchHit\x12\x1d\n\x02id\x18\x01 \x01(\x0b\x32\x11.vectordb.PointID\x12\r\n\x05score\x18\x02 \x01(\x02\"3\n\x0eSearchResponse\x12!\n\x04hits\x18\x01 \x03(\x0b\x32\x13.vectordb.SearchHit\"\x1d\n\x0b\x44\x65nseVector\x12\x0e\n\x06values\x18\x01 \x03(\x02\"2\n\x0cSparseVector\x12\x11\n\ttoken_ids\x18\x01 \x03(\r\x12\x0f\n\x07weights\x18\x02 \x03(\x02\"\xa6\x01\n\x05Point\x12\x1d\n\x02id\x18\x01 \x01(\x0b\x32\x11.vectordb.PointID\x12\"\n\x07payload\x18\x02 \x01(\x0b\x32\x11.vectordb.Payload\x12+\n\x0c\x64\x65nse_vector\x18\x03 \x01(\x0b\x32\x15.vectordb.DenseVector\x12-\n\rsparse_vector\x18\x04 \x01(\x0b\x32\x16.vectordb.SparseVector\"%\n\x07PointID\x12\x1a\n\x02id\x18\x01 \x01(\x0b\x32\x0e.vectordb.UUID\"G\n\x07Payload\x12+\n\x0c\x63ontent_type\x18\x01 \x01(\x0e\x32\x15.vectordb.ContentType\x12\x0f\n\x07\x63ontent\x18\x02 \x01(\t\"@\n\x18InsertPointsBatchRequest\x12$\n\x06points\x18\x01 \x03(\x0b\x32\x14.vectordb.PointInput\";\n\x19InsertPointsBatchResponse\x12\x1e\n\x03ids\x18\x01 \x03(\x0b\x32\x11.vectordb.PointID\">\n\x12SearchBatchRequest\x12(\n\x07queries\x18\x01 \x03(\x0b\x32\x17.vectordb.SearchRequest\"@\n\x13SearchBatchResponse\x12)\n\x07results\x18\x01 \x03(\x0b\x32\x18.vectordb.SearchResponse*C\n\nSimilarity\x12\r\n\tEuclidean\x10\x00\x12\r\n\tManhattan\x10\x01\x12\x0b\n\x07Hamming\x10\x02\x12\n\n\x06\x43osine\x10\x03*\"\n\x0b\x43ontentType\x12\t\n\x05Image\x10\x00\x12\x08\n\x04Text\x10\x01\x32\x9f\x03\n\x08VectorDB\x12\x38\n\x0bInsertPoint\x12\x14.vectordb.PointInput\x1a\x11.vectordb.PointID\"\x00\x12:\n\x0b\x44\x65letePoint\x12\x11.vectordb.PointID\x1a\x16.google.protobuf.Empty\"\x00\x12\x30\n\x08GetPoint\x12\x11.vectordb.PointID\x1a\x0f.vectordb.Point\"\x00\x12=\n\x06Search\x12\x17.vectordb.SearchRequest\x1a\x18.vectordb.SearchResponse\"\x00\x12^\n\x11InsertPointsBatch\x12\".vectordb.InsertPointsBatchRequest\x1a#.vectordb.InsertPointsBatchResponse\"\x00\x12L\n\x0bSearchBatch\x12\x1c.vectordb.SearchBatchRequest\x1a\x1d.vectordb.SearchBatchResponse\"\x00\x62\x06proto3') _globals = globals() _builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, _globals) _builder.BuildTopDescriptorsAndMessages(DESCRIPTOR, 'vector_db_pb2', _globals) if not _descriptor._USE_C_DESCRIPTORS: DESCRIPTOR._loaded_options = None - _globals['_SIMILARITY']._serialized_start=929 - _globals['_SIMILARITY']._serialized_end=996 - _globals['_CONTENTTYPE']._serialized_start=998 - _globals['_CONTENTTYPE']._serialized_end=1032 + _globals['_SIMILARITY']._serialized_start=1229 + _globals['_SIMILARITY']._serialized_end=1296 + _globals['_CONTENTTYPE']._serialized_start=1298 + _globals['_CONTENTTYPE']._serialized_end=1332 _globals['_UUID']._serialized_start=58 _globals['_UUID']._serialized_end=79 - _globals['_INSERTVECTORREQUEST']._serialized_start=81 - _globals['_INSERTVECTORREQUEST']._serialized_end=177 - _globals['_SEARCHREQUEST']._serialized_start=180 - _globals['_SEARCHREQUEST']._serialized_end=309 - _globals['_SEARCHRESPONSE']._serialized_start=311 - _globals['_SEARCHRESPONSE']._serialized_end=388 - _globals['_DENSEVECTOR']._serialized_start=390 - _globals['_DENSEVECTOR']._serialized_end=419 - _globals['_POINT']._serialized_start=421 - _globals['_POINT']._serialized_end=534 - _globals['_POINTID']._serialized_start=536 - _globals['_POINTID']._serialized_end=573 - _globals['_PAYLOAD']._serialized_start=575 - _globals['_PAYLOAD']._serialized_end=646 - _globals['_INSERTVECTORSBATCHREQUEST']._serialized_start=648 - _globals['_INSERTVECTORSBATCHREQUEST']._serialized_end=723 - _globals['_INSERTVECTORSBATCHRESPONSE']._serialized_start=725 - _globals['_INSERTVECTORSBATCHRESPONSE']._serialized_end=785 - _globals['_SEARCHPOINTSBATCHREQUEST']._serialized_start=787 - _globals['_SEARCHPOINTSBATCHREQUEST']._serialized_end=855 - _globals['_SEARCHPOINTSBATCHRESPONSE']._serialized_start=857 - _globals['_SEARCHPOINTSBATCHRESPONSE']._serialized_end=927 - _globals['_VECTORDB']._serialized_start=1035 - _globals['_VECTORDB']._serialized_end=1487 + _globals['_POINTINPUT']._serialized_start=82 + _globals['_POINTINPUT']._serialized_end=253 + _globals['_DENSEQUERY']._serialized_start=255 + _globals['_DENSEQUERY']._serialized_end=372 + _globals['_SEARCHREQUEST']._serialized_start=374 + _globals['_SEARCHREQUEST']._serialized_end=494 + _globals['_SEARCHHIT']._serialized_start=496 + _globals['_SEARCHHIT']._serialized_end=553 + _globals['_SEARCHRESPONSE']._serialized_start=555 + _globals['_SEARCHRESPONSE']._serialized_end=606 + _globals['_DENSEVECTOR']._serialized_start=608 + _globals['_DENSEVECTOR']._serialized_end=637 + _globals['_SPARSEVECTOR']._serialized_start=639 + _globals['_SPARSEVECTOR']._serialized_end=689 + _globals['_POINT']._serialized_start=692 + _globals['_POINT']._serialized_end=858 + _globals['_POINTID']._serialized_start=860 + _globals['_POINTID']._serialized_end=897 + _globals['_PAYLOAD']._serialized_start=899 + _globals['_PAYLOAD']._serialized_end=970 + _globals['_INSERTPOINTSBATCHREQUEST']._serialized_start=972 + _globals['_INSERTPOINTSBATCHREQUEST']._serialized_end=1036 + _globals['_INSERTPOINTSBATCHRESPONSE']._serialized_start=1038 + _globals['_INSERTPOINTSBATCHRESPONSE']._serialized_end=1097 + _globals['_SEARCHBATCHREQUEST']._serialized_start=1099 + _globals['_SEARCHBATCHREQUEST']._serialized_end=1161 + _globals['_SEARCHBATCHRESPONSE']._serialized_start=1163 + _globals['_SEARCHBATCHRESPONSE']._serialized_end=1227 + _globals['_VECTORDB']._serialized_start=1335 + _globals['_VECTORDB']._serialized_end=1750 # @@protoc_insertion_point(module_scope) diff --git a/client/python/vortexdb/grpc/vector_db_pb2_grpc.py b/client/python/vortexdb/grpc/vector_db_pb2_grpc.py index 1442f69..cb43ba5 100644 --- a/client/python/vortexdb/grpc/vector_db_pb2_grpc.py +++ b/client/python/vortexdb/grpc/vector_db_pb2_grpc.py @@ -35,9 +35,9 @@ def __init__(self, channel): Args: channel: A grpc.Channel. """ - self.InsertVector = channel.unary_unary( - '/vectordb.VectorDB/InsertVector', - request_serializer=vector__db__pb2.InsertVectorRequest.SerializeToString, + self.InsertPoint = channel.unary_unary( + '/vectordb.VectorDB/InsertPoint', + request_serializer=vector__db__pb2.PointInput.SerializeToString, response_deserializer=vector__db__pb2.PointID.FromString, _registered_method=True) self.DeletePoint = channel.unary_unary( @@ -50,61 +50,61 @@ def __init__(self, channel): request_serializer=vector__db__pb2.PointID.SerializeToString, response_deserializer=vector__db__pb2.Point.FromString, _registered_method=True) - self.SearchPoints = channel.unary_unary( - '/vectordb.VectorDB/SearchPoints', + self.Search = channel.unary_unary( + '/vectordb.VectorDB/Search', request_serializer=vector__db__pb2.SearchRequest.SerializeToString, response_deserializer=vector__db__pb2.SearchResponse.FromString, _registered_method=True) - self.InsertVectorsBatch = channel.unary_unary( - '/vectordb.VectorDB/InsertVectorsBatch', - request_serializer=vector__db__pb2.InsertVectorsBatchRequest.SerializeToString, - response_deserializer=vector__db__pb2.InsertVectorsBatchResponse.FromString, + self.InsertPointsBatch = channel.unary_unary( + '/vectordb.VectorDB/InsertPointsBatch', + request_serializer=vector__db__pb2.InsertPointsBatchRequest.SerializeToString, + response_deserializer=vector__db__pb2.InsertPointsBatchResponse.FromString, _registered_method=True) - self.SearchPointsBatch = channel.unary_unary( - '/vectordb.VectorDB/SearchPointsBatch', - request_serializer=vector__db__pb2.SearchPointsBatchRequest.SerializeToString, - response_deserializer=vector__db__pb2.SearchPointsBatchResponse.FromString, + self.SearchBatch = channel.unary_unary( + '/vectordb.VectorDB/SearchBatch', + request_serializer=vector__db__pb2.SearchBatchRequest.SerializeToString, + response_deserializer=vector__db__pb2.SearchBatchResponse.FromString, _registered_method=True) class VectorDBServicer: """Missing associated documentation comment in .proto file.""" - def InsertVector(self, request, context): - """Insert a vector with a payload and return the assigned PointID + def InsertPoint(self, request, context): + """Insert a point with a dense and/or sparse vector and return its PointID """ context.set_code(grpc.StatusCode.UNIMPLEMENTED) context.set_details('Method not implemented!') raise NotImplementedError('Method not implemented!') def DeletePoint(self, request, context): - """Delete a vector by its PointID + """Delete a point by its PointID """ context.set_code(grpc.StatusCode.UNIMPLEMENTED) context.set_details('Method not implemented!') raise NotImplementedError('Method not implemented!') def GetPoint(self, request, context): - """Get a vector and its payload by PointID + """Get a point with its vectors and payload by PointID """ context.set_code(grpc.StatusCode.UNIMPLEMENTED) context.set_details('Method not implemented!') raise NotImplementedError('Method not implemented!') - def SearchPoints(self, request, context): - """Search for the k nearest vectors to a target vector given a distance function + def Search(self, request, context): + """Search for the nearest points to a dense or sparse query """ context.set_code(grpc.StatusCode.UNIMPLEMENTED) context.set_details('Method not implemented!') raise NotImplementedError('Method not implemented!') - def InsertVectorsBatch(self, request, context): + def InsertPointsBatch(self, request, context): """Missing associated documentation comment in .proto file.""" context.set_code(grpc.StatusCode.UNIMPLEMENTED) context.set_details('Method not implemented!') raise NotImplementedError('Method not implemented!') - def SearchPointsBatch(self, request, context): + def SearchBatch(self, request, context): """Missing associated documentation comment in .proto file.""" context.set_code(grpc.StatusCode.UNIMPLEMENTED) context.set_details('Method not implemented!') @@ -113,9 +113,9 @@ def SearchPointsBatch(self, request, context): def add_VectorDBServicer_to_server(servicer, server): rpc_method_handlers = { - 'InsertVector': grpc.unary_unary_rpc_method_handler( - servicer.InsertVector, - request_deserializer=vector__db__pb2.InsertVectorRequest.FromString, + 'InsertPoint': grpc.unary_unary_rpc_method_handler( + servicer.InsertPoint, + request_deserializer=vector__db__pb2.PointInput.FromString, response_serializer=vector__db__pb2.PointID.SerializeToString, ), 'DeletePoint': grpc.unary_unary_rpc_method_handler( @@ -128,20 +128,20 @@ def add_VectorDBServicer_to_server(servicer, server): request_deserializer=vector__db__pb2.PointID.FromString, response_serializer=vector__db__pb2.Point.SerializeToString, ), - 'SearchPoints': grpc.unary_unary_rpc_method_handler( - servicer.SearchPoints, + 'Search': grpc.unary_unary_rpc_method_handler( + servicer.Search, request_deserializer=vector__db__pb2.SearchRequest.FromString, response_serializer=vector__db__pb2.SearchResponse.SerializeToString, ), - 'InsertVectorsBatch': grpc.unary_unary_rpc_method_handler( - servicer.InsertVectorsBatch, - request_deserializer=vector__db__pb2.InsertVectorsBatchRequest.FromString, - response_serializer=vector__db__pb2.InsertVectorsBatchResponse.SerializeToString, + 'InsertPointsBatch': grpc.unary_unary_rpc_method_handler( + servicer.InsertPointsBatch, + request_deserializer=vector__db__pb2.InsertPointsBatchRequest.FromString, + response_serializer=vector__db__pb2.InsertPointsBatchResponse.SerializeToString, ), - 'SearchPointsBatch': grpc.unary_unary_rpc_method_handler( - servicer.SearchPointsBatch, - request_deserializer=vector__db__pb2.SearchPointsBatchRequest.FromString, - response_serializer=vector__db__pb2.SearchPointsBatchResponse.SerializeToString, + 'SearchBatch': grpc.unary_unary_rpc_method_handler( + servicer.SearchBatch, + request_deserializer=vector__db__pb2.SearchBatchRequest.FromString, + response_serializer=vector__db__pb2.SearchBatchResponse.SerializeToString, ), } generic_handler = grpc.method_handlers_generic_handler( @@ -155,7 +155,7 @@ class VectorDB: """Missing associated documentation comment in .proto file.""" @staticmethod - def InsertVector(request, + def InsertPoint(request, target, options=(), channel_credentials=None, @@ -168,8 +168,8 @@ def InsertVector(request, return grpc.experimental.unary_unary( request, target, - '/vectordb.VectorDB/InsertVector', - vector__db__pb2.InsertVectorRequest.SerializeToString, + '/vectordb.VectorDB/InsertPoint', + vector__db__pb2.PointInput.SerializeToString, vector__db__pb2.PointID.FromString, options, channel_credentials, @@ -236,7 +236,7 @@ def GetPoint(request, _registered_method=True) @staticmethod - def SearchPoints(request, + def Search(request, target, options=(), channel_credentials=None, @@ -249,7 +249,7 @@ def SearchPoints(request, return grpc.experimental.unary_unary( request, target, - '/vectordb.VectorDB/SearchPoints', + '/vectordb.VectorDB/Search', vector__db__pb2.SearchRequest.SerializeToString, vector__db__pb2.SearchResponse.FromString, options, @@ -263,7 +263,7 @@ def SearchPoints(request, _registered_method=True) @staticmethod - def InsertVectorsBatch(request, + def InsertPointsBatch(request, target, options=(), channel_credentials=None, @@ -276,9 +276,9 @@ def InsertVectorsBatch(request, return grpc.experimental.unary_unary( request, target, - '/vectordb.VectorDB/InsertVectorsBatch', - vector__db__pb2.InsertVectorsBatchRequest.SerializeToString, - vector__db__pb2.InsertVectorsBatchResponse.FromString, + '/vectordb.VectorDB/InsertPointsBatch', + vector__db__pb2.InsertPointsBatchRequest.SerializeToString, + vector__db__pb2.InsertPointsBatchResponse.FromString, options, channel_credentials, insecure, @@ -290,7 +290,7 @@ def InsertVectorsBatch(request, _registered_method=True) @staticmethod - def SearchPointsBatch(request, + def SearchBatch(request, target, options=(), channel_credentials=None, @@ -303,9 +303,9 @@ def SearchPointsBatch(request, return grpc.experimental.unary_unary( request, target, - '/vectordb.VectorDB/SearchPointsBatch', - vector__db__pb2.SearchPointsBatchRequest.SerializeToString, - vector__db__pb2.SearchPointsBatchResponse.FromString, + '/vectordb.VectorDB/SearchBatch', + vector__db__pb2.SearchBatchRequest.SerializeToString, + vector__db__pb2.SearchBatchResponse.FromString, options, channel_credentials, insecure, diff --git a/client/python/vortexdb/models.py b/client/python/vortexdb/models.py index 0ba1e6c..1e96808 100644 --- a/client/python/vortexdb/models.py +++ b/client/python/vortexdb/models.py @@ -1,3 +1,4 @@ +import math from dataclasses import dataclass from enum import Enum from typing import List @@ -99,49 +100,191 @@ def to_proto(self) -> vector_db_pb2.Payload: ) +@dataclass(frozen=True) +class SparseVector: + token_ids: List[int] + weights: List[float] + + def __post_init__(self): + if not isinstance(self.token_ids, (list, tuple)) or not isinstance( + self.weights, (list, tuple) + ): + raise TypeError("SparseVector expects lists of token_ids and weights") + + if len(self.token_ids) != len(self.weights): + raise ValueError( + "SparseVector token_ids and weights must have equal length" + ) + + if not self.token_ids: + raise ValueError("SparseVector cannot be empty") + + for t in self.token_ids: + if not isinstance(t, int) or isinstance(t, bool) or t < 0: + raise TypeError("SparseVector token_ids must be non-negative integers") + + for w in self.weights: + if not isinstance(w, (int, float)): + raise TypeError("SparseVector weights must be numeric (int or float)") + if not math.isfinite(w): + raise ValueError("SparseVector weights must be finite") + if w < 0: + raise ValueError("SparseVector weights must not be negative") + + if any(a >= b for a, b in zip(self.token_ids, self.token_ids[1:])): + raise ValueError("SparseVector token_ids must be strictly increasing") + + object.__setattr__(self, "token_ids", list(self.token_ids)) + object.__setattr__(self, "weights", [float(w) for w in self.weights]) + + def to_proto(self) -> vector_db_pb2.SparseVector: + return vector_db_pb2.SparseVector( + token_ids=self.token_ids, weights=self.weights + ) + + +@dataclass(frozen=True) +class PointInput: + dense_vector: DenseVector | None = None + sparse_vector: SparseVector | None = None + payload: Payload | None = None + point_id: str | None = None + + def __post_init__(self): + if self.dense_vector is None and self.sparse_vector is None: + raise ValueError("PointInput needs a dense_vector, a sparse_vector or both") + + if self.dense_vector is not None and not isinstance( + self.dense_vector, DenseVector + ): + raise TypeError( + "dense_vector must be a DenseVector. Use: DenseVector([1.0, 2.0, 3.0])" + ) + + if self.sparse_vector is not None and not isinstance( + self.sparse_vector, SparseVector + ): + raise TypeError("sparse_vector must be a SparseVector") + + if self.payload is not None and not isinstance(self.payload, Payload): + raise TypeError("payload must be a Payload") + + def to_proto(self) -> vector_db_pb2.PointInput: + return vector_db_pb2.PointInput( + id=( + vector_db_pb2.PointID(id=vector_db_pb2.UUID(value=self.point_id)) + if self.point_id is not None + else None + ), + dense_vector=self.dense_vector.to_proto() if self.dense_vector else None, + sparse_vector=self.sparse_vector.to_proto() if self.sparse_vector else None, + payload=self.payload.to_proto() if self.payload else None, + ) + + @dataclass(frozen=True) class Point: id: str - vector: DenseVector - payload: Payload + dense_vector: DenseVector | None + sparse_vector: SparseVector | None + payload: Payload | None @staticmethod def from_proto(proto: vector_db_pb2.Point) -> "Point": - payload = proto.payload - if payload is None: - payload_obj = Payload.text("") - else: - payload_obj = Payload( - content_type=ContentType.from_proto(payload.content_type), - content=payload.content, + dense_vector = None + if proto.HasField("dense_vector"): + dense_vector = DenseVector(list(proto.dense_vector.values)) + + sparse_vector = None + if proto.HasField("sparse_vector"): + sparse_vector = SparseVector( + list(proto.sparse_vector.token_ids), + list(proto.sparse_vector.weights), + ) + + payload = None + if proto.HasField("payload"): + payload = Payload( + content_type=ContentType.from_proto(proto.payload.content_type), + content=proto.payload.content, ) return Point( id=proto.id.id.value, - vector=DenseVector(list(proto.vector.values)), - payload=payload_obj, + dense_vector=dense_vector, + sparse_vector=sparse_vector, + payload=payload, ) def pretty(self) -> str: + dense = self.dense_vector.values if self.dense_vector else None + sparse = ( + dict(zip(self.sparse_vector.token_ids, self.sparse_vector.weights)) + if self.sparse_vector + else None + ) + payload_type = self.payload.content_type.name if self.payload else None + payload = self.payload.content if self.payload else None return ( f"\nPoint:\n id = {self.id},\n" - f" vector_dim = {len(self.vector.values)},\n" - f" vector = {self.vector},\n" - f" payload_type = {self.payload.content_type.name},\n" - f" payload = '{self.payload.content}'" + f" dense_vector = {dense},\n" + f" sparse_vector = {sparse},\n" + f" payload_type = {payload_type},\n" + f" payload = {payload!r}" ) -# I added this because using tuples will get messy if we increase fields in a search query @dataclass(frozen=True) -class SearchQuery: +class DenseQuery: vector: DenseVector similarity: Similarity + ef: int | None = None + + def __post_init__(self): + if not isinstance(self.vector, DenseVector): + raise TypeError( + "vector must be a DenseVector. Use: DenseVector([1.0, 2.0, 3.0])" + ) + if not isinstance(self.similarity, Similarity): + raise TypeError("similarity must be a Similarity enum") + if self.ef is not None and not isinstance(self.ef, int): + raise TypeError("ef must be an int") + + def to_proto(self) -> vector_db_pb2.DenseQuery: + return vector_db_pb2.DenseQuery( + vector=self.vector.to_proto(), + similarity=self.similarity.to_proto(), + ef=self.ef, + ) + + +# I added this because using tuples will get messy if we increase fields in a search query +@dataclass(frozen=True) +class SearchQuery: limit: int + query: DenseQuery | SparseVector + + def __post_init__(self): + if not isinstance(self.limit, int): + raise TypeError("limit must be an int") + if not isinstance(self.query, (DenseQuery, SparseVector)): + raise TypeError("query must be a DenseQuery or a SparseVector") def to_proto(self) -> vector_db_pb2.SearchRequest: + if isinstance(self.query, DenseQuery): + return vector_db_pb2.SearchRequest( + limit=self.limit, dense=self.query.to_proto() + ) return vector_db_pb2.SearchRequest( - query_vector=self.vector.to_proto(), - similarity=self.similarity.to_proto(), - limit=self.limit, + limit=self.limit, sparse=self.query.to_proto() ) + + +@dataclass(frozen=True) +class SearchHit: + id: str + score: float + + @staticmethod + def from_proto(proto: vector_db_pb2.SearchHit) -> "SearchHit": + return SearchHit(id=proto.id.id.value, score=proto.score) diff --git a/client/python/vortexdb/protoutils.py b/client/python/vortexdb/protoutils.py index c6aae69..d176f25 100644 --- a/client/python/vortexdb/protoutils.py +++ b/client/python/vortexdb/protoutils.py @@ -1,106 +1,62 @@ from vortexdb.grpc import vector_db_pb2 -from vortexdb.models import DenseVector, Payload, Similarity +from vortexdb.models import PointInput, SearchHit, SearchQuery -def build_insert_request( - *, - vector: DenseVector, - payload: Payload, -) -> vector_db_pb2.InsertVectorRequest: - return vector_db_pb2.InsertVectorRequest( - vector=vector.to_proto(), - payload=payload.to_proto(), - ) +def build_insert_request(point: PointInput) -> vector_db_pb2.PointInput: + if not isinstance(point, PointInput): + raise TypeError("point must be a PointInput") + + return point.to_proto() def build_batch_insert_request( *, - items: list[tuple[DenseVector, Payload]], -) -> vector_db_pb2.InsertVectorsBatchRequest: - if not isinstance(items, (list, tuple)): - raise TypeError("Items must be a list of (DenseVector, Payload) tuples") - - if not items: - raise ValueError("Items cannot be empty") + points: list[PointInput], +) -> vector_db_pb2.InsertPointsBatchRequest: + if not isinstance(points, (list, tuple)): + raise TypeError("points must be a list of PointInput") - requests = [] + if not points: + raise ValueError("points cannot be empty") - for i, pair in enumerate(items): - if not isinstance(pair, (list, tuple)) or len(pair) != 2: - raise TypeError(f"items[{i}] must be a tuple of (DenseVector, Payload)") + for i, point in enumerate(points): + if not isinstance(point, PointInput): + raise TypeError(f"points[{i}] must be a PointInput") - vector, payload = pair - if not isinstance(vector, DenseVector): - raise TypeError( - f"items[{i}][0] must be a DenseVectorUse: DenseVector([1.0, 2.0, 3.0])" - ) - - if not isinstance(payload, Payload): - raise TypeError(f"items[{i}][1] must be Payload") - - requests.append(build_insert_request(vector=vector, payload=payload)) - - return vector_db_pb2.InsertVectorsBatchRequest(vectors=requests) + return vector_db_pb2.InsertPointsBatchRequest( + points=[point.to_proto() for point in points] + ) def build_point_id_request(point_id: str) -> vector_db_pb2.PointID: return vector_db_pb2.PointID(id=vector_db_pb2.UUID(value=point_id)) -def build_search_request( - *, - vector: DenseVector, - similarity: Similarity, - limit: int, - ef: int | None = None, -) -> vector_db_pb2.SearchRequest: - return vector_db_pb2.SearchRequest( - query_vector=vector.to_proto(), - similarity=similarity.to_proto(), - limit=limit, - ef=ef or 0, - ) +def build_search_request(query: SearchQuery) -> vector_db_pb2.SearchRequest: + if not isinstance(query, SearchQuery): + raise TypeError("query must be a SearchQuery") + + return query.to_proto() def build_batch_search_request( *, - queries: list[tuple[DenseVector, Similarity, int]], - ef: int | None = None, -) -> vector_db_pb2.SearchPointsBatchRequest: + queries: list[SearchQuery], +) -> vector_db_pb2.SearchBatchRequest: if not isinstance(queries, (list, tuple)): - raise TypeError( - "Queries must be a list of (DenseVector, Similarity, Limit (int)) tuples" - ) + raise TypeError("queries must be a list of SearchQuery") if not queries: - raise ValueError("Queries cannot be empty") - - requests = [] - - for i, trio in enumerate(queries): - if not isinstance(trio, (list, tuple)) or len(trio) != 3: - raise TypeError( - f"queries[{i}] must be a tuple of (DenseVector, Similarity, Limit(int))" - ) - - vector, similarity, limit = trio - if not isinstance(vector, DenseVector): - raise TypeError( - f"queries[{i}][0] must be a DenseVector" - "Use: DenseVector([1.0, 2.0, 3.0])" - ) - if not isinstance(similarity, Similarity): - raise TypeError(f"queries[{i}][1] must be Similarity") - if not isinstance(limit, int): - raise TypeError(f"queries[{i}][2] must be an integer value") - - requests.append( - vector_db_pb2.SearchRequest( - query_vector=vector.to_proto(), - similarity=similarity.to_proto(), - limit=limit, - ef=ef, - ) - ) - - return vector_db_pb2.SearchPointsBatchRequest(queries=requests) + raise ValueError("queries cannot be empty") + + for i, query in enumerate(queries): + if not isinstance(query, SearchQuery): + raise TypeError(f"queries[{i}] must be a SearchQuery") + + return vector_db_pb2.SearchBatchRequest( + queries=[query.to_proto() for query in queries] + ) + + +def parse_search_response(response: vector_db_pb2.SearchResponse) -> list[SearchHit]: + return [SearchHit.from_proto(hit) for hit in response.hits] diff --git a/crates/api/src/error.rs b/crates/api/src/error.rs index 622e58b..f0cecaa 100644 --- a/crates/api/src/error.rs +++ b/crates/api/src/error.rs @@ -1,4 +1,4 @@ -use defs::{Dimension, PointId}; +use defs::{Dimension, PointId, SparseError}; use snafu::prelude::*; #[derive(Debug, Snafu)] @@ -21,6 +21,15 @@ pub enum ApiError { #[snafu(display("Point {id} not found"))] PointNotFound { id: PointId }, + #[snafu(display("Point {id} has neither a dense nor a sparse vector"))] + MissingVector { id: PointId }, + + #[snafu(display("Point {id} already exists"))] + PointAlreadyExists { id: PointId }, + + #[snafu(display("Invalid sparse vector: {source}"))] + InvalidSparseVector { source: SparseError }, + #[snafu(display("Invalid search limit: {limit}"))] InvalidSearchLimit { limit: usize }, diff --git a/crates/api/src/lib.rs b/crates/api/src/lib.rs index 5377dde..5d399bf 100644 --- a/crates/api/src/lib.rs +++ b/crates/api/src/lib.rs @@ -1,7 +1,9 @@ -use defs::{DbError, Dimension, IndexedVector, SearchQueryInput, Similarity, SnapshottableDb}; -use defs::{DenseVector, Payload, Point, PointId, PointInput}; +use defs::{DbError, Dimension, IndexedVector, Similarity, SnapshottableDb}; +use defs::{Point, PointId, PointInput, QueryVector, SearchHit, SearchQuery}; use index::hnsw::{HnswConfig, HnswIndex}; +use index::inverted::InvertedIndex; use index::kd_tree::{KDTree, KDTreeConfig}; +use std::collections::HashSet; use std::path::{Path, PathBuf}; use tempfile::tempdir; // use std::sync::atomic::{AtomicU64, Ordering}; @@ -9,13 +11,14 @@ use std::sync::{Arc, RwLock}; use index::flat::index::FlatIndex; use index::{IndexType, VectorIndex}; +use snafu::ResultExt; use snapshot::Snapshot; -use storage::{StorageEngine, StorageType, VectorPage, create_storage_engine}; +use storage::{PointPage, StorageEngine, StorageType, create_storage_engine}; use uuid::Uuid; pub mod error; -pub use error::{ApiError, Result}; +pub use error::{ApiError, InvalidSparseVectorSnafu, Result}; // static NEXT_ID: AtomicU64 = AtomicU64::new(1); @@ -27,9 +30,12 @@ fn generate_point_id() -> PointId { Uuid::new_v4() } +const MAX_SEARCH_LIMIT: usize = 10_000; + pub struct VectorDb { storage: Arc, index: Arc>, // Using a RwLock instead of Mutex to improve concurrency + sparse_index: RwLock, dimension: usize, } @@ -42,59 +48,73 @@ impl VectorDb { Self { storage, index, + sparse_index: RwLock::new(InvertedIndex::new()), dimension, } } - //TODO: Make this an atomic operation - pub fn insert(&self, vector: DenseVector, payload: Payload) -> Result { - if vector.len() != self.dimension { + fn validate_point(&self, point: &PointInput, id: PointId) -> Result<()> { + if point.dense_vector.is_none() && point.sparse_vector.is_none() { + return Err(ApiError::MissingVector { id }); + } + + if let Some(ref v) = point.dense_vector + && v.len() != self.dimension + { return Err(ApiError::DimensionMismatch { expected: self.dimension, - got: vector.len(), + got: v.len(), }); } - // Generate a new point id - let point_id = generate_point_id(); - self.storage - .insert_point(point_id, Some(vector.clone()), Some(payload))?; - // Get write lock on the index - let mut index = self.index.write().map_err(|_| ApiError::LockError)?; - index.insert(IndexedVector { - vector, - id: point_id, - })?; + if let Some(ref sv) = point.sparse_vector { + sv.validate().context(InvalidSparseVectorSnafu)?; + } - Ok(point_id) + Ok(()) } + pub fn insert(&self, point: PointInput) -> Result { + let mut ids = self.insert_batch(vec![point])?; + Ok(ids.pop().expect("insert_batch returns one id per point")) + } + + //TODO: Make this an atomic operation pub fn insert_batch(&self, points: Vec) -> Result> { + // Validate the whole batch before writing anything let mut ids = Vec::with_capacity(points.len()); + let mut seen = HashSet::with_capacity(points.len()); - for point in points { - let id = point.id.unwrap_or_else(Uuid::new_v4); - let vector = point.vector; - let payload = point.payload; + for point in &points { + let id = point.id.unwrap_or_else(generate_point_id); + self.validate_point(point, id)?; - if let Some(ref v) = vector - && v.len() != self.dimension - { - return Err(ApiError::DimensionMismatch { - expected: self.dimension, - got: v.len(), - }); + if !seen.insert(id) || (point.id.is_some() && self.storage.contains_point(id)?) { + return Err(ApiError::PointAlreadyExists { id }); } + ids.push(id); + } + + for (&id, point) in ids.iter().zip(points) { + let PointInput { + dense_vector, + sparse_vector, + payload, + .. + } = point; - self.storage.insert_point(id, vector.clone(), payload)?; + self.storage + .insert_point(id, dense_vector.clone(), sparse_vector.clone(), payload)?; - if let Some(v) = vector { - let indexed = IndexedVector { id, vector: v }; + if let Some(vector) = dense_vector { let mut index = self.index.write().map_err(|_| ApiError::LockError)?; - index.insert(indexed)?; + index.insert(IndexedVector { id, vector })?; + } + if let Some(sv) = sparse_vector { + let mut sparse_index = + self.sparse_index.write().map_err(|_| ApiError::LockError)?; + sparse_index.insert(id, &sv); } - - ids.push(id); } Ok(ids) @@ -102,71 +122,70 @@ impl VectorDb { //TODO: Make this an atomic operation pub fn delete(&self, id: PointId) -> Result { - // Remove from storage + let Some(point) = self.storage.get_point(id)? else { + return Ok(false); + }; + self.storage.delete_point(id)?; - // Remove from index - let mut index = self.index.write().map_err(|_| ApiError::LockError)?; - let point_found = index.delete(id)?; - Ok(point_found) + + if point.dense_vector.is_some() { + let mut index = self.index.write().map_err(|_| ApiError::LockError)?; + index.delete(id)?; + } + if let Some(sv) = point.sparse_vector { + let mut sparse_index = self.sparse_index.write().map_err(|_| ApiError::LockError)?; + sparse_index.delete(id, &sv); + } + + Ok(true) } pub fn get(&self, id: PointId) -> Result> { - // Search for the Point with given id in storage - let payload = self.storage.get_payload(id)?; - let vector = self.storage.get_vector(id)?; - if payload.is_some() || vector.is_some() { - Ok(Some(Point { - id, - payload, - vector, - })) - } else { - Ok(None) - } + Ok(self.storage.get_point(id)?) } - pub fn search(&self, query: SearchQueryInput) -> Result> { - // Validate search limit - if query.limit == 0 { + pub fn search(&self, query: SearchQuery) -> Result> { + if query.limit == 0 || query.limit > MAX_SEARCH_LIMIT { return Err(ApiError::InvalidSearchLimit { limit: query.limit }); } - // Validate query dimension - if query.vector.len() != self.dimension { - return Err(ApiError::DimensionMismatch { - expected: self.dimension, - got: query.vector.len(), - }); - } - - // Use vector index to find similar vectors - let index = self.index.read().map_err(|_| ApiError::LockError)?; - - let vectors = - index.search_with_ef(query.vector, query.similarity, query.limit, query.ef)?; + let results = match query.query { + QueryVector::Dense(dense) => { + if dense.vector.len() != self.dimension { + return Err(ApiError::DimensionMismatch { + expected: self.dimension, + got: dense.vector.len(), + }); + } + let index = self.index.read().map_err(|_| ApiError::LockError)?; + index.search_with_ef(dense.vector, dense.similarity, query.limit, dense.ef)? + } + QueryVector::Sparse(sparse_vector) => { + sparse_vector.validate().context(InvalidSparseVectorSnafu)?; + let sparse_index = self.sparse_index.read().map_err(|_| ApiError::LockError)?; + sparse_index.search(&sparse_vector, query.limit) + } + }; - Ok(vectors) + Ok(results + .into_iter() + .map(|(id, score)| SearchHit { id, score }) + .collect()) } - pub fn search_batch(&self, queries: Vec) -> Result>> { - let mut results = Vec::with_capacity(queries.len()); - let index = self.index.read().unwrap(); - - for query in queries { - let found = - index.search_with_ef(query.vector, query.similarity, query.limit, query.ef)?; - results.push(found); - } - - Ok(results) + pub fn search_batch(&self, queries: Vec) -> Result>> { + queries + .into_iter() + .map(|query| self.search(query)) + .collect() } - pub fn list(&self, offset: PointId, limit: usize) -> Result> { - let page = self.storage.list_vectors(offset, limit)?; + pub fn list(&self, offset: PointId, limit: usize) -> Result> { + let page = self.storage.list_points(offset, limit)?; Ok(page) } - // populates the current index with vectors from the storage + // populates both indexes with the points in storage pub fn build_index(&self) -> Result { // start from the minimal UUID and fetch in bounded batches and insert let mut offset = Uuid::nil(); @@ -174,14 +193,26 @@ impl VectorDb { let mut inserted: usize = 0; let mut index = self.index.write().map_err(|_| ApiError::LockError)?; + let mut sparse_index = self.sparse_index.write().map_err(|_| ApiError::LockError)?; - while let Some((batch, next_offset)) = self.storage.list_vectors(offset, page_size)? { + while let Some((batch, next_offset)) = self.storage.list_points(offset, page_size)? { if batch.is_empty() || next_offset == offset { break; } - for (id, vector) in batch { - index.insert(IndexedVector { id, vector })?; + for Point { + id, + dense_vector, + sparse_vector, + .. + } in batch + { + if let Some(vector) = dense_vector { + index.insert(IndexedVector { id, vector })?; + } + if let Some(sv) = sparse_vector { + sparse_index.insert(id, &sv); + } inserted += 1; } @@ -190,6 +221,30 @@ impl VectorDb { Ok(inserted) } + + // populates only the sparse index, used after a snapshot restore where the dense index comes from the snapshot + fn build_sparse_index(&self) -> Result<()> { + let mut offset = Uuid::nil(); + let page_size: usize = 1000; + + let mut sparse_index = self.sparse_index.write().map_err(|_| ApiError::LockError)?; + + while let Some((batch, next_offset)) = self.storage.list_points(offset, page_size)? { + if batch.is_empty() || next_offset == offset { + break; + } + + for point in batch { + if let Some(sv) = point.sparse_vector { + sparse_index.insert(point.id, &sv); + } + } + + offset = next_offset; + } + + Ok(()) + } } impl SnapshottableDb for VectorDb { @@ -249,7 +304,10 @@ pub fn restore_from_snapshot(config: &DbRestoreConfig) -> Result Result { diff --git a/crates/api/src/tests.rs b/crates/api/src/tests.rs index f99d77b..36db56b 100644 --- a/crates/api/src/tests.rs +++ b/crates/api/src/tests.rs @@ -4,7 +4,7 @@ use super::*; use std::sync::Mutex; -use defs::ContentType; +use defs::{ContentType, DenseQuery, Payload, SparseVector}; use snapshot::{engine::SnapshotEngine, registry::local::LocalRegistry}; use tempfile::{TempDir, tempdir}; @@ -34,6 +34,26 @@ fn test_payload(content: &str) -> Payload { } } +fn dense_point(vector: Vec, payload: Payload) -> PointInput { + PointInput { + id: None, + dense_vector: Some(vector), + sparse_vector: None, + payload: Some(payload), + } +} + +fn dense_query(vector: Vec, similarity: Similarity, limit: usize) -> SearchQuery { + SearchQuery { + limit, + query: QueryVector::Dense(DenseQuery { + vector, + similarity, + ef: None, + }), + } +} + #[test] fn test_insert_and_get() { let (db, _temp_dir) = create_test_db(); @@ -44,13 +64,15 @@ fn test_insert_and_get() { }; // Test insert - let id = db.insert(vector.clone(), payload.clone()).unwrap(); + let id = db + .insert(dense_point(vector.clone(), payload.clone())) + .unwrap(); assert!(id != Uuid::nil()); // Test get let point = db.get(id).unwrap().unwrap(); assert_eq!(point.id, id); - assert_eq!(point.vector.as_ref().unwrap(), &vector); + assert_eq!(point.dense_vector.as_ref().unwrap(), &vector); assert_eq!(point.payload.as_ref().unwrap(), &payload); assert_eq!( point.payload.as_ref().unwrap().content_type, @@ -65,11 +87,13 @@ fn test_insert_and_get_with_in_memory_storage() { let vector = vec![1.0, 2.0, 3.0]; let payload = test_payload("Test content"); - let id = db.insert(vector.clone(), payload.clone()).unwrap(); + let id = db + .insert(dense_point(vector.clone(), payload.clone())) + .unwrap(); let point = db.get(id).unwrap().unwrap(); assert_eq!(point.id, id); - assert_eq!(point.vector, Some(vector)); + assert_eq!(point.dense_vector, Some(vector)); assert_eq!(point.payload, Some(payload)); } @@ -83,11 +107,11 @@ fn test_dimension_mismatch() { content: "tester".to_string(), }; - let res1 = db.insert(v1, payload.clone()); + let res1 = db.insert(dense_point(v1, payload.clone())); assert!(res1.is_ok()); // Insert vector of dimension 2 != 3 - let res2 = db.insert(v2, payload); + let res2 = db.insert(dense_point(v2, payload)); assert!(res2.is_err()); match res2.unwrap_err() { ApiError::DimensionMismatch { expected, got } => { @@ -108,7 +132,7 @@ fn test_delete() { }; // Insert a point - let id = db.insert(vector, payload).unwrap(); + let id = db.insert(dense_point(vector, payload)).unwrap(); // try deleting a point that does not exist let found = db.delete(Uuid::new_v4()); @@ -138,25 +162,20 @@ fn test_search() { content_type: ContentType::Text, content: format!("Test content {vector:?}"), }; - let id = db.insert(vector, payload).unwrap(); + let id = db.insert(dense_point(vector, payload)).unwrap(); ids.push(id); } // Search for the closest vector to [1.0, 0.1, 0.1] let query = vec![1.0, 0.1, 0.1]; let results = db - .search(SearchQueryInput { - vector: query, - similarity: Similarity::Cosine, - limit: 1, - ef: None, - }) + .search(dense_query(query, Similarity::Cosine, 1)) .unwrap(); assert_eq!(results.len(), 1); - assert_eq!(results[0].0, ids[0]); // The first vector should be closest + assert_eq!(results[0].id, ids[0]); // The first vector should be closest let expected = 1.0 - 1.0 / 1.02f32.sqrt(); - assert!((results[0].1 - expected).abs() < 1e-6); + assert!((results[0].score - expected).abs() < 1e-6); } #[test] @@ -168,13 +187,13 @@ fn test_search_limit() { for i in 0..5 { let vector = vec![i as f32, 0.0, 0.0]; let id = db - .insert( + .insert(dense_point( vector, Payload { content_type: ContentType::Text, content: format!("Test content {i}"), }, - ) + )) .unwrap(); ids.push(id); } @@ -182,12 +201,7 @@ fn test_search_limit() { // Search with limit 3 let query = vec![0.0, 0.0, 0.0]; let results = db - .search(SearchQueryInput { - vector: query, - similarity: Similarity::Euclidean, - limit: 3, - ef: None, - }) + .search(dense_query(query, Similarity::Euclidean, 3)) .unwrap(); assert_eq!(results.len(), 3); @@ -198,12 +212,7 @@ fn test_search_zero_limit() { let (db, _temp_dir) = create_test_db(); let query = vec![1.0, 2.0, 3.0]; - let result = db.search(SearchQueryInput { - vector: query, - similarity: Similarity::Cosine, - limit: 0, - ef: None, - }); + let result = db.search(dense_query(query, Similarity::Cosine, 0)); assert!(result.is_err()); match result.unwrap_err() { @@ -214,6 +223,24 @@ fn test_search_zero_limit() { } } +#[test] +fn test_search_limit_above_max() { + let (db, _temp_dir) = create_test_db(); + + assert!( + db.search(dense_query(vec![1.0, 2.0, 3.0], Similarity::Cosine, 10_000)) + .is_ok() + ); + assert!(matches!( + db.search(dense_query(vec![1.0, 2.0, 3.0], Similarity::Cosine, 10_001)), + Err(ApiError::InvalidSearchLimit { limit: 10_001 }) + )); + assert!(matches!( + db.search(sparse_query(vec![1], vec![1.0], usize::MAX)), + Err(ApiError::InvalidSearchLimit { .. }) + )); +} + #[test] fn test_empty_database() { let (db, _temp_dir) = create_test_db(); @@ -223,12 +250,7 @@ fn test_empty_database() { let query = vec![1.0, 2.0, 3.0]; let results = db - .search(SearchQueryInput { - vector: query, - similarity: Similarity::Cosine, - limit: 10, - ef: None, - }) + .search(dense_query(query, Similarity::Cosine, 10)) .unwrap(); assert_eq!(results.len(), 0); } @@ -242,13 +264,13 @@ fn test_list_vectors() { let i = i as f32; let vector = vec![i, i + 1.0, i + 2.0]; let id = db - .insert( + .insert(dense_point( vector, Payload { content_type: ContentType::Text, content: format!("Test content {i}"), }, - ) + )) .unwrap(); ids.push(id); } @@ -272,13 +294,13 @@ fn test_build_index() { for i in 0..10 { let i = i as f32; let vector = vec![i, i + 1.0, i + 2.0]; - db.insert( + db.insert(dense_point( vector, Payload { content_type: ContentType::Text, content: format!("Test content {i}"), }, - ) + )) .unwrap(); } @@ -296,23 +318,23 @@ fn test_create_and_load_snapshot() { let v3 = vec![6.0, 7.0, 8.0]; let id1 = old_db - .insert( + .insert(dense_point( v1.clone(), Payload { content_type: ContentType::Text, content: "test".to_string(), }, - ) + )) .unwrap(); let id2 = old_db - .insert( + .insert(dense_point( v2.clone(), Payload { content_type: ContentType::Text, content: "test".to_string(), }, - ) + )) .unwrap(); let temp_snapshot_dir = tempdir().unwrap(); @@ -320,13 +342,13 @@ fn test_create_and_load_snapshot() { // insert v3 after snapshot let id3 = old_db - .insert( + .insert(dense_point( v3.clone(), Payload { content_type: ContentType::Text, content: "test".to_string(), }, - ) + )) .unwrap(); let reload_config = DbRestoreConfig { @@ -342,8 +364,8 @@ fn test_create_and_load_snapshot() { assert!(loaded_db.get(id3).unwrap_or(None).is_none()); // v3 was inserted after snapshot was taken // vector restore check - assert!(loaded_db.get(id1).unwrap().unwrap().vector.unwrap() == v1); - assert!(loaded_db.get(id2).unwrap().unwrap().vector.unwrap() == v2); + assert!(loaded_db.get(id1).unwrap().unwrap().dense_vector.unwrap() == v1); + assert!(loaded_db.get(id2).unwrap().unwrap().dense_vector.unwrap() == v2); } #[test] @@ -354,13 +376,19 @@ fn test_create_and_load_snapshot_with_in_memory_storage() { let v2 = vec![3.0, 4.0, 5.0]; let v3 = vec![6.0, 7.0, 8.0]; - let id1 = old_db.insert(v1.clone(), test_payload("one")).unwrap(); - let id2 = old_db.insert(v2.clone(), test_payload("two")).unwrap(); + let id1 = old_db + .insert(dense_point(v1.clone(), test_payload("one"))) + .unwrap(); + let id2 = old_db + .insert(dense_point(v2.clone(), test_payload("two"))) + .unwrap(); let temp_snapshot_dir = tempdir().unwrap(); let snapshot_path = old_db.create_snapshot(temp_snapshot_dir.path()).unwrap(); - let id3 = old_db.insert(v3, test_payload("three")).unwrap(); + let id3 = old_db + .insert(dense_point(v3, test_payload("three"))) + .unwrap(); let reload_config = DbRestoreConfig { data_path: temp_dir.path().to_path_buf(), @@ -369,8 +397,8 @@ fn test_create_and_load_snapshot_with_in_memory_storage() { let loaded_db = restore_from_snapshot(&reload_config).unwrap(); - assert_eq!(loaded_db.get(id1).unwrap().unwrap().vector, Some(v1)); - assert_eq!(loaded_db.get(id2).unwrap().unwrap().vector, Some(v2)); + assert_eq!(loaded_db.get(id1).unwrap().unwrap().dense_vector, Some(v1)); + assert_eq!(loaded_db.get(id2).unwrap().unwrap().dense_vector, Some(v2)); assert!(loaded_db.get(id3).unwrap().is_none()); } @@ -400,13 +428,13 @@ fn test_snapshot_engine() { let id = db .lock() .unwrap() - .insert( + .insert(dense_point( vector.clone(), Payload { content_type: ContentType::Text, content: format!("{}", i), }, - ) + )) .unwrap(); inserted_ids.push(id); } @@ -433,7 +461,14 @@ fn test_snapshot_engine() { // test if point is present assert!(db.get(inserted_ids[j]).unwrap_or(None).is_some()); // test vector restore - assert!(db.get(inserted_ids[j]).unwrap().unwrap().vector.unwrap() == test_vectors[j]); + assert!( + db.get(inserted_ids[j]) + .unwrap() + .unwrap() + .dense_vector + .unwrap() + == test_vectors[j] + ); // test payload restore assert!( db.get(inserted_ids[j]) @@ -451,3 +486,309 @@ fn test_snapshot_engine() { std::mem::drop(db); } } + +fn sparse_point(id: PointId, token_ids: Vec, weights: Vec) -> PointInput { + PointInput { + id: Some(id), + dense_vector: None, + sparse_vector: Some(SparseVector { token_ids, weights }), + payload: Some(test_payload("sparse")), + } +} + +fn sparse_query(token_ids: Vec, weights: Vec, limit: usize) -> SearchQuery { + SearchQuery { + limit, + query: QueryVector::Sparse(SparseVector { token_ids, weights }), + } +} + +fn hit_ids(hits: &[SearchHit]) -> Vec { + hits.iter().map(|hit| hit.id).collect() +} + +#[test] +fn test_sparse_insert_get_and_search() { + let (db, _temp_dir) = create_test_db(); + let id1 = Uuid::from_u128(1); + let id2 = Uuid::from_u128(2); + db.insert_batch(vec![ + sparse_point(id1, vec![1, 2], vec![1.0, 2.0]), + sparse_point(id2, vec![2, 3], vec![1.0, 3.0]), + ]) + .unwrap(); + + let point = db.get(id1).unwrap().unwrap(); + assert_eq!(point.dense_vector, None); + assert_eq!(point.sparse_vector.unwrap().token_ids, vec![1, 2]); + + let hits = db + .search(sparse_query(vec![2, 3], vec![1.0, 1.0], 10)) + .unwrap(); + assert_eq!( + hits, + vec![ + SearchHit { + id: id2, + score: 4.0 + }, + SearchHit { + id: id1, + score: 2.0 + } + ] + ); +} + +#[test] +fn test_insert_mixed_batch() { + let (db, _temp_dir) = create_test_db(); + let dense_id = Uuid::from_u128(1); + let sparse_id = Uuid::from_u128(2); + let both_id = Uuid::from_u128(3); + db.insert_batch(vec![ + PointInput { + id: Some(dense_id), + ..dense_point(vec![1.0, 0.0, 0.0], test_payload("dense")) + }, + sparse_point(sparse_id, vec![5], vec![1.0]), + PointInput { + id: Some(both_id), + dense_vector: Some(vec![0.0, 1.0, 0.0]), + sparse_vector: Some(SparseVector { + token_ids: vec![5], + weights: vec![2.0], + }), + payload: None, + }, + ]) + .unwrap(); + + let both = db.get(both_id).unwrap().unwrap(); + assert!(both.dense_vector.is_some() && both.sparse_vector.is_some()); + + let dense_hits = db + .search(dense_query(vec![1.0, 0.0, 0.0], Similarity::Cosine, 10)) + .unwrap(); + assert_eq!(hit_ids(&dense_hits), vec![dense_id, both_id]); + + let sparse_hits = db.search(sparse_query(vec![5], vec![1.0], 10)).unwrap(); + assert_eq!(hit_ids(&sparse_hits), vec![both_id, sparse_id]); +} + +#[test] +fn test_insert_rejects_point_without_vectors() { + let (db, _temp_dir) = create_test_db(); + let result = db.insert(PointInput { + id: None, + dense_vector: None, + sparse_vector: None, + payload: Some(test_payload("empty")), + }); + assert!(matches!(result, Err(ApiError::MissingVector { .. }))); +} + +#[test] +fn test_sparse_insert_rejects_invalid_vector() { + let (db, _temp_dir) = create_test_db(); + let result = db.insert_batch(vec![sparse_point( + Uuid::new_v4(), + vec![2, 1], + vec![1.0, 1.0], + )]); + assert!(matches!(result, Err(ApiError::InvalidSparseVector { .. }))); +} + +#[test] +fn test_sparse_weights_must_not_be_negative() { + let (db, _temp_dir) = create_test_db(); + let result = db.insert_batch(vec![sparse_point( + Uuid::new_v4(), + vec![1, 2], + vec![1.0, -0.5], + )]); + assert!(matches!(result, Err(ApiError::InvalidSparseVector { .. }))); + + assert!( + db.insert_batch(vec![sparse_point( + Uuid::new_v4(), + vec![1, 2], + vec![0.0, 1.0] + )]) + .is_ok() + ); + assert!(matches!( + db.search(sparse_query(vec![1], vec![-1.0], 5)), + Err(ApiError::InvalidSparseVector { .. }) + )); +} + +#[test] +fn test_insert_rejects_existing_point_id() { + let (db, _temp_dir) = create_test_db(); + let id = Uuid::new_v4(); + db.insert_batch(vec![sparse_point(id, vec![1], vec![1.0])]) + .unwrap(); + + let result = db.insert_batch(vec![sparse_point(id, vec![2], vec![1.0])]); + assert!(matches!(result, Err(ApiError::PointAlreadyExists { .. }))); +} + +#[test] +fn test_insert_rejects_duplicate_id_within_batch() { + let (db, _temp_dir) = create_test_db(); + let id = Uuid::new_v4(); + let result = db.insert_batch(vec![ + sparse_point(id, vec![1], vec![1.0]), + sparse_point(id, vec![2], vec![1.0]), + ]); + + assert!(matches!(result, Err(ApiError::PointAlreadyExists { .. }))); + assert!(db.get(id).unwrap().is_none()); +} + +#[test] +fn test_failed_batch_writes_nothing() { + let (db, _temp_dir) = create_test_db(); + let id1 = Uuid::from_u128(1); + let id2 = Uuid::from_u128(2); + let result = db.insert_batch(vec![ + sparse_point(id1, vec![1], vec![1.0]), + sparse_point(id2, vec![2], vec![1.0]), + dense_point(vec![1.0, 2.0], test_payload("wrong dimension")), + ]); + + assert!(matches!(result, Err(ApiError::DimensionMismatch { .. }))); + assert!(db.get(id1).unwrap().is_none()); + assert!(db.get(id2).unwrap().is_none()); + assert!( + db.search(sparse_query(vec![1, 2], vec![1.0, 1.0], 5)) + .unwrap() + .is_empty() + ); +} + +#[test] +fn test_sparse_search_validates_query() { + let (db, _temp_dir) = create_test_db(); + assert!(matches!( + db.search(sparse_query(vec![1], vec![1.0], 0)), + Err(ApiError::InvalidSearchLimit { .. }) + )); + assert!(matches!( + db.search(sparse_query(vec![], vec![], 5)), + Err(ApiError::InvalidSparseVector { .. }) + )); +} + +#[test] +fn test_search_batch_mixes_dense_and_sparse() { + let (db, _temp_dir) = create_test_db(); + let dense_id = db + .insert(dense_point(vec![1.0, 0.0, 0.0], test_payload("dense"))) + .unwrap(); + let sparse_id = Uuid::from_u128(2); + db.insert(sparse_point(sparse_id, vec![2], vec![2.0])) + .unwrap(); + + let results = db + .search_batch(vec![ + dense_query(vec![1.0, 0.0, 0.0], Similarity::Cosine, 5), + sparse_query(vec![2], vec![1.0], 5), + ]) + .unwrap(); + assert_eq!(hit_ids(&results[0]), vec![dense_id]); + assert_eq!( + results[1], + vec![SearchHit { + id: sparse_id, + score: 2.0 + }] + ); +} + +#[test] +fn test_search_batch_rejects_invalid_query() { + let (db, _temp_dir) = create_test_db(); + let result = db.search_batch(vec![ + sparse_query(vec![1], vec![1.0], 5), + dense_query(vec![1.0], Similarity::Cosine, 5), + ]); + assert!(matches!(result, Err(ApiError::DimensionMismatch { .. }))); +} + +#[test] +fn test_delete_removes_point_from_both_indexes() { + let (db, _temp_dir) = create_test_db(); + let id = Uuid::new_v4(); + db.insert(PointInput { + id: Some(id), + dense_vector: Some(vec![1.0, 0.0, 0.0]), + sparse_vector: Some(SparseVector { + token_ids: vec![1, 2], + weights: vec![1.0, 1.0], + }), + payload: None, + }) + .unwrap(); + + assert!(db.delete(id).unwrap()); + assert!(db.get(id).unwrap().is_none()); + assert!( + db.search(sparse_query(vec![1, 2], vec![1.0, 1.0], 5)) + .unwrap() + .is_empty() + ); + assert!( + db.search(dense_query(vec![1.0, 0.0, 0.0], Similarity::Cosine, 5)) + .unwrap() + .is_empty() + ); + assert!(!db.delete(id).unwrap()); +} + +#[test] +fn test_sparse_index_rebuilt_on_reopen() { + let temp_dir = tempdir().unwrap(); + let config = || DbConfig { + storage_type: StorageType::RocksDb, + index_type: IndexType::Flat, + data_path: temp_dir.path().to_path_buf(), + dimension: 3, + similarity: Similarity::Cosine, + hnsw_config: HnswConfig::default(), + kd_tree_config: KDTreeConfig::default(), + }; + let id = Uuid::new_v4(); + + let db = init_api(config()).unwrap(); + db.insert(sparse_point(id, vec![3], vec![1.5])).unwrap(); + std::mem::drop(db); + + let reopened = init_api(config()).unwrap(); + let hits = reopened + .search(sparse_query(vec![3], vec![1.0], 5)) + .unwrap(); + assert_eq!(hits, vec![SearchHit { id, score: 1.5 }]); +} + +#[test] +fn test_sparse_index_rebuilt_on_snapshot_restore() { + let (old_db, temp_dir) = create_test_db(); + let id = Uuid::new_v4(); + old_db.insert(sparse_point(id, vec![3], vec![1.5])).unwrap(); + + let temp_snapshot_dir = tempdir().unwrap(); + let snapshot_path = old_db.create_snapshot(temp_snapshot_dir.path()).unwrap(); + std::mem::drop(old_db); + + let loaded_db = restore_from_snapshot(&DbRestoreConfig { + data_path: temp_dir.path().to_path_buf(), + snapshot_path, + }) + .unwrap(); + let hits = loaded_db + .search(sparse_query(vec![3], vec![1.0], 5)) + .unwrap(); + assert_eq!(hits, vec![SearchHit { id, score: 1.5 }]); +} diff --git a/crates/defs/Cargo.toml b/crates/defs/Cargo.toml index 81d8dd9..64d8944 100644 --- a/crates/defs/Cargo.toml +++ b/crates/defs/Cargo.toml @@ -7,7 +7,6 @@ edition.workspace = true license.workspace = true [dependencies] -axum.workspace = true serde.workspace = true snafu.workspace = true uuid.workspace = true diff --git a/crates/defs/src/error.rs b/crates/defs/src/error.rs index fa0e852..a4db0e2 100644 --- a/crates/defs/src/error.rs +++ b/crates/defs/src/error.rs @@ -24,7 +24,6 @@ pub enum DbError { PointNotFound { id: PointId }, } -use axum::{http::StatusCode, response::IntoResponse}; use snafu::Snafu; #[derive(Debug, Snafu)] @@ -36,20 +35,23 @@ pub enum ServerError { Serve { source: io::Error }, } -#[derive(Debug)] -pub enum AppError { - ServerError(ServerError), - Api(String), -} - -impl IntoResponse for AppError { - fn into_response(self) -> axum::response::Response { - let (status, message) = match self { - AppError::Api(msg) => (StatusCode::BAD_REQUEST, msg), - AppError::ServerError(e) => (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()), - }; - (status, message).into_response() - } -} // Error type for server pub type BoxError = Box; + +#[derive(Debug, Snafu)] +pub enum SparseError { + #[snafu(display("token_ids length {ids} does not match weights length {weights}"))] + LengthMismatch { ids: usize, weights: usize }, + + #[snafu(display("sparse vector must not be empty"))] + Empty, + + #[snafu(display("token_ids must be strictly increasing, violated at index {index}"))] + UnsortedTokenIds { index: usize }, + + #[snafu(display("weight at index {index} is not finite"))] + NonFiniteWeight { index: usize }, + + #[snafu(display("weight at index {index} is negative"))] + NegativeWeight { index: usize }, +} diff --git a/crates/defs/src/types.rs b/crates/defs/src/types.rs index 7e67399..5bb2f23 100644 --- a/crates/defs/src/types.rs +++ b/crates/defs/src/types.rs @@ -2,6 +2,8 @@ use serde::{Deserialize, Serialize}; use std::cmp::Ordering; use uuid::Uuid; +use crate::SparseError; + pub type PointId = Uuid; /// Type of vector element. @@ -11,16 +13,47 @@ pub type ElementByte = u8; pub type Dimension = usize; -// Dense Vector and Vector are considered same -// Sparse vector implementation not supported yet. Refer lib/sparse/src/common/sparse_vector.rs +pub type TokenID = u32; + pub type DenseVector = Vec; -pub type Magic = [u8; 4]; +#[derive(Serialize, Deserialize, Debug, Clone, PartialEq)] +pub struct SparseVector { + pub token_ids: Vec, + pub weights: Vec, +} + +impl SparseVector { + pub fn validate(&self) -> Result<(), SparseError> { + if self.token_ids.len() != self.weights.len() { + return Err(SparseError::LengthMismatch { + ids: self.token_ids.len(), + weights: self.weights.len(), + }); + } + + if self.token_ids.is_empty() { + return Err(SparseError::Empty); + } -pub enum StoredVector { - Dense(DenseVector), + if let Some(i) = self.token_ids.windows(2).position(|w| w[0] >= w[1]) { + return Err(SparseError::UnsortedTokenIds { index: i + 1 }); + } + + if let Some(i) = self.weights.iter().position(|w| !w.is_finite()) { + return Err(SparseError::NonFiniteWeight { index: i }); + } + + if let Some(i) = self.weights.iter().position(|&w| w < 0.0) { + return Err(SparseError::NegativeWeight { index: i }); + } + + Ok(()) + } } +pub type Magic = [u8; 4]; + #[derive(Serialize, Deserialize, Clone, Copy, Debug, Default, PartialEq)] pub enum ContentType { #[default] @@ -37,14 +70,16 @@ pub struct Payload { #[derive(Serialize, Deserialize, Clone, Debug, PartialEq)] pub struct PointInput { pub id: Option, - pub vector: Option, + pub dense_vector: Option, + pub sparse_vector: Option, pub payload: Option, } #[derive(Serialize, Deserialize, Clone, Debug, PartialEq)] pub struct Point { pub id: PointId, - pub vector: Option, + pub dense_vector: Option, + pub sparse_vector: Option, pub payload: Option, } @@ -77,22 +112,39 @@ pub struct BatchInsertResponse { // For batch search #[derive(Debug, Clone, Serialize, Deserialize)] pub struct BatchSearchRequest { - pub queries: Vec, + pub queries: Vec, } -#[derive(Debug, Clone, Serialize, Deserialize)] -pub struct SearchQueryInput { +#[derive(Serialize, Deserialize, Debug, Clone)] +pub struct DenseQuery { pub vector: DenseVector, pub similarity: Similarity, - pub limit: usize, #[serde(default)] pub ef: Option, } +#[derive(Serialize, Deserialize, Debug, Clone, PartialEq)] +pub struct SearchHit { + pub id: PointId, + pub score: f32, +} + +#[derive(Serialize, Deserialize, Debug, Clone)] +#[serde(rename_all = "snake_case")] +pub enum QueryVector { + Dense(DenseQuery), + Sparse(SparseVector), +} + +#[derive(Serialize, Deserialize, Debug, Clone)] +pub struct SearchQuery { + pub limit: usize, + pub query: QueryVector, +} + #[derive(Clone, Serialize, Deserialize, Debug)] pub struct SearchResponse { - pub results: Vec, - pub scores: Vec, + pub hits: Vec, } #[derive(Debug, Clone, Serialize, Deserialize)] diff --git a/crates/grpc/proto/vector-db.proto b/crates/grpc/proto/vector-db.proto index ecca463..c959d06 100644 --- a/crates/grpc/proto/vector-db.proto +++ b/crates/grpc/proto/vector-db.proto @@ -9,50 +9,71 @@ message UUID { } service VectorDB { - //Insert a vector with a payload and return the assigned PointID - rpc InsertVector(InsertVectorRequest) returns (PointID) {} + //Insert a point with a dense and/or sparse vector and return its PointID + rpc InsertPoint(PointInput) returns (PointID) {} - //Delete a vector by its PointID + //Delete a point by its PointID rpc DeletePoint(PointID) returns (google.protobuf.Empty) {} - //Get a vector and its payload by PointID + //Get a point with its vectors and payload by PointID rpc GetPoint(PointID) returns (Point) {} - //Search for the k nearest vectors to a target vector given a distance function - rpc SearchPoints(SearchRequest) returns (SearchResponse) {} + //Search for the nearest points to a dense or sparse query + rpc Search(SearchRequest) returns (SearchResponse) {} - rpc InsertVectorsBatch(InsertVectorsBatchRequest) returns (InsertVectorsBatchResponse) {} -rpc SearchPointsBatch(SearchPointsBatchRequest) returns (SearchPointsBatchResponse) {} + rpc InsertPointsBatch(InsertPointsBatchRequest) returns (InsertPointsBatchResponse) {} + rpc SearchBatch(SearchBatchRequest) returns (SearchBatchResponse) {} } -message InsertVectorRequest { +message PointInput { + PointID id = 1; + DenseVector dense_vector = 2; + SparseVector sparse_vector = 3; + Payload payload = 4; +} + + +message DenseQuery { DenseVector vector = 1; - Payload payload = 2; + Similarity similarity = 2; + optional uint64 ef = 3; } message SearchRequest { - DenseVector query_vector = 1; - Similarity similarity = 2; - uint64 limit = 3; - uint64 ef = 4; + uint64 limit = 1; + oneof query { + DenseQuery dense = 2; + SparseVector sparse = 3; + } +} + + +message SearchHit { + PointID id = 1; + float score = 2; } message SearchResponse { - repeated PointID result_point_ids = 1; - repeated float scores = 2; + repeated SearchHit hits = 1; } message DenseVector { repeated float values = 1; } +message SparseVector { + repeated uint32 token_ids = 1; + repeated float weights = 2; +} + message Point { PointID id = 1; Payload payload = 2; - DenseVector vector = 3; + DenseVector dense_vector = 3; + SparseVector sparse_vector = 4; } message PointID { @@ -76,18 +97,18 @@ message Payload { string content = 2; } -message InsertVectorsBatchRequest { - repeated InsertVectorRequest vectors = 1; +message InsertPointsBatchRequest { + repeated PointInput points = 1; } -message InsertVectorsBatchResponse { +message InsertPointsBatchResponse { repeated PointID ids = 1; } -message SearchPointsBatchRequest { +message SearchBatchRequest { repeated SearchRequest queries = 1; } -message SearchPointsBatchResponse { +message SearchBatchResponse { repeated SearchResponse results = 1; } diff --git a/crates/grpc/src/error.rs b/crates/grpc/src/error.rs index faf02b4..0d1a360 100644 --- a/crates/grpc/src/error.rs +++ b/crates/grpc/src/error.rs @@ -57,6 +57,15 @@ impl From for GrpcError { message: format!("invalid search limit: {}", limit), }, ApiError::InitializationFailed { reason } => GrpcError::Internal { message: reason }, + ApiError::PointAlreadyExists { id } => GrpcError::AlreadyExists { + message: format!("point already exists: {}", id), + }, + ApiError::MissingVector { id } => GrpcError::InvalidArgument { + message: format!("point has neither a dense nor a sparse vector: {}", id), + }, + ApiError::InvalidSparseVector { source } => GrpcError::InvalidArgument { + message: format!("invalid sparse vector: {}", source), + }, } } } diff --git a/crates/grpc/src/service.rs b/crates/grpc/src/service.rs index 9807bf2..73e9b0d 100644 --- a/crates/grpc/src/service.rs +++ b/crates/grpc/src/service.rs @@ -6,13 +6,13 @@ use crate::interceptors; use crate::service::vectordb::{ContentType, Uuid}; use crate::utils::log_rpc; use crate::{constants::SIMILARITY_PROTOBUFF_MAP, utils::ServerEndpoint}; -use defs::SearchQueryInput; use tonic::{Request, Response, Status, service::InterceptorLayer, transport::Server}; use tracing::{Level, event}; use uuid::Uuid as UuidCrate; use vectordb::{ - DenseVector, InsertVectorRequest, InsertVectorsBatchRequest, InsertVectorsBatchResponse, Point, - PointId, SearchPointsBatchRequest, SearchPointsBatchResponse, SearchRequest, SearchResponse, + DenseVector, InsertPointsBatchRequest, InsertPointsBatchResponse, Point, PointId, PointInput, + SearchBatchRequest, SearchBatchResponse, SearchHit, SearchRequest, SearchResponse, + SparseVector, search_request, vector_db_server::{VectorDb, VectorDbServer}, }; @@ -31,240 +31,214 @@ impl VectorDBService { } } +fn parse_point_id(point_id: PointId) -> Result { + let value = point_id.id.unwrap_or_default().value; + UuidCrate::from_str(&value) + .map_err(|_| Status::invalid_argument(format!("invalid point id: {}", value))) +} + +fn to_proto_point_id(id: UuidCrate) -> PointId { + PointId { + id: Some(Uuid { + value: id.to_string(), + }), + } +} + +fn from_proto_payload(payload: vectordb::Payload) -> defs::Payload { + let content_type = + match ContentType::try_from(payload.content_type).unwrap_or(ContentType::Text) { + ContentType::Text => defs::ContentType::Text, + ContentType::Image => defs::ContentType::Image, + }; + defs::Payload { + content_type, + content: payload.content, + } +} + +fn to_proto_payload(payload: defs::Payload) -> vectordb::Payload { + let content_type = match payload.content_type { + defs::ContentType::Text => ContentType::Text, + defs::ContentType::Image => ContentType::Image, + }; + vectordb::Payload { + content_type: content_type as i32, + content: payload.content, + } +} + +fn from_proto_sparse(sparse_vector: SparseVector) -> defs::SparseVector { + defs::SparseVector { + token_ids: sparse_vector.token_ids, + weights: sparse_vector.weights, + } +} + +fn from_proto_point_input(point: PointInput) -> Result { + Ok(defs::PointInput { + id: point.id.map(parse_point_id).transpose()?, + dense_vector: point.dense_vector.map(|v| v.values), + sparse_vector: point.sparse_vector.map(from_proto_sparse), + payload: point.payload.map(from_proto_payload), + }) +} + +fn to_proto_point(point: defs::Point) -> Point { + Point { + id: Some(to_proto_point_id(point.id)), + dense_vector: point.dense_vector.map(|values| DenseVector { values }), + sparse_vector: point.sparse_vector.map(|sv| SparseVector { + token_ids: sv.token_ids, + weights: sv.weights, + }), + payload: point.payload.map(to_proto_payload), + } +} + +fn from_proto_search_request(request: SearchRequest) -> Result { + let query = match request + .query + .ok_or_else(|| Status::invalid_argument("query is required"))? + { + search_request::Query::Dense(dense) => { + let similarity = SIMILARITY_PROTOBUFF_MAP + .get(dense.similarity as usize) + .ok_or_else(|| Status::invalid_argument("Invalid similarity"))?; + let vector = dense + .vector + .ok_or_else(|| Status::invalid_argument("dense query vector is required"))?; + defs::QueryVector::Dense(defs::DenseQuery { + vector: vector.values, + similarity: *similarity, + ef: dense.ef.map(|ef| ef as usize), + }) + } + search_request::Query::Sparse(sparse_vector) => { + defs::QueryVector::Sparse(from_proto_sparse(sparse_vector)) + } + }; + + Ok(defs::SearchQuery { + limit: request.limit as usize, + query, + }) +} + +fn to_proto_search_response(hits: Vec) -> SearchResponse { + SearchResponse { + hits: hits + .into_iter() + .map(|hit| SearchHit { + id: Some(to_proto_point_id(hit.id)), + score: hit.score, + }) + .collect(), + } +} + #[tonic::async_trait] impl VectorDb for VectorDBService { - async fn insert_vector( + async fn insert_point( &self, - request: Request, + request: Request, ) -> Result, Status> { - log_rpc("insert_vector", self.logging); + log_rpc("insert_point", self.logging); interceptors::require_write_role(&request)?; - let inner_request = request.into_inner(); - - let dense_vector = inner_request.vector; - if dense_vector.is_none() { - return Err(Status::invalid_argument("dense_vector is empty")); - } + let point = from_proto_point_input(request.into_inner())?; + let id = self + .vector_db + .insert(point) + .map_err(|e| Status::from(GrpcError::from(e)))?; - // TODO: Implement payload handling once its defined - // fetch payload and default to empty struct otherwise - - let payload = inner_request.payload.unwrap_or_default(); - let payload_content = payload.content; - let payload_type = - match ContentType::try_from(payload.content_type).unwrap_or(ContentType::Text) { - ContentType::Text => defs::ContentType::Text, - ContentType::Image => defs::ContentType::Image, - }; - - let point_id = self.vector_db.insert( - dense_vector.unwrap().values, - defs::Payload { - content_type: payload_type, - content: payload_content, - }, - ); - - let res = point_id.map_err(|e| Status::from(crate::error::GrpcError::from(e)))?; - - Ok(Response::new(PointId { - id: Some(Uuid { - value: res.to_string(), - }), - })) + Ok(Response::new(to_proto_point_id(id))) } async fn get_point(&self, request: Request) -> Result, Status> { log_rpc("get_point", self.logging); - let inner_request = request.into_inner(); - - let point_id = inner_request.id.unwrap_or_default().value; - let point_opt = self + let point_id = parse_point_id(request.into_inner())?; + let point = self .vector_db - .get(UuidCrate::from_str(&point_id).unwrap()) - .map_err(|e| Status::from(crate::error::GrpcError::from(e)))?; - - // return error if not found - let point = point_opt.ok_or(Status::not_found(format!("point not found: {}", point_id)))?; - - let payload = point.payload.map(|p| vectordb::Payload { - content_type: match p.content_type { - defs::ContentType::Text => ContentType::Text as i32, - defs::ContentType::Image => ContentType::Image as i32, - }, - content: p.content, - }); - - Ok(Response::new(Point { - id: Some(PointId { - id: Some(Uuid { - value: point.id.to_string(), - }), - }), - vector: Some(DenseVector { - values: point.vector.unwrap_or_default(), - }), - payload, - })) + .get(point_id) + .map_err(|e| Status::from(GrpcError::from(e)))? + .ok_or_else(|| Status::not_found(format!("point not found: {}", point_id)))?; + + Ok(Response::new(to_proto_point(point))) } - async fn search_points( + async fn search( &self, request: Request, ) -> Result, Status> { - log_rpc("search_points", self.logging); + log_rpc("search", self.logging); - let search_request = request.into_inner(); - - // extract request arguments - let query_vect = search_request - .query_vector - .ok_or(Status::invalid_argument("Invalid query_vector"))?; - let similarity = SIMILARITY_PROTOBUFF_MAP - .get(search_request.similarity as usize) - .ok_or(Status::invalid_argument("Invalid similarity"))?; - let limit = search_request.limit; - - if limit == 0 { - return Err(Status::invalid_argument("Limit must be greater than zero")); - } - - let result_point_ids = self + let query = from_proto_search_request(request.into_inner())?; + let hits = self .vector_db - .search(SearchQueryInput { - vector: query_vect.values, - similarity: *similarity, - limit: limit as usize, - ef: (search_request.ef > 0).then_some(search_request.ef as usize), - }) - .map_err(|e| Status::from(crate::error::GrpcError::from(e)))?; - - // create a mapped vector of PointIds - let (result, scores) = result_point_ids - .into_iter() - .map(|(id, score)| { - ( - PointId { - id: Some(Uuid { - value: id.to_string(), - }), - }, - score, - ) - }) - .unzip(); + .search(query) + .map_err(|e| Status::from(GrpcError::from(e)))?; - Ok(Response::new(SearchResponse { - result_point_ids: result, - scores, - })) + Ok(Response::new(to_proto_search_response(hits))) } async fn delete_point(&self, request: Request) -> Result, Status> { log_rpc("delete_point", self.logging); interceptors::require_write_role(&request)?; - let point_id = request.into_inner().id.unwrap_or_default().value; + let point_id = parse_point_id(request.into_inner())?; - match self - .vector_db - .delete(UuidCrate::from_str(&point_id).unwrap()) - { - Ok(found) => { - if found { - Ok(Response::new(())) - } else { - Err(Status::not_found("Point not found")) - } - } - Err(e) => Err(Status::from(crate::error::GrpcError::from(e))), + match self.vector_db.delete(point_id) { + Ok(true) => Ok(Response::new(())), + Ok(false) => Err(Status::not_found("Point not found")), + Err(e) => Err(Status::from(GrpcError::from(e))), } } - async fn insert_vectors_batch( + async fn insert_points_batch( &self, - request: tonic::Request, - ) -> Result, tonic::Status> { + request: Request, + ) -> Result, Status> { interceptors::require_write_role(&request)?; - let req = request.into_inner(); - let mut ids = Vec::with_capacity(req.vectors.len()); - - for vec in req.vectors { - let payload = vec.payload.map(|p| defs::Payload { - content_type: match ContentType::try_from(p.content_type) - .unwrap_or(ContentType::Text) - { - ContentType::Text => defs::ContentType::Text, - ContentType::Image => defs::ContentType::Image, - }, - content: p.content, - }); - - let id = self - .vector_db - .insert( - vec.vector.unwrap_or_default().values, - payload.unwrap_or_default(), - ) - .map_err(|e| tonic::Status::internal(e.to_string()))?; - - ids.push(PointId { - id: Some(Uuid { - value: id.to_string(), - }), - }); - } - Ok(tonic::Response::new(InsertVectorsBatchResponse { ids })) + let points = request + .into_inner() + .points + .into_iter() + .map(from_proto_point_input) + .collect::, Status>>()?; + + let ids = self + .vector_db + .insert_batch(points) + .map_err(|e| Status::from(GrpcError::from(e)))?; + + Ok(Response::new(InsertPointsBatchResponse { + ids: ids.into_iter().map(to_proto_point_id).collect(), + })) } - async fn search_points_batch( + async fn search_batch( &self, - request: tonic::Request, - ) -> Result, tonic::Status> { - let req = request.into_inner(); - let mut results = Vec::with_capacity(req.queries.len()); + request: Request, + ) -> Result, Status> { + let queries = request + .into_inner() + .queries + .into_iter() + .map(from_proto_search_request) + .collect::, Status>>()?; - for query in req.queries { - let similarity = SIMILARITY_PROTOBUFF_MAP - .get(query.similarity as usize) - .ok_or(tonic::Status::invalid_argument("Invalid similarity"))?; - - let ids = self - .vector_db - .search(SearchQueryInput { - vector: query - .query_vector - .ok_or(tonic::Status::invalid_argument("missing query_vector"))? - .values, - similarity: *similarity, - limit: query.limit as usize, - ef: (query.ef > 0).then_some(query.ef as usize), - }) - .map_err(|e| tonic::Status::from(GrpcError::from(e)))?; - - let (result_point_ids, scores) = ids - .into_iter() - .map(|(id, score)| { - ( - PointId { - id: Some(Uuid { - value: id.to_string(), - }), - }, - score, - ) - }) - .unzip(); - results.push(SearchResponse { - result_point_ids, - scores, - }); - } + let results = self + .vector_db + .search_batch(queries) + .map_err(|e| Status::from(GrpcError::from(e)))? + .into_iter() + .map(to_proto_search_response) + .collect(); - Ok(tonic::Response::new(SearchPointsBatchResponse { results })) + Ok(Response::new(SearchBatchResponse { results })) } } pub async fn run_server( diff --git a/crates/grpc/src/tests.rs b/crates/grpc/src/tests.rs index 006649b..6fc0e1b 100644 --- a/crates/grpc/src/tests.rs +++ b/crates/grpc/src/tests.rs @@ -1,7 +1,8 @@ use crate::constants::AUTHORIZATION_HEADER_KEY; use crate::service::vectordb::vector_db_client::VectorDbClient; use crate::service::vectordb::{ - ContentType, DenseVector, InsertVectorRequest, Payload, PointId, SearchRequest, + ContentType, DenseQuery, DenseVector, InsertPointsBatchRequest, Payload, PointId, PointInput, + SearchBatchRequest, SearchRequest, SparseVector, Uuid, search_request, }; use crate::service::{VectorDBService, run_server}; use crate::utils::ServerEndpoint; @@ -79,6 +80,35 @@ async fn create_test_client( Ok(VectorDbClient::new(channel)) } +fn dense_search(values: Vec, similarity: i32, limit: u64) -> SearchRequest { + SearchRequest { + limit, + query: Some(search_request::Query::Dense(DenseQuery { + vector: Some(DenseVector { values }), + similarity, + ef: None, + })), + } +} + +fn sparse_search(token_ids: Vec, weights: Vec, limit: u64) -> SearchRequest { + SearchRequest { + limit, + query: Some(search_request::Query::Sparse(SparseVector { + token_ids, + weights, + })), + } +} + +fn point_id(value: &str) -> PointId { + PointId { + id: Some(Uuid { + value: value.to_string(), + }), + } +} + #[tokio::test] async fn test_grpc_server_start() { let (server_addr, _temp_dir) = start_test_server().await.unwrap(); @@ -87,15 +117,16 @@ async fn test_grpc_server_start() { // insert a test vector let test_vec = vec![1.0, 2.0, 3.0]; - let mut request = tonic::Request::new(InsertVectorRequest { - vector: Some(DenseVector { + let mut request = tonic::Request::new(PointInput { + dense_vector: Some(DenseVector { values: test_vec.clone(), }), payload: Some(Payload::default()), + ..Default::default() }); append_test_auth_header(&mut request, TEST_AUTH_BEARER_TOKEN); - let status = client.insert_vector(request).await.is_ok(); + let status = client.insert_point(request).await.is_ok(); assert!(status); } @@ -107,18 +138,19 @@ async fn test_insert_vector_rpc() { // insert a test vector let test_vec = vec![1.0, 2.0, 3.0]; - let mut request = tonic::Request::new(InsertVectorRequest { - vector: Some(DenseVector { + let mut request = tonic::Request::new(PointInput { + dense_vector: Some(DenseVector { values: test_vec.clone(), }), payload: Some(Payload { content_type: ContentType::Text as i32, content: "test".to_string(), }), + ..Default::default() }); append_test_auth_header(&mut request, TEST_AUTH_BEARER_TOKEN); - let resp = client.insert_vector(request).await; + let resp = client.insert_point(request).await; // check if request is successful assert!(resp.is_ok()); @@ -133,7 +165,7 @@ async fn test_insert_vector_rpc() { // check if request is successful assert!(resp.is_ok()); let point = resp.unwrap().into_inner(); - assert_eq!(point.vector.unwrap().values, test_vec); + assert_eq!(point.dense_vector.unwrap().values, test_vec); // payload assertions let payload = point.payload.unwrap(); @@ -141,15 +173,16 @@ async fn test_insert_vector_rpc() { assert_eq!(payload.content, "test"); // insert a new vector with mismatched dimensions - let mut request = tonic::Request::new(InsertVectorRequest { - vector: Some(DenseVector { + let mut request = tonic::Request::new(PointInput { + dense_vector: Some(DenseVector { values: vec![1.0, 2.0], }), payload: Some(Payload::default()), + ..Default::default() }); append_test_auth_header(&mut request, TEST_AUTH_BEARER_TOKEN); - let resp = client.insert_vector(request).await; + let resp = client.insert_point(request).await; // request must fail assert!(resp.is_err()); @@ -162,15 +195,16 @@ async fn test_delete_vector_rpc() { // insert a test vector let test_vec = vec![1.0, 2.0, 3.0]; - let mut request = tonic::Request::new(InsertVectorRequest { - vector: Some(DenseVector { + let mut request = tonic::Request::new(PointInput { + dense_vector: Some(DenseVector { values: test_vec.clone(), }), payload: Some(Payload::default()), + ..Default::default() }); append_test_auth_header(&mut request, TEST_AUTH_BEARER_TOKEN); - let resp = client.insert_vector(request).await; + let resp = client.insert_point(request).await; // check if request is successful assert!(resp.is_ok()); @@ -204,15 +238,16 @@ async fn test_search_vector_rpc() { // insert a test vector let test_vec = vec![1.0, 2.0, 3.0]; - let mut request = tonic::Request::new(InsertVectorRequest { - vector: Some(DenseVector { + let mut request = tonic::Request::new(PointInput { + dense_vector: Some(DenseVector { values: test_vec.clone(), }), payload: Some(Payload::default()), + ..Default::default() }); append_test_auth_header(&mut request, TEST_AUTH_BEARER_TOKEN); - let resp = client.insert_vector(request).await; + let resp = client.insert_point(request).await; // check if request is successful assert!(resp.is_ok()); @@ -221,28 +256,21 @@ async fn test_search_vector_rpc() { let query_vec = vec![2.0, 2.0, 2.0]; // search for the vector - let mut request = tonic::Request::new(SearchRequest { - query_vector: Some(DenseVector { - values: query_vec.clone(), - }), - similarity: 0, // euclidean distance - limit: 1, - ef: 0, - }); + let mut request = tonic::Request::new(dense_search(query_vec.clone(), 0, 1)); append_test_auth_header(&mut request, TEST_AUTH_BEARER_TOKEN); - let resp = client.search_points(request).await; + let resp = client.search(request).await; // check if request is successful assert!(resp.is_ok()); let result = resp.unwrap().into_inner(); // 1 vector has to be returned - assert_eq!(result.result_point_ids.len(), 1); + assert_eq!(result.hits.len(), 1); // check if the returned point id matches the inserted point id - assert_eq!(result.result_point_ids[0], PointId { id: point.id }); - assert_eq!(result.scores, vec![2f32.sqrt()]); + assert_eq!(result.hits[0].id, Some(PointId { id: point.id })); + assert_eq!(result.hits[0].score, 2f32.sqrt()); } #[tokio::test] @@ -252,15 +280,16 @@ async fn test_unauthorized_rpc() { // insert a test vector let test_vec = vec![1.0, 2.0, 3.0]; - let mut request = tonic::Request::new(InsertVectorRequest { - vector: Some(DenseVector { + let mut request = tonic::Request::new(PointInput { + dense_vector: Some(DenseVector { values: test_vec.clone(), }), payload: Some(Payload::default()), + ..Default::default() }); append_test_auth_header(&mut request, "43121"); - let resp = client.insert_vector(request).await; + let resp = client.insert_point(request).await; // request must fail assert!(resp.is_err()); @@ -271,15 +300,16 @@ async fn test_readonly_key_cannot_write() { let (server_addr, _temp_dir) = start_test_server().await.unwrap(); let mut client = create_test_client(server_addr).await.unwrap(); - let mut request = tonic::Request::new(InsertVectorRequest { - vector: Some(DenseVector { + let mut request = tonic::Request::new(PointInput { + dense_vector: Some(DenseVector { values: vec![1.0, 2.0, 3.0], }), payload: Some(Payload::default()), + ..Default::default() }); append_test_auth_header(&mut request, TEST_READONLY_BEARER_TOKEN); - let resp = client.insert_vector(request).await; + let resp = client.insert_point(request).await; assert_eq!(resp.unwrap_err().code(), tonic::Code::PermissionDenied); } @@ -289,17 +319,95 @@ async fn test_readonly_key_can_search() { let (server_addr, _temp_dir) = start_test_server().await.unwrap(); let mut client = create_test_client(server_addr).await.unwrap(); - let mut request = tonic::Request::new(SearchRequest { - query_vector: Some(DenseVector { - values: vec![1.0, 2.0, 3.0], - }), - similarity: 0, - limit: 1, - ef: 0, - }); + let mut request = tonic::Request::new(dense_search(vec![1.0, 2.0, 3.0], 0, 1)); append_test_auth_header(&mut request, TEST_READONLY_BEARER_TOKEN); - let resp = client.search_points(request).await; + let resp = client.search(request).await; assert!(resp.is_ok()); } + +#[tokio::test] +async fn test_sparse_and_mixed_batch_rpc() { + let (server_addr, _temp_dir) = start_test_server().await.unwrap(); + let mut client = create_test_client(server_addr).await.unwrap(); + let sparse_id = "00000000-0000-0000-0000-000000000002"; + + let mut request = tonic::Request::new(InsertPointsBatchRequest { + points: vec![ + PointInput { + dense_vector: Some(DenseVector { + values: vec![1.0, 0.0, 0.0], + }), + ..Default::default() + }, + PointInput { + id: Some(point_id(sparse_id)), + sparse_vector: Some(SparseVector { + token_ids: vec![1, 4], + weights: vec![0.5, 2.0], + }), + ..Default::default() + }, + ], + }); + append_test_auth_header(&mut request, TEST_AUTH_BEARER_TOKEN); + let ids = client + .insert_points_batch(request) + .await + .unwrap() + .into_inner() + .ids; + assert_eq!(ids[1], point_id(sparse_id)); + + let mut request = tonic::Request::new(point_id(sparse_id)); + append_test_auth_header(&mut request, TEST_AUTH_BEARER_TOKEN); + let point = client.get_point(request).await.unwrap().into_inner(); + assert!(point.dense_vector.is_none()); + assert_eq!(point.sparse_vector.unwrap().token_ids, vec![1, 4]); + + let mut request = tonic::Request::new(SearchBatchRequest { + queries: vec![ + dense_search(vec![1.0, 0.0, 0.0], 3, 5), + sparse_search(vec![4], vec![1.0], 5), + ], + }); + append_test_auth_header(&mut request, TEST_AUTH_BEARER_TOKEN); + let results = client + .search_batch(request) + .await + .unwrap() + .into_inner() + .results; + assert_eq!(results[0].hits[0].id, Some(ids[0].clone())); + assert_eq!(results[1].hits.len(), 1); + assert_eq!(results[1].hits[0].id, Some(point_id(sparse_id))); + assert_eq!(results[1].hits[0].score, 2.0); +} + +#[tokio::test] +async fn test_invalid_requests_return_invalid_argument() { + let (server_addr, _temp_dir) = start_test_server().await.unwrap(); + let mut client = create_test_client(server_addr).await.unwrap(); + + let mut request = tonic::Request::new(point_id("not-a-uuid")); + append_test_auth_header(&mut request, TEST_AUTH_BEARER_TOKEN); + let status = client.get_point(request).await.unwrap_err(); + assert_eq!(status.code(), tonic::Code::InvalidArgument); + + let mut request = tonic::Request::new(SearchRequest { + limit: 1, + query: None, + }); + append_test_auth_header(&mut request, TEST_AUTH_BEARER_TOKEN); + let status = client.search(request).await.unwrap_err(); + assert_eq!(status.code(), tonic::Code::InvalidArgument); + + let mut request = tonic::Request::new(PointInput { + payload: Some(Payload::default()), + ..Default::default() + }); + append_test_auth_header(&mut request, TEST_AUTH_BEARER_TOKEN); + let status = client.insert_point(request).await.unwrap_err(); + assert_eq!(status.code(), tonic::Code::InvalidArgument); +} diff --git a/crates/http/src/handler.rs b/crates/http/src/handler.rs index e8e171d..eb03e0f 100644 --- a/crates/http/src/handler.rs +++ b/crates/http/src/handler.rs @@ -5,8 +5,8 @@ use axum::{ http::StatusCode, }; use defs::{ - AppError, BatchInsertRequest, BatchInsertResponse, BatchSearchRequest, BatchSearchResponse, - DenseVector, Payload, Point, PointId, SearchQueryInput, SearchResponse, + BatchInsertRequest, BatchInsertResponse, BatchSearchRequest, BatchSearchResponse, Point, + PointId, PointInput, SearchQuery, SearchResponse, }; use index::error::IndexError; use serde::{Deserialize, Serialize}; @@ -15,12 +15,6 @@ use tracing::error; use crate::AppState; -#[derive(Deserialize, Debug)] -pub struct InsertRequest { - pub vector: DenseVector, - pub payload: Payload, -} - #[derive(Serialize, Deserialize)] pub struct InsertResponse { pub point_id: PointId, @@ -36,9 +30,9 @@ pub async fn health_handler() -> &'static str { pub async fn insert_point_handler( State(app_state): State, - Json(request): Json, + Json(request): Json, ) -> Result<(StatusCode, Json), (StatusCode, String)> { - match app_state.db.insert(request.vector, request.payload) { + match app_state.db.insert(request) { Ok(point_id) => { let response = InsertResponse { point_id }; Ok((StatusCode::CREATED, Json(response))) @@ -53,11 +47,11 @@ pub async fn insert_point_handler( pub async fn batch_insert_handler( State(state): State, Json(request): Json, -) -> Result, AppError> { - let ids = state - .db - .insert_batch(request.points) - .map_err(|e| AppError::Api(e.to_string()))?; +) -> Result, (StatusCode, String)> { + let ids = state.db.insert_batch(request.points).map_err(|e| { + error!("Failed to insert batch: {:?}", e); + api_error_to_response(&e) + })?; Ok(Json(BatchInsertResponse { inserted: ids.len(), @@ -94,14 +88,10 @@ pub async fn delete_point_handler( pub async fn search_points_handler( State(app_state): State, - Json(request): Json, + Json(request): Json, ) -> Result, (StatusCode, String)> { match app_state.db.search(request) { - Ok(results) => { - let (results, scores) = results.into_iter().unzip(); - let response = SearchResponse { results, scores }; - Ok(Json(response)) - } + Ok(hits) => Ok(Json(SearchResponse { hits })), Err(e) => { error!("Failed to search points: {:?}", e); Err(api_error_to_response(&e)) @@ -112,16 +102,16 @@ pub async fn search_points_handler( pub async fn batch_search_handler( State(state): State, Json(request): Json, -) -> Result, AppError> { +) -> Result, (StatusCode, String)> { let results = state .db .search_batch(request.queries) - .map_err(|e| AppError::Api(e.to_string()))? + .map_err(|e| { + error!("Failed to search batch: {:?}", e); + api_error_to_response(&e) + })? .into_iter() - .map(|found| { - let (results, scores) = found.into_iter().unzip(); - SearchResponse { results, scores } - }) + .map(|hits| SearchResponse { hits }) .collect(); Ok(Json(BatchSearchResponse { results })) @@ -133,6 +123,9 @@ fn api_error_to_response(err: &ApiError) -> (StatusCode, String) { ApiError::PointNotFound { .. } => (StatusCode::NOT_FOUND, err.to_string()), ApiError::DimensionMismatch { .. } => (StatusCode::BAD_REQUEST, err.to_string()), ApiError::InvalidSearchLimit { .. } => (StatusCode::BAD_REQUEST, err.to_string()), + ApiError::InvalidSparseVector { .. } => (StatusCode::BAD_REQUEST, err.to_string()), + ApiError::PointAlreadyExists { .. } => (StatusCode::CONFLICT, err.to_string()), + ApiError::MissingVector { .. } => (StatusCode::BAD_REQUEST, err.to_string()), ApiError::LockError => ( StatusCode::SERVICE_UNAVAILABLE, "Server is busy, try again".to_string(), diff --git a/crates/http/src/tests.rs b/crates/http/src/tests.rs index 204ed8b..0e28373 100644 --- a/crates/http/src/tests.rs +++ b/crates/http/src/tests.rs @@ -52,7 +52,7 @@ async fn in_memory_storage_http_smoke_test() { .post("/points") .add_header("api-key", API_KEY) .json(&json!({ - "vector": [1.0, 0.0, 0.0], + "dense_vector": [1.0, 0.0, 0.0], "payload": { "content_type": "Text", "content": "smoke-test" @@ -70,21 +70,19 @@ async fn in_memory_storage_http_smoke_test() { get_response.assert_status_ok(); let point_body: serde_json::Value = get_response.json(); assert_eq!(point_body["payload"]["content"], "smoke-test"); - assert_eq!(point_body["vector"], json!([1.0, 0.0, 0.0])); + assert_eq!(point_body["dense_vector"], json!([1.0, 0.0, 0.0])); let search_response = server .post("/points/search") .add_header("api-key", API_KEY) .json(&json!({ - "vector": [1.0, 0.0, 0.0], - "similarity": "Cosine", - "limit": 1 + "limit": 1, + "query": {"dense": {"vector": [1.0, 0.0, 0.0], "similarity": "Cosine"}} })) .await; search_response.assert_status_ok(); let search_body: serde_json::Value = search_response.json(); - assert_eq!(search_body["results"], json!([point_id])); - assert_eq!(search_body["scores"], json!([0.0])); + assert_eq!(search_body["hits"], json!([{"id": point_id, "score": 0.0}])); let delete_response = server .delete(&format!("/points/{point_id}")) @@ -115,7 +113,7 @@ async fn protected_routes_reject_missing_key() { .assert_status(StatusCode::UNAUTHORIZED); server .post("/points") - .json(&json!({"vector": [1.0, 0.0, 0.0], "payload": {"content_type": "Text", "content": "x"}})) + .json(&json!({"dense_vector": [1.0, 0.0, 0.0], "payload": {"content_type": "Text", "content": "x"}})) .await .assert_status(StatusCode::UNAUTHORIZED); } @@ -137,7 +135,7 @@ async fn readonly_key_can_read_but_not_write() { let insert_response = server .post("/points") .add_header("api-key", API_KEY) - .json(&json!({"vector": [1.0, 0.0, 0.0], "payload": {"content_type": "Text", "content": "x"}})) + .json(&json!({"dense_vector": [1.0, 0.0, 0.0], "payload": {"content_type": "Text", "content": "x"}})) .await; let point_id = insert_response.json::()["point_id"] .as_str() @@ -153,7 +151,7 @@ async fn readonly_key_can_read_but_not_write() { server .post("/points/search") .add_header("api-key", READONLY_API_KEY) - .json(&json!({"vector": [1.0, 0.0, 0.0], "similarity": "Cosine", "limit": 1})) + .json(&json!({"limit": 1, "query": {"dense": {"vector": [1.0, 0.0, 0.0], "similarity": "Cosine"}}})) .await .assert_status_ok(); @@ -166,7 +164,74 @@ async fn readonly_key_can_read_but_not_write() { server .post("/points") .add_header("api-key", READONLY_API_KEY) - .json(&json!({"vector": [1.0, 0.0, 0.0], "payload": {"content_type": "Text", "content": "y"}})) + .json(&json!({"dense_vector": [1.0, 0.0, 0.0], "payload": {"content_type": "Text", "content": "y"}})) .await .assert_status(StatusCode::FORBIDDEN); } + +#[tokio::test] +async fn sparse_and_dense_batch_round_trip() { + let server = test_server(); + let sparse_id = "00000000-0000-0000-0000-000000000002"; + + let insert_response = server + .post("/points/batch") + .add_header("api-key", API_KEY) + .json(&json!({"points": [ + {"dense_vector": [1.0, 0.0, 0.0]}, + {"id": sparse_id, "sparse_vector": {"token_ids": [1, 4], "weights": [0.5, 2.0]}} + ]})) + .await; + insert_response.assert_status_ok(); + let dense_id = insert_response.json::()["ids"][0].clone(); + + let search_response = server + .post("/points/search/batch") + .add_header("api-key", API_KEY) + .json(&json!({"queries": [ + {"limit": 5, "query": {"dense": {"vector": [1.0, 0.0, 0.0], "similarity": "Cosine"}}}, + {"limit": 5, "query": {"sparse": {"token_ids": [4], "weights": [1.0]}}} + ]})) + .await; + search_response.assert_status_ok(); + let body: serde_json::Value = search_response.json(); + assert_eq!(body["results"][0]["hits"][0]["id"], dense_id); + assert_eq!( + body["results"][1]["hits"], + json!([{"id": sparse_id, "score": 2.0}]) + ); + + let point: serde_json::Value = server + .get(&format!("/points/{sparse_id}")) + .add_header("api-key", API_KEY) + .await + .json(); + assert_eq!(point["dense_vector"], json!(null)); + assert_eq!(point["sparse_vector"]["token_ids"], json!([1, 4])); +} + +#[tokio::test] +async fn batch_errors_use_api_status_codes() { + let server = test_server(); + let id = "00000000-0000-0000-0000-000000000003"; + let point = json!({"id": id, "dense_vector": [1.0, 0.0, 0.0]}); + + server + .post("/points/batch") + .add_header("api-key", API_KEY) + .json(&json!({"points": [point]})) + .await + .assert_status_ok(); + server + .post("/points/batch") + .add_header("api-key", API_KEY) + .json(&json!({"points": [point]})) + .await + .assert_status(StatusCode::CONFLICT); + server + .post("/points") + .add_header("api-key", API_KEY) + .json(&json!({"payload": {"content_type": "Text", "content": "no vectors"}})) + .await + .assert_status(StatusCode::BAD_REQUEST); +} diff --git a/crates/index/src/inverted/index.rs b/crates/index/src/inverted/index.rs new file mode 100644 index 0000000..15f2447 --- /dev/null +++ b/crates/index/src/inverted/index.rs @@ -0,0 +1,77 @@ +use std::{ + cmp::Reverse, + collections::{BinaryHeap, HashMap, hash_map::Entry}, +}; + +use defs::{OrdF32, PointId, SparseVector, TokenID}; + +use crate::inverted::types::PostingList; + +pub struct InvertedIndex { + pub(super) postings: HashMap, +} + +impl Default for InvertedIndex { + fn default() -> Self { + Self::new() + } +} + +impl InvertedIndex { + pub fn new() -> Self { + InvertedIndex { + postings: HashMap::new(), + } + } + + pub fn insert(&mut self, point_id: PointId, vector: &SparseVector) { + for (&token, &weight) in vector.token_ids.iter().zip(&vector.weights) { + self.postings + .entry(token) + .or_default() + .insert(point_id, weight); + } + } + + pub fn search(&self, query: &SparseVector, k: usize) -> Vec<(PointId, f32)> { + if k == 0 { + return Vec::new(); + } + let mut scores: HashMap = HashMap::new(); + for (&token, &query_weight) in query.token_ids.iter().zip(&query.weights) { + let Some(list) = self.postings.get(&token) else { + continue; + }; + for (point_id, weight) in list.iter() { + *scores.entry(point_id).or_default() += query_weight * weight; + } + } + + let mut heap: BinaryHeap<(Reverse, PointId)> = + BinaryHeap::with_capacity(k.min(scores.len()) + 1); + for (point_id, score) in scores { + heap.push((Reverse(OrdF32(score)), point_id)); + if heap.len() > k { + heap.pop(); + } + } + + heap.into_sorted_vec() + .into_iter() + .map(|(Reverse(score), point_id)| (point_id, score.0)) + .collect() + } + + pub fn delete(&mut self, point_id: PointId, vector: &SparseVector) -> bool { + let mut removed = false; + for &token in &vector.token_ids { + if let Entry::Occupied(mut entry) = self.postings.entry(token) { + removed |= entry.get_mut().remove(point_id); + if entry.get().is_empty() { + entry.remove(); + } + } + } + removed + } +} diff --git a/crates/index/src/inverted/mod.rs b/crates/index/src/inverted/mod.rs new file mode 100644 index 0000000..19ff5d8 --- /dev/null +++ b/crates/index/src/inverted/mod.rs @@ -0,0 +1,6 @@ +pub mod index; +pub mod types; +pub use index::InvertedIndex; + +#[cfg(test)] +mod tests; diff --git a/crates/index/src/inverted/tests.rs b/crates/index/src/inverted/tests.rs new file mode 100644 index 0000000..0347567 --- /dev/null +++ b/crates/index/src/inverted/tests.rs @@ -0,0 +1,92 @@ +use super::index::InvertedIndex; +use super::types::PostingList; +use defs::{SparseVector, TokenID}; +use uuid::Uuid; + +const ID1: Uuid = Uuid::from_u128(1); +const ID2: Uuid = Uuid::from_u128(2); +const ID3: Uuid = Uuid::from_u128(3); + +fn sparse(token_ids: &[TokenID], weights: &[f32]) -> SparseVector { + SparseVector { + token_ids: token_ids.to_vec(), + weights: weights.to_vec(), + } +} + +fn sample_index() -> InvertedIndex { + let mut index = InvertedIndex::new(); + index.insert(ID1, &sparse(&[1, 2], &[1.0, 2.0])); + index.insert(ID2, &sparse(&[2, 3], &[1.0, 3.0])); + index.insert(ID3, &sparse(&[4], &[5.0])); + index +} + +#[test] +fn test_posting_list_remove() { + let mut list = PostingList::default(); + list.insert(ID1, 1.0); + list.insert(ID2, 2.0); + + assert!(list.remove(ID1)); + assert_eq!(list.iter().collect::>(), vec![(ID2, 2.0)]); + assert!(!list.remove(ID1)); + + assert!(list.remove(ID2)); + assert!(list.is_empty()); +} + +#[test] +fn test_search_scores_by_dot_product() { + let index = sample_index(); + let results = index.search(&sparse(&[2, 3], &[1.0, 1.0]), 10); + assert_eq!(results, vec![(ID2, 4.0), (ID1, 2.0)]); +} + +#[test] +fn test_search_respects_limit() { + let index = sample_index(); + let results = index.search(&sparse(&[2, 3], &[1.0, 1.0]), 1); + assert_eq!(results, vec![(ID2, 4.0)]); +} + +#[test] +fn test_search_k_zero_returns_empty() { + let index = sample_index(); + assert!(index.search(&sparse(&[2], &[1.0]), 0).is_empty()); +} + +#[test] +fn test_search_unknown_token_returns_empty() { + let index = sample_index(); + assert!(index.search(&sparse(&[99], &[1.0]), 10).is_empty()); +} + +#[test] +fn test_search_empty_index_returns_empty() { + let index = InvertedIndex::new(); + assert!(index.search(&sparse(&[1], &[1.0]), 10).is_empty()); +} + +#[test] +fn test_delete_removes_point_from_results() { + let mut index = sample_index(); + + assert!(index.delete(ID2, &sparse(&[2, 3], &[1.0, 3.0]))); + + let results = index.search(&sparse(&[2, 3], &[1.0, 1.0]), 10); + assert_eq!(results, vec![(ID1, 2.0)]); + + assert!(index.postings.contains_key(&2)); + assert!(!index.postings.contains_key(&3)); +} + +#[test] +fn test_delete_missing_point_returns_false() { + let mut index = sample_index(); + let vector = sparse(&[4], &[5.0]); + + assert!(index.delete(ID3, &vector)); + assert!(!index.delete(ID3, &vector)); + assert!(index.search(&vector, 10).is_empty()); +} diff --git a/crates/index/src/inverted/types.rs b/crates/index/src/inverted/types.rs new file mode 100644 index 0000000..5a75620 --- /dev/null +++ b/crates/index/src/inverted/types.rs @@ -0,0 +1,34 @@ +use defs::{Element, PointId}; + +#[derive(Default, Debug, Clone)] +pub struct PostingList { + entries: Vec<(PointId, Element)>, +} + +impl PostingList { + pub fn insert(&mut self, point_id: PointId, weight: Element) { + self.entries.push((point_id, weight)); + } + + pub fn remove(&mut self, point_id: PointId) -> bool { + match self.entries.iter().position(|&(id, _)| id == point_id) { + Some(pos) => { + self.entries.swap_remove(pos); + true + } + None => false, + } + } + + pub fn iter(&self) -> impl Iterator + '_ { + self.entries.iter().copied() + } + + pub fn len(&self) -> usize { + self.entries.len() + } + + pub fn is_empty(&self) -> bool { + self.entries.is_empty() + } +} diff --git a/crates/index/src/lib.rs b/crates/index/src/lib.rs index 48452e3..215876d 100644 --- a/crates/index/src/lib.rs +++ b/crates/index/src/lib.rs @@ -7,6 +7,7 @@ use serde::{Deserialize, Serialize}; use storage::StorageEngine; pub mod flat; pub mod hnsw; +pub mod inverted; pub mod kd_tree; pub trait VectorIndex: Send + Sync + SerializableIndex { diff --git a/crates/storage/src/in_memory.rs b/crates/storage/src/in_memory.rs index 3ee86fb..cf06c08 100644 --- a/crates/storage/src/in_memory.rs +++ b/crates/storage/src/in_memory.rs @@ -1,8 +1,8 @@ use crate::StorageType; use crate::error::StorageError; -use crate::{StorageEngine, VectorPage, checkpoint::StorageCheckpoint}; +use crate::{PointPage, StorageEngine, checkpoint::StorageCheckpoint}; use bincode::{deserialize_from, serialize_into}; -use defs::{DenseVector, Payload, Point, PointId}; +use defs::{DenseVector, Payload, Point, PointId, SparseVector}; use std::collections::BTreeMap; use std::fs::File; use std::io::{BufReader, BufWriter, Read, Write}; @@ -39,6 +39,7 @@ impl StorageEngine for MemoryStorage { &self, id: PointId, vector: Option, + sparse_vector: Option, payload: Option, ) -> Result<(), StorageError> { let mut points = self @@ -49,7 +50,8 @@ impl StorageEngine for MemoryStorage { id, Point { id, - vector, + dense_vector: vector, + sparse_vector, payload, }, ); @@ -70,25 +72,25 @@ impl StorageEngine for MemoryStorage { points.remove(&id); Ok(()) } - fn get_payload(&self, id: PointId) -> Result, StorageError> { + fn get_point(&self, id: PointId) -> Result, StorageError> { let points = self .points .read() .map_err(|_| StorageError::InMemoryLock {})?; - Ok(points.get(&id).and_then(|point| point.payload.clone())) + Ok(points.get(&id).cloned()) } fn get_vector(&self, id: PointId) -> Result, StorageError> { let points = self .points .read() .map_err(|_| StorageError::InMemoryLock {})?; - Ok(points.get(&id).and_then(|point| point.vector.clone())) + Ok(points.get(&id).and_then(|point| point.dense_vector.clone())) } - fn list_vectors( + fn list_points( &self, offset: PointId, limit: usize, - ) -> Result, StorageError> { + ) -> Result, StorageError> { if limit < 1 { return Ok(None); } @@ -97,18 +99,12 @@ impl StorageEngine for MemoryStorage { .points .read() .map_err(|_| StorageError::InMemoryLock {})?; - let mut result = Vec::with_capacity(limit); - let mut last_id = offset; - - for (id, point) in points.range((Excluded(offset), Unbounded)) { - if let Some(vector) = &point.vector { - last_id = *id; - result.push((*id, vector.clone())); - if result.len() == limit { - break; - } - } - } + let result: Vec = points + .range((Excluded(offset), Unbounded)) + .take(limit) + .map(|(_, point)| point.clone()) + .collect(); + let last_id = result.last().map_or(offset, |point| point.id); Ok(Some((result, last_id))) } diff --git a/crates/storage/src/in_memory/tests.rs b/crates/storage/src/in_memory/tests.rs index 9d8510e..7ef34e5 100644 --- a/crates/storage/src/in_memory/tests.rs +++ b/crates/storage/src/in_memory/tests.rs @@ -23,7 +23,9 @@ fn test_insert_and_get_vector() { let vector = Some(vec![0.1, 0.2, 0.3]); let payload = Some(test_payload("Test")); - storage.insert_point(id, vector.clone(), payload).unwrap(); + storage + .insert_point(id, vector.clone(), None, payload) + .unwrap(); assert_eq!(storage.get_vector(id).unwrap(), vector); } @@ -34,9 +36,17 @@ fn test_insert_and_get_payload() { let id = Uuid::new_v4(); let payload = Some(test_payload("Test")); - storage.insert_point(id, None, payload.clone()).unwrap(); + storage + .insert_point(id, None, None, payload.clone()) + .unwrap(); - assert_eq!(storage.get_payload(id).unwrap(), payload); + assert_eq!( + storage + .get_point(id) + .unwrap() + .and_then(|point| point.payload), + payload + ); } #[test] @@ -47,18 +57,29 @@ fn test_contains_and_delete_point() { assert!(!storage.contains_point(id).unwrap()); storage - .insert_point(id, Some(vec![0.4, 0.5, 0.6]), Some(test_payload("Test"))) + .insert_point( + id, + Some(vec![0.4, 0.5, 0.6]), + None, + Some(test_payload("Test")), + ) .unwrap(); assert!(storage.contains_point(id).unwrap()); storage.delete_point(id).unwrap(); assert!(!storage.contains_point(id).unwrap()); assert_eq!(storage.get_vector(id).unwrap(), None); - assert_eq!(storage.get_payload(id).unwrap(), None); + assert_eq!( + storage + .get_point(id) + .unwrap() + .and_then(|point| point.payload), + None + ); } #[test] -fn test_list_vectors_respects_offset_limit_and_skips_payload_only_points() { +fn test_list_points_respects_offset_and_limit() { let storage = create_test_storage(); let ids = [ Uuid::from_u128(1), @@ -66,37 +87,70 @@ fn test_list_vectors_respects_offset_limit_and_skips_payload_only_points() { Uuid::from_u128(3), Uuid::from_u128(4), ]; + let sparse_vector = SparseVector { + token_ids: vec![7, 9], + weights: vec![0.5, 1.5], + }; storage - .insert_point(ids[0], Some(vec![1.0, 1.1]), Some(test_payload("one"))) + .insert_point(ids[0], Some(vec![1.0, 1.1]), None, None) .unwrap(); storage - .insert_point(ids[1], None, Some(test_payload("payload-only"))) + .insert_point(ids[1], None, Some(sparse_vector.clone()), None) .unwrap(); storage - .insert_point(ids[2], Some(vec![3.0, 3.1]), Some(test_payload("three"))) + .insert_point(ids[2], None, None, Some(test_payload("payload-only"))) .unwrap(); storage - .insert_point(ids[3], Some(vec![4.0, 4.1]), Some(test_payload("four"))) + .insert_point(ids[3], Some(vec![4.0, 4.1]), None, None) .unwrap(); - let (first_page, next_offset) = storage.list_vectors(Uuid::nil(), 2).unwrap().unwrap(); - assert_eq!( - first_page, - vec![(ids[0], vec![1.0, 1.1]), (ids[2], vec![3.0, 3.1])] - ); - assert_eq!(next_offset, ids[2]); + let (first_page, next_offset) = storage.list_points(Uuid::nil(), 2).unwrap().unwrap(); + let first_ids: Vec = first_page.iter().map(|point| point.id).collect(); + assert_eq!(first_ids, vec![ids[0], ids[1]]); + assert_eq!(first_page[1].sparse_vector, Some(sparse_vector)); + assert_eq!(next_offset, ids[1]); + + let (second_page, next_offset) = storage.list_points(next_offset, 2).unwrap().unwrap(); + let second_ids: Vec = second_page.iter().map(|point| point.id).collect(); + assert_eq!(second_ids, vec![ids[2], ids[3]]); + assert_eq!(next_offset, ids[3]); - let (second_page, next_offset) = storage.list_vectors(next_offset, 2).unwrap().unwrap(); - assert_eq!(second_page, vec![(ids[3], vec![4.0, 4.1])]); + let (empty_page, next_offset) = storage.list_points(next_offset, 2).unwrap().unwrap(); + assert!(empty_page.is_empty()); assert_eq!(next_offset, ids[3]); } #[test] -fn test_list_vectors_with_zero_limit_returns_none() { +fn test_list_points_with_zero_limit_returns_none() { let storage = create_test_storage(); - assert_eq!(storage.list_vectors(Uuid::nil(), 0).unwrap(), None); + assert_eq!(storage.list_points(Uuid::nil(), 0).unwrap(), None); +} + +#[test] +fn test_get_point_returns_all_fields() { + let storage = create_test_storage(); + let id = Uuid::new_v4(); + let sparse_vector = SparseVector { + token_ids: vec![1], + weights: vec![2.0], + }; + + storage + .insert_point( + id, + Some(vec![0.1, 0.2]), + Some(sparse_vector.clone()), + Some(test_payload("both")), + ) + .unwrap(); + + let point = storage.get_point(id).unwrap().unwrap(); + assert_eq!(point.dense_vector, Some(vec![0.1, 0.2])); + assert_eq!(point.sparse_vector, Some(sparse_vector)); + assert_eq!(point.payload, Some(test_payload("both"))); + assert_eq!(storage.get_point(Uuid::new_v4()).unwrap(), None); } #[test] @@ -110,6 +164,7 @@ fn test_create_and_restore_checkpoint() { .insert_point( id_before_checkpoint, Some(vec![0.1, 0.2, 0.3]), + None, Some(test_payload("before")), ) .unwrap(); @@ -119,6 +174,7 @@ fn test_create_and_restore_checkpoint() { .insert_point( id_after_checkpoint, Some(vec![0.4, 0.5, 0.6]), + None, Some(test_payload("after")), ) .unwrap(); @@ -128,7 +184,10 @@ fn test_create_and_restore_checkpoint() { assert!(storage.contains_point(id_before_checkpoint).unwrap()); assert!(!storage.contains_point(id_after_checkpoint).unwrap()); assert_eq!( - storage.get_payload(id_before_checkpoint).unwrap(), + storage + .get_point(id_before_checkpoint) + .unwrap() + .and_then(|point| point.payload), Some(test_payload("before")) ); } @@ -140,7 +199,12 @@ fn test_checkpoint_writes_header() { let id = Uuid::new_v4(); storage - .insert_point(id, Some(vec![0.1, 0.2, 0.3]), Some(test_payload("point"))) + .insert_point( + id, + Some(vec![0.1, 0.2, 0.3]), + None, + Some(test_payload("point")), + ) .unwrap(); let checkpoint = storage.checkpoint_at(temp_dir.path()).unwrap(); diff --git a/crates/storage/src/lib.rs b/crates/storage/src/lib.rs index 31f8170..9bd97b3 100644 --- a/crates/storage/src/lib.rs +++ b/crates/storage/src/lib.rs @@ -1,11 +1,11 @@ use crate::rocks_db::RocksDbStorage; -use defs::{DenseVector, Payload, PointId}; +use defs::{DenseVector, Payload, Point, PointId, SparseVector}; use serde::{Deserialize, Serialize}; use std::path::{Path, PathBuf}; use std::sync::Arc; pub mod checkpoint; -pub type VectorPage = (Vec<(PointId, DenseVector)>, PointId); +pub type PointPage = (Vec, PointId); pub mod error; pub mod in_memory; @@ -18,13 +18,14 @@ pub trait StorageEngine: Send + Sync { &self, id: PointId, vector: Option, + sparse_vector: Option, payload: Option, ) -> Result<()>; + fn get_point(&self, id: PointId) -> Result>; fn get_vector(&self, id: PointId) -> Result>; - fn get_payload(&self, id: PointId) -> Result>; fn delete_point(&self, id: PointId) -> Result<()>; fn contains_point(&self, id: PointId) -> Result; - fn list_vectors(&self, offset: PointId, limit: usize) -> Result>; + fn list_points(&self, offset: PointId, limit: usize) -> Result>; fn checkpoint_at(&self, path: &Path) -> Result; fn restore_checkpoint(&mut self, checkpoint: &checkpoint::StorageCheckpoint) -> Result<()>; diff --git a/crates/storage/src/rocks_db.rs b/crates/storage/src/rocks_db.rs index 3fd2828..a691a34 100644 --- a/crates/storage/src/rocks_db.rs +++ b/crates/storage/src/rocks_db.rs @@ -5,9 +5,9 @@ use crate::error::{ self, RocksDbCheckpointIoSnafu, RocksDbCheckpointMsgSnafu, RocksDbCheckpointSnafu, RocksDbFlushSnafu, RocksDbInitializationSnafu, StorageError, }; -use crate::{StorageEngine, StorageType, VectorPage}; +use crate::{PointPage, StorageEngine, StorageType}; use bincode::{deserialize, serialize}; -use defs::{DenseVector, Payload, Point, PointId}; +use defs::{DenseVector, Payload, Point, PointId, SparseVector}; use flate2::{Compression, read::GzDecoder, write::GzEncoder}; use rocksdb::{DB, Error, Options}; use snafu::{OptionExt, ResultExt}; @@ -70,12 +70,14 @@ impl StorageEngine for RocksDbStorage { &self, id: PointId, vector: Option, + sparse_vector: Option, payload: Option, ) -> Result<(), StorageError> { let key = id.to_string(); let point = Point { id, - vector, + dense_vector: vector, + sparse_vector, payload, }; let value = serialize(&point).context(error::SerializationSnafu { id })?; @@ -122,7 +124,7 @@ impl StorageEngine for RocksDbStorage { Ok(()) } - fn get_payload(&self, id: PointId) -> Result, StorageError> { + fn get_point(&self, id: PointId) -> Result, StorageError> { let key = id.to_string(); let Some(value_serialized) = self .db @@ -134,35 +136,21 @@ impl StorageEngine for RocksDbStorage { return Ok(None); // This should not return error but rather give None }; - let value = + let point = deserialize::(&value_serialized).context(error::DeserializationSnafu { id })?; - Ok(value.payload) + Ok(Some(point)) } fn get_vector(&self, id: PointId) -> Result, StorageError> { - let key = id.to_string(); - let Some(value_serialized) = self - .db - .as_ref() - .ok_or(StorageError::RocksDbInitialization {})? - .get(key) - .context(error::RocksDbReadSnafu { id })? - else { - return Ok(None); // This should not return error but rather give None - }; - - let value = - deserialize::(&value_serialized).context(error::DeserializationSnafu { id })?; - - Ok(value.vector) + Ok(self.get_point(id)?.and_then(|point| point.dense_vector)) } - fn list_vectors( + fn list_points( &self, offset: PointId, limit: usize, - ) -> Result, StorageError> { + ) -> Result, StorageError> { if limit < 1 { return Ok(None); } @@ -189,12 +177,10 @@ impl StorageEngine for RocksDbStorage { continue; } - if let Some(vec) = point.vector { - last_id = id; - result.push((id, vec)); - if result.len() == limit { - break; - } + last_id = id; + result.push(point); + if result.len() == limit { + break; } } Ok(Some((result, last_id))) diff --git a/crates/storage/src/rocks_db/tests.rs b/crates/storage/src/rocks_db/tests.rs index 4673f2f..593da50 100644 --- a/crates/storage/src/rocks_db/tests.rs +++ b/crates/storage/src/rocks_db/tests.rs @@ -28,11 +28,57 @@ fn test_insert_and_get_vector() { content: "Test".to_string(), }); - assert!(db.insert_point(id, vector.clone(), payload).is_ok()); + assert!(db.insert_point(id, vector.clone(), None, payload).is_ok()); let result = db.get_vector(id).unwrap(); assert_eq!(result, vector); } +#[test] +fn test_insert_and_get_sparse_vector() { + let (db, _temp_dir) = create_test_db(); + let id = Uuid::new_v4(); + let sparse_vector = Some(SparseVector { + token_ids: vec![1, 5, 9], + weights: vec![0.5, 1.0, 2.0], + }); + + db.insert_point(id, None, sparse_vector.clone(), None) + .unwrap(); + assert_eq!( + db.get_point(id) + .unwrap() + .and_then(|point| point.sparse_vector), + sparse_vector + ); + assert_eq!(db.get_vector(id).unwrap(), None); +} + +#[test] +fn test_list_points_pages_dense_and_sparse_points() { + let (db, _temp_dir) = create_test_db(); + let ids = [Uuid::from_u128(1), Uuid::from_u128(2), Uuid::from_u128(3)]; + let sparse_vector = SparseVector { + token_ids: vec![2, 4], + weights: vec![1.0, 0.5], + }; + + db.insert_point(ids[0], Some(vec![0.1, 0.2]), None, None) + .unwrap(); + db.insert_point(ids[1], None, Some(sparse_vector.clone()), None) + .unwrap(); + db.insert_point(ids[2], Some(vec![0.3, 0.4]), None, None) + .unwrap(); + + let (first_page, next_offset) = db.list_points(Uuid::nil(), 2).unwrap().unwrap(); + let first_ids: Vec = first_page.iter().map(|point| point.id).collect(); + assert_eq!(first_ids, vec![ids[0], ids[1]]); + assert_eq!(first_page[1].sparse_vector, Some(sparse_vector)); + + let (second_page, _) = db.list_points(next_offset, 2).unwrap().unwrap(); + let second_ids: Vec = second_page.iter().map(|point| point.id).collect(); + assert_eq!(second_ids, vec![ids[2]]); +} + #[test] fn test_insert_and_get_payload() { let (db, _temp_dir) = create_test_db(); @@ -44,8 +90,8 @@ fn test_insert_and_get_payload() { let vector = None; // Move payload into insert_point and recreate expected for comparison - assert!(db.insert_point(id, vector, payload).is_ok()); - let result = db.get_payload(id).unwrap(); + assert!(db.insert_point(id, vector, None, payload).is_ok()); + let result = db.get_point(id).unwrap().and_then(|point| point.payload); let expected = Some(Payload { content_type: ContentType::Text, content: "Test".to_string(), @@ -65,7 +111,7 @@ fn test_contains_point() { assert!(!db.contains_point(id).unwrap()); let vector = Some(vec![0.4, 0.5, 0.6]); - db.insert_point(id, vector, payload).unwrap(); + db.insert_point(id, vector, None, payload).unwrap(); assert!(db.contains_point(id).unwrap()); } @@ -81,7 +127,7 @@ fn test_delete_point() { let vector = vec![0.7, 0.8, 0.9]; - db.insert_point(id, Some(vector), payload).unwrap(); + db.insert_point(id, Some(vector), None, payload).unwrap(); assert!(db.contains_point(id).unwrap()); @@ -89,7 +135,10 @@ fn test_delete_point() { assert!(!db.contains_point(id).unwrap()); assert_eq!(db.get_vector(id).unwrap(), None); - assert_eq!(db.get_payload(id).unwrap(), None); + assert_eq!( + db.get_point(id).unwrap().and_then(|point| point.payload), + None + ); } #[test] @@ -105,7 +154,10 @@ fn test_get_nonexistent_payload() { let (db, _temp_dir) = create_test_db(); let id = Uuid::new_v4(); - assert_eq!(db.get_payload(id).unwrap(), None); + assert_eq!( + db.get_point(id).unwrap().and_then(|point| point.payload), + None + ); } #[test] fn test_error_context_preservation() { @@ -141,7 +193,7 @@ fn test_create_and_load_checkpoint() { }); assert!( - db.insert_point(id1, vector.clone(), payload.clone()) + db.insert_point(id1, vector.clone(), None, payload.clone()) .is_ok() ); @@ -150,7 +202,7 @@ fn test_create_and_load_checkpoint() { .expect("Failed to create checkpoint"); assert!( - db.insert_point(id2, vector.clone(), payload.clone()) + db.insert_point(id2, vector.clone(), None, payload.clone()) .is_ok() ); diff --git a/crates/tui/src/app/events.rs b/crates/tui/src/app/events.rs index c2bf708..23abedd 100644 --- a/crates/tui/src/app/events.rs +++ b/crates/tui/src/app/events.rs @@ -1,7 +1,7 @@ use super::{App, AppState, ModalType, VectorListItem}; use crate::constants::VECTOR_LIST_LIMIT; use crossterm::event::{Event, KeyCode, KeyEvent}; -use defs::{ContentType, Payload, SearchQueryInput, Similarity}; +use defs::{ContentType, DenseQuery, Payload, PointInput, QueryVector, SearchQuery, Similarity}; use std::io; use std::path::PathBuf; use uuid::Uuid; @@ -311,12 +311,14 @@ fn execute_modal_action(app: &mut App) -> io::Result<()> { } }; - let ids = db - .search(SearchQueryInput { - vector: query, - similarity: Similarity::Cosine, + let hits = db + .search(SearchQuery { limit: k, - ef: None, + query: QueryVector::Dense(DenseQuery { + vector: query, + similarity: Similarity::Cosine, + ef: None, + }), }) .map_err(to_io)?; @@ -326,9 +328,9 @@ fn execute_modal_action(app: &mut App) -> io::Result<()> { app.vector_list_post_restore = false; app.vector_list_selected_index = 0; - for (id, _) in ids.into_iter() { - if let Some(point) = db.get(id).map_err(to_io)? - && let Some(vector) = point.vector + for hit in hits { + if let Some(point) = db.get(hit.id).map_err(to_io)? + && let Some(vector) = point.dense_vector { app.vector_list_items.push(VectorListItem { id: point.id, @@ -373,7 +375,14 @@ fn execute_modal_action(app: &mut App) -> io::Result<()> { content: trimmed_text.to_string(), }; - match db.insert(vector, payload).map_err(to_io) { + let point = PointInput { + id: None, + dense_vector: Some(vector), + sparse_vector: None, + payload: Some(payload), + }; + + match db.insert(point).map_err(to_io) { Ok(new_id) => app.modal.show_success(format!( "Text embedding inserted (id={new_id}, {dims} dims)." )), @@ -424,7 +433,14 @@ fn execute_modal_action(app: &mut App) -> io::Result<()> { content: trimmed_path.to_string(), }; - match db.insert(vector, payload).map_err(to_io) { + let point = PointInput { + id: None, + dense_vector: Some(vector), + sparse_vector: None, + payload: Some(payload), + }; + + match db.insert(point).map_err(to_io) { Ok(new_id) => app.modal.show_success(format!( "Image embedding inserted (id={new_id}, {dims} dims)." )), @@ -530,28 +546,29 @@ fn fetch_next_vector_page(app: &mut App) -> io::Result { let response = db.list(offset, VECTOR_LIST_LIMIT).map_err(to_io)?; - let Some((vectors, next_offset)) = response else { + let Some((points, next_offset)) = response else { app.vector_list_next_offset = None; return Ok(false); }; - if vectors.is_empty() { + if points.is_empty() { app.vector_list_next_offset = None; return Ok(false); } - let vector_count = vectors.len(); + let point_count = points.len(); - for (id, vector) in vectors.into_iter() { - let payload = db.get(id).map_err(to_io)?.and_then(|p| p.payload); - app.vector_list_items.push(VectorListItem { - id, - vector, - payload, - }); + for point in points { + if let Some(vector) = point.dense_vector { + app.vector_list_items.push(VectorListItem { + id: point.id, + vector, + payload: point.payload, + }); + } } - if vector_count == VECTOR_LIST_LIMIT { + if point_count == VECTOR_LIST_LIMIT { app.vector_list_next_offset = Some(next_offset); } else { app.vector_list_next_offset = None; @@ -616,7 +633,7 @@ fn show_vector_info(app: &mut App, id: Uuid) -> io::Result<()> { } let point = point_opt.unwrap(); - let vector = point.vector.unwrap_or_default(); + let vector = point.dense_vector.unwrap_or_default(); let payload = point.payload; app.vector_detail = Some(VectorListItem { diff --git a/docs/api-reference/grpc.mdx b/docs/api-reference/grpc.mdx index bb544b2..65c483e 100644 --- a/docs/api-reference/grpc.mdx +++ b/docs/api-reference/grpc.mdx @@ -25,10 +25,10 @@ grpcurl -plaintext localhost:50051 list All gRPC calls require the `authorization` header with your API key: ```bash --H "authorization: your-api-key" +-H "authorization: Bearer your-api-key" ``` -Valid keys come from the JSON file pointed to by the `VORTEXDB_KEYS_FILE` environment variable, shared with the HTTP server. `readonly` keys can call `GetPoint`, `SearchPoints`, and `SearchPointsBatch`; `readwrite` keys can additionally call `InsertVector`, `InsertVectorsBatch`, and `DeletePoint`. A `readonly` key calling a write RPC gets a `PERMISSION_DENIED` status. +Valid keys come from the JSON file pointed to by the `VORTEXDB_KEYS_FILE` environment variable, shared with the HTTP server. `readonly` keys can call `GetPoint`, `Search`, and `SearchBatch`; `readwrite` keys can additionally call `InsertPoint`, `InsertPointsBatch`, and `DeletePoint`. A `readonly` key calling a write RPC gets a `PERMISSION_DENIED` status. --- @@ -39,12 +39,12 @@ syntax = "proto3"; package vectordb; service VectorDB { - rpc InsertVector(InsertVectorRequest) returns (PointID); - rpc InsertVectorsBatch(InsertVectorsBatchRequest) returns (InsertVectorsBatchResponse); + rpc InsertPoint(PointInput) returns (PointID); + rpc InsertPointsBatch(InsertPointsBatchRequest) returns (InsertPointsBatchResponse); rpc DeletePoint(PointID) returns (google.protobuf.Empty); rpc GetPoint(PointID) returns (Point); - rpc SearchPoints(SearchRequest) returns (SearchResponse); - rpc SearchPointsBatch(SearchPointsBatchRequest) returns (SearchPointsBatchResponse); + rpc Search(SearchRequest) returns (SearchResponse); + rpc SearchBatch(SearchBatchRequest) returns (SearchBatchResponse); } ``` @@ -52,55 +52,50 @@ service VectorDB { ## Methods -### InsertVector +### InsertPoint -Insert a vector with its associated payload. +Insert a point with a dense vector, a sparse vector, or both. At least one is required. - - The vector to insert. Must match the configured `DIMENSION`. + + Optional id. Generated when unset. An id that already exists returns `ALREADY_EXISTS`. - - Metadata associated with the vector. + + Must match the configured `DIMENSION`. -**Request:** + + Non-empty, strictly increasing `token_ids` with one finite, non-negative weight each. + -```protobuf -message InsertVectorRequest { - DenseVector vector = 1; - Payload payload = 2; -} -``` + + Optional metadata. + -**Response:** +**Request:** ```protobuf -message PointID { - UUID id = 1; +message PointInput { + PointID id = 1; + DenseVector dense_vector = 2; + SparseVector sparse_vector = 3; + Payload payload = 4; } ``` +**Response:** `PointID` of the created point. + **Example:** ```bash grpcurl -plaintext \ - -H "authorization: secret" \ + -H "authorization: Bearer your-api-key" \ -d '{ - "vector": {"values": [0.1, 0.2, 0.3, 0.4]}, - "payload": {"content_type": 1, "content": "Hello world"} + "dense_vector": {"values": [0.1, 0.2, 0.3, 0.4]}, + "sparse_vector": {"token_ids": [3, 17, 42], "weights": [0.5, 1.2, 0.8]}, + "payload": {"content_type": "Text", "content": "Hello, VortexDB!"} }' \ - localhost:50051 vectordb.VectorDB/InsertVector -``` - -**Response:** - -```json -{ - "id": { - "value": "550e8400-e29b-41d4-a716-446655440000" - } -} + localhost:50051 vectordb.VectorDB/InsertPoint ``` --- @@ -109,125 +104,66 @@ grpcurl -plaintext \ Retrieve a point by its ID. - - The unique identifier of the point. - - -**Request:** - -```protobuf -message PointID { - UUID id = 1; -} -``` +**Request:** `PointID` **Response:** ```protobuf message Point { PointID id = 1; - Payload payload = 2; - DenseVector vector = 3; + Payload payload = 2; // unset if the point has no payload + DenseVector dense_vector = 3; // unset if the point has no dense vector + SparseVector sparse_vector = 4; // unset if the point has no sparse vector } ``` -**Example:** +Returns `NOT_FOUND` if the point does not exist. ```bash grpcurl -plaintext \ - -H "authorization: secret" \ + -H "authorization: Bearer your-api-key" \ -d '{"id": {"value": "550e8400-e29b-41d4-a716-446655440000"}}' \ localhost:50051 vectordb.VectorDB/GetPoint ``` -**Response:** - -```json -{ - "id": { - "id": { - "value": "550e8400-e29b-41d4-a716-446655440000" - } - }, - "payload": { - "contentType": "Text", - "content": "Hello world" - }, - "vector": { - "values": [0.1, 0.2, 0.3, 0.4] - } -} -``` - --- ### DeletePoint -Delete a point by its ID. - - - The unique identifier of the point to delete. - - -**Request:** - -```protobuf -message PointID { - UUID id = 1; -} -``` +Delete a point by its ID. The point is removed from both the dense and the sparse index. -**Response:** +**Request:** `PointID` -```protobuf -google.protobuf.Empty -``` - -**Example:** +**Response:** `google.protobuf.Empty` ```bash grpcurl -plaintext \ - -H "authorization: secret" \ + -H "authorization: Bearer your-api-key" \ -d '{"id": {"value": "550e8400-e29b-41d4-a716-446655440000"}}' \ localhost:50051 vectordb.VectorDB/DeletePoint ``` -**Response:** - -```json -{} -``` - --- -### SearchPoints - -Search for the k nearest neighbors to a query vector. - - - The vector to search with. Must match the configured `DIMENSION`. - - - - The distance metric to use. - - - - Maximum number of results to return. - +### Search - - Search breadth for HNSW. Larger values trade speed for accuracy. Defaults to the server's `HNSW_EF` setting. - +Find the nearest points to a dense or a sparse query. Exactly one of `dense` or `sparse` must be set. **Request:** ```protobuf message SearchRequest { - DenseVector query_vector = 1; + uint64 limit = 1; + oneof query { + DenseQuery dense = 2; + SparseVector sparse = 3; + } +} + +message DenseQuery { + DenseVector vector = 1; Similarity similarity = 2; - uint64 limit = 3; - uint64 ef = 4; + optional uint64 ef = 3; // HNSW search breadth, ignored by other indexes } ``` @@ -235,105 +171,85 @@ message SearchRequest { ```protobuf message SearchResponse { - repeated PointID result_point_ids = 1; + repeated SearchHit hits = 1; // best first +} + +message SearchHit { + PointID id = 1; + float score = 2; } ``` -**Example:** +Dense queries return a distance as the score (lower is closer). Sparse queries return a dot product (higher is closer). ```bash +# Dense grpcurl -plaintext \ - -H "authorization: secret" \ - -d '{ - "query_vector": {"values": [0.1, 0.2, 0.3, 0.4]}, - "similarity": 3, - "limit": 5, - "ef": 200 - }' \ - localhost:50051 vectordb.VectorDB/SearchPoints -``` + -H "authorization: Bearer your-api-key" \ + -d '{"limit": 5, "dense": {"vector": {"values": [0.1, 0.2, 0.3, 0.4]}, "similarity": "Cosine"}}' \ + localhost:50051 vectordb.VectorDB/Search -**Response:** - -```json -{ - "resultPointIds": [ - {"id": {"value": "550e8400-e29b-41d4-a716-446655440000"}}, - {"id": {"value": "6ba7b810-9dad-11d1-80b4-00c04fd430c8"}} - ] -} +# Sparse +grpcurl -plaintext \ + -H "authorization: Bearer your-api-key" \ + -d '{"limit": 5, "sparse": {"token_ids": [3, 17], "weights": [1.0, 0.5]}}' \ + localhost:50051 vectordb.VectorDB/Search ``` --- -### InsertVectorsBatch +### InsertPointsBatch -Insert multiple vectors in a single request. - -**Request:** +Insert multiple points in one call. The whole batch is validated before anything is written, so one invalid point rejects the batch. ```protobuf -message InsertVectorsBatchRequest { - repeated InsertVectorRequest vectors = 1; +message InsertPointsBatchRequest { + repeated PointInput points = 1; } -``` -**Response:** - -```protobuf -message InsertVectorsBatchResponse { - repeated PointID ids = 1; +message InsertPointsBatchResponse { + repeated PointID ids = 1; // in input order } ``` -**Example:** - ```bash grpcurl -plaintext \ - -H "authorization: secret" \ + -H "authorization: Bearer your-api-key" \ -d '{ - "vectors": [ - {"vector": {"values": [0.1, 0.2, 0.3]}, "payload": {"content_type": 1, "content": "doc one"}}, - {"vector": {"values": [0.4, 0.5, 0.6]}, "payload": {"content_type": 1, "content": "doc two"}} + "points": [ + {"dense_vector": {"values": [0.1, 0.2, 0.3, 0.4]}, "payload": {"content_type": "Text", "content": "doc one"}}, + {"sparse_vector": {"token_ids": [2, 9], "weights": [1.0, 0.25]}} ] }' \ - localhost:50051 vectordb.VectorDB/InsertVectorsBatch + localhost:50051 vectordb.VectorDB/InsertPointsBatch ``` --- -### SearchPointsBatch - -Search against multiple query vectors in a single request. +### SearchBatch -**Request:** +Run multiple queries in one call. Dense and sparse queries can be mixed. ```protobuf -message SearchPointsBatchRequest { +message SearchBatchRequest { repeated SearchRequest queries = 1; } -``` - -**Response:** -```protobuf -message SearchPointsBatchResponse { - repeated SearchResponse results = 1; +message SearchBatchResponse { + repeated SearchResponse results = 1; // one per query, in order } ``` -**Example:** - ```bash grpcurl -plaintext \ - -H "authorization: secret" \ + -H "authorization: Bearer your-api-key" \ -d '{ "queries": [ - {"query_vector": {"values": [0.1, 0.2, 0.3]}, "similarity": 3, "limit": 2}, - {"query_vector": {"values": [0.4, 0.5, 0.6]}, "similarity": 0, "limit": 2} + {"limit": 2, "dense": {"vector": {"values": [0.1, 0.2, 0.3, 0.4]}, "similarity": "Cosine"}}, + {"limit": 2, "sparse": {"token_ids": [3, 17], "weights": [1.0, 0.5]}} ] }' \ - localhost:50051 vectordb.VectorDB/SearchPointsBatch + localhost:50051 vectordb.VectorDB/SearchBatch ``` --- @@ -344,33 +260,32 @@ grpcurl -plaintext \ ```protobuf message UUID { - string value = 1; // UUID v4 string + string value = 1; // UUID string } ``` -### DenseVector +### PointID ```protobuf -message DenseVector { - repeated float values = 1; // Vector components +message PointID { + UUID id = 1; } ``` -### Point +### DenseVector ```protobuf -message Point { - PointID id = 1; // Unique identifier - Payload payload = 2; // Associated metadata - DenseVector vector = 3; // Vector values +message DenseVector { + repeated float values = 1; } ``` -### PointID +### SparseVector ```protobuf -message PointID { - UUID id = 1; +message SparseVector { + repeated uint32 token_ids = 1; // strictly increasing + repeated float weights = 2; // one per token id } ``` @@ -378,8 +293,8 @@ message PointID { ```protobuf message Payload { - ContentType content_type = 1; // Type of content - string content = 2; // Content string + ContentType content_type = 1; + string content = 2; } ``` @@ -414,8 +329,10 @@ Type of payload content. | gRPC Code | Name | Description | |-----------|------|-------------| | `0` | `OK` | Success | -| `3` | `INVALID_ARGUMENT` | Invalid request (e.g., wrong dimensions) | +| `3` | `INVALID_ARGUMENT` | Invalid request (wrong dimension, invalid sparse vector, no vector, missing query, malformed id) | | `5` | `NOT_FOUND` | Point does not exist | +| `6` | `ALREADY_EXISTS` | A point with this id already exists | +| `7` | `PERMISSION_DENIED` | `readonly` key called a write RPC | | `13` | `INTERNAL` | Server error | | `16` | `UNAUTHENTICATED` | Invalid or missing API key | @@ -426,28 +343,32 @@ Type of payload content. ### Python ```python -from vortexdb import VortexDB, DenseVector, Payload, Similarity +from vortexdb import ( + VortexDB, DenseQuery, DenseVector, Payload, PointInput, SearchQuery, Similarity, SparseVector, +) with VortexDB(grpc_url="localhost:50051", api_key="secret") as db: # Insert point_id = db.insert( - vector=DenseVector([0.1, 0.2, 0.3, 0.4]), - payload=Payload.text("Hello") + dense_vector=DenseVector([0.1, 0.2, 0.3, 0.4]), + sparse_vector=SparseVector([3, 17], [0.5, 1.2]), + payload=Payload.text("Hello"), ) # Batch insert - ids = db.batch_insert(items=[ - (DenseVector([0.1, 0.2, 0.3]), Payload.text("doc one")), - (DenseVector([0.4, 0.5, 0.6]), Payload.text("doc two")), + ids = db.batch_insert(points=[ + PointInput(dense_vector=DenseVector([0.1, 0.2, 0.3, 0.4]), payload=Payload.text("doc one")), + PointInput(sparse_vector=SparseVector([2, 9], [1.0, 0.25])), ]) - # Search with ef parameter - results = db.search( - vector=DenseVector([0.1, 0.2, 0.3, 0.4]), - similarity=Similarity.COSINE, + # Dense search with ef parameter + hits = db.search(query=SearchQuery( limit=5, - ef=200, - ) + query=DenseQuery(DenseVector([0.1, 0.2, 0.3, 0.4]), Similarity.COSINE, ef=200), + )) + + # Sparse search + hits = db.search(query=SearchQuery(limit=5, query=SparseVector([3], [1.0]))) ``` ### Generate Clients diff --git a/docs/api-reference/http.mdx b/docs/api-reference/http.mdx index 26ae86f..8e14626 100644 --- a/docs/api-reference/http.mdx +++ b/docs/api-reference/http.mdx @@ -61,12 +61,31 @@ Vector Database server is running! ### Insert Point - - Array of floating-point numbers representing the vector. Must match the configured `DIMENSION`. +A point holds a dense vector, a sparse vector, or both. At least one is required. + + + Optional UUID for the point. Generated when omitted. Inserting an id that already exists returns `409`. + + + + Array of floating-point numbers. Must match the configured `DIMENSION`. + + + + Sparse vector, for example from a SPLADE or BM25 encoder. + + + + Non-empty, strictly increasing integer token ids + + + Finite weights, one per token id + + - - Metadata object associated with the vector. + + Optional metadata associated with the point. @@ -87,7 +106,8 @@ Vector Database server is running! curl -X POST http://localhost:3000/points \ -H "Content-Type: application/json" \ -d '{ - "vector": [0.1, 0.2, 0.3, 0.4], + "dense_vector": [0.1, 0.2, 0.3, 0.4], + "sparse_vector": {"token_ids": [3, 17, 42], "weights": [0.5, 1.2, 0.8]}, "payload": { "content_type": "Text", "content": "Hello, VortexDB!" @@ -106,21 +126,26 @@ curl -X POST http://localhost:3000/points \ | Status | Description | |--------|-------------| -| `400 Bad Request` | Invalid JSON or missing fields | +| `400 Bad Request` | Invalid JSON, no vector, wrong dense dimension, or invalid sparse vector | +| `409 Conflict` | A point with this id already exists | | `500 Internal Server Error` | Server error during insertion | --- ### Batch Insert -Insert multiple vectors in a single request. +Insert multiple points in a single request. The whole batch is validated before anything is written, so one invalid point rejects the batch. - - Array of insert objects, each with `vector` and `payload` fields (same shape as single insert). + + Array of point objects (same shape as single insert). - - Array of UUIDs for each created point, in the same order as the input. + + Number of points inserted. + + + + UUIDs of the created points, in the same order as the input. @@ -128,22 +153,23 @@ Insert multiple vectors in a single request. curl -X POST http://localhost:3000/points/batch \ -H "Content-Type: application/json" \ -d '{ - "vectors": [ + "points": [ { - "vector": [0.1, 0.2, 0.3, 0.4], + "dense_vector": [0.1, 0.2, 0.3, 0.4], "payload": {"content_type": "Text", "content": "Document one"} }, { - "vector": [0.5, 0.6, 0.7, 0.8], + "sparse_vector": {"token_ids": [2, 9], "weights": [1.0, 0.25]}, "payload": {"content_type": "Text", "content": "Document two"} } ] }' ``` -```json Response (201 Created) +```json Response (200 OK) { - "point_ids": [ + "inserted": 2, + "ids": [ "550e8400-e29b-41d4-a716-446655440000", "6ba7b810-9dad-11d1-80b4-00c04fd430c8" ] @@ -155,7 +181,8 @@ curl -X POST http://localhost:3000/points/batch \ | Status | Description | |--------|-------------| -| `400 Bad Request` | Invalid JSON or missing fields | +| `400 Bad Request` | Invalid JSON or an invalid point | +| `409 Conflict` | An id already exists or is repeated in the batch | | `500 Internal Server Error` | Server error during insertion | --- @@ -172,12 +199,16 @@ Retrieve a point by its ID. The point's UUID. - - The stored vector values. + + The stored dense vector, or `null`. - - The associated payload metadata. + + The stored sparse vector, or `null`. + + + + The associated payload, or `null`. @@ -188,7 +219,8 @@ curl http://localhost:3000/points/550e8400-e29b-41d4-a716-446655440000 ```json Response (200 OK) { "id": "550e8400-e29b-41d4-a716-446655440000", - "vector": [0.1, 0.2, 0.3, 0.4], + "dense_vector": [0.1, 0.2, 0.3, 0.4], + "sparse_vector": {"token_ids": [3, 17, 42], "weights": [0.5, 1.2, 0.8]}, "payload": { "content_type": "Text", "content": "Hello, VortexDB!" @@ -208,7 +240,7 @@ curl http://localhost:3000/points/550e8400-e29b-41d4-a716-446655440000 ### Delete Point -Delete a point by its ID. +Delete a point by its ID. The point is removed from both the dense and the sparse index. UUID of the point to delete. @@ -234,41 +266,53 @@ curl -X DELETE http://localhost:3000/points/550e8400-e29b-41d4-a716-446655440000 ### Search Points -Search for the k nearest neighbors to a query vector. +Find the nearest points to a dense or a sparse query. `query` must contain exactly one of `dense` or `sparse`. - - Query vector. Must match the configured `DIMENSION`. + + Maximum number of results to return. - - Distance metric: `"Euclidean"`, `"Manhattan"`, `"Hamming"`, or `"Cosine"` - + + The query vector. - - Maximum number of results to return. + + + Dense query: `vector` (must match `DIMENSION`), `similarity` (`"Euclidean"`, `"Manhattan"`, `"Hamming"` or `"Cosine"`), and optional `ef` (HNSW search breadth, ignored by other indexes). + + + Sparse query: `token_ids` and `weights`, with the same rules as insert. + + - - Array of point IDs ordered by similarity (closest first). + + Array of `{id, score}` objects, best first. Dense queries return a distance (lower is closer). Sparse queries return a dot product (higher is closer). -```bash Request +```bash Dense curl -X POST http://localhost:3000/points/search \ -H "Content-Type: application/json" \ -d '{ - "vector": [0.1, 0.2, 0.3, 0.4], - "similarity": "Cosine", - "limit": 5 + "limit": 5, + "query": {"dense": {"vector": [0.1, 0.2, 0.3, 0.4], "similarity": "Cosine"}} + }' +``` + +```bash Sparse +curl -X POST http://localhost:3000/points/search \ + -H "Content-Type: application/json" \ + -d '{ + "limit": 5, + "query": {"sparse": {"token_ids": [3, 17], "weights": [1.0, 0.5]}} }' ``` ```json Response (200 OK) { - "results": [ - "550e8400-e29b-41d4-a716-446655440000", - "6ba7b810-9dad-11d1-80b4-00c04fd430c8", - "f47ac10b-58cc-4372-a567-0e02b2c3d479" + "hits": [ + {"id": "550e8400-e29b-41d4-a716-446655440000", "score": 0.012}, + {"id": "6ba7b810-9dad-11d1-80b4-00c04fd430c8", "score": 0.087} ] } ``` @@ -278,21 +322,21 @@ curl -X POST http://localhost:3000/points/search \ | Status | Description | |--------|-------------| -| `400 Bad Request` | Invalid JSON or missing fields | +| `400 Bad Request` | Invalid JSON, zero limit, wrong dense dimension, or invalid sparse vector | | `500 Internal Server Error` | Server error during search | --- ### Batch Search -Search against multiple query vectors in a single request. +Run multiple queries in a single request. Dense and sparse queries can be mixed. - Array of search query objects, each with `vector`, `similarity`, and `limit`. + Array of search objects (same shape as single search). - Array of result arrays, one per input query, each ordered by similarity. + One `{hits}` object per input query, in order. @@ -301,8 +345,8 @@ curl -X POST http://localhost:3000/points/search/batch \ -H "Content-Type: application/json" \ -d '{ "queries": [ - {"vector": [0.1, 0.2, 0.3, 0.4], "similarity": "Cosine", "limit": 2}, - {"vector": [0.5, 0.6, 0.7, 0.8], "similarity": "Euclidean", "limit": 2} + {"limit": 2, "query": {"dense": {"vector": [0.1, 0.2, 0.3, 0.4], "similarity": "Cosine"}}}, + {"limit": 2, "query": {"sparse": {"token_ids": [3, 17], "weights": [1.0, 0.5]}}} ] }' ``` @@ -310,8 +354,8 @@ curl -X POST http://localhost:3000/points/search/batch \ ```json Response (200 OK) { "results": [ - ["550e8400-...", "6ba7b810-..."], - ["f47ac10b-...", "9a1b2c3d-..."] + {"hits": [{"id": "550e8400-...", "score": 0.012}, {"id": "6ba7b810-...", "score": 0.087}]}, + {"hits": [{"id": "f47ac10b-...", "score": 1.5}]} ] } ``` @@ -321,7 +365,7 @@ curl -X POST http://localhost:3000/points/search/batch \ | Status | Description | |--------|-------------| -| `400 Bad Request` | Invalid JSON or missing fields | +| `400 Bad Request` | Invalid JSON or an invalid query | | `500 Internal Server Error` | Server error during search | --- @@ -356,7 +400,7 @@ curl http://localhost:3000/health POINT_ID=$(curl -s -X POST http://localhost:3000/points \ -H "Content-Type: application/json" \ -d '{ - "vector": [0.1, 0.2, 0.3, 0.4], + "dense_vector": [0.1, 0.2, 0.3, 0.4], "payload": {"content_type": "Text", "content": "First document"} }' | jq -r '.point_id') @@ -369,9 +413,8 @@ curl http://localhost:3000/points/$POINT_ID curl -X POST http://localhost:3000/points/search \ -H "Content-Type: application/json" \ -d '{ - "vector": [0.15, 0.25, 0.35, 0.45], - "similarity": "Cosine", - "limit": 10 + "limit": 10, + "query": {"dense": {"vector": [0.15, 0.25, 0.35, 0.45], "similarity": "Cosine"}} }' # 5. Delete the point diff --git a/docs/api-reference/overview.mdx b/docs/api-reference/overview.mdx index 57d29a7..2a293bc 100644 --- a/docs/api-reference/overview.mdx +++ b/docs/api-reference/overview.mdx @@ -58,7 +58,7 @@ The HTTP API requires an `api-key` header on every request under `/points`. `/` curl -X POST "http://localhost:3000/points/search" \ -H "api-key: your-api-key" \ -H "Content-Type: application/json" \ - -d '{"vector": [0.1, 0.2, 0.3], "similarity": "Cosine", "limit": 5}' + -d '{"limit": 5, "query": {"dense": {"vector": [0.1, 0.2, 0.3], "similarity": "Cosine"}}}' ``` Keys come from the same `VORTEXDB_KEYS_FILE` used by gRPC (see [gRPC API](/api-reference/grpc)), but the HTTP API additionally enforces each key's role: `readonly` keys can fetch and search points; `readwrite` keys can also insert, batch-insert, and delete. A `readonly` key used against a write route gets `403 Forbidden`; a missing or unrecognized key gets `401 Unauthorized`. @@ -69,12 +69,12 @@ Both APIs support the same core operations: | Operation | gRPC Method | HTTP Endpoint | |-----------|-------------|---------------| -| Insert vector | `InsertVector` | `POST /points` | -| Batch insert | `InsertVectorsBatch` | `POST /points/batch` | +| Insert point | `InsertPoint` | `POST /points` | +| Batch insert | `InsertPointsBatch` | `POST /points/batch` | | Get point | `GetPoint` | `GET /points/:id` | | Delete point | `DeletePoint` | `DELETE /points/:id` | -| Search vectors | `SearchPoints` | `POST /points/search` | -| Batch search | `SearchPointsBatch` | `POST /points/search/batch` | +| Search (dense or sparse) | `Search` | `POST /points/search` | +| Batch search | `SearchBatch` | `POST /points/search/batch` | | Health check | - | `GET /health` | ## Error Handling @@ -86,6 +86,8 @@ Both APIs support the same core operations: | `0` | OK | Success | | `3` | INVALID_ARGUMENT | Invalid request parameters | | `5` | NOT_FOUND | Point not found | +| `6` | ALREADY_EXISTS | Point id already exists | +| `7` | PERMISSION_DENIED | `readonly` key called a write RPC | | `13` | INTERNAL | Server error | | `16` | UNAUTHENTICATED | Invalid or missing API key | @@ -98,13 +100,14 @@ Both APIs support the same core operations: | `204` | No Content (for delete) | | `400` | Bad Request | | `404` | Not Found | +| `409` | Conflict (point id already exists) | | `500` | Internal Server Error | ## Data Types -### Vector +### Vectors -A dense vector of floating-point values: +A point holds a dense vector, a sparse vector, or both. A dense vector is a list of floats whose length matches `DIMENSION`: @@ -117,12 +120,36 @@ A dense vector of floating-point values: ```json { - "vector": [0.1, 0.2, 0.3, 0.4] + "dense_vector": [0.1, 0.2, 0.3, 0.4] } ``` +A sparse vector stores only its non-zero entries, for example from a SPLADE or BM25 encoder. Token ids must be strictly increasing and each weight must be finite and non-negative: + + + + ```protobuf + message SparseVector { + repeated uint32 token_ids = 1; + repeated float weights = 2; + } + ``` + + + ```json + { + "sparse_vector": {"token_ids": [3, 17, 42], "weights": [0.5, 1.2, 0.8]} + } + ``` + + + +### Search Scores + +Every search hit carries a raw score. Dense queries return a distance for the chosen similarity, so lower is closer. Sparse queries return a dot product, so higher is closer. Hits are always ordered best first. + ### Payload Metadata attached to vectors: diff --git a/docs/concepts/architecture.mdx b/docs/concepts/architecture.mdx index 5d31837..cf0e212 100644 --- a/docs/concepts/architecture.mdx +++ b/docs/concepts/architecture.mdx @@ -38,10 +38,12 @@ The **gRPC layer** is the primary high-performance interface: ```protobuf service VectorDB { - rpc InsertVector(InsertVectorRequest) returns (PointID); + rpc InsertPoint(PointInput) returns (PointID); + rpc InsertPointsBatch(InsertPointsBatchRequest) returns (InsertPointsBatchResponse); rpc DeletePoint(PointID) returns (google.protobuf.Empty); rpc GetPoint(PointID) returns (Point); - rpc SearchPoints(SearchRequest) returns (SearchResponse); + rpc Search(SearchRequest) returns (SearchResponse); + rpc SearchBatch(SearchBatchRequest) returns (SearchBatchResponse); } ``` diff --git a/docs/getting-started/quickstart.mdx b/docs/getting-started/quickstart.mdx index 8c9622c..3f290a1 100644 --- a/docs/getting-started/quickstart.mdx +++ b/docs/getting-started/quickstart.mdx @@ -16,7 +16,7 @@ Make sure VortexDB is running (see [Installation](/getting-started/installation) ```python - from vortexdb import VortexDB, DenseVector, Payload, Similarity + from vortexdb import VortexDB, DenseVector, Payload # Connect to VortexDB db = VortexDB( @@ -26,7 +26,7 @@ Make sure VortexDB is running (see [Installation](/getting-started/installation) # Insert a vector with a text payload point_id = db.insert( - vector=DenseVector([0.1, 0.2, 0.3, 0.4]), + dense_vector=DenseVector([0.1, 0.2, 0.3, 0.4]), payload=Payload.text("Hello, VortexDB!") ) print(f"Inserted point: {point_id}") @@ -38,9 +38,10 @@ Make sure VortexDB is running (see [Installation](/getting-started/installation) ```bash curl -X POST http://localhost:3000/points \ + -H "api-key: your-secure-password" \ -H "Content-Type: application/json" \ -d '{ - "vector": [0.1, 0.2, 0.3, 0.4], + "dense_vector": [0.1, 0.2, 0.3, 0.4], "payload": { "content_type": "Text", "content": "Hello, VortexDB!" @@ -59,12 +60,12 @@ Make sure VortexDB is running (see [Installation](/getting-started/installation) ```protobuf // Using grpcurl grpcurl -plaintext \ - -H "authorization: your-secure-password" \ + -H "authorization: Bearer your-secure-password" \ -d '{ - "vector": {"values": [0.1, 0.2, 0.3, 0.4]}, - "payload": {"content_type": 1, "content": "Hello, VortexDB!"} + "dense_vector": {"values": [0.1, 0.2, 0.3, 0.4]}, + "payload": {"content_type": "Text", "content": "Hello, VortexDB!"} }' \ - localhost:50051 vectordb.VectorDB/InsertVector + localhost:50051 vectordb.VectorDB/InsertPoint ``` @@ -74,23 +75,24 @@ Make sure VortexDB is running (see [Installation](/getting-started/installation) ```python - from vortexdb import VortexDB, DenseVector, Similarity + from vortexdb import VortexDB, DenseQuery, DenseVector, SearchQuery, Similarity db = VortexDB( grpc_url="localhost:50051", api_key="your-secure-password" ) - # Search for 5 most similar vectors using cosine similarity - results = db.search( - vector=DenseVector([0.1, 0.2, 0.3, 0.4]), - similarity=Similarity.COSINE, - limit=5 + # Search for the 5 closest vectors using cosine distance + hits = db.search( + query=SearchQuery( + limit=5, + query=DenseQuery(DenseVector([0.1, 0.2, 0.3, 0.4]), Similarity.COSINE), + ) ) - print(f"Found {len(results)} similar vectors:") - for point_id in results: - print(f" - {point_id}") + print(f"Found {len(hits)} similar vectors:") + for hit in hits: + print(f" - {hit.id} (distance {hit.score:.3f})") db.close() ``` @@ -98,20 +100,20 @@ Make sure VortexDB is running (see [Installation](/getting-started/installation) ```bash curl -X POST http://localhost:3000/points/search \ + -H "api-key: your-secure-password" \ -H "Content-Type: application/json" \ -d '{ - "vector": [0.1, 0.2, 0.3, 0.4], - "similarity": "Cosine", - "limit": 5 + "limit": 5, + "query": {"dense": {"vector": [0.1, 0.2, 0.3, 0.4], "similarity": "Cosine"}} }' ``` Response: ```json { - "results": [ - "550e8400-e29b-41d4-a716-446655440000", - "6ba7b810-9dad-11d1-80b4-00c04fd430c8" + "hits": [ + {"id": "550e8400-e29b-41d4-a716-446655440000", "score": 0.0}, + {"id": "6ba7b810-9dad-11d1-80b4-00c04fd430c8", "score": 0.087} ] } ``` @@ -129,21 +131,26 @@ Make sure VortexDB is running (see [Installation](/getting-started/installation) if point: print(point.pretty()) # Output: - # Point ID: 550e8400-e29b-41d4-a716-446655440000 - # Vector: [0.1, 0.2, 0.3, 0.4] - # Payload: Hello, VortexDB! + # Point: + # id = 550e8400-e29b-41d4-a716-446655440000, + # dense_vector = [0.1, 0.2, 0.3, 0.4], + # sparse_vector = None, + # payload_type = TEXT, + # payload = 'Hello, VortexDB!' ``` ```bash - curl http://localhost:3000/points/550e8400-e29b-41d4-a716-446655440000 + curl http://localhost:3000/points/550e8400-e29b-41d4-a716-446655440000 \ + -H "api-key: your-secure-password" ``` Response: ```json { "id": "550e8400-e29b-41d4-a716-446655440000", - "vector": [0.1, 0.2, 0.3, 0.4], + "dense_vector": [0.1, 0.2, 0.3, 0.4], + "sparse_vector": null, "payload": { "content_type": "Text", "content": "Hello, VortexDB!" @@ -164,7 +171,8 @@ Make sure VortexDB is running (see [Installation](/getting-started/installation) ```bash - curl -X DELETE http://localhost:3000/points/550e8400-e29b-41d4-a716-446655440000 + curl -X DELETE http://localhost:3000/points/550e8400-e29b-41d4-a716-446655440000 \ + -H "api-key: your-secure-password" ``` @@ -174,7 +182,7 @@ Make sure VortexDB is running (see [Installation](/getting-started/installation) Here's a complete example using the Python SDK with context manager: ```python -from vortexdb import VortexDB, DenseVector, Payload, Similarity +from vortexdb import VortexDB, DenseQuery, DenseVector, Payload, SearchQuery, Similarity # Using context manager for automatic cleanup with VortexDB(grpc_url="localhost:50051", api_key="secret") as db: @@ -188,22 +196,23 @@ with VortexDB(grpc_url="localhost:50051", api_key="secret") as db: point_ids = [] for vec, text in vectors: pid = db.insert( - vector=DenseVector(vec), + dense_vector=DenseVector(vec), payload=Payload.text(text) ) point_ids.append(pid) print(f"Inserted: {text} -> {pid}") # Search for vectors similar to the first one - results = db.search( - vector=DenseVector([0.15, 0.25, 0.35, 0.45]), - similarity=Similarity.COSINE, - limit=2 + hits = db.search( + query=SearchQuery( + limit=2, + query=DenseQuery(DenseVector([0.15, 0.25, 0.35, 0.45]), Similarity.COSINE), + ) ) print(f"\nTop 2 similar vectors:") - for pid in results: - point = db.get(point_id=pid) + for hit in hits: + point = db.get(point_id=hit.id) print(f" - {point.payload.content}") ``` diff --git a/docs/openapi.yaml b/docs/openapi.yaml index f1b3a69..1e55279 100644 --- a/docs/openapi.yaml +++ b/docs/openapi.yaml @@ -61,16 +61,21 @@ paths: tags: - Points summary: Insert a point - description: Insert a new vector with its associated payload + description: | + Insert a point with a dense vector, a sparse vector, or both, plus an optional payload. + The id is generated when omitted. operationId: insertPoint requestBody: required: true content: application/json: schema: - $ref: '#/components/schemas/InsertRequest' + $ref: '#/components/schemas/PointInput' example: - vector: [0.1, 0.2, 0.3, 0.4] + dense_vector: [0.1, 0.2, 0.3, 0.4] + sparse_vector: + token_ids: [3, 17, 42] + weights: [0.5, 1.2, 0.8] payload: content_type: Text content: Hello, VortexDB! @@ -84,26 +89,33 @@ paths: example: point_id: 550e8400-e29b-41d4-a716-446655440000 '400': - description: Invalid request + description: Invalid request (no vector, wrong dense dimension, or invalid sparse vector) + content: + text/plain: + schema: + type: string + example: "Dimension mismatch: expected 4, got 3" + '409': + description: A point with this id already exists content: text/plain: schema: type: string - example: Invalid vector dimensions '500': description: Server error content: text/plain: schema: type: string - example: Failed to insert point /points/batch: post: tags: - Points summary: Batch insert points - description: Insert multiple vectors in a single request + description: | + Insert multiple points in a single request. The whole batch is validated + before anything is written, so one invalid point rejects the batch. operationId: batchInsertPoints requestBody: required: true @@ -112,24 +124,25 @@ paths: schema: $ref: '#/components/schemas/BatchInsertRequest' example: - vectors: - - vector: [0.1, 0.2, 0.3, 0.4] + points: + - dense_vector: [0.1, 0.2, 0.3, 0.4] payload: content_type: Text content: Document one - - vector: [0.5, 0.6, 0.7, 0.8] - payload: - content_type: Text - content: Document two + - id: 6ba7b810-9dad-11d1-80b4-00c04fd430c8 + sparse_vector: + token_ids: [2, 9] + weights: [1.0, 0.25] responses: - '201': + '200': description: Points created successfully content: application/json: schema: $ref: '#/components/schemas/BatchInsertResponse' example: - point_ids: + inserted: 2 + ids: - 550e8400-e29b-41d4-a716-446655440000 - 6ba7b810-9dad-11d1-80b4-00c04fd430c8 '400': @@ -138,6 +151,12 @@ paths: text/plain: schema: type: string + '409': + description: A point id already exists or is repeated in the batch + content: + text/plain: + schema: + type: string '500': description: Server error content: @@ -170,7 +189,8 @@ paths: $ref: '#/components/schemas/Point' example: id: 550e8400-e29b-41d4-a716-446655440000 - vector: [0.1, 0.2, 0.3, 0.4] + dense_vector: [0.1, 0.2, 0.3, 0.4] + sparse_vector: null payload: content_type: Text content: Hello, VortexDB! @@ -193,7 +213,7 @@ paths: tags: - Points summary: Delete a point - description: Delete a point by its ID + description: Delete a point by its ID. The point is removed from both the dense and the sparse index. operationId: deletePoint parameters: - name: point_id @@ -219,19 +239,32 @@ paths: post: tags: - Points - summary: Search for similar vectors - description: Find the k nearest neighbors to a query vector + summary: Search points + description: | + Find the nearest points to a dense or a sparse query. + `query` must contain exactly one of `dense` or `sparse`. operationId: searchPoints requestBody: required: true content: application/json: schema: - $ref: '#/components/schemas/SearchRequest' - example: - vector: [0.1, 0.2, 0.3, 0.4] - similarity: Cosine - limit: 5 + $ref: '#/components/schemas/SearchQuery' + examples: + dense: + value: + limit: 5 + query: + dense: + vector: [0.1, 0.2, 0.3, 0.4] + similarity: Cosine + sparse: + value: + limit: 5 + query: + sparse: + token_ids: [3, 17] + weights: [1.0, 0.5] responses: '200': description: Search results @@ -240,16 +273,18 @@ paths: schema: $ref: '#/components/schemas/SearchResponse' example: - results: - - 550e8400-e29b-41d4-a716-446655440000 - - 6ba7b810-9dad-11d1-80b4-00c04fd430c8 + hits: + - id: 550e8400-e29b-41d4-a716-446655440000 + score: 0.012 + - id: 6ba7b810-9dad-11d1-80b4-00c04fd430c8 + score: 0.087 '400': description: Invalid request content: text/plain: schema: type: string - example: Invalid vector dimensions + example: "Dimension mismatch: expected 4, got 3" '500': description: Server error content: @@ -262,8 +297,8 @@ paths: post: tags: - Points - summary: Batch search for similar vectors - description: Search against multiple query vectors in a single request + summary: Batch search points + description: Run multiple dense and/or sparse queries in a single request operationId: batchSearchPoints requestBody: required: true @@ -273,12 +308,16 @@ paths: $ref: '#/components/schemas/BatchSearchRequest' example: queries: - - vector: [0.1, 0.2, 0.3, 0.4] - similarity: Cosine - limit: 2 - - vector: [0.5, 0.6, 0.7, 0.8] - similarity: Euclidean - limit: 2 + - limit: 2 + query: + dense: + vector: [0.1, 0.2, 0.3, 0.4] + similarity: Cosine + - limit: 2 + query: + sparse: + token_ids: [3, 17] + weights: [1.0, 0.5] responses: '200': description: Batch search results @@ -288,9 +327,12 @@ paths: $ref: '#/components/schemas/BatchSearchResponse' example: results: - - - 550e8400-e29b-41d4-a716-446655440000 - - 6ba7b810-9dad-11d1-80b4-00c04fd430c8 - - - f47ac10b-58cc-4372-a567-0e02b2c3d479 + - hits: + - id: 550e8400-e29b-41d4-a716-446655440000 + score: 0.012 + - hits: + - id: f47ac10b-58cc-4372-a567-0e02b2c3d479 + score: 1.5 '400': description: Invalid request content: @@ -311,9 +353,31 @@ components: items: type: number format: float - description: Array of floating-point values representing the vector + description: Array of floating-point values; its length must match the database dimension example: [0.1, 0.2, 0.3, 0.4] + SparseVector: + type: object + required: + - token_ids + - weights + properties: + token_ids: + type: array + items: + type: integer + format: uint32 + description: Non-empty, strictly increasing token ids + example: [3, 17, 42] + weights: + type: array + items: + type: number + format: float + minimum: 0 + description: Finite, non-negative weights, one per token id + example: [0.5, 1.2, 0.8] + ContentType: type: string enum: @@ -354,27 +418,37 @@ components: type: object required: - id - - vector - - payload properties: id: type: string format: uuid description: Unique identifier for the point example: 550e8400-e29b-41d4-a716-446655440000 - vector: - $ref: '#/components/schemas/DenseVector' + dense_vector: + allOf: + - $ref: '#/components/schemas/DenseVector' + nullable: true + sparse_vector: + allOf: + - $ref: '#/components/schemas/SparseVector' + nullable: true payload: - $ref: '#/components/schemas/Payload' + allOf: + - $ref: '#/components/schemas/Payload' + nullable: true - InsertRequest: + PointInput: type: object - required: - - vector - - payload + description: At least one of `dense_vector` or `sparse_vector` is required properties: - vector: + id: + type: string + format: uuid + description: Optional id; generated when omitted + dense_vector: $ref: '#/components/schemas/DenseVector' + sparse_vector: + $ref: '#/components/schemas/SparseVector' payload: $ref: '#/components/schemas/Payload' @@ -392,57 +466,98 @@ components: BatchInsertRequest: type: object required: - - vectors + - points properties: - vectors: + points: type: array items: - $ref: '#/components/schemas/InsertRequest' - description: Array of insert objects + $ref: '#/components/schemas/PointInput' BatchInsertResponse: type: object required: - - point_ids + - inserted + - ids properties: - point_ids: + inserted: + type: integer + description: Number of points inserted + ids: type: array items: type: string format: uuid description: UUIDs of the created points, in input order - SearchRequest: + DenseQuery: type: object required: - vector - similarity - - limit properties: vector: $ref: '#/components/schemas/DenseVector' similarity: $ref: '#/components/schemas/Similarity' + ef: + type: integer + nullable: true + description: HNSW search breadth; ignored by other indexes + + QueryVector: + type: object + description: Exactly one of `dense` or `sparse` + oneOf: + - required: [dense] + properties: + dense: + $ref: '#/components/schemas/DenseQuery' + - required: [sparse] + properties: + sparse: + $ref: '#/components/schemas/SparseVector' + + SearchQuery: + type: object + required: + - limit + - query + properties: limit: type: integer minimum: 1 + maximum: 10000 description: Maximum number of results to return example: 5 + query: + $ref: '#/components/schemas/QueryVector' + + SearchHit: + type: object + required: + - id + - score + properties: + id: + type: string + format: uuid + score: + type: number + format: float + description: | + Dense queries return a distance (lower is closer). + Sparse queries return a dot product (higher is closer). SearchResponse: type: object required: - - results + - hits properties: - results: + hits: type: array items: - type: string - format: uuid - description: Array of point IDs ordered by similarity (closest first) - example: - - 550e8400-e29b-41d4-a716-446655440000 - - 6ba7b810-9dad-11d1-80b4-00c04fd430c8 + $ref: '#/components/schemas/SearchHit' + description: Hits ordered best first BatchSearchRequest: type: object @@ -452,8 +567,7 @@ components: queries: type: array items: - $ref: '#/components/schemas/SearchRequest' - description: Array of search query objects + $ref: '#/components/schemas/SearchQuery' BatchSearchResponse: type: object @@ -463,8 +577,5 @@ components: results: type: array items: - type: array - items: - type: string - format: uuid - description: One result array per input query + $ref: '#/components/schemas/SearchResponse' + description: One response per input query, in order diff --git a/docs/sdk/examples.mdx b/docs/sdk/examples.mdx index 0ba249c..4168f3a 100644 --- a/docs/sdk/examples.mdx +++ b/docs/sdk/examples.mdx @@ -10,29 +10,30 @@ Ready-to-run code examples demonstrating VortexDB Python SDK usage. ## Basic Usage ```python -from vortexdb import VortexDB, DenseVector, Payload, Similarity +from vortexdb import ( + VortexDB, DenseQuery, DenseVector, Payload, PointInput, SearchQuery, Similarity, SparseVector, +) with VortexDB(grpc_url="localhost:50051", api_key="secret") as db: # Insert point_id = db.insert( - vector=DenseVector([0.1, 0.2, 0.3]), + dense_vector=DenseVector([0.1, 0.2, 0.3]), payload=Payload.text("hello world"), ) print(f"Inserted: {point_id}") - # Batch insert - ids = db.batch_insert(items=[ - (DenseVector([0.1, 0.2, 0.3]), Payload.text("doc one")), - (DenseVector([0.4, 0.5, 0.6]), Payload.text("doc two")), + # Batch insert: dense, sparse, or both on one point + ids = db.batch_insert(points=[ + PointInput(dense_vector=DenseVector([0.1, 0.2, 0.3]), payload=Payload.text("doc one")), + PointInput(sparse_vector=SparseVector([4, 8], [1.0, 0.5]), payload=Payload.text("doc two")), ]) - # Search - results = db.search( - vector=DenseVector([0.1, 0.2, 0.3]), - similarity=Similarity.COSINE, - limit=3, - ) - print(f"Found {len(results)} results") + # Dense search: score is a distance, lower is closer + hits = db.search(query=SearchQuery(3, DenseQuery(DenseVector([0.1, 0.2, 0.3]), Similarity.COSINE))) + print(f"Found {len(hits)} results") + + # Sparse search: score is a dot product, higher is closer + hits = db.search(query=SearchQuery(3, SparseVector([8], [1.0]))) ``` --- @@ -42,7 +43,7 @@ with VortexDB(grpc_url="localhost:50051", api_key="secret") as db: Using sentence-transformers for text embedding: ```python -from vortexdb import VortexDB, DenseVector, Payload, Similarity +from vortexdb import VortexDB, DenseQuery, DenseVector, Payload, SearchQuery, Similarity from sentence_transformers import SentenceTransformer model = SentenceTransformer('all-MiniLM-L6-v2') @@ -58,15 +59,14 @@ def embed(text: str) -> DenseVector: with VortexDB(grpc_url="localhost:50051", api_key="secret") as db: for doc in documents: - db.insert(vector=embed(doc), payload=Payload.text(doc)) + db.insert(dense_vector=embed(doc), payload=Payload.text(doc)) - results = db.search( - vector=embed("AI and programming"), - similarity=Similarity.COSINE, + hits = db.search(query=SearchQuery( limit=2, - ) - for pid in results: - point = db.get(point_id=pid) + query=DenseQuery(embed("AI and programming"), Similarity.COSINE), + )) + for hit in hits: + point = db.get(point_id=hit.id) print(f" {point.payload.content}") ``` @@ -75,16 +75,22 @@ with VortexDB(grpc_url="localhost:50051", api_key="secret") as db: ## Batch Processing ```python -from vortexdb import VortexDB, DenseVector, Payload, Similarity +from vortexdb import ( + VortexDB, DenseQuery, DenseVector, Payload, PointInput, SearchQuery, Similarity, SparseVector, +) with VortexDB(grpc_url="localhost:50051", api_key="secret") as db: - items = [(DenseVector([i * 0.1 for _ in range(3)]), Payload.text(f"doc {i}")) for i in range(100)] - ids = db.batch_insert(items=items) + points = [ + PointInput(dense_vector=DenseVector([(i + 1) * 0.1] * 3), payload=Payload.text(f"doc {i}")) + for i in range(100) + ] + ids = db.batch_insert(points=points) - # Batch search + # Batch search, dense and sparse queries can be mixed queries = [ - (DenseVector([0.1, 0.2, 0.3]), Similarity.COSINE, 3), - (DenseVector([0.4, 0.5, 0.6]), Similarity.EUCLIDEAN, 3), + SearchQuery(3, DenseQuery(DenseVector([0.1, 0.2, 0.3]), Similarity.COSINE)), + SearchQuery(3, DenseQuery(DenseVector([0.4, 0.5, 0.6]), Similarity.EUCLIDEAN)), + SearchQuery(3, SparseVector([1, 2], [1.0, 1.0])), ] batch_results = db.batch_search(queries=queries) for i, res in enumerate(batch_results): @@ -97,7 +103,7 @@ with VortexDB(grpc_url="localhost:50051", api_key="secret") as db: ```python import pytest -from vortexdb import VortexDB, DenseVector, Payload, Similarity +from vortexdb import VortexDB, DenseQuery, DenseVector, Payload, SearchQuery, Similarity @pytest.fixture def db(): @@ -108,7 +114,7 @@ def db(): class TestVortexDB: def test_insert_and_get(self, db): point_id = db.insert( - vector=DenseVector([0.1, 0.2, 0.3, 0.4]), + dense_vector=DenseVector([0.1, 0.2, 0.3, 0.4]), payload=Payload.text("Test document"), ) point = db.get(point_id=point_id) @@ -118,15 +124,14 @@ class TestVortexDB: def test_search(self, db): point_id = db.insert( - vector=DenseVector([1.0, 2.0, 3.0]), + dense_vector=DenseVector([1.0, 2.0, 3.0, 4.0]), payload=Payload.text("target"), ) - results = db.search( - vector=DenseVector([1.0, 2.0, 3.0]), - similarity=Similarity.COSINE, + hits = db.search(query=SearchQuery( limit=10, - ) - assert point_id in results + query=DenseQuery(DenseVector([1.0, 2.0, 3.0, 4.0]), Similarity.COSINE), + )) + assert point_id in [hit.id for hit in hits] db.delete(point_id=point_id) ``` diff --git a/docs/sdk/reference.mdx b/docs/sdk/reference.mdx index ad1e6e2..b522495 100644 --- a/docs/sdk/reference.mdx +++ b/docs/sdk/reference.mdx @@ -69,17 +69,23 @@ db = VortexDB() #### insert -Insert a vector with its payload into the database. +Insert a point with a dense vector, a sparse vector, or both. At least one vector is required. ```python def insert( self, *, - vector: DenseVector, - payload: Payload, + dense_vector: DenseVector | None = None, + sparse_vector: SparseVector | None = None, + payload: Payload | None = None, + point_id: str | None = None, ) -> str ``` + + Optional UUID for the point. Generated by the server when omitted. + + UUID of the created point. @@ -88,7 +94,8 @@ def insert( ```python point_id = db.insert( - vector=DenseVector([0.1, 0.2, 0.3, 0.4]), + dense_vector=DenseVector([0.1, 0.2, 0.3, 0.4]), + sparse_vector=SparseVector([3, 17], [0.5, 1.2]), payload=Payload.text("My document"), ) ``` @@ -97,13 +104,13 @@ point_id = db.insert( #### batch_insert -Insert multiple vectors in a single request. +Insert multiple points in a single request. The whole batch is validated before anything is written. ```python def batch_insert( self, *, - items: list[tuple[DenseVector, Payload]], + points: list[PointInput], ) -> list[str] ``` @@ -114,9 +121,9 @@ def batch_insert( **Example:** ```python -ids = db.batch_insert(items=[ - (DenseVector([0.1, 0.2, 0.3]), Payload.text("doc one")), - (DenseVector([0.4, 0.5, 0.6]), Payload.text("doc two")), +ids = db.batch_insert(points=[ + PointInput(dense_vector=DenseVector([0.1, 0.2, 0.3]), payload=Payload.text("doc one")), + PointInput(sparse_vector=SparseVector([2, 9], [1.0, 0.25])), ]) ``` @@ -142,86 +149,68 @@ def get( ```python point = db.get(point_id="550e8400-e29b-41d4-a716-446655440000") -if point: - print(f"Vector: {point.vector.to_list()}") - print(f"Payload: {point.payload.content}") +if point and point.dense_vector: + print(f"Dense vector: {point.dense_vector.to_list()}") ``` --- #### search -Search for the k nearest neighbors to a query vector. +Find the nearest points to a dense or a sparse query. ```python def search( self, *, - vector: DenseVector | None = None, - similarity: Similarity | None = None, - limit: int | None = None, - query: SearchQuery | None = None, - ef: int | None = None, -) -> List[str] + query: SearchQuery, +) -> list[SearchHit] ``` - - A `SearchQuery` object bundling vector, similarity, and limit. Use this or pass individual args. - - - - Search breadth for HNSW. Uses server default if not set. - - - - List of point IDs ordered by similarity (closest first). + + Hits ordered best first. For dense queries `score` is a distance (lower is closer). For sparse queries it is a dot product (higher is closer). **Example:** ```python -# Using a SearchQuery -query = SearchQuery(DenseVector([0.1, 0.2, 0.3, 0.4]), Similarity.COSINE, 10) -results = db.search(query=query) - -# Using individual args with ef -results = db.search( - vector=DenseVector([0.1, 0.2, 0.3, 0.4]), - similarity=Similarity.COSINE, +# Dense search, with an optional HNSW ef +hits = db.search(query=SearchQuery( limit=10, - ef=200, -) + query=DenseQuery(DenseVector([0.1, 0.2, 0.3, 0.4]), Similarity.COSINE, ef=200), +)) + +# Sparse search +hits = db.search(query=SearchQuery(limit=10, query=SparseVector([3, 17], [1.0, 0.5]))) + +for hit in hits: + print(hit.id, hit.score) ``` --- #### batch_search -Search against multiple query vectors in a single request. +Run multiple queries in a single request. Dense and sparse queries can be mixed. ```python def batch_search( self, *, - queries, - similarity: Similarity | None = None, - limit: int | None = None, - ef: int | None = None, -) -> List[List[str]] + queries: list[SearchQuery], +) -> list[list[SearchHit]] ``` -Accepts `List[SearchQuery]`, `List[(DenseVector, Similarity, int)]`, or bare `List[DenseVector]` with global `similarity` and `limit`. - - - One result list per input query. + + One hit list per input query, in order. **Example:** ```python results = db.batch_search(queries=[ - SearchQuery(DenseVector([0.1, 0.2, 0.3]), Similarity.COSINE, 5), - (DenseVector([0.4, 0.5, 0.6]), Similarity.EUCLIDEAN, 3), + SearchQuery(5, DenseQuery(DenseVector([0.1, 0.2, 0.3]), Similarity.COSINE)), + SearchQuery(3, SparseVector([4, 8], [1.0, 2.0])), ]) ``` @@ -272,7 +261,7 @@ The client supports the context manager protocol for automatic cleanup: ```python with VortexDB(grpc_url="localhost:50051", api_key="secret") as db: point_id = db.insert( - vector=DenseVector([0.1, 0.2, 0.3]), + dense_vector=DenseVector([0.1, 0.2, 0.3]), payload=Payload.text("Hello"), ) # Connection automatically closed @@ -409,12 +398,48 @@ Payload(content_type: ContentType, content: str) --- +## SparseVector + +A sparse vector that stores only its non-zero entries, for example from a SPLADE or BM25 encoder. + +```python +from vortexdb import SparseVector + +v = SparseVector([3, 17, 42], [0.5, 1.2, 0.8]) +``` + +| Param | Type | Description | +|-------|------|-------------| +| `token_ids` | `list[int]` | Non-empty, non-negative, strictly increasing | +| `weights` | `list[float]` | Finite, non-negative numbers, one per token id | + +Raises `TypeError` for wrong types and `ValueError` for empty, mismatched, unsorted, non-finite, or negative values. + +--- + +## PointInput + +A point to insert with `batch_insert`. Takes the same arguments as `insert`, and at least one vector is required. + +```python +from vortexdb import PointInput + +PointInput( + dense_vector: DenseVector | None = None, + sparse_vector: SparseVector | None = None, + payload: Payload | None = None, + point_id: str | None = None, +) +``` + +--- + ## Point -A point returned from the database (vector + payload + ID). +A point returned from the database. ```python -from vortexdb.models import Point +from vortexdb import Point ``` ### Properties @@ -422,8 +447,9 @@ from vortexdb.models import Point | Property | Type | Description | |----------|------|-------------| | `id` | `str` | Point UUID | -| `vector` | `DenseVector` | The vector values | -| `payload` | `Payload` | Associated metadata | +| `dense_vector` | `DenseVector \| None` | The dense vector, if stored | +| `sparse_vector` | `SparseVector \| None` | The sparse vector, if stored | +| `payload` | `Payload \| None` | Associated metadata, if stored | ### Methods @@ -441,10 +467,12 @@ def pretty(self) -> str point = db.get(point_id="...") print(point.pretty()) # Output: -# Point ID: 550e8400-e29b-41d4-a716-446655440000 -# Vector: [0.1, 0.2, 0.3, 0.4] -# Payload Type: Text -# Payload Content: My document +# Point: +# id = 550e8400-e29b-41d4-a716-446655440000, +# dense_vector = [0.1, 0.2, 0.3, 0.4], +# sparse_vector = None, +# payload_type = TEXT, +# payload = 'My document' ``` --- @@ -471,12 +499,8 @@ from vortexdb import Similarity ```python from vortexdb import Similarity -# Use in search -results = db.search( - vector=DenseVector([0.1, 0.2, 0.3]), - similarity=Similarity.COSINE, - limit=5, -) +# Use in a dense query +query = DenseQuery(DenseVector([0.1, 0.2, 0.3]), Similarity.COSINE) ``` --- @@ -500,20 +524,34 @@ from vortexdb.models import ContentType ## SearchQuery -Bundles a search's parameters into a single object. +A search request: a result limit plus either a `DenseQuery` or a `SparseVector`. ```python -from vortexdb import SearchQuery +from vortexdb import DenseQuery, SearchQuery, SparseVector -query = SearchQuery(DenseVector([0.1, 0.2, 0.3]), Similarity.COSINE, 10) -results = db.search(query=query) +dense = SearchQuery(limit=10, query=DenseQuery(DenseVector([0.1, 0.2, 0.3]), Similarity.COSINE)) +sparse = SearchQuery(limit=10, query=SparseVector([3, 17], [1.0, 0.5])) ``` +| Param | Type | Description | +|-------|------|-------------| +| `limit` | `int` | Max results, between 1 and 10000 | +| `query` | `DenseQuery \| SparseVector` | The query vector | + +### DenseQuery + | Param | Type | Description | |-------|------|-------------| | `vector` | `DenseVector` | Query vector | | `similarity` | `Similarity` | Distance metric | -| `limit` | `int` | Max results | +| `ef` | `int \| None` | HNSW search breadth. Uses the server default if not set | + +### SearchHit + +| Property | Type | Description | +|----------|------|-------------| +| `id` | `str` | Point UUID | +| `score` | `float` | Distance for dense queries (lower is closer), dot product for sparse queries (higher is closer) | --- @@ -527,6 +565,7 @@ from vortexdb import ( AuthenticationError, NotFoundError, InvalidArgumentError, + AlreadyExistsError, TimeoutError, ServiceUnavailableError, InternalServerError, @@ -541,6 +580,7 @@ from vortexdb import ( | `AuthenticationError` | Invalid or missing API key | | `NotFoundError` | Requested resource not found | | `InvalidArgumentError` | Invalid input parameters | +| `AlreadyExistsError` | A point with that id already exists | | `TimeoutError` | Request timed out | | `ServiceUnavailableError` | Server is unavailable | | `InternalServerError` | Server-side error | @@ -597,21 +637,18 @@ config = VortexDBConfig.from_env( The SDK is fully typed. Import types for type hints: ```python -from typing import List, Optional -from vortexdb import VortexDB, DenseVector, Payload, Similarity -from vortexdb.models import Point, ContentType +from vortexdb import VortexDB, DenseQuery, DenseVector, Point, SearchQuery, Similarity def search_documents( db: VortexDB, - query_vector: List[float], + query_vector: list[float], limit: int = 10, -) -> List[Optional[Point]]: - results = db.search( - vector=DenseVector(query_vector), - similarity=Similarity.COSINE, +) -> list[Point | None]: + hits = db.search(query=SearchQuery( limit=limit, - ) - return [db.get(point_id=pid) for pid in results] + query=DenseQuery(DenseVector(query_vector), Similarity.COSINE), + )) + return [db.get(point_id=hit.id) for hit in hits] ``` ## Next Steps