From 9a53a5da90f6313c81fd221f5beeecbad5c479f4 Mon Sep 17 00:00:00 2001 From: William Fu-Hinthorn <13333726+hinthornw@users.noreply.github.com> Date: Tue, 17 Sep 2024 10:12:06 -0700 Subject: [PATCH] [JS] Anno crud --- js/src/client.ts | 213 ++++++++++++++++++++++++++++++++++++++++++ js/src/schemas.ts | 28 ++++++ js/src/utils/_uuid.ts | 9 +- 3 files changed, 248 insertions(+), 2 deletions(-) diff --git a/js/src/client.ts b/js/src/client.ts index 5a57c3651..c362af893 100644 --- a/js/src/client.ts +++ b/js/src/client.ts @@ -32,6 +32,8 @@ import { TracerSession, TracerSessionResult, ValueType, + AnnotationQueue, + RunWithAnnotationQueueInfo, } from "./schemas.js"; import { convertLangChainMessageToExample, @@ -3103,6 +3105,217 @@ export class Client { return results; } + /** + * API for managing annotation queues + */ + + /** + * List the annotation queues on the LangSmith API. + * @param options - The options for listing annotation queues + * @param options.queueIds - The IDs of the queues to filter by + * @param options.name - The name of the queue to filter by + * @param options.nameContains - The substring that the queue name should contain + * @param options.limit - The maximum number of queues to return + * @returns An iterator of AnnotationQueue objects + */ + public async *listAnnotationQueues( + options: { + queueIds?: string[]; + name?: string; + nameContains?: string; + limit?: number; + } = {} + ): AsyncIterableIterator { + const { queueIds, name, nameContains, limit } = options; + const params: Record = { + ids: queueIds?.map((id, i) => assertUuid(id, `queueIds[${i}]`)), + name, + name_contains: nameContains, + limit: limit !== undefined ? Math.min(limit, 100) : 100, + }; + + let count = 0; + for await (const queue of this._getPaginated( + "/annotation-queues", + params as URLSearchParams + )) { + yield queue as unknown as AnnotationQueue; + count++; + if (limit !== undefined && count >= limit) break; + } + } + + /** + * Create an annotation queue on the LangSmith API. + * @param options - The options for creating an annotation queue + * @param options.name - The name of the annotation queue + * @param options.description - The description of the annotation queue + * @param options.queueId - The ID of the annotation queue + * @returns The created AnnotationQueue object + */ + public async createAnnotationQueue(options: { + name: string; + description?: string; + queueId?: string; + }): Promise { + const { name, description, queueId } = options; + const body = { + name, + description, + id: queueId || uuid.v4(), + }; + + const response = await this.caller.call( + _getFetchImplementation(), + `${this.apiUrl}/annotation-queues`, + { + method: "POST", + headers: { ...this.headers, "Content-Type": "application/json" }, + body: JSON.stringify( + Object.fromEntries( + Object.entries(body).filter(([_, v]) => v !== undefined) + ) + ), + signal: AbortSignal.timeout(this.timeout_ms), + ...this.fetchOptions, + } + ); + await raiseForStatus(response, "create annotation queue"); + const data = await response.json(); + return data as AnnotationQueue; + } + + /** + * Read an annotation queue with the specified queue ID. + * @param queueId - The ID of the annotation queue to read + * @returns The AnnotationQueue object + */ + public async readAnnotationQueue(queueId: string): Promise { + // TODO: Replace when actual endpoint is added + const queue = await this.listAnnotationQueues({ + queueIds: [queueId], + }).next(); + if (queue.done) { + throw new Error(`Annotation queue with ID ${queueId} not found`); + } + return queue.value; + } + + /** + * Update an annotation queue with the specified queue ID. + * @param queueId - The ID of the annotation queue to update + * @param options - The options for updating the annotation queue + * @param options.name - The new name for the annotation queue + * @param options.description - The new description for the annotation queue + */ + public async updateAnnotationQueue( + queueId: string, + options: { + name: string; + description?: string; + } + ): Promise { + const { name, description } = options; + const response = await this.caller.call( + _getFetchImplementation(), + `${this.apiUrl}/annotation-queues/${assertUuid(queueId, "queueId")}`, + { + method: "PATCH", + headers: { ...this.headers, "Content-Type": "application/json" }, + body: JSON.stringify({ name, description }), + signal: AbortSignal.timeout(this.timeout_ms), + ...this.fetchOptions, + } + ); + await raiseForStatus(response, "update annotation queue"); + } + + /** + * Delete an annotation queue with the specified queue ID. + * @param queueId - The ID of the annotation queue to delete + */ + public async deleteAnnotationQueue(queueId: string): Promise { + const response = await this.caller.call( + _getFetchImplementation(), + `${this.apiUrl}/annotation-queues/${assertUuid(queueId, "queueId")}`, + { + method: "DELETE", + headers: { ...this.headers, Accept: "application/json" }, + signal: AbortSignal.timeout(this.timeout_ms), + ...this.fetchOptions, + } + ); + await raiseForStatus(response, "delete annotation queue"); + } + + /** + * Add runs to an annotation queue with the specified queue ID. + * @param queueId - The ID of the annotation queue + * @param runIds - The IDs of the runs to be added to the annotation queue + */ + public async addRunsToAnnotationQueue( + queueId: string, + runIds: string[] + ): Promise { + const response = await this.caller.call( + _getFetchImplementation(), + `${this.apiUrl}/annotation-queues/${assertUuid(queueId, "queueId")}/runs`, + { + method: "POST", + headers: { ...this.headers, "Content-Type": "application/json" }, + body: JSON.stringify( + runIds.map((id, i) => assertUuid(id, `runIds[${i}]`).toString()) + ), + signal: AbortSignal.timeout(this.timeout_ms), + ...this.fetchOptions, + } + ); + await raiseForStatus(response, "add runs to annotation queue"); + } + + /** + * List runs from an annotation queue with the specified queue ID. + * @param queueId - The ID of the annotation queue + * @param limit - The maximum number of runs to return + * @returns An iterator of RunWithAnnotationQueueInfo objects + */ + public async *listRunsFromAnnotationQueue( + queueId: string, + limit?: number + ): AsyncIterableIterator { + const path = `/annotation-queues/${assertUuid(queueId, "queueId")}/runs`; + const limit_ = limit !== undefined ? Math.min(limit, 100) : 100; + + let offset = 0; + let count = 0; + while (true) { + const url = `${this.apiUrl}${path}?offset=${offset}&limit=${limit_}`; + const response = await this.caller.call(_getFetchImplementation(), url, { + method: "GET", + headers: this.headers, + signal: AbortSignal.timeout(this.timeout_ms), + ...this.fetchOptions, + }); + await raiseForStatus(response, `Failed to fetch ${path}`); + const data = await response.json(); + const runs = data.runs as RunWithAnnotationQueueInfo[]; + + if (runs.length === 0) { + break; + } + + for (const run of runs) { + yield run; + count++; + if (limit !== undefined && count >= limit) { + return; + } + } + + offset += runs.length; + } + } + protected async _currentTenantIsOwner(owner: string): Promise { const settings = await this._getSettings(); return owner == "-" || settings.tenant_handle === owner; diff --git a/js/src/schemas.ts b/js/src/schemas.ts index 4d73f29aa..7275f6d39 100644 --- a/js/src/schemas.ts +++ b/js/src/schemas.ts @@ -474,3 +474,31 @@ export interface LangSmithSettings { created_at: string; tenant_handle?: string; } + +export interface AnnotationQueue { + /** The unique identifier of the annotation queue. */ + id: string; + + /** The name of the annotation queue. */ + name: string; + + /** An optional description of the annotation queue. */ + description?: string; + + /** The timestamp when the annotation queue was created. */ + created_at: string; + + /** The timestamp when the annotation queue was last updated. */ + updated_at: string; + + /** The ID of the tenant associated with the annotation queue. */ + tenant_id: string; +} + +export interface RunWithAnnotationQueueInfo extends BaseRun { + /** The last time this run was reviewed. */ + last_reviewed_time?: string; + + /** The time this run was added to the queue. */ + added_at?: string; +} diff --git a/js/src/utils/_uuid.ts b/js/src/utils/_uuid.ts index 714235131..51d71f020 100644 --- a/js/src/utils/_uuid.ts +++ b/js/src/utils/_uuid.ts @@ -1,7 +1,12 @@ import * as uuid from "uuid"; -export function assertUuid(str: string): void { +export function assertUuid(str: string, which?: string): string { if (!uuid.validate(str)) { - throw new Error(`Invalid UUID: ${str}`); + const msg = + which !== undefined + ? `Invalid UUID for ${which}: ${str}` + : `Invalid UUID: ${str}`; + throw new Error(msg); } + return str; }