[PLT-4347] Route embedding metadata operations through GraphQL (#2074)

Author: abacchilb

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

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:
         """

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
 

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