diff --git a/libs/labelbox/src/labelbox/adv_client.py b/libs/labelbox/src/labelbox/adv_client.py index 20766a9a4..52190df2c 100644 --- a/libs/labelbox/src/labelbox/adv_client.py +++ b/libs/labelbox/src/labelbox/adv_client.py @@ -1,7 +1,7 @@ import io import json import logging -from typing import Any, Callable, Dict, List, Optional +from typing import Any, Callable, Dict, Optional from urllib.parse import urlparse import requests @@ -17,19 +17,6 @@ def __init__(self, endpoint: str, api_key: str): self.api_key = api_key self.session = self._create_session() - def create_embedding(self, name: str, dims: int) -> Dict[str, Any]: - data = {"name": name, "dims": dims} - return self._request("POST", "/adv/v1/embeddings", data) - - def delete_embedding(self, id: str): - return self._request("DELETE", f"/adv/v1/embeddings/{id}") - - def get_embedding(self, id: str) -> Dict[str, Any]: - return self._request("GET", f"/adv/v1/embeddings/{id}") - - def get_embeddings(self) -> List[Dict[str, Any]]: - return self._request("GET", "/adv/v1/embeddings").get("results", []) - def import_vectors_from_file(self, id: str, file_path: str, callback=None): self._send_ndjson( f"/adv/v1/embeddings/{id}/_import_ndjson", file_path, callback diff --git a/libs/labelbox/src/labelbox/client.py b/libs/labelbox/src/labelbox/client.py index a00d2c7f0..082d5db01 100644 --- a/libs/labelbox/src/labelbox/client.py +++ b/libs/labelbox/src/labelbox/client.py @@ -2150,8 +2150,21 @@ def create_embedding(self, name: str, dims: int) -> Embedding: Returns: A new Embedding object. """ - data = self._adv_client.create_embedding(name, dims) - return Embedding(self._adv_client, **data) + mutation = """ + mutation CreateEmbeddingPyApi($data: CreateEmbeddingInput!) { + createEmbedding(data: $data) { + id + name + dims + custom + } + } + """ + data = self.execute( + mutation, + {"data": {"name": name, "dims": dims}}, + )["createEmbedding"] + return Embedding(self, **data) def get_embeddings(self) -> List[Embedding]: """ @@ -2160,8 +2173,18 @@ def get_embeddings(self) -> List[Embedding]: Returns: A list of embedding objects. """ - results = self._adv_client.get_embeddings() - return [Embedding(self._adv_client, **data) for data in results] + query = """ + query GetEmbeddingsPyApi { + embeddings { + id + name + dims + custom + } + } + """ + results = self.execute(query)["embeddings"] + return [Embedding(self, **data) for data in results] def get_embedding_by_id(self, id: str) -> Embedding: """ @@ -2173,8 +2196,34 @@ def get_embedding_by_id(self, id: str) -> Embedding: Returns: The embedding object. """ - data = self._adv_client.get_embedding(id) - return Embedding(self._adv_client, **data) + for embedding in self.get_embeddings(): + if embedding.id == id: + return embedding + raise ResourceNotFoundError(Embedding, dict(id=id)) + + def delete_embedding(self, id: str): + """ + Delete a custom embedding through the GraphQL API. + + Args: + id: The embedding ID. + """ + mutation = """ + mutation DeleteEmbeddingPyApi($data: DeleteEmbeddingInput!) { + deleteEmbedding(data: $data) + } + """ + return self.execute(mutation, {"data": {"id": id}})["deleteEmbedding"] + + def import_vectors_from_file(self, id: str, file_path: str, callback=None): + """Upload embedding vectors directly to ADV.""" + return self._adv_client.import_vectors_from_file( + id, file_path, callback + ) + + def get_imported_vector_count(self, id: str) -> int: + """Return an embedding's imported vector count directly from ADV.""" + return self._adv_client.get_imported_vector_count(id) def get_embedding_by_name(self, name: str) -> Embedding: """ diff --git a/libs/labelbox/src/labelbox/schema/embedding.py b/libs/labelbox/src/labelbox/schema/embedding.py index dd5224c7e..3a1ae199c 100644 --- a/libs/labelbox/src/labelbox/schema/embedding.py +++ b/libs/labelbox/src/labelbox/schema/embedding.py @@ -1,9 +1,21 @@ -from typing import Optional, Callable, Dict, Any, List +from typing import Any, Callable, Dict, List, Optional, Protocol -from labelbox.adv_client import AdvClient from pydantic import BaseModel, PrivateAttr +class EmbeddingClient(Protocol): + def delete_embedding(self, id: str): ... + + def import_vectors_from_file( + self, + id: str, + file_path: str, + callback: Optional[Callable[[Dict[str, Any]], None]] = None, + ): ... + + def get_imported_vector_count(self, id: str) -> int: ... + + class EmbeddingVector(BaseModel): """ A Vector Embedding for Custom Embedding. @@ -43,9 +55,9 @@ class Embedding(BaseModel): name: str custom: bool dims: int - _client: AdvClient = PrivateAttr() + _client: EmbeddingClient = PrivateAttr() - def __init__(self, client: AdvClient, **data): + def __init__(self, client: EmbeddingClient, **data): super().__init__(**data) self._client = client diff --git a/libs/labelbox/tests/unit/test_client.py b/libs/labelbox/tests/unit/test_client.py index 3f3d51566..bbdf0857c 100644 --- a/libs/labelbox/tests/unit/test_client.py +++ b/libs/labelbox/tests/unit/test_client.py @@ -1,4 +1,10 @@ +from unittest.mock import Mock + +import pytest +from lbox.exceptions import ResourceNotFoundError + from labelbox.client import Client +from labelbox.schema.embedding import Embedding # @patch.dict(os.environ, {'LABELBOX_API_KEY': 'bar'}) @@ -14,3 +20,104 @@ def test_headers(): def test_enable_experimental(): client = Client(api_key="api_key", enable_experimental=True) assert client.enable_experimental + + +def test_create_embedding_uses_graphql(): + client = Client(api_key="api_key") + client.execute = Mock( + return_value={ + "createEmbedding": { + "id": "embedding-id", + "name": "custom", + "dims": 8, + "custom": True, + } + } + ) + + embedding = client.create_embedding("custom", 8) + + assert embedding.id == "embedding-id" + query, variables = client.execute.call_args.args + assert "createEmbedding" in query + assert variables == {"data": {"name": "custom", "dims": 8}} + + +def test_get_embeddings_uses_graphql(): + client = Client(api_key="api_key") + client.execute = Mock( + return_value={ + "embeddings": [ + { + "id": "embedding-id", + "name": "custom", + "dims": 8, + "custom": True, + } + ] + } + ) + + embeddings = client.get_embeddings() + + assert [embedding.id for embedding in embeddings] == ["embedding-id"] + assert "embeddings" in client.execute.call_args.args[0] + + +def test_get_embedding_by_id_filters_graphql_results(): + client = Client(api_key="api_key") + client.get_embeddings = Mock( + return_value=[ + Embedding( + client, + id="embedding-id", + name="custom", + dims=8, + custom=True, + ) + ] + ) + + assert client.get_embedding_by_id("embedding-id").name == "custom" + + with pytest.raises(ResourceNotFoundError): + client.get_embedding_by_id("missing") + + +def test_embedding_delete_uses_graphql(): + client = Client(api_key="api_key") + client.execute = Mock(return_value={"deleteEmbedding": True}) + embedding = Embedding( + client, + id="embedding-id", + name="custom", + dims=8, + custom=True, + ) + + embedding.delete() + + query, variables = client.execute.call_args.args + assert "deleteEmbedding" in query + assert variables == {"data": {"id": "embedding-id"}} + + +def test_embedding_vector_operations_remain_on_adv(): + client = Client(api_key="api_key") + client._adv_client.import_vectors_from_file = Mock() + client._adv_client.get_imported_vector_count = Mock(return_value=12) + callback = Mock() + embedding = Embedding( + client, + id="embedding-id", + name="custom", + dims=8, + custom=True, + ) + + embedding.import_vectors_from_file("vectors.ndjson", callback) + + client._adv_client.import_vectors_from_file.assert_called_once_with( + "embedding-id", "vectors.ndjson", callback + ) + assert embedding.get_imported_vector_count() == 12