feat(python/sdk): add ability to purge LeMUR requests (#2046)

GitOrigin-RevId: 4f596802d9bc99741304bc51d5dc1bca9fb2d2b7

Author: mez171

README.md

@@ -260,6 +260,35 @@ print(result.response)
 
 </details>
 
+<details>
+  <summary>Delete data previously sent to LeMUR</summary>
+
+```python
+import assemblyai as aai
+
+# Create a transcript and a corresponding LeMUR request that may contain senstive information.
+transcriber = aai.Transcriber()
+transcript_group = transcriber.transcribe_group(
+  [
+    "https://example.org/customer1.mp3",
+  ],
+)
+
+result = transcript_group.lemur.summarize(
+  context="Customers providing sensitive, personally identifiable information",
+  answer_format="TLDR"
+)
+
+# Get the request ID from the LeMUR response
+request_id = result.request_id
+
+# Now we can delete the data about this request
+deletion_result = aai.Lemur.purge_request_data(request_id)
+print(deletion_result)
+```
+
+</details>
+
 ---
 
 ### **Audio Intelligence Examples**

assemblyai/__init__.py

@@ -20,6 +20,8 @@
     LemurActionItemsResponse,
     LemurError,
     LemurModel,
+    LemurPurgeRequest,
+    LemurPurgeResponse,
     LemurQuestion,
     LemurQuestionAnswer,
     LemurQuestionResponse,
@@ -81,6 +83,8 @@
     "LemurActionItemsResponse",
     "LemurError",
     "LemurModel",
+    "LemurPurgeRequest",
+    "LemurPurgeResponse",
     "LemurSource",
     "LemurSourceType",
     "LemurTranscriptSource",

assemblyai/api.py

@@ -7,7 +7,8 @@
 
 ENDPOINT_TRANSCRIPT = "/v2/transcript"
 ENDPOINT_UPLOAD = "/v2/upload"
-ENDPOINT_LEMUR = "/lemur/v3/generate"
+ENDPOINT_LEMUR_BASE = "/lemur/v3"
+ENDPOINT_LEMUR = f"{ENDPOINT_LEMUR_BASE}/generate"
 
 
 def _get_error_message(response: httpx.Response) -> str:
@@ -310,3 +311,21 @@ def lemur_task(
         )
 
     return types.LemurTaskResponse.parse_obj(response.json())
+
+
+def lemur_purge_request_data(
+    client: httpx.Client,
+    request: types.LemurPurgeRequest,
+    http_timeout: Optional[float],
+) -> types.LemurPurgeRequest:
+    response = client.delete(
+        f"{ENDPOINT_LEMUR_BASE}/{request.request_id}",
+        timeout=http_timeout,
+    )
+
+    if response.status_code != httpx.codes.ok:
+        raise types.LemurError(
+            f"Failed to purge LeMUR request data for provided request ID: {request.request_id}. Error: {_get_error_message(response)}"
+        )
+
+    return types.LemurPurgeResponse.parse_obj(response.json())

assemblyai/lemur.py

@@ -283,3 +283,26 @@ def task(
             timeout=timeout,
             temperature=temperature,
         )
+
+    @classmethod
+    def purge_request_data(
+        cls,
+        request_id: str,
+        timeout: Optional[float] = None,
+    ) -> types.LemurPurgeResponse:
+        """
+        Purge sent LeMUR request data that was previously sent.
+
+        Args:
+            request_id: The request ID that was returned to you from the original LeMUR request that should be purged.
+
+        Returns: A response saying whether the LeMUR request data was successfully purged.
+        """
+
+        return api.lemur_purge_request_data(
+            client=_client.Client.get_default().http_client,
+            request=types.LemurPurgeRequest(
+                request_id=request_id,
+            ),
+            http_timeout=timeout,
+        )

assemblyai/types.py

@@ -1829,6 +1829,25 @@ class LemurActionItemsResponse(BaseModel):
     "The action items for your LeMUR request"
 
 
+class LemurPurgeRequest(BaseModel):
+    request_id: str
+
+
+class LemurPurgeResponse(BaseModel):
+    """
+    The result of your LeMUR purge request.
+    """
+
+    request_id: str
+    "The unique identifier of the LeMUR purge request"
+
+    request_id_to_purge: str
+    "The unique identifier of the LeMUR request nneds to be purged"
+
+    deleted: bool
+    "The result of the LeMUR purge request"
+
+
 class RealtimeMessageTypes(str, Enum):
     """
     The type of message received from the real-time API

tests/unit/factories.py

@@ -193,6 +193,15 @@ class Meta:
     response = factory.Faker("text")
 
 
+class LemurPurgeResponse(factory.Factory):
+    class Meta:
+        model = types.LemurPurgeResponse
+
+    request_id = factory.Faker("uuid4")
+    request_id_to_purge = factory.Faker("uuid4")
+    deleted = True
+
+
 class WordSearchMatchFactory(factory.Factory):
     class Meta:
         model = types.WordSearchMatch

tests/unit/test_lemur.py

@@ -5,7 +5,10 @@
 from pytest_httpx import HTTPXMock
 
 import assemblyai as aai
-from assemblyai.api import ENDPOINT_LEMUR
+from assemblyai.api import (
+    ENDPOINT_LEMUR,
+    ENDPOINT_LEMUR_BASE,
+)
 from tests.unit import factories
 
 aai.settings.api_key = "test"
@@ -330,3 +333,59 @@ def test_lemur_ask_coach_fails(httpx_mock: HTTPXMock):
 
     # check whether we mocked everything
     assert len(httpx_mock.get_requests()) == 1
+
+
+def test_lemur_purge_request_data_succeeds(httpx_mock: HTTPXMock):
+    """
+    Tests whether LeMUR request purging succeeds.
+    """
+
+    # create a mock response of a LemurPurgeResponse
+    mock_lemur_purge_response = factories.generate_dict_factory(
+        factories.LemurPurgeResponse
+    )()
+
+    mock_request_id: str = str(uuid.uuid4())
+
+    # mock the specific endpoints
+    httpx_mock.add_response(
+        url=f"{aai.settings.base_url}{ENDPOINT_LEMUR_BASE}/{mock_request_id}",
+        status_code=httpx.codes.OK,
+        method="DELETE",
+        json=mock_lemur_purge_response,
+    )
+
+    # mimic the usage of the SDK
+    result = aai.Lemur.purge_request_data(request_id=mock_request_id)
+
+    # check the response
+    assert isinstance(result, aai.LemurPurgeResponse)
+
+    # check whether we mocked everything
+    assert len(httpx_mock.get_requests()) == 1
+
+
+def test_lemur_purge_request_data_fails(httpx_mock: HTTPXMock):
+    """
+    Tests whether LeMUR request purging fails.
+    """
+
+    # create a mock response of a LemurPurgeResponse
+    mock_lemur_purge_response = factories.generate_dict_factory(
+        factories.LemurPurgeResponse
+    )()
+
+    mock_request_id: str = str(uuid.uuid4())
+
+    # mock the specific endpoints
+    httpx_mock.add_response(
+        url=f"{aai.settings.base_url}{ENDPOINT_LEMUR_BASE}/{mock_request_id}",
+        status_code=httpx.codes.INTERNAL_SERVER_ERROR,
+        method="DELETE",
+        json=mock_lemur_purge_response,
+    )
+
+    with pytest.raises(aai.LemurError) as error:
+        aai.Lemur.purge_request_data(mock_request_id)
+
+    assert len(httpx_mock.get_requests()) == 1