Make DatasetRecords input_data type more strict, add tests

Yun-Kim · Yun-Kim · commit 50f602ec4782 · 2025-07-14T19:34:55.000-04:00
diff --git a/ddtrace/llmobs/_constants.py b/ddtrace/llmobs/_constants.py
@@ -96,4 +96,3 @@
 PROXY_REQUEST = "llmobs.proxy_request"
 
 EXPERIMENT_ID_KEY = "_ml_obs.experiment_id"
-EXPERIMENT_EXPECTED_OUTPUT_KEY = "_ml_obs.meta.input.expected_output"
diff --git a/ddtrace/llmobs/_experiment.py b/ddtrace/llmobs/_experiment.py
@@ -1,9 +1,8 @@
-from typing import TYPE_CHECKING
 from concurrent.futures import ThreadPoolExecutor
 from copy import deepcopy
 import sys
 import time
-import traceback
+from typing import TYPE_CHECKING
 from typing import Any
 from typing import Callable
 from typing import Dict
@@ -19,7 +18,6 @@
 from ddtrace.constants import ERROR_MSG
 from ddtrace.constants import ERROR_STACK
 from ddtrace.constants import ERROR_TYPE
-from ddtrace.llmobs._constants import EXPERIMENT_EXPECTED_OUTPUT_KEY
 
 
 if TYPE_CHECKING:
@@ -29,12 +27,10 @@
 
 JSONType = Union[str, int, float, bool, None, List["JSONType"], Dict[str, "JSONType"]]
 NonNoneJSONType = Union[str, int, float, bool, List[JSONType], Dict[str, JSONType]]
-API_PROCESSING_TIME_SLEEP = 6  # median events processor processing time in seconds
-FLUSH_EVERY = 10  # default number of records to process before flushing
 
 
 class DatasetRecord(TypedDict):
-    input_data: NonNoneJSONType
+    input_data: Dict[str, NonNoneJSONType]
     expected_output: JSONType
     metadata: Dict[str, Any]
     record_id: NotRequired[Optional[str]]
@@ -96,7 +92,7 @@ def __init__(
         project_name: str,
         description: str = "",
         tags: Optional[List[str]] = None,
-        config: Optional[Dict[str, Any]] = None,
+        config: Optional[Dict[str, JSONType]] = None,
         _llmobs_instance: Optional["LLMObs"] = None,
     ) -> None:
         self.name = name
