client: Add 'process_bulk' and 'redact' parameters

Add client support for the 'process_bulk' task, which allows providing a
list of documents to be annotated, and add missing configuration options
to the 'redact' method.

Signed-off-by: Phoevos Kalemkeris <[email protected]>

Author: phoevos

client/cogstack_model_gateway_client/client.py

@@ -1,5 +1,6 @@
 import asyncio
 import json
+from collections.abc import Iterable
 from functools import wraps
 
 import httpx
@@ -46,6 +47,7 @@ async def submit_task(
         model_name: str = None,
         task: str = None,
         data=None,
+        json=None,
         files=None,
         params=None,
         headers=None,
@@ -57,7 +59,9 @@ async def submit_task(
         if not model_name:
             raise ValueError("Please provide a model name or set a default model for the client.")
         url = f"{self.base_url}/models/{model_name}/tasks/{task}"
-        resp = await self._client.post(url, data=data, files=files, params=params, headers=headers)
+        resp = await self._client.post(
+            url, data=data, json=json, files=files, params=params, headers=headers
+        )
         resp.raise_for_status()
         task_info = resp.json()
         if wait_for_completion:
@@ -84,18 +88,51 @@ async def process(
             return_result=return_result,
         )
 
+    async def process_bulk(
+        self,
+        texts: list[str],
+        model_name: str = None,
+        wait_for_completion: bool = True,
+        return_result: bool = True,
+    ):
+        """Generate annotations for a list of texts."""
+        return await self.submit_task(
+            model_name=model_name,
+            task="process_bulk",
+            json=texts,
+            headers={"Content-Type": "application/json"},
+            wait_for_completion=wait_for_completion,
+            return_result=return_result,
+        )
+
     async def redact(
         self,
         text: str,
+        concepts_to_keep: Iterable[str] = None,
+        warn_on_no_redaction: bool = None,
+        mask: str = None,
+        hash: bool = None,
         model_name: str = None,
         wait_for_completion: bool = True,
         return_result: bool = True,
     ):
         """Redact sensitive information from the provided text."""
+        params = {
+            k: v
+            for k, v in {
+                "concepts_to_keep": concepts_to_keep,
+                "warn_on_no_redaction": warn_on_no_redaction,
+                "mask": mask,
+                "hash": hash,
+            }.items()
+            if v is not None
+        } or None
+
         return await self.submit_task(
             model_name=model_name,
             task="redact",
             data=text,
+            params=params,
             headers={"Content-Type": "text/plain"},
             wait_for_completion=wait_for_completion,
             return_result=return_result,
@@ -225,6 +262,9 @@ def submit_task(self, *args, **kwargs):
     def process(self, *args, **kwargs):
         return self._loop.run_until_complete(self._client.process(*args, **kwargs))
 
+    def process_bulk(self, *args, **kwargs):
+        return self._loop.run_until_complete(self._client.process_bulk(*args, **kwargs))
+
     def redact(self, *args, **kwargs):
         return self._loop.run_until_complete(self._client.redact(*args, **kwargs))
 

tests/unit/client/test_client.py

@@ -76,6 +76,7 @@ async def test_submit_task_success(mock_httpx_async_client):
     mock_client_instance.post.assert_awaited_once_with(
         "http://test-gateway.com/models/my_model/tasks/process",
         data="some text",
+        json=None,
         files=None,
         params=None,
         headers=None,
@@ -100,6 +101,7 @@ async def test_submit_task_with_default_model(mock_httpx_async_client):
     mock_client_instance.post.assert_awaited_once_with(
         "http://test-gateway.com/models/default_model/tasks/process",
         data="some text",
+        json=None,
         files=None,
         params=None,
         headers=None,
@@ -176,16 +178,35 @@ async def test_process_method(mocker):
         )
 
 
+@pytest.mark.asyncio
+async def test_process_bulk_method(mocker):
+    """Test the process_bulk method."""
+    async with GatewayClient(base_url="http://test-gateway.com") as client:
+        mock_submit_task = mocker.patch.object(client, "submit_task", new=AsyncMock())
+        await client.process_bulk(texts=["text1", "text2"], model_name="bulk_model")
+        mock_submit_task.assert_awaited_once_with(
+            model_name="bulk_model",
+            task="process_bulk",
+            json=["text1", "text2"],
+            headers={"Content-Type": "application/json"},
+            wait_for_completion=True,
+            return_result=True,
+        )
+
+
 @pytest.mark.asyncio
 async def test_redact_method(mocker):
     """Test the redact method."""
     async with GatewayClient(base_url="http://test-gateway.com") as client:
         mock_submit_task = mocker.patch.object(client, "submit_task", new=AsyncMock())
-        await client.redact(text="sensitive text", model_name="deid_model")
+        await client.redact(
+            text="sensitive text", concepts_to_keep=["label1", "label2"], model_name="deid_model"
+        )
         mock_submit_task.assert_awaited_once_with(
             model_name="deid_model",
             task="redact",
             data="sensitive text",
+            params={"concepts_to_keep": ["label1", "label2"]},
             headers={"Content-Type": "text/plain"},
             wait_for_completion=True,
             return_result=True,