Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat(js): add readExamplesWithRuns method for fetching examples w/ runs #1410

Open
wants to merge 5 commits into
base: main
Choose a base branch
from
Open
Changes from 1 commit
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
81 changes: 76 additions & 5 deletions js/src/client.ts
Original file line number Diff line number Diff line change
Expand Up @@ -697,7 +697,7 @@ export class Client implements LangSmithTracingClientInterface {
path: string,
body: RecordStringAny | null = null,
requestMethod = "POST",
dataKey = "runs"
dataKey?: string
): AsyncIterable<T[]> {
const bodyParams = body ? { ...body } : {};
while (true) {
Expand All @@ -716,10 +716,13 @@ export class Client implements LangSmithTracingClientInterface {
if (!responseBody) {
break;
}
if (!responseBody[dataKey]) {
break;
if (dataKey) {
if (!responseBody[dataKey]) {
break;
}
yield responseBody[dataKey];
}
yield responseBody[dataKey];
yield responseBody;
tendev-eoghan marked this conversation as resolved.
Show resolved Hide resolved
const cursors = responseBody.cursors;
if (!cursors) {
break;
Expand Down Expand Up @@ -1544,7 +1547,9 @@ export class Client implements LangSmithTracingClientInterface {
let runsYielded = 0;
for await (const runs of this._getCursorPaginatedList<Run>(
"/runs/query",
body
body,
"POST",
"runs"
)) {
if (limit) {
if (runsYielded >= limit) {
Expand Down Expand Up @@ -2779,6 +2784,72 @@ export class Client implements LangSmithTracingClientInterface {
return example;
}

/**
* Reads examples with runs from a dataset.
*
* @param {string} datasetId - The ID of the dataset.
* @param {Object} options - The options for reading examples.
* @param {string[]} options.sessionIds - The session IDs.
* @param {string} [options.comparativeExperimentId] - The comparative experiment ID.
* @param {KVMap} [options.filters] - The filters to apply.
* @param {boolean} [options.preview=false] - Whether to preview the examples.
* @param {number} [options.offset=0] - The offset for pagination.
* @param {number} [options.limit=20] - The limit for pagination.
* @returns {AsyncIterable<Example>} An async iterable of examples.
* @throws {Error} If the dataset ID or any session ID is invalid.
*/
public async *readExamplesWithRuns(
datasetId: string,
{
sessionIds,
comparativeExperimentId,
filters,
preview = false,
offset = 0,
limit = 20,
}: {
sessionIds: string[];
tendev-eoghan marked this conversation as resolved.
Show resolved Hide resolved
comparativeExperimentId?: string;
filters?: KVMap;
preview?: boolean;
offset?: number;
limit?: number;
}
): AsyncIterable<Example> {
assertUuid(datasetId);
const endpoint = `/datasets/${datasetId}/runs`;
sessionIds.forEach((id) => assertUuid(id));
const body: RecordStringAny = {
session_ids: sessionIds,
comparative_experiment_id: comparativeExperimentId,
filters,
preview,
offset,
limit,
};

let examplesYielded = 0;
for await (const exampleWithRuns of this._getCursorPaginatedList<Example>(
endpoint,
body
)) {
if (limit) {
if (examplesYielded >= limit) {
tendev-eoghan marked this conversation as resolved.
Show resolved Hide resolved
break;
}
if (exampleWithRuns.length + examplesYielded > limit) {
const newRuns = exampleWithRuns.slice(0, limit - examplesYielded);
yield* newRuns;
break;
}
examplesYielded += exampleWithRuns.length;
yield* exampleWithRuns;
} else {
yield* exampleWithRuns;
}
}
}

public async *listExamples({
datasetId,
datasetName,
Expand Down