chore(llmobs): fix deletion conflicts with insertions (#14255)

gary-huang · web-flow · commit 2aa770070de7 · 2025-08-11T14:11:21.000-04:00
consider a case where a dataset is created, and the local state is being manipulated, specifically: 1. records are inserted / appended 2. appended records are then deleted 3. then a push occurs when this happens, the POST to batch_update will have the records in the records to be inserted, and the deleted_ids keys will contain empty strings for the new records, leading to backend error before <img width="934" height="923" alt="image" src="https://github.com/user-attachments/assets/72ce1c0c-2246-4e60-ab8c-ef39764b1b6e" /> leads to[ this backend error](https://app.datadoghq.com/apm/trace/689634c90000000023a901545cda196b?colorLegendSort=time&colsTraces=service%2Cresource_name%2C%40duration%2C%40_span.count%2C%40_duration.by_service&graphType=flamegraph&shouldShowLegend=true&spanID=2939228585138612840&spanViewType=errors&timeHint=1754674377308.8823&traceQuery=) [ after](https://dddev.datadoghq.com/llm/datasets/e9d31c52-f28e-4c9c-ab5a-eaa846e6d0a5?page=1) <img width="977" height="901" alt="image" src="https://github.com/user-attachments/assets/c407e0dc-2a7f-4170-8116-d6cadcd6d5fc" /> ## Checklist - [x] PR author has checked that all the criteria below are met - The PR description includes an overview of the change - The PR description articulates the motivation for the change - The change includes tests OR the PR description describes a testing strategy - The PR description notes risks associated with the change, if any - Newly-added code is easy to change - The change follows the [library release note guidelines](https://ddtrace.readthedocs.io/en/stable/releasenotes.html) - The change includes or references documentation updates if necessary - Backport labels are set (if [applicable](https://ddtrace.readthedocs.io/en/latest/contributing.html#backporting)) ## Reviewer Checklist - [x] Reviewer has checked that all the criteria below are met - Title is accurate - All changes are related to the pull request's stated goal - Avoids breaking [API](https://ddtrace.readthedocs.io/en/stable/versioning.html#interfaces) changes - Testing strategy adequately addresses listed risks - Newly-added code is easy to change - Release note makes sense to a user of the library - If necessary, author has acknowledged and discussed the performance implications of this PR as reported in the benchmarks PR comment - Backport labels are set in a manner that is consistent with the [release branch maintenance policy](https://ddtrace.readthedocs.io/en/latest/contributing.html#backporting)
diff --git a/ddtrace/llmobs/_experiment.py b/ddtrace/llmobs/_experiment.py
@@ -14,6 +14,7 @@
 from typing import Union
 from typing import cast
 from typing import overload
+import uuid
 
 from typing_extensions import NotRequired
 
@@ -95,7 +96,7 @@ class Dataset:
     _records: List[DatasetRecord]
     _version: int
     _dne_client: "LLMObsExperimentsClient"
-    _new_records: List[DatasetRecordRaw]
+    _new_records_by_record_id: Dict[str, DatasetRecordRaw]
     _updated_record_ids_to_new_fields: Dict[str, UpdatableDatasetRecord]
     _deleted_record_ids: List[str]
 
@@ -114,7 +115,7 @@ def __init__(
         self._version = version
         self._dne_client = _dne_client
         self._records = records
-        self._new_records = []
+        self._new_records_by_record_id = {}
         self._updated_record_ids_to_new_fields = {}
         self._deleted_record_ids = []
 
@@ -136,16 +137,16 @@ def push(self) -> None:
 
         updated_records = list(self._updated_record_ids_to_new_fields.values())
         new_version, new_record_ids = self._dne_client.dataset_batch_update(
-            self._id, self._new_records, updated_records, self._deleted_record_ids
+            self._id, list(self._new_records_by_record_id.values()), updated_records, self._deleted_record_ids
         )
 
         # attach record ids to newly created records
-        for record, record_id in zip(self._new_records, new_record_ids):
+        for record, record_id in zip(self._new_records_by_record_id.values(), new_record_ids):
             record["record_id"] = record_id  # type: ignore
 
         # FIXME: we don't get version numbers in responses to deletion requests
         self._version = new_version if new_version != -1 else self._version + 1
-        self._new_records = []
+        self._new_records_by_record_id = {}
         self._deleted_record_ids = []
         self._updated_record_ids_to_new_fields = {}
 
@@ -164,19 +165,34 @@ def update(self, index: int, record: DatasetRecordRaw) -> None:
         self._records[index] = {**self._records[index], **record, "record_id": record_id}
 
     def append(self, record: DatasetRecordRaw) -> None:
-        r: DatasetRecord = {**record, "record_id": ""}
+        record_id: str = uuid.uuid4().hex
+        # this record ID will be discarded after push, BE will generate a new one, this is just
+        # for tracking new records locally before the push
+        r: DatasetRecord = {**record, "record_id": record_id}
         # keep the same reference in both lists to enable us to update the record_id after push
-        self._new_records.append(r)
+        self._new_records_by_record_id[record_id] = r
         self._records.append(r)
 
     def delete(self, index: int) -> None:
         record_id = self._records[index]["record_id"]
-        self._deleted_record_ids.append(record_id)
+        should_append_to_be_deleted = True
+
         del self._records[index]
 
+        if record_id is None or record_id == "":
+            logger.warning("encountered unexpected record_id on deletion %s", record_id)
+            return
+
         if record_id in self._updated_record_ids_to_new_fields:
             del self._updated_record_ids_to_new_fields[record_id]
 
+        if record_id in self._new_records_by_record_id:
+            del self._new_records_by_record_id[record_id]
+            should_append_to_be_deleted = False
+
+        if should_append_to_be_deleted:
+            self._deleted_record_ids.append(record_id)
+
     @property
     def url(self) -> str:
         # FIXME: will not work for subdomain orgs
diff --git a/tests/llmobs/llmobs_cassettes/datadog/datadog_api_unstable_llm-obs_v1_datasets_cfdb514f-92b0-404f-9b3f-020f82b0cd51_batch_update_post_01ab4862.yaml b/tests/llmobs/llmobs_cassettes/datadog/datadog_api_unstable_llm-obs_v1_datasets_cfdb514f-92b0-404f-9b3f-020f82b0cd51_batch_update_post_01ab4862.yaml
@@ -0,0 +1,48 @@
+interactions:
+- request:
+    body: '{"data": {"type": "datasets", "id": "cfdb514f-92b0-404f-9b3f-020f82b0cd51",
+      "attributes": {"insert_records": [{"input": {"prompt": "What is the capital
+      of Sweden?"}, "expected_output": null, "metadata": null}], "update_records":
+      [], "delete_records": ["0c84833f-c2a2-43e0-bbbe-5dee1336c5f3"]}}}'
+    headers:
+      Accept:
+      - '*/*'
+      ? !!python/object/apply:multidict._multidict.istr
+      - Accept-Encoding
+      : - identity
+      Connection:
+      - keep-alive
+      Content-Length:
+      - '294'
+      ? !!python/object/apply:multidict._multidict.istr
+      - Content-Type
+      : - application/json
+      User-Agent:
+      - python-requests/2.32.3
+    method: POST
+    uri: https://api.datadoghq.com/api/unstable/llm-obs/v1/datasets/cfdb514f-92b0-404f-9b3f-020f82b0cd51/batch_update
+  response:
+    body:
+      string: '{"data":[{"id":"e6c19140-44dd-41ae-8617-adb52e2a79c1","type":"datasets","attributes":{"author":{"id":"de473b30-eb9f-11e9-a77a-c7405862b8bd"},"created_at":"2025-08-08T20:50:03.226626591Z","dataset_id":"cfdb514f-92b0-404f-9b3f-020f82b0cd51","input":{"prompt":"What
+        is the capital of Sweden?"},"updated_at":"2025-08-08T20:50:03.226626591Z","version":2}}]}'
+    headers:
+      content-length:
+      - '352'
+      content-security-policy:
+      - frame-ancestors 'self'; report-uri https://logs.browser-intake-datadoghq.com/api/v2/logs?dd-api-key=pube4f163c23bbf91c16b8f57f56af9fc58&dd-evp-origin=content-security-policy&ddsource=csp-report&ddtags=site%3Adatadoghq.com
+      content-type:
+      - application/vnd.api+json
+      date:
+      - Fri, 08 Aug 2025 20:50:03 GMT
+      strict-transport-security:
+      - max-age=31536000; includeSubDomains; preload
+      vary:
+      - Accept-Encoding
+      x-content-type-options:
+      - nosniff
+      x-frame-options:
+      - SAMEORIGIN
+    status:
+      code: 200
+      message: OK
+version: 1
diff --git a/tests/llmobs/llmobs_cassettes/datadog/datadog_api_unstable_llm-obs_v1_datasets_cfdb514f-92b0-404f-9b3f-020f82b0cd51_batch_update_post_a50f2e77.yaml b/tests/llmobs/llmobs_cassettes/datadog/datadog_api_unstable_llm-obs_v1_datasets_cfdb514f-92b0-404f-9b3f-020f82b0cd51_batch_update_post_a50f2e77.yaml
@@ -0,0 +1,50 @@
+interactions:
+- request:
+    body: '{"data": {"type": "datasets", "id": "cfdb514f-92b0-404f-9b3f-020f82b0cd51",
+      "attributes": {"insert_records": [{"input": {"prompt": "What is the capital
+      of France?"}, "expected_output": {"answer": "Paris"}, "metadata": null}, {"input":
+      {"prompt": "What is the capital of Italy?"}, "expected_output": {"answer": "Rome"},
+      "metadata": null}], "update_records": [], "delete_records": []}}}'
+    headers:
+      Accept:
+      - '*/*'
+      ? !!python/object/apply:multidict._multidict.istr
+      - Accept-Encoding
+      : - identity
+      Connection:
+      - keep-alive
+      Content-Length:
+      - '384'
+      ? !!python/object/apply:multidict._multidict.istr
+      - Content-Type
+      : - application/json
+      User-Agent:
+      - python-requests/2.32.3
+    method: POST
+    uri: https://api.datadoghq.com/api/unstable/llm-obs/v1/datasets/cfdb514f-92b0-404f-9b3f-020f82b0cd51/batch_update
+  response:
+    body:
+      string: '{"data":[{"id":"0c84833f-c2a2-43e0-bbbe-5dee1336c5f3","type":"datasets","attributes":{"author":{"id":"de473b30-eb9f-11e9-a77a-c7405862b8bd"},"created_at":"2025-08-08T20:49:58.695339032Z","dataset_id":"cfdb514f-92b0-404f-9b3f-020f82b0cd51","expected_output":{"answer":"Paris"},"input":{"prompt":"What
+        is the capital of France?"},"updated_at":"2025-08-08T20:49:58.695339032Z","version":1}},{"id":"0a37e076-a220-427b-babe-74b9b7ca74a1","type":"datasets","attributes":{"author":{"id":"de473b30-eb9f-11e9-a77a-c7405862b8bd"},"created_at":"2025-08-08T20:49:58.695339032Z","dataset_id":"cfdb514f-92b0-404f-9b3f-020f82b0cd51","expected_output":{"answer":"Rome"},"input":{"prompt":"What
+        is the capital of Italy?"},"updated_at":"2025-08-08T20:49:58.695339032Z","version":1}}]}'
+    headers:
+      content-length:
+      - '766'
+      content-security-policy:
+      - frame-ancestors 'self'; report-uri https://logs.browser-intake-datadoghq.com/api/v2/logs?dd-api-key=pube4f163c23bbf91c16b8f57f56af9fc58&dd-evp-origin=content-security-policy&ddsource=csp-report&ddtags=site%3Adatadoghq.com
+      content-type:
+      - application/vnd.api+json
+      date:
+      - Fri, 08 Aug 2025 20:49:59 GMT
+      strict-transport-security:
+      - max-age=31536000; includeSubDomains; preload
+      vary:
+      - Accept-Encoding
+      x-content-type-options:
+      - nosniff
+      x-frame-options:
+      - SAMEORIGIN
+    status:
+      code: 200
+      message: OK
+version: 1
diff --git a/tests/llmobs/llmobs_cassettes/datadog/datadog_api_unstable_llm-obs_v1_datasets_cfdb514f-92b0-404f-9b3f-020f82b0cd51_records_get_8c03b3af.yaml b/tests/llmobs/llmobs_cassettes/datadog/datadog_api_unstable_llm-obs_v1_datasets_cfdb514f-92b0-404f-9b3f-020f82b0cd51_records_get_8c03b3af.yaml
@@ -0,0 +1,47 @@
+interactions:
+- request:
+    body: null
+    headers:
+      Accept:
+      - '*/*'
+      ? !!python/object/apply:multidict._multidict.istr
+      - Accept-Encoding
+      : - identity
+      Connection:
+      - keep-alive
+      ? !!python/object/apply:multidict._multidict.istr
+      - Content-Length
+      : - '0'
+      ? !!python/object/apply:multidict._multidict.istr
+      - Content-Type
+      : - application/json
+      User-Agent:
+      - python-requests/2.32.3
+    method: GET
+    uri: https://api.datadoghq.com/api/unstable/llm-obs/v1/datasets/cfdb514f-92b0-404f-9b3f-020f82b0cd51/records
+  response:
+    body:
+      string: '{"data":[{"id":"e6c19140-44dd-41ae-8617-adb52e2a79c1","type":"datasets","attributes":{"author":{"id":"de473b30-eb9f-11e9-a77a-c7405862b8bd"},"created_at":"2025-08-08T20:50:03.226626Z","dataset_id":"cfdb514f-92b0-404f-9b3f-020f82b0cd51","input":{"prompt":"What
+        is the capital of Sweden?"},"updated_at":"2025-08-08T20:50:03.226626Z"}},{"id":"0a37e076-a220-427b-babe-74b9b7ca74a1","type":"datasets","attributes":{"author":{"id":"de473b30-eb9f-11e9-a77a-c7405862b8bd"},"created_at":"2025-08-08T20:49:58.695339Z","dataset_id":"cfdb514f-92b0-404f-9b3f-020f82b0cd51","expected_output":{"answer":"Rome"},"input":{"prompt":"What
+        is the capital of Italy?"},"updated_at":"2025-08-08T20:49:58.695339Z"}}],"meta":{"after":""}}'
+    headers:
+      content-length:
+      - '713'
+      content-security-policy:
+      - frame-ancestors 'self'; report-uri https://logs.browser-intake-datadoghq.com/api/v2/logs?dd-api-key=pube4f163c23bbf91c16b8f57f56af9fc58&dd-evp-origin=content-security-policy&ddsource=csp-report&ddtags=site%3Adatadoghq.com
+      content-type:
+      - application/vnd.api+json
+      date:
+      - Fri, 08 Aug 2025 20:50:07 GMT
+      strict-transport-security:
+      - max-age=31536000; includeSubDomains; preload
+      vary:
+      - Accept-Encoding
+      x-content-type-options:
+      - nosniff
+      x-frame-options:
+      - SAMEORIGIN
+    status:
+      code: 200
+      message: OK
+version: 1
diff --git a/tests/llmobs/llmobs_cassettes/datadog/datadog_api_unstable_llm-obs_v1_datasets_delete_post_19f1f445.yaml b/tests/llmobs/llmobs_cassettes/datadog/datadog_api_unstable_llm-obs_v1_datasets_delete_post_19f1f445.yaml
@@ -0,0 +1,46 @@
+interactions:
+- request:
+    body: '{"data": {"type": "datasets", "attributes": {"type": "soft", "dataset_ids":
+      ["cfdb514f-92b0-404f-9b3f-020f82b0cd51"]}}}'
+    headers:
+      Accept:
+      - '*/*'
+      ? !!python/object/apply:multidict._multidict.istr
+      - Accept-Encoding
+      : - identity
+      Connection:
+      - keep-alive
+      Content-Length:
+      - '119'
+      ? !!python/object/apply:multidict._multidict.istr
+      - Content-Type
+      : - application/json
+      User-Agent:
+      - python-requests/2.32.3
+    method: POST
+    uri: https://api.datadoghq.com/api/unstable/llm-obs/v1/datasets/delete
+  response:
+    body:
+      string: '{"data":[{"id":"cfdb514f-92b0-404f-9b3f-020f82b0cd51","type":"datasets","attributes":{"author":{"id":"de473b30-eb9f-11e9-a77a-c7405862b8bd"},"created_at":"2025-08-08T20:49:45.077961Z","current_version":2,"deleted_at":"2025-08-08T20:50:07.174644Z","description":"A
+        test dataset","name":"test-dataset-test_dataset_delete_after_append[test_dataset_records0]","updated_at":"2025-08-08T20:50:03.656733Z"}}]}'
+    headers:
+      content-length:
+      - '402'
+      content-security-policy:
+      - frame-ancestors 'self'; report-uri https://logs.browser-intake-datadoghq.com/api/v2/logs?dd-api-key=pube4f163c23bbf91c16b8f57f56af9fc58&dd-evp-origin=content-security-policy&ddsource=csp-report&ddtags=site%3Adatadoghq.com
+      content-type:
+      - application/vnd.api+json
+      date:
+      - Fri, 08 Aug 2025 20:50:07 GMT
+      strict-transport-security:
+      - max-age=31536000; includeSubDomains; preload
+      vary:
+      - Accept-Encoding
+      x-content-type-options:
+      - nosniff
+      x-frame-options:
+      - SAMEORIGIN
+    status:
+      code: 200
+      message: OK
+version: 1
diff --git a/tests/llmobs/llmobs_cassettes/datadog/datadog_api_unstable_llm-obs_v1_datasets_filter_name__test-dataset-test_dataset_delete_after_append_test_dataset_records0__get_7ff2210e.yaml b/tests/llmobs/llmobs_cassettes/datadog/datadog_api_unstable_llm-obs_v1_datasets_filter_name__test-dataset-test_dataset_delete_after_append_test_dataset_records0__get_7ff2210e.yaml
@@ -0,0 +1,46 @@
+interactions:
+- request:
+    body: null
+    headers:
+      Accept:
+      - '*/*'
+      ? !!python/object/apply:multidict._multidict.istr
+      - Accept-Encoding
+      : - identity
+      Connection:
+      - keep-alive
+      ? !!python/object/apply:multidict._multidict.istr
+      - Content-Length
+      : - '0'
+      ? !!python/object/apply:multidict._multidict.istr
+      - Content-Type
+      : - application/json
+      User-Agent:
+      - python-requests/2.32.3
+    method: GET
+    uri: https://api.datadoghq.com/api/unstable/llm-obs/v1/datasets?filter%5Bname%5D=test-dataset-test_dataset_delete_after_append%5Btest_dataset_records0%5D
+  response:
+    body:
+      string: '{"data":[{"id":"cfdb514f-92b0-404f-9b3f-020f82b0cd51","type":"datasets","attributes":{"author":{"id":"de473b30-eb9f-11e9-a77a-c7405862b8bd"},"created_at":"2025-08-08T20:49:45.077961Z","current_version":2,"description":"A
+        test dataset","name":"test-dataset-test_dataset_delete_after_append[test_dataset_records0]","updated_at":"2025-08-08T20:50:03.656733Z"}}],"meta":{"after":""}}'
+    headers:
+      content-length:
+      - '379'
+      content-security-policy:
+      - frame-ancestors 'self'; report-uri https://logs.browser-intake-datadoghq.com/api/v2/logs?dd-api-key=pube4f163c23bbf91c16b8f57f56af9fc58&dd-evp-origin=content-security-policy&ddsource=csp-report&ddtags=site%3Adatadoghq.com
+      content-type:
+      - application/vnd.api+json
+      date:
+      - Fri, 08 Aug 2025 20:50:06 GMT
+      strict-transport-security:
+      - max-age=31536000; includeSubDomains; preload
+      vary:
+      - Accept-Encoding
+      x-content-type-options:
+      - nosniff
+      x-frame-options:
+      - SAMEORIGIN
+    status:
+      code: 200
+      message: OK
+version: 1
diff --git a/tests/llmobs/llmobs_cassettes/datadog/datadog_api_unstable_llm-obs_v1_datasets_post_f9798e4e.yaml b/tests/llmobs/llmobs_cassettes/datadog/datadog_api_unstable_llm-obs_v1_datasets_post_f9798e4e.yaml
@@ -0,0 +1,46 @@
+interactions:
+- request:
+    body: '{"data": {"type": "datasets", "attributes": {"name": "test-dataset-test_dataset_delete_after_append[test_dataset_records0]",
+      "description": "A test dataset"}}}'
+    headers:
+      Accept:
+      - '*/*'
+      ? !!python/object/apply:multidict._multidict.istr
+      - Accept-Encoding
+      : - identity
+      Connection:
+      - keep-alive
+      Content-Length:
+      - '159'
+      ? !!python/object/apply:multidict._multidict.istr
+      - Content-Type
+      : - application/json
+      User-Agent:
+      - python-requests/2.32.3
+    method: POST
+    uri: https://api.datadoghq.com/api/unstable/llm-obs/v1/datasets
+  response:
+    body:
+      string: '{"data":{"id":"cfdb514f-92b0-404f-9b3f-020f82b0cd51","type":"datasets","attributes":{"author":{"id":"de473b30-eb9f-11e9-a77a-c7405862b8bd"},"created_at":"2025-08-08T20:49:45.077961Z","current_version":0,"description":"A
+        test dataset","name":"test-dataset-test_dataset_delete_after_append[test_dataset_records0]","updated_at":"2025-08-08T20:49:45.077961Z"}}}'
+    headers:
+      content-length:
+      - '357'
+      content-security-policy:
+      - frame-ancestors 'self'; report-uri https://logs.browser-intake-datadoghq.com/api/v2/logs?dd-api-key=pube4f163c23bbf91c16b8f57f56af9fc58&dd-evp-origin=content-security-policy&ddsource=csp-report&ddtags=site%3Adatadoghq.com
+      content-type:
+      - application/vnd.api+json
+      date:
+      - Fri, 08 Aug 2025 20:49:58 GMT
+      strict-transport-security:
+      - max-age=31536000; includeSubDomains; preload
+      vary:
+      - Accept-Encoding
+      x-content-type-options:
+      - nosniff
+      x-frame-options:
+      - SAMEORIGIN
+    status:
+      code: 200
+      message: OK
+version: 1
diff --git a/tests/llmobs/test_experiments.py b/tests/llmobs/test_experiments.py
