Skip to content

Commit 7147dfc

Browse files
tonyredondotonyredondo
authored
Merge branch 'main' into feat/civisibility-read-cache
2 parents 3e19d5d + ee0a579 commit 7147dfc

5 files changed

Lines changed: 259 additions & 179 deletions

File tree

‎internal/llmobs/transport/dne.go‎

Lines changed: 50 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -212,6 +212,21 @@ type (
212212
CreateExperimentResponse = Response[ExperimentView]
213213
)
214214

215+
// DatasetRecordItemV2 is the wire format for a single item in the v2 dataset-records list response.
216+
// The v2 endpoint returns a flat object (id, input, expected_output, metadata at the top level),
217+
// unlike the unstable endpoint which wraps fields under "attributes".
218+
type DatasetRecordItemV2 struct {
219+
ID string `json:"id"`
220+
Input any `json:"input"`
221+
ExpectedOutput any `json:"expected_output"`
222+
Metadata any `json:"metadata"`
223+
}
224+
225+
type GetDatasetRecordsResponseV2 struct {
226+
Data []DatasetRecordItemV2 `json:"data"`
227+
Meta ResponseMeta `json:"meta"`
228+
}
229+
215230
func (c *Transport) GetDatasetByName(ctx context.Context, name, projectID string) (*DatasetView, error) {
216231
q := url.Values{}
217232
q.Set("filter[name]", name)
@@ -353,60 +368,83 @@ func (c *Transport) BatchUpdateDataset(
353368
}
354369

355370
// GetDatasetRecordsPage fetches a single page of records for the given dataset.
371+
// projectID is the LLM Observability project UUID that owns the dataset.
372+
// version, when non-nil, requests a specific historical snapshot; nil fetches the latest version.
356373
// Returns the records, the cursor for the next page (empty string if no more pages), and any error.
357-
func (c *Transport) GetDatasetRecordsPage(ctx context.Context, datasetID, cursor string) ([]DatasetRecordView, string, error) {
358-
method := http.MethodGet
359-
recordsPath := fmt.Sprintf("%s/datasets/%s/records", endpointPrefixDNE, url.PathEscape(datasetID))
374+
func (c *Transport) GetDatasetRecordsPage(ctx context.Context, projectID, datasetID, cursor string, version *int) ([]DatasetRecordView, string, error) {
375+
recordsPath := fmt.Sprintf("%s/%s/datasets/%s/records", endpointPrefixDNEStable,
376+
url.PathEscape(projectID), url.PathEscape(datasetID))
360377

378+
q := url.Values{}
379+
if version != nil {
380+
q.Set("filter[version]", fmt.Sprintf("%d", *version))
381+
}
361382
if cursor != "" {
362-
recordsPath = fmt.Sprintf("%s?page[cursor]=%s", recordsPath, url.QueryEscape(cursor))
383+
q.Set("page[cursor]", cursor)
384+
}
385+
if len(q) > 0 {
386+
recordsPath = recordsPath + "?" + q.Encode()
363387
}
364388

365-
result, err := c.jsonRequest(ctx, method, recordsPath, subdomainDNE, nil, getDatasetRecordsTimeout)
389+
result, err := c.jsonRequest(ctx, http.MethodGet, recordsPath, subdomainDNE, nil, getDatasetRecordsTimeout)
366390
if err != nil {
367391
return nil, "", err
368392
}
369393
if result.statusCode != http.StatusOK {
370394
return nil, "", fmt.Errorf("unexpected status %d: %s", result.statusCode, string(result.body))
371395
}
372396

373-
var recordsResp GetDatasetRecordsResponse
397+
var recordsResp GetDatasetRecordsResponseV2
374398
if err := json.Unmarshal(result.body, &recordsResp); err != nil {
375399
return nil, "", fmt.Errorf("failed to decode json response: %w", err)
376400
}
377401

378402
records := make([]DatasetRecordView, 0, len(recordsResp.Data))
379403
for _, r := range recordsResp.Data {
380-
rec := r.Attributes
381-
rec.ID = r.ID
382-
records = append(records, rec)
404+
records = append(records, DatasetRecordView{
405+
ID: r.ID,
406+
Input: r.Input,
407+
ExpectedOutput: r.ExpectedOutput,
408+
Metadata: r.Metadata,
409+
})
383410
}
384411

385412
return records, recordsResp.Meta.After, nil
386413
}
387414

388415
// GetDatasetWithRecords fetches the given Dataset and all its records from DataDog.
416+
// version, when non-nil, requests a specific historical snapshot; nil fetches the latest version.
389417
// This eagerly fetches all pages of records.
390-
func (c *Transport) GetDatasetWithRecords(ctx context.Context, name, projectID string) (*DatasetView, []DatasetRecordView, error) {
418+
func (c *Transport) GetDatasetWithRecords(ctx context.Context, name, projectID string, version *int) (*DatasetView, []DatasetRecordView, error) {
391419
// 1) Fetch dataset by name
392420
ds, err := c.GetDatasetByName(ctx, name, projectID)
393421
if err != nil {
394422
return nil, nil, err
395423
}
396424

397-
// 2) Fetch all records with pagination support
425+
// 2) Fetch all records with pagination support.
426+
// The v2 records endpoint has no per-record version field, so stamp the effective
427+
// version (requested snapshot or the dataset's current version) onto every record.
428+
effectiveVersion := ds.CurrentVersion
429+
if version != nil {
430+
effectiveVersion = *version
431+
}
432+
398433
var allRecords []DatasetRecordView
399434
nextCursor := ""
400435
pageNum := 0
401436

402437
for {
403438
log.Debug("llmobs/transport: fetching dataset records page %d", pageNum)
404439

405-
records, cursor, err := c.GetDatasetRecordsPage(ctx, ds.ID, nextCursor)
440+
records, cursor, err := c.GetDatasetRecordsPage(ctx, projectID, ds.ID, nextCursor, version)
406441
if err != nil {
407442
return nil, nil, fmt.Errorf("get dataset records failed on page %d: %w", pageNum, err)
408443
}
409444

445+
for i := range records {
446+
records[i].Version = effectiveVersion
447+
}
410448
allRecords = append(allRecords, records...)
411449

412450
nextCursor = cursor

‎internal/llmobs/transport/transport.go‎

Lines changed: 5 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -37,8 +37,9 @@ const (
3737
endpointEvalMetric = "/api/intake/llm-obs/v2/eval-metric"
3838
endpointLLMSpan = "/api/v2/llmobs"
3939

40-
endpointPrefixEVPProxy = "/evp_proxy/v2"
41-
endpointPrefixDNE = "/api/unstable/llm-obs/v1"
40+
endpointPrefixEVPProxy = "/evp_proxy/v2"
41+
endpointPrefixDNE = "/api/unstable/llm-obs/v1"
42+
endpointPrefixDNEStable = "/api/v2/llm-obs/v1"
4243

4344
subdomainLLMSpan = "llmobs-intake"
4445
subdomainEvalMetric = "api"
@@ -213,8 +214,8 @@ func (c *Transport) request(ctx context.Context, method, path, subdomain string,
213214
req.Header.Set(headerEVPSubdomain, subdomain)
214215
}
215216

216-
// Set headers for datasets and experiments endpoints
217-
if strings.HasPrefix(path, endpointPrefixDNE) {
217+
// Set headers for datasets and experiments endpoints (both unstable and stable v2 paths)
218+
if strings.HasPrefix(path, endpointPrefixDNE) || strings.HasPrefix(path, endpointPrefixDNEStable) {
218219
if c.agentless && c.appKey != "" {
219220
// In agentless mode, set the app key header if available
220221
req.Header.Set("DD-APPLICATION-KEY", c.appKey)

‎llmobs/dataset/dataset.go‎

Lines changed: 9 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -324,7 +324,7 @@ func Pull(ctx context.Context, name string, opts ...PullOption) (*Dataset, error
324324
return nil, fmt.Errorf("failed to get or create project: %w", err)
325325
}
326326

327-
dsResp, recordsResp, err := ll.Transport.GetDatasetWithRecords(ctx, name, project.ID)
327+
dsResp, recordsResp, err := ll.Transport.GetDatasetWithRecords(ctx, name, project.ID, cfg.version)
328328
if err != nil {
329329
return nil, fmt.Errorf("failed to get dataset: %w", err)
330330
}
@@ -339,12 +339,19 @@ func Pull(ctx context.Context, name string, opts ...PullOption) (*Dataset, error
339339
version: rec.Version,
340340
})
341341
}
342+
// When pulling a specific historical version, report that version so that
343+
// experiment.Run registers the run against the correct dataset snapshot
344+
// rather than the latest current_version returned by GetDatasetByName.
345+
dsVersion := dsResp.CurrentVersion
346+
if cfg.version != nil {
347+
dsVersion = *cfg.version
348+
}
342349
ds := &Dataset{
343350
id: dsResp.ID,
344351
name: dsResp.Name,
345352
description: dsResp.Description,
346353
records: records,
347-
version: dsResp.CurrentVersion,
354+
version: dsVersion,
348355
}
349356
return ds, nil
350357
}

0 commit comments

Comments
 (0)