Add Label functionallity (#14)

.vscode/settings.json

@@ -6,6 +6,7 @@
   "cSpell.words": [
     "anonymization",
     "knowledgeboxes",
+    "labelsets",
     "Serviceaccount",
     "tqdm"
   ]

nuclia/sdk/kb.py

@@ -1,9 +1,12 @@
+from typing import List, Optional
 from warnings import warn
 
+from nucliadb_models.labels import KnowledgeBoxLabels, Label, LabelSet, LabelSetKind
 from nucliadb_models.resource import Resource, ResourceList
 
 from nuclia.data import get_auth
 from nuclia.decorators import kb, pretty
+from nuclia.lib.kb import NucliaDBClient
 from nuclia.sdk.auth import NucliaAuth
 from nuclia.sdk.logger import logger
 from nuclia.sdk.resource import NucliaResource
@@ -23,12 +26,85 @@ def __init__(self):
         self.resource = NucliaResource()
 
     @kb
-    def list(self, **kwargs):
-        ndb = kwargs["ndb"]
+    def list(self, **kwargs) -> None:
+        ndb: NucliaDBClient = kwargs["ndb"]
         data: ResourceList = ndb.ndb.list_resources(kbid=ndb.kbid)
         for resource in data.resources:
             print(f"{resource.id} {resource.icon:30} {resource.title}")
 
+    @kb
+    def get_labelset(
+        self,
+        *,
+        labelset: str,
+        **kwargs,
+    ) -> LabelSet:
+        ndb: NucliaDBClient = kwargs["ndb"]
+        return ndb.ndb.get_labelset(kbid=ndb.kbid, labelset=labelset)
+
+    @kb
+    def add_labelset(
+        self,
+        *,
+        labelset: str,
+        kind: LabelSetKind = LabelSetKind.RESOURCES,
+        multiple: bool = True,
+        title: Optional[str] = None,
+        color: Optional[str] = None,
+        labels: Optional[List[str]] = None,
+        **kwargs,
+    ):
+        ndb: NucliaDBClient = kwargs["ndb"]
+        if labels is None:
+            labels_list = []
+        else:
+            labels_list = [Label(title=label) for label in labels]
+
+        labelset_obj = LabelSet(
+            title=title if title is not None else labelset,
+            color=color if color is not None else "blue",
+            labels=labels_list,
+            kind=[kind],
+            multiple=multiple,
+        )
+
+        ndb.ndb.set_labelset(kbid=ndb.kbid, labelset=labelset, content=labelset_obj)
+
+    @kb
+    def list_labelsets(self, **kwargs) -> KnowledgeBoxLabels:
+        ndb: NucliaDBClient = kwargs["ndb"]
+        data: KnowledgeBoxLabels = ndb.ndb.get_labelsets(kbid=ndb.kbid)
+        return data
+
+    @kb
+    def del_labelset(self, *, labelset: str, **kwargs):
+        ndb: NucliaDBClient = kwargs["ndb"]
+        ndb.ndb.delete_labelset(kbid=ndb.kbid, labelset=labelset)
+
+    @kb
+    def add_label(
+        self,
+        *,
+        labelset: str,
+        label: str,
+        text: Optional[str] = None,
+        uri: Optional[str] = None,
+        **kwargs,
+    ):
+        ndb: NucliaDBClient = kwargs["ndb"]
+        labelset_obj: LabelSet = ndb.ndb.get_labelset(kbid=ndb.kbid, labelset=labelset)
+        label_obj = Label(title=label, text=text, uri=uri)
+        labelset_obj.labels.append(label_obj)
+        ndb.ndb.set_labelset(kbid=ndb.kbid, labelset=labelset, content=labelset_obj)
+
+    @kb
+    def del_label(self, *, labelset: str, label: str, **kwargs):
+        ndb: NucliaDBClient = kwargs["ndb"]
+        labelset_obj: LabelSet = ndb.ndb.get_labelset(kbid=ndb.kbid, labelset=labelset)
+        label_to_delete = next(x for x in labelset_obj.labels if x.title == label)
+        labelset_obj.labels.remove(label_to_delete)
+        ndb.ndb.set_labelset(kbid=ndb.kbid, labelset=labelset, content=labelset_obj)
+
     @kb
     @pretty
     def get_resource_by_id(self, *, rid: str, **kwargs) -> Resource:
@@ -37,7 +113,7 @@ def get_resource_by_id(self, *, rid: str, **kwargs) -> Resource:
             DeprecationWarning,
         )
         logger.warning("get_resource_by_slug is deprecated, use resource.get instead")
-        ndb = kwargs["ndb"]
+        ndb: NucliaDBClient = kwargs["ndb"]
         return ndb.ndb.get_resource_by_id(
             kbid=ndb.kbid, rid=rid, query_params={"show": "values"}
         )
@@ -50,7 +126,7 @@ def get_resource_by_slug(self, *, slug: str, **kwargs) -> Resource:
             DeprecationWarning,
         )
         logger.warning("get_resource_by_slug is deprecated, use resource.get instead")
-        ndb = kwargs["ndb"]
+        ndb: NucliaDBClient = kwargs["ndb"]
         return ndb.ndb.get_resource_by_slug(
             kbid=ndb.kbid, slug=slug, query_params={"show": "values"}
         )
@@ -59,5 +135,5 @@ def get_resource_by_slug(self, *, slug: str, **kwargs) -> Resource:
     def delete(self, *, rid: str, **kwargs):
         warn("delete is deprecated, use resource.delete instead", DeprecationWarning)
         logger.warning("delete is deprecated, use resource.delete instead")
-        ndb = kwargs["ndb"]
+        ndb: NucliaDBClient = kwargs["ndb"]
         ndb.ndb.delete_resource(kbid=ndb.kbid, rid=rid)

nuclia/tests/test_labels.py

@@ -0,0 +1,25 @@
+from nuclia.sdk.kb import NucliaKB
+
+
+def test_labels(testing_config):
+    nkb = NucliaKB()
+    nkb.add_label(labelset="labelset1", label="label1")
+    nkb.add_labelset(labelset="labelset1")
+    nkb.add_label(labelset="labelset1", label="label1")
+    nkb.add_label(labelset="labelset1", label="label2")
+    nkb.del_labelset(labelset="labelset2")
+    labelset = nkb.get_labelset(labelset="labelset1")
+    assert len(labelset.labels) == 2
+
+    nkb.del_label(labelset="labelset1", label="label2")
+
+    labelset = nkb.get_labelset(labelset="labelset1")
+    assert len(labelset.labels) == 1
+    assert labelset.labels[0].title == "label1"
+
+    labelsets = nkb.list_labelsets()
+    assert len(labelsets.labelsets) == 1
+    assert "labelset1" in labelsets.labelsets
+    nkb.del_labelset(labelset="labelset1")
+    labelsets = nkb.list_labelsets()
+    assert len(labelsets.labelsets) == 0

requirements.txt

@@ -3,6 +3,6 @@ pydantic<2.0.0
 pyyaml==5.3.1
 requests
 gnureadline
-nucliadb_sdk>=2.17.0
-nucliadb_models>=2.17.0
+nucliadb_sdk>=2.18.0
+nucliadb_models>=2.18.0
 tqdm