@@ -105,7 +101,7 @@ def __init__(
         self._evaluators = evaluators
         self._description = description
         self._tags = tags or []
-        self._config: Dict[str, Any] = config or {}
+        self._config: Dict[str, JSONType] = config or {}
         self._llmobs_instance = _llmobs_instance
 
         if not project_name:
@@ -141,24 +137,28 @@ def run(self, jobs: int = 1, raise_errors: bool = False, sample_size: Optional[i
         self._run_evaluators(task_results, raise_errors=raise_errors)
         return
 
-    def _process_record(self, idx_record: Tuple[int, DatasetRecord]) -> Dict[str, Any]:
+    def _process_record(self, idx_record: Tuple[int, DatasetRecord]) -> Dict[str, JSONType]:
+        if not self._llmobs_instance or not self._llmobs_instance.enabled:
+            return {}
         idx, record = idx_record
         start_ns = time.time_ns()
         with self._llmobs_instance._experiment(name=self._task.__name__, experiment_id=self._id) as span:
             span_context = self._llmobs_instance.export_span(span=span)
-            span_id = span_context.get("span_id", "")
-            trace_id = span_context.get("trace_id", "")
+            if span_context:
+                span_id = span_context.get("span_id", "")
+                trace_id = span_context.get("trace_id", "")
+            else:
+                span_id, trace_id = "", ""
             input_data = record["input_data"]
             record_id = record.get("record_id", "")
             expected_output = record["expected_output"]
             tags = {"dataset_id": self._dataset._id, "dataset_record_id": record_id, "experiment_id": self._id}
             output_data = None
             try:
-                output_data = self._task(input_data)  # FIXME: support config?
+                output_data = self._task(input_data, self._config)
             except Exception:
                 span.set_exc_info(*sys.exc_info())
             self._llmobs_instance.annotate(span, input_data=input_data, output_data=output_data, tags=tags)
-            span._set_ctx_item(EXPERIMENT_EXPECTED_OUTPUT_KEY, expected_output)  # FIXME: should we be doing this here?
             return {
                 "idx": idx,
                 "output": output_data,
@@ -178,11 +178,21 @@ def _process_record(self, idx_record: Tuple[int, DatasetRecord]) -> Dict[str, An
                 },
             }
 
-    def _run_task(self, jobs: int, raise_errors: bool = False, sample_size: Optional[int] = None) -> List[Any]:
+    def _run_task(
+        self, jobs: int, raise_errors: bool = False, sample_size: Optional[int] = None
+    ) -> List[Dict[str, JSONType]]:
+        if not self._llmobs_instance or not self._llmobs_instance.enabled:
+            return []
         if sample_size is not None and sample_size < len(self._dataset):
-            subset_data = [deepcopy(record) for record in self._dataset._data[:sample_size]]
+            subset_records = [deepcopy(record) for record in self._dataset._records[:sample_size]]
             subset_name = "[Test subset of {} records] {}".format(sample_size, self._dataset.name)
-            subset_dataset = Dataset(name=subset_name, dataset_id=self._dataset._id, data=subset_data)
+            subset_dataset = Dataset(
+                name=subset_name,
+                dataset_id=self._dataset._id,
+                records=subset_records,
+                description=self._dataset.description,
+                version=self._dataset._version,
+            )
         else:
             subset_dataset = self._dataset
         task_results = []
@@ -191,8 +201,7 @@ def _run_task(self, jobs: int, raise_errors: bool = False, sample_size: Optional
                 task_results.append(result)
                 if raise_errors and result["error"]["message"]:
                     raise RuntimeError("Error on record {}: {}".format(result["idx"], result["error"]["message"]))
-        self._llmobs_instance.flush()
-        time.sleep(API_PROCESSING_TIME_SLEEP)
+        self._llmobs_instance.flush()  # Ensure spans get submitted in serverless environments
         return task_results
 
     def _run_evaluators(self, task_results, raise_errors: bool = False) -> None:
diff --git a/ddtrace/llmobs/_llmobs.py b/ddtrace/llmobs/_llmobs.py
@@ -630,8 +630,8 @@ def experiment(
             raise TypeError("task must be a callable function.")
         sig = inspect.signature(task)
         params = sig.parameters
-        if "input_data" not in params:
-            raise TypeError("Task function must have an 'input_data' parameter.")
+        if "input_data" not in params or "config" not in params:
+            raise TypeError("Task function must have 'input_data' and 'config' parameters.")
         if not isinstance(dataset, Dataset):
             raise TypeError("Dataset must be an LLMObs Dataset object.")
         if not evaluators or not all(callable(evaluator) for evaluator in evaluators):
diff --git a/ddtrace/llmobs/_writer.py b/ddtrace/llmobs/_writer.py
@@ -346,7 +346,7 @@ def dataset_create_with_records(self, name: str, description: str, records: List
     def dataset_batch_update(self, dataset_id: str, records: List[DatasetRecord]) -> int:
         rs: JSONType = [
             {
-                "input": r["input_data"],
+                "input": cast(Dict[str, JSONType], r["input_data"]),
                 "expected_output": r["expected_output"],
                 "metadata": r.get("metadata", {}),
                 "record_id": r.get("record_id", None),
diff --git a/tests/llmobs/llmobs_cassettes/datadog/datadog_api_unstable_llm-obs_v1_datasets_3bf4897d-e6aa-43a3-8d9c-5097b1f85177_batch_update_post_2d58a82a.yaml b/tests/llmobs/llmobs_cassettes/datadog/datadog_api_unstable_llm-obs_v1_datasets_3bf4897d-e6aa-43a3-8d9c-5097b1f85177_batch_update_post_2d58a82a.yaml
@@ -0,0 +1,50 @@
+interactions:
+- request:
+    body: '{"data": {"type": "datasets", "attributes": {"insert_records": [{"input":
+      {"prompt": "What is the capital of France?"}, "expected_output": {"answer":
+      "Paris"}, "metadata": {}, "record_id": null}, {"input": {"prompt": "What is
+      the capital of Canada?"}, "expected_output": {"answer": "Ottawa"}, "metadata":
+      {}, "record_id": null}]}}}'
+    headers:
+      Accept:
+      - '*/*'
+      ? !!python/object/apply:multidict._multidict.istr
+      - Accept-Encoding
+      : - identity
+      Connection:
+      - keep-alive
+      Content-Length:
+      - '331'
+      ? !!python/object/apply:multidict._multidict.istr
+      - Content-Type
+      : - application/json
+      User-Agent:
+      - python-requests/2.32.4
+    method: POST
+    uri: https://api.datadoghq.com/api/unstable/llm-obs/v1/datasets/3bf4897d-e6aa-43a3-8d9c-5097b1f85177/batch_update
+  response:
+    body:
+      string: '{"data":[{"id":"3eda96b0-5590-4886-8633-34154e381dc3","type":"datasets","attributes":{"author":{"id":"df7d11c9-da50-11ed-af19-2e9f609a4672"},"created_at":"2025-07-14T23:18:56.69063271Z","dataset_id":"3bf4897d-e6aa-43a3-8d9c-5097b1f85177","expected_output":{"answer":"Paris"},"input":{"prompt":"What
+        is the capital of France?"},"metadata":{},"updated_at":"2025-07-14T23:18:56.69063271Z","version":1}},{"id":"54fd2188-bdae-47d5-bf76-4c5d8c9fba9f","type":"datasets","attributes":{"author":{"id":"df7d11c9-da50-11ed-af19-2e9f609a4672"},"created_at":"2025-07-14T23:18:56.69063271Z","dataset_id":"3bf4897d-e6aa-43a3-8d9c-5097b1f85177","expected_output":{"answer":"Ottawa"},"input":{"prompt":"What
+        is the capital of Canada?"},"metadata":{},"updated_at":"2025-07-14T23:18:56.69063271Z","version":1}}]}'
+    headers:
+      content-length:
+      - '793'
+      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:
+      - Mon, 14 Jul 2025 23:18:56 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_67c7b6cc-ce98-481e-ab9b-e4925564826c_batch_update_post_a1f44751.yaml b/tests/llmobs/llmobs_cassettes/datadog/datadog_api_unstable_llm-obs_v1_datasets_67c7b6cc-ce98-481e-ab9b-e4925564826c_batch_update_post_a1f44751.yaml
@@ -0,0 +1,47 @@
+interactions:
+- request:
+    body: '{"data": {"type": "datasets", "attributes": {"insert_records": [{"input":
+      {"prompt": "What is the capital of France?"}, "expected_output": {"answer":
+      "Paris"}, "metadata": {}, "record_id": null}]}}}'
+    headers:
+      Accept:
+      - '*/*'
+      ? !!python/object/apply:multidict._multidict.istr
+      - Accept-Encoding
+      : - identity
+      Connection:
+      - keep-alive
+      Content-Length:
+      - '198'
+      ? !!python/object/apply:multidict._multidict.istr
+      - Content-Type
+      : - application/json
+      User-Agent:
+      - python-requests/2.32.4
+    method: POST
+    uri: https://api.datadoghq.com/api/unstable/llm-obs/v1/datasets/67c7b6cc-ce98-481e-ab9b-e4925564826c/batch_update
+  response:
+    body:
+      string: '{"data":[{"id":"ddffedb3-cfa9-459c-80d1-cdfcb7062ec9","type":"datasets","attributes":{"author":{"id":"df7d11c9-da50-11ed-af19-2e9f609a4672"},"created_at":"2025-07-14T23:19:03.665305678Z","dataset_id":"67c7b6cc-ce98-481e-ab9b-e4925564826c","expected_output":{"answer":"Paris"},"input":{"prompt":"What
+        is the capital of France?"},"metadata":{},"updated_at":"2025-07-14T23:19:03.665305678Z","version":1}}]}'
+    headers:
+      content-length:
+      - '403'
+      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:
+      - Mon, 14 Jul 2025 23:19: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_delete_post_264c9f32.yaml b/tests/llmobs/llmobs_cassettes/datadog/datadog_api_unstable_llm-obs_v1_datasets_delete_post_264c9f32.yaml
@@ -0,0 +1,46 @@
+interactions:
+- request:
+    body: '{"data": {"type": "datasets", "attributes": {"type": "soft", "dataset_ids":
+      ["67c7b6cc-ce98-481e-ab9b-e4925564826c"]}}}'
+    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.4
+    method: POST
+    uri: https://api.datadoghq.com/api/unstable/llm-obs/v1/datasets/delete
+  response:
+    body:
+      string: '{"data":[{"id":"67c7b6cc-ce98-481e-ab9b-e4925564826c","type":"datasets","attributes":{"author":{"id":"df7d11c9-da50-11ed-af19-2e9f609a4672"},"created_at":"2025-07-14T23:19:03.530317Z","current_version":1,"deleted_at":"2025-07-14T23:19:09.996669Z","description":"A
+        test dataset","name":"test-dataset-test_experiment_run_task_error[test_dataset_records0]","updated_at":"2025-07-14T23:19:03.816944Z"}}]}'
+    headers:
+      content-length:
+      - '400'
+      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:
+      - Mon, 14 Jul 2025 23:19: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
diff --git a/tests/llmobs/llmobs_cassettes/datadog/datadog_api_unstable_llm-obs_v1_datasets_delete_post_d28fa230.yaml b/tests/llmobs/llmobs_cassettes/datadog/datadog_api_unstable_llm-obs_v1_datasets_delete_post_d28fa230.yaml
@@ -0,0 +1,46 @@
+interactions:
+- request:
+    body: '{"data": {"type": "datasets", "attributes": {"type": "soft", "dataset_ids":
+      ["3bf4897d-e6aa-43a3-8d9c-5097b1f85177"]}}}'
+    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.4
+    method: POST
+    uri: https://api.datadoghq.com/api/unstable/llm-obs/v1/datasets/delete
+  response:
+    body:
+      string: '{"data":[{"id":"3bf4897d-e6aa-43a3-8d9c-5097b1f85177","type":"datasets","attributes":{"author":{"id":"df7d11c9-da50-11ed-af19-2e9f609a4672"},"created_at":"2025-07-14T23:18:56.592831Z","current_version":1,"deleted_at":"2025-07-14T23:19:03.156151Z","description":"A
+        test dataset","name":"test-dataset-test_experiment_run_task[test_dataset_records0]","updated_at":"2025-07-14T23:18:56.822333Z"}}]}'
+    headers:
+      content-length:
+      - '394'
+      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:
+      - Mon, 14 Jul 2025 23:19: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_post_35e295b1.yaml b/tests/llmobs/llmobs_cassettes/datadog/datadog_api_unstable_llm-obs_v1_datasets_post_35e295b1.yaml
@@ -0,0 +1,46 @@
+interactions:
+- request:
+    body: '{"data": {"type": "datasets", "attributes": {"name": "test-dataset-test_experiment_run_task[test_dataset_records0]",
+      "description": "A test dataset"}}}'
+    headers:
+      Accept:
+      - '*/*'
+      ? !!python/object/apply:multidict._multidict.istr
+      - Accept-Encoding
+      : - identity
+      Connection:
+      - keep-alive
+      Content-Length:
+      - '151'
+      ? !!python/object/apply:multidict._multidict.istr
+      - Content-Type
+      : - application/json
+      User-Agent:
+      - python-requests/2.32.4
+    method: POST
+    uri: https://api.datadoghq.com/api/unstable/llm-obs/v1/datasets
+  response:
+    body:
+      string: '{"data":{"id":"3bf4897d-e6aa-43a3-8d9c-5097b1f85177","type":"datasets","attributes":{"author":{"id":"df7d11c9-da50-11ed-af19-2e9f609a4672"},"created_at":"2025-07-14T23:18:56.592831479Z","current_version":0,"description":"A
+        test dataset","name":"test-dataset-test_experiment_run_task[test_dataset_records0]","updated_at":"2025-07-14T23:18:56.592831479Z"}}}'
+    headers:
+      content-length:
+      - '355'
+      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:
+      - Mon, 14 Jul 2025 23:18:56 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_3f098e52.yaml b/tests/llmobs/llmobs_cassettes/datadog/datadog_api_unstable_llm-obs_v1_datasets_post_3f098e52.yaml
diff --git a/tests/llmobs/test_experiments.py b/tests/llmobs/test_experiments.py

Original file line number	Diff line number	Diff line change
`@@ -96,4 +96,3 @@`
`96`	`96`	`PROXY_REQUEST = "llmobs.proxy_request"`
`97`	`97`
`98`	`98`	`EXPERIMENT_ID_KEY = "_ml_obs.experiment_id"`
`99`		`-EXPERIMENT_EXPECTED_OUTPUT_KEY = "_ml_obs.meta.input.expected_output"`
Original file line number	Diff line number	Diff line change
`@@ -346,7 +346,7 @@ def dataset_create_with_records(self, name: str, description: str, records: List`
`346`	`346`	`def dataset_batch_update(self, dataset_id: str, records: List[DatasetRecord]) -> int:`
`347`	`347`	`rs: JSONType = [`
`348`	`348`	`{`
`349`		`- "input": r["input_data"],`
	`349`	`+ "input": cast(Dict[str, JSONType], r["input_data"]),`
`350`	`350`	`"expected_output": r["expected_output"],`
`351`	`351`	`"metadata": r.get("metadata", {}),`
`352`	`352`	`"record_id": r.get("record_id", None),`