chore(llmobs): fix issue where non updated fields would lose state locally (#14254)

when updating, the previous behavior would lose state of the fields that
weren't updated, for example - if only expected_output is updated, then
the local state would lose track of input_data and metadata, but remote
state still has the old state

this PR fixes that and maintains the state of the fields that weren't
updated

[before
example](https://app.datadoghq.com/llm/datasets/421c182d-14e8-4df7-956a-1056633a601f?page=1)
<img width="819" height="765" alt="image"
src="https://github.com/user-attachments/assets/7db8115b-7eac-45a4-aa33-02d6b93ad719"
/>
<img width="1335" height="404" alt="image"
src="https://github.com/user-attachments/assets/ea533503-422f-4f49-b79f-49166bb3fdc2"
/>


[after](https://app.datadoghq.com/llm/datasets/e660b3f3-d57e-4eff-a223-2fa5ac6cdcd6?page=1)
<img width="748" height="766" alt="image"
src="https://github.com/user-attachments/assets/23c3bf39-fc82-4bdb-af99-520c01b8a44f"
/>



## 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)

Author: gary-huang

ddtrace/llmobs/_experiment.py

@@ -15,6 +15,8 @@
 from typing import cast
 from typing import overload
 
+from typing_extensions import NotRequired
+
 import ddtrace
 from ddtrace import config
 from ddtrace.constants import ERROR_MSG
@@ -46,6 +48,13 @@ class DatasetRecordRaw(TypedDict):
     metadata: Dict[str, Any]
 
 
+class UpdatableDatasetRecord(TypedDict):
+    input_data: NotRequired[DatasetRecordInputType]
+    expected_output: NotRequired[JSONType]
+    metadata: NotRequired[Dict[str, Any]]
+    record_id: str
+
+
 class DatasetRecord(DatasetRecordRaw):
     record_id: str
 
@@ -87,7 +96,7 @@ class Dataset:
     _version: int
     _dne_client: "LLMObsExperimentsClient"
     _new_records: List[DatasetRecordRaw]
-    _updated_record_ids: List[str]
+    _updated_record_ids_to_new_fields: Dict[str, UpdatableDatasetRecord]
     _deleted_record_ids: List[str]
 
     def __init__(
@@ -106,7 +115,7 @@ def __init__(
         self._dne_client = _dne_client
         self._records = records
         self._new_records = []
-        self._updated_record_ids = []
+        self._updated_record_ids_to_new_fields = {}
         self._deleted_record_ids = []
 
     def push(self) -> None:
@@ -125,7 +134,7 @@ def push(self) -> None:
                 )
             )
 
-        updated_records = [r for r in self._records if "record_id" in r and r["record_id"] in self._updated_record_ids]
+        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
         )
@@ -138,7 +147,7 @@ def push(self) -> None:
         self._version = new_version if new_version != -1 else self._version + 1
         self._new_records = []
         self._deleted_record_ids = []
-        self._updated_record_ids = []
+        self._updated_record_ids_to_new_fields = {}
 
     def update(self, index: int, record: DatasetRecordRaw) -> None:
         if all(k not in record for k in ("input_data", "expected_output", "metadata")):
@@ -147,8 +156,12 @@ def update(self, index: int, record: DatasetRecordRaw) -> None:
                 "input_data, expected_output, or metadata to update"
             )
         record_id = self._records[index]["record_id"]
-        self._updated_record_ids.append(record_id)
-        self._records[index] = {**record, "record_id": record_id}
+        self._updated_record_ids_to_new_fields[record_id] = {
+            **self._updated_record_ids_to_new_fields.get(record_id, {"record_id": record_id}),
+            **record,
+            "record_id": record_id,
+        }
+        self._records[index] = {**self._records[index], **record, "record_id": record_id}
 
     def append(self, record: DatasetRecordRaw) -> None:
         r: DatasetRecord = {**record, "record_id": ""}
@@ -161,6 +174,9 @@ def delete(self, index: int) -> None:
         self._deleted_record_ids.append(record_id)
         del self._records[index]
 
+        if record_id in self._updated_record_ids_to_new_fields:
+            del self._updated_record_ids_to_new_fields[record_id]
+
     @property
     def url(self) -> str:
         # FIXME: will not work for subdomain orgs

ddtrace/llmobs/_writer.py

@@ -45,6 +45,7 @@
 from ddtrace.llmobs._experiment import DatasetRecord
 from ddtrace.llmobs._experiment import DatasetRecordRaw
 from ddtrace.llmobs._experiment import JSONType
+from ddtrace.llmobs._experiment import UpdatableDatasetRecord
 from ddtrace.llmobs._utils import safe_json
 from ddtrace.settings._agent import config as agent_config
 
@@ -355,7 +356,7 @@ def dataset_create(self, name: str, description: str) -> Dataset:
         return Dataset(name, dataset_id, [], description, curr_version, _dne_client=self)
 
     @staticmethod
-    def _get_record_json(record: Union[DatasetRecord, DatasetRecordRaw], is_update: bool) -> JSONType:
+    def _get_record_json(record: Union[UpdatableDatasetRecord, DatasetRecordRaw], is_update: bool) -> JSONType:
         # for now, if a user wants to "erase" the value of expected_output, they are expected to
         # set expected_output to None, and we serialize that as empty string to indicate this to BE
         expected_output: JSONType = None
@@ -383,7 +384,7 @@ def dataset_batch_update(
         self,
         dataset_id: str,
         insert_records: List[DatasetRecordRaw],
-        update_records: List[DatasetRecord],
+        update_records: List[UpdatableDatasetRecord],
         delete_record_ids: List[str],
     ) -> Tuple[int, List[str]]:
         irs: JSONType = [self._get_record_json(r, False) for r in insert_records]

tests/llmobs/llmobs_cassettes/datadog/datadog_api_unstable_llm-obs_v1_datasets_0a3ce469-bbb5-4d8e-8f9d-86024dba131f_batch_update_post_8f4153f3.yaml

@@ -0,0 +1,51 @@
+interactions:
+- request:
+    body: '{"data": {"type": "datasets", "id": "0a3ce469-bbb5-4d8e-8f9d-86024dba131f",
+      "attributes": {"insert_records": [], "update_records": [{"input": {"prompt":
+      "What is the capital of Germany?"}, "expected_output": {"answer": "Berlin"},
+      "metadata": null, "id": "61e3dfe8-ab26-4952-84a0-ef6717a6508c"}, {"input": {"prompt":
+      "What is the capital of Mexico?"}, "expected_output": "", "metadata": {"difficulty":
+      "easy"}, "id": "fe4a6696-deab-4410-8667-dc6b32d1c8c2"}], "delete_records": []}}}'
+    headers:
+      Accept:
+      - '*/*'
+      ? !!python/object/apply:multidict._multidict.istr
+      - Accept-Encoding
+      : - identity
+      Connection:
+      - keep-alive
+      Content-Length:
+      - '481'
+      ? !!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/0a3ce469-bbb5-4d8e-8f9d-86024dba131f/batch_update
+  response:
+    body:
+      string: '{"data":[{"id":"3bace74a-68a9-4a5a-ad70-7af6aba25b57","type":"datasets","attributes":{"author":{"id":"de473b30-eb9f-11e9-a77a-c7405862b8bd"},"created_at":"2025-08-08T05:14:12.868054731Z","dataset_id":"0a3ce469-bbb5-4d8e-8f9d-86024dba131f","expected_output":{"answer":"Berlin"},"input":{"prompt":"What
+        is the capital of Germany?"},"updated_at":"2025-08-08T05:14:12.868054805Z","version":2}},{"id":"0f4490c1-6a55-4688-b695-e46d9e172671","type":"datasets","attributes":{"author":{"id":"de473b30-eb9f-11e9-a77a-c7405862b8bd"},"created_at":"2025-08-08T05:14:12.868072454Z","dataset_id":"0a3ce469-bbb5-4d8e-8f9d-86024dba131f","expected_output":"","input":{"prompt":"What
+        is the capital of Mexico?"},"metadata":{"difficulty":"easy"},"updated_at":"2025-08-08T05:14:12.868072528Z","version":2}}]}'
+    headers:
+      content-length:
+      - '787'
+      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 05:14:12 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

tests/llmobs/llmobs_cassettes/datadog/datadog_api_unstable_llm-obs_v1_datasets_0a3ce469-bbb5-4d8e-8f9d-86024dba131f_batch_update_post_aa4b1918.yaml

@@ -0,0 +1,50 @@
+interactions:
+- request:
+    body: '{"data": {"type": "datasets", "id": "0a3ce469-bbb5-4d8e-8f9d-86024dba131f",
+      "attributes": {"insert_records": [{"input": {"prompt": "What is the capital
+      of France?"}, "expected_output": {"answer": "Paris"}, "metadata": null}, {"input":
+      {"prompt": "What is the capital of China?"}, "expected_output": {"answer": "Beijing"},
+      "metadata": null}], "update_records": [], "delete_records": []}}}'
+    headers:
+      Accept:
+      - '*/*'
+      ? !!python/object/apply:multidict._multidict.istr
+      - Accept-Encoding
+      : - identity
+      Connection:
+      - keep-alive
+      Content-Length:
+      - '387'
+      ? !!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/0a3ce469-bbb5-4d8e-8f9d-86024dba131f/batch_update
+  response:
+    body:
+      string: '{"data":[{"id":"61e3dfe8-ab26-4952-84a0-ef6717a6508c","type":"datasets","attributes":{"author":{"id":"de473b30-eb9f-11e9-a77a-c7405862b8bd"},"created_at":"2025-08-08T05:14:09.951986612Z","dataset_id":"0a3ce469-bbb5-4d8e-8f9d-86024dba131f","expected_output":{"answer":"Paris"},"input":{"prompt":"What
+        is the capital of France?"},"updated_at":"2025-08-08T05:14:09.951986612Z","version":1}},{"id":"fe4a6696-deab-4410-8667-dc6b32d1c8c2","type":"datasets","attributes":{"author":{"id":"de473b30-eb9f-11e9-a77a-c7405862b8bd"},"created_at":"2025-08-08T05:14:09.951986612Z","dataset_id":"0a3ce469-bbb5-4d8e-8f9d-86024dba131f","expected_output":{"answer":"Beijing"},"input":{"prompt":"What
+        is the capital of China?"},"updated_at":"2025-08-08T05:14:09.951986612Z","version":1}}]}'
+    headers:
+      content-length:
+      - '769'
+      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 05:14:10 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

tests/llmobs/llmobs_cassettes/datadog/datadog_api_unstable_llm-obs_v1_datasets_0a3ce469-bbb5-4d8e-8f9d-86024dba131f_records_get_8a97b03b.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/0a3ce469-bbb5-4d8e-8f9d-86024dba131f/records
+  response:
+    body:
+      string: '{"data":[{"id":"3bace74a-68a9-4a5a-ad70-7af6aba25b57","type":"datasets","attributes":{"author":{"id":"de473b30-eb9f-11e9-a77a-c7405862b8bd"},"created_at":"2025-08-08T05:14:12.868054Z","dataset_id":"0a3ce469-bbb5-4d8e-8f9d-86024dba131f","expected_output":{"answer":"Berlin"},"input":{"prompt":"What
+        is the capital of Germany?"},"updated_at":"2025-08-08T05:14:12.868054Z"}},{"id":"0f4490c1-6a55-4688-b695-e46d9e172671","type":"datasets","attributes":{"author":{"id":"de473b30-eb9f-11e9-a77a-c7405862b8bd"},"created_at":"2025-08-08T05:14:12.868072Z","dataset_id":"0a3ce469-bbb5-4d8e-8f9d-86024dba131f","expected_output":"","input":{"prompt":"What
+        is the capital of Mexico?"},"metadata":{"difficulty":"easy"},"updated_at":"2025-08-08T05:14:12.868072Z"}}],"meta":{"after":""}}'
+    headers:
+      content-length:
+      - '771'
+      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 05:14:17 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

tests/llmobs/llmobs_cassettes/datadog/datadog_api_unstable_llm-obs_v1_datasets_1ebc4710-c5a1-4487-b56f-25ecefcc1bf1_batch_update_post_0f800e23.yaml

@@ -0,0 +1,46 @@
+interactions:
+- request:
+    body: '{"data": {"type": "datasets", "id": "1ebc4710-c5a1-4487-b56f-25ecefcc1bf1",
+      "attributes": {"insert_records": [], "update_records": [], "delete_records":
+      ["f681d2a5-01aa-4369-b760-5eed65ae632d"]}}}'
+    headers:
+      Accept:
+      - '*/*'
+      ? !!python/object/apply:multidict._multidict.istr
+      - Accept-Encoding
+      : - identity
+      Connection:
+      - keep-alive
+      Content-Length:
+      - '196'
+      ? !!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/1ebc4710-c5a1-4487-b56f-25ecefcc1bf1/batch_update
+  response:
+    body:
+      string: '{"data":[]}'
+    headers:
+      content-length:
+      - '11'
+      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 05:56:40 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

tests/llmobs/llmobs_cassettes/datadog/datadog_api_unstable_llm-obs_v1_datasets_1ebc4710-c5a1-4487-b56f-25ecefcc1bf1_batch_update_post_27a9e507.yaml

@@ -0,0 +1,50 @@
+interactions:
+- request:
+    body: '{"data": {"type": "datasets", "id": "1ebc4710-c5a1-4487-b56f-25ecefcc1bf1",
+      "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/1ebc4710-c5a1-4487-b56f-25ecefcc1bf1/batch_update
+  response:
+    body:
+      string: '{"data":[{"id":"f681d2a5-01aa-4369-b760-5eed65ae632d","type":"datasets","attributes":{"author":{"id":"de473b30-eb9f-11e9-a77a-c7405862b8bd"},"created_at":"2025-08-08T05:56:35.363359262Z","dataset_id":"1ebc4710-c5a1-4487-b56f-25ecefcc1bf1","expected_output":{"answer":"Paris"},"input":{"prompt":"What
+        is the capital of France?"},"updated_at":"2025-08-08T05:56:35.363359262Z","version":1}},{"id":"f28cabd6-e1c6-4b00-9888-4dea090b13d3","type":"datasets","attributes":{"author":{"id":"de473b30-eb9f-11e9-a77a-c7405862b8bd"},"created_at":"2025-08-08T05:56:35.363359262Z","dataset_id":"1ebc4710-c5a1-4487-b56f-25ecefcc1bf1","expected_output":{"answer":"Rome"},"input":{"prompt":"What
+        is the capital of Italy?"},"updated_at":"2025-08-08T05:56:35.363359262Z","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 05:56:35 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