fix(llmobs): stop failing dataset pushes that delete and insert together (#5180)

### What does this PR do?

Fixes `Dataset.Push` returning an error for any push that deletes and inserts in the same batch, even though the backend applied the change. Deleting every record and appending a replacement set, which is the natural way to overwrite a dataset, always failed:

```
push dataset: received a different number of new records than what it was sent (want: 7, got :0)
```

Inserted records now carry the id the client already generates in `Append`, and that id survives the push, so callers can keep using it to update or delete the record later.

### Motivation

The client asserted that the batch_update response contains exactly `len(insert) + len(update)` entries. The response actually contains one entry per affected record, deletions included, so 7 inserts plus 7 deletes came back as 14. The assertion failed, the new record ids were dropped, and the push reported failure after the server had already committed it. A delete-only batch would have failed the same check from the other direction.

The same assumption also made the code read `data[len(update):]` as the inserted records, which on a mixed batch is a slice of deletion entries.

Neither of the other clients for this endpoint validates that count. The Python tracer takes the ids straight from the response and matches them by lookup, never by position, and `ddeval` reads only the new version and ignores the rest. Both send a client-generated id on inserts, which is what removes the need to correlate the response at all; this change does the same. `Append` already minted a UUID for local tracking, so the id simply stops being thrown away.

Verified against the API: a dataset created through this path comes back with exactly the ids the client sent, and a mixed delete-plus-insert push now succeeds.

### Reviewer's Checklist

- [ ] Changed code has unit tests for its functionality at or near 100% coverage.
- [ ] [System-Tests](https://github.com/DataDog/system-tests/) covering this feature have been added and enabled with the va.b.c-dev version tag.
- [ ] There is a benchmark for any new code, or changes to existing code.
- [ ] If this interacts with the agent in a new way, a system test has been added.
- [ ] New code is free of linting errors. You can check this by running `make lint` locally.
- [ ] New code doesn't break existing tests. You can check this by running `make test` locally.
- [ ] Add an appropriate team label so this PR gets put in the right place for the release notes.
- [ ] All generated files are up to date. You can check this by running `make generate` locally.
- [ ] Non-trivial go.mod changes, e.g. adding new modules, are reviewed by @DataDog/dd-trace-go-guild. Make sure all nested modules are up to date by running `make fix-modules` locally.


Co-authored-by: rodrigo.arguello <rodrigo.arguello@datadoghq.com>

Author: rarguelloF

internal/llmobs/transport/dne.go

@@ -67,9 +67,13 @@ type ExperimentView struct {
 }
 
 type DatasetRecordCreate struct {
-	Input          any `json:"input,omitempty"`
-	ExpectedOutput any `json:"expected_output,omitempty"`
-	Metadata       any `json:"metadata,omitempty"`
+	// ID is supplied by the client and persisted by the backend as the record's
+	// id. Sending it means the caller already knows the id of everything it
+	// inserted and does not have to recover it from the response.
+	ID             string `json:"id,omitempty"`
+	Input          any    `json:"input,omitempty"`
+	ExpectedOutput any    `json:"expected_output,omitempty"`
+	Metadata       any    `json:"metadata,omitempty"`
 }
 
 type DatasetRecordUpdate struct {
@@ -303,7 +307,7 @@ func (c *Transport) BatchUpdateDataset(
 	insert []DatasetRecordCreate,
 	update []DatasetRecordUpdate,
 	delete []string,
-) (int, []string, error) {
+) (int, error) {
 	path := fmt.Sprintf("%s/datasets/%s/batch_update", endpointPrefixDNE, url.PathEscape(datasetID))
 	method := http.MethodPost
 	body := BatchUpdateDatasetRequest{
@@ -320,37 +324,30 @@ func (c *Transport) BatchUpdateDataset(
 
 	result, err := c.jsonRequest(ctx, method, path, subdomainDNE, body, payloadLimits)
 	if err != nil {
-		return -1, nil, err
+		return -1, err
 	}
 	if result.statusCode != http.StatusOK {
-		return -1, nil, fmt.Errorf("unexpected status %d: %s", result.statusCode, string(result.body))
+		return -1, fmt.Errorf("unexpected status %d: %s", result.statusCode, string(result.body))
 	}
 
 	var resp BatchUpdateDatasetResponse
 	if err := json.Unmarshal(result.body, &resp); err != nil {
-		return -1, nil, fmt.Errorf("failed to decode json response: %w", err)
-	}
-
-	// FIXME: we don't get version numbers in responses to deletion requests
-	// TODO(rarguelloF): the backend could return a better response here...
-	var (
-		newDatasetVersion = -1
-		newRecordIDs      []string
-	)
-	if len(resp.Data) > 0 {
-		if resp.Data[0].Attributes.Version > 0 {
-			newDatasetVersion = resp.Data[0].Attributes.Version
-		}
-	}
-	if len(resp.Data) == len(insert)+len(update) {
-		// new records are at the end of the slice
-		for _, rec := range resp.Data[len(update):] {
-			newRecordIDs = append(newRecordIDs, rec.ID)
-		}
-	} else {
-		log.Warn("llmobs/internal/transport: BatchUpdateDataset: expected %d records in response, got %d", len(insert)+len(update), len(resp.Data))
-	}
-	return newDatasetVersion, newRecordIDs, nil
+		return -1, fmt.Errorf("failed to decode json response: %w", err)
+	}
+
+	// The response carries one entry per affected record, deletions included:
+	// those come back with deleted_at and ttl set. Nothing here correlates
+	// entries with what was sent, because inserts already carry a client-supplied
+	// id, so the only thing worth reading is the new version. Every entry in a
+	// batch shares it, so the first will do.
+	//
+	// A deletes-only batch returns no version, in which case the caller falls
+	// back to incrementing.
+	newDatasetVersion := -1
+	if len(resp.Data) > 0 && resp.Data[0].Attributes.Version > 0 {
+		newDatasetVersion = resp.Data[0].Attributes.Version
+	}
+	return newDatasetVersion, nil
 }
 
 // GetDatasetRecordsPage fetches a single page of records for the given dataset.

llmobs/dataset/dataset.go

@@ -379,8 +379,8 @@ func (d *Dataset) Append(records ...Record) {
 	d.initialize()
 
 	for _, rec := range records {
-		// This id will be discarded after push, since the backend will generate a new one.
-		// It is used for tracking new records locally before the push.
+		// The id tracks the record locally before the push, and Push sends it as
+		// the record's id, so it stays valid afterwards.
 		id := uuid.New().String()
 		rec.id = id
 
@@ -484,11 +484,13 @@ func (d *Dataset) Push(ctx context.Context) error {
 	}
 
 	// Build slices for inserts, updates, and deletes from the pending maps.
-	insertOldIDs := make([]string, 0, len(d.appendRecords))
+	// Inserts carry the id Append already minted, and the backend persists it.
+	// That is what makes the response uninteresting: nothing has to be matched
+	// back to what was sent.
 	insert := make([]transport.DatasetRecordCreate, 0, len(d.appendRecords))
 	for id, rec := range d.appendRecords {
-		insertOldIDs = append(insertOldIDs, id)
 		insert = append(insert, transport.DatasetRecordCreate{
+			ID:             id,
 			Input:          rec.Input,
 			ExpectedOutput: rec.ExpectedOutput,
 			Metadata:       rec.Metadata,
@@ -522,7 +524,6 @@ func (d *Dataset) Push(ctx context.Context) error {
 		for i := 0; i < len(insert); i += chunkSize {
 			end := min(i+chunkSize, len(insert))
 			chunkInsert := insert[i:end]
-			chunkOldIDs := insertOldIDs[i:end]
 			chunkNum := (i / chunkSize) + 1
 
 			log.Debug("llmobs: uploading dataset chunk %d/%d (%d records)", chunkNum, numBatches, len(chunkInsert))
@@ -535,17 +536,11 @@ func (d *Dataset) Push(ctx context.Context) error {
 				chunkDel = del
 			}
 
-			newVersion, newRecordIDs, err := ll.Transport.BatchUpdateDataset(ctx, d.id, chunkInsert, chunkUpdate, chunkDel)
+			newVersion, err := ll.Transport.BatchUpdateDataset(ctx, d.id, chunkInsert, chunkUpdate, chunkDel)
 			if err != nil {
 				return fmt.Errorf("failed to batch update dataset (chunk %d): %w", chunkNum, err)
 			}
 			log.Debug("llmobs: successfully uploaded dataset chunk %d/%d", chunkNum, numBatches)
-			if len(chunkOldIDs) != len(newRecordIDs) {
-				return fmt.Errorf("received a different number of new records than what it was sent (want: %d, got: %d)", len(chunkOldIDs), len(newRecordIDs))
-			}
-			for j, newID := range newRecordIDs {
-				d.appendRecords[chunkOldIDs[j]].id = newID
-			}
 			if newVersion > 0 {
 				lastVersion = newVersion
 			}
@@ -565,29 +560,17 @@ func (d *Dataset) Push(ctx context.Context) error {
 	// Small delta: a single batch_update request is sufficient.
 	log.Debug("llmobs: dataset delta is %d bytes, using batch update", deltaSize)
 
-	// newRecordIDs should go in the same order
-	newVersion, newRecordIDs, err := ll.Transport.BatchUpdateDataset(ctx, d.id, insert, update, del)
+	newVersion, err := ll.Transport.BatchUpdateDataset(ctx, d.id, insert, update, del)
 	if err != nil {
 		return fmt.Errorf("failed to batch update dataset: %w", err)
 	}
 
-	// TODO(rarguelloF): migrate to new backend response format so this is not necessary
-	if len(insertOldIDs) != len(newRecordIDs) {
-		return fmt.Errorf("received a different number of new records than what it was sent (want: %d, got :%d)", len(insertOldIDs), len(newRecordIDs))
-	}
-
-	// FIXME(rarguelloF): we don't get version numbers in responses to deletion requests
+	// A batch containing only deletions comes back without a version.
 	if newVersion > 0 {
 		d.version = newVersion
 	} else {
 		d.version++
 	}
-
-	// update the inserted records with the new IDs generated by the backend
-	for i, newID := range newRecordIDs {
-		oldID := insertOldIDs[i]
-		d.appendRecords[oldID].id = newID
-	}
 	d.appendRecords = make(map[string]*Record)
 	d.updateRecords = make(map[string]*RecordUpdate)
 	d.deleteRecords = make(map[string]struct{})

llmobs/dataset/dataset_test.go

@@ -1559,21 +1559,35 @@ func handleMockDatasetBatchUpdate(w http.ResponseWriter, r *http.Request) {
 		})
 	}
 
-	// Add inserted records (these get new IDs)
-	for i, insertRec := range attrs.InsertRecords {
-		newID := fmt.Sprintf("new-record-id-%d", i+1)
+	// Inserted records keep the id the client supplied, which is what the
+	// backend does with it.
+	for _, insertRec := range attrs.InsertRecords {
 		response.Data = append(response.Data, llmobstransport.ResponseData[llmobstransport.DatasetRecordView]{
-			ID:   newID,
+			ID:   insertRec.ID,
 			Type: "dataset_records",
 			Attributes: llmobstransport.DatasetRecordView{
-				ID:             newID,
+				ID:             insertRec.ID,
 				Input:          insertRec.Input,
 				ExpectedOutput: insertRec.ExpectedOutput,
 				Version:        2,
 			},
 		})
 	}
 
+	// Deletions come back as entries too, carrying the new version. Emitting
+	// them is what makes this mock able to reproduce a response longer than
+	// insert+update.
+	for _, deletedID := range attrs.DeleteRecords {
+		response.Data = append(response.Data, llmobstransport.ResponseData[llmobstransport.DatasetRecordView]{
+			ID:   deletedID,
+			Type: "dataset_records",
+			Attributes: llmobstransport.DatasetRecordView{
+				ID:      deletedID,
+				Version: 2,
+			},
+		})
+	}
+
 	respData, _ := json.Marshal(response)
 	w.Header().Set("Content-Type", "application/json")
 	w.WriteHeader(http.StatusOK)
@@ -1716,3 +1730,47 @@ func TestLargeDatasetPushChunking(t *testing.T) {
 			"each batch_update request body (%d bytes) must be within the %d byte limit to avoid EVP proxy rejection", size, batchUpdateThreshold)
 	}
 }
+
+// A push that deletes and inserts in the same batch gets back one entry per
+// affected record, deletions included, so the response is longer than
+// insert+update. Counting it and requiring a match failed the push even though
+// the backend had applied it.
+func TestDatasetPushDeleteAndInsertTogether(t *testing.T) {
+	testTracer(t)
+	ctx := context.Background()
+
+	ds, err := Create(ctx, "test-dataset", []Record{
+		{Input: map[string]any{"question": "keep"}, ExpectedOutput: "1"},
+		{Input: map[string]any{"question": "drop"}, ExpectedOutput: "2"},
+	})
+	require.NoError(t, err)
+
+	ds.Delete(1)
+	ds.Append(Record{Input: map[string]any{"question": "added"}, ExpectedOutput: "3"})
+
+	require.NoError(t, ds.Push(ctx))
+	assert.Equal(t, 2, ds.Len())
+}
+
+// Records keep the id they were given locally, because Push sends it and the
+// backend persists it. Callers hold onto these ids to update or delete later,
+// so they have to stay valid across a push.
+func TestDatasetPushKeepsClientRecordIDs(t *testing.T) {
+	testTracer(t)
+	ctx := context.Background()
+
+	ds, err := Create(ctx, "test-dataset", nil)
+	require.NoError(t, err)
+
+	ds.Append(Record{Input: map[string]any{"question": "one"}, ExpectedOutput: "1"})
+	rec, ok := ds.Record(0)
+	require.True(t, ok)
+	idBefore := rec.ID()
+	require.NotEmpty(t, idBefore)
+
+	require.NoError(t, ds.Push(ctx))
+
+	rec, ok = ds.Record(0)
+	require.True(t, ok)
+	assert.Equal(t, idBefore, rec.ID(), "the record id must survive the push")
+}