Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
14 changes: 13 additions & 1 deletion eval/dataset_api.go
Original file line number Diff line number Diff line change
Expand Up @@ -18,7 +18,7 @@ type DatasetAPI[I, R any] struct {

// DatasetQueryOpts contains options for querying datasets.
type DatasetQueryOpts struct {
// Name is the dataset name (requires project context)
// Name is the dataset name (requires ProjectID or ProjectName)
Name string

// ID is the dataset ID
Expand All @@ -29,6 +29,12 @@ type DatasetQueryOpts struct {

// Limit specifies the maximum number of records to return (0 = unlimited)
Limit int

// ProjectID is the project ID (used with Name to disambiguate datasets)
ProjectID string

// ProjectName is the project name (used with Name to disambiguate datasets)
ProjectName string
}

// Get loads a dataset by ID and returns a Dataset iterator.
Expand Down Expand Up @@ -65,6 +71,12 @@ func (d *DatasetAPI[I, R]) Query(ctx context.Context, opts DatasetQueryOpts) (Da
if opts.Version != "" {
queryParams.Version = opts.Version
}
if opts.ProjectID != "" {
queryParams.ProjectID = opts.ProjectID
}
if opts.ProjectName != "" {
queryParams.ProjectName = opts.ProjectName
}

response, err := d.api.Datasets().Query(ctx, queryParams)
if err != nil {
Expand Down
72 changes: 72 additions & 0 deletions eval/dataset_api_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -224,6 +224,78 @@ func TestDatasetAPI_TypeSafety(t *testing.T) {
assert.NotNil(t, datasetAPI)
}

// TestDatasetAPI_QueryByProject tests querying a dataset by name with ProjectName and ProjectID
func TestDatasetAPI_QueryByProject(t *testing.T) {
t.Parallel()

ctx := context.Background()

// Create API client with VCR support
apiClient := createIntegrationTestAPIClient(t)

// Create a different project (not the default integration test project)
testProjectName := integrationTestProject + "-dataset-query"
project, err := apiClient.Projects().Create(ctx, projects.CreateParams{Name: testProjectName})
require.NoError(t, err)

// Use fixed name for VCR determinism
datasetName := "test-dataset-by-project"
dataset, err := apiClient.Datasets().Create(ctx, datasets.CreateParams{
ProjectID: project.ID,
Name: datasetName,
Description: "Test dataset for querying by name and project",
})
require.NoError(t, err)
defer func() {
_ = apiClient.Datasets().Delete(ctx, dataset.ID)
}()

// Insert test data
events := []datasets.Event{
{
Input: map[string]interface{}{
"question": "What is Go?",
},
Expected: map[string]interface{}{
"answer": "A programming language",
},
},
}

err = apiClient.Datasets().InsertEvents(ctx, dataset.ID, events)
require.NoError(t, err)

datasetAPI := &DatasetAPI[testDatasetInput, testDatasetOutput]{
api: apiClient,
}

// Test Query with Name and ProjectName
casesByName, err := datasetAPI.Query(ctx, DatasetQueryOpts{
Name: datasetName,
ProjectName: testProjectName,
})
require.NoError(t, err)
require.NotNil(t, casesByName)

testCase, err := casesByName.Next()
require.NoError(t, err)
assert.Equal(t, "What is Go?", testCase.Input.Question)
assert.Equal(t, "A programming language", testCase.Expected.Answer)

// Test Query with Name and ProjectID
casesByID, err := datasetAPI.Query(ctx, DatasetQueryOpts{
Name: datasetName,
ProjectID: project.ID,
})
require.NoError(t, err)
require.NotNil(t, casesByID)

testCase2, err := casesByID.Next()
require.NoError(t, err)
assert.Equal(t, "What is Go?", testCase2.Input.Question)
assert.Equal(t, "A programming language", testCase2.Expected.Answer)
}

// TestDatasetAPI_PopulatesDatasetFields tests that dataset iterator populates ID, XactID, and Created
func TestDatasetAPI_PopulatesDatasetFields(t *testing.T) {
t.Parallel()
Expand Down
Loading