feat: added maestro functions

Author: MikeSquall

examples/studio/maestro/maestro-run-create-and-poll.ts

@@ -0,0 +1,21 @@
+import { AI21 } from 'ai21';
+
+const TIMEOUT = 20000;
+const INTERVAL = 1500;
+
+async function main() {
+  const client = new AI21({ apiKey: process.env.AI21_API_KEY });
+
+  const response = await client.beta.maestro.runs.create_and_poll(
+    {
+      input: 'Hello, how are you? tell me a short story about a wizard',
+    },
+    {
+      timeout: TIMEOUT,
+      interval: INTERVAL,
+    },
+  );
+  console.log(response);
+}
+
+main().catch(console.error);

examples/studio/maestro/maestro-run.ts

@@ -0,0 +1,15 @@
+import { AI21 } from 'ai21';
+
+async function main() {
+  const client = new AI21({ apiKey: process.env.AI21_API_KEY });
+
+  const { id } = await client.beta.maestro.runs.create({
+    input: 'Hello, how are you? tell me a short story about a wizard',
+  });
+
+  const response = await client.beta.maestro.runs.get(id);
+
+  console.log(response);
+}
+
+main().catch(console.error);

src/AI21.ts

@@ -4,7 +4,7 @@ import { AI21Error, MissingAPIKeyError } from './errors';
 import { APIClient } from './APIClient';
 import { Headers } from './types';
 import * as Runtime from './runtime';
-import { Chat, ConversationalRag, Library } from './resources';
+import { Beta, Chat, ConversationalRag, Library } from './resources';
 
 export interface ClientOptions {
   baseURL?: string | undefined;
@@ -67,6 +67,7 @@ export class AI21 extends APIClient {
   chat: Chat = new Chat(this);
   conversationalRag: ConversationalRag = new ConversationalRag(this);
   library: Library = new Library(this);
+  beta: Beta = new Beta(this);
 
   // eslint-disable-next-line @typescript-eslint/no-unused-vars
   protected override authHeaders(_: Types.FinalRequestOptions): Types.Headers {

src/resources/chat/completions.ts

@@ -2,7 +2,7 @@ import * as Models from '../../types';
 import { APIResource } from '../../APIResource';
 import { Stream } from '../../streaming';
 
-const deprecatedModels = ['jamba-1.5-mini', 'jamba-1.5-large'];
+const deprecatedModels = ['jamba-1.6-mini', 'jamba-1.6-large'];
 
 export class Completions extends APIResource {
   create(

src/resources/index.ts

@@ -1,3 +1,4 @@
 export { Chat, Completions } from './chat';
 export { ConversationalRag } from './rag';
 export { Library } from './library';
+export { Maestro, Beta, Runs } from './maestro';

src/resources/maestro/beta.ts

@@ -0,0 +1,6 @@
+import { APIResource } from '../../APIResource';
+import { Maestro } from './maestro';
+
+export class Beta extends APIResource {
+  maestro: Maestro = new Maestro(this.client);
+}

src/resources/maestro/index.ts

@@ -0,0 +1,3 @@
+export { Beta } from './beta';
+export { Maestro } from './maestro';
+export { Runs } from './runs';

src/resources/maestro/maestro.ts

@@ -0,0 +1,6 @@
+import { APIResource } from '../../APIResource';
+import { Runs } from './runs';
+
+export class Maestro extends APIResource {
+  runs: Runs = new Runs(this.client);
+}

src/resources/maestro/runs.ts

@@ -0,0 +1,55 @@
+import { APIResource } from '../../APIResource';
+import { MaestroRunRequest, MaestroRunResponse, RequestOptions, MaestroRunRequestOptions } from '../../types';
+
+const MAESTRO_PATH = '/maestro/runs';
+const DEFAULT_TIMEOUT = 30000;
+const DEFAULT_INTERVAL = 1000;
+
+export class Runs extends APIResource {
+  async create(body: MaestroRunRequest): Promise<MaestroRunResponse> {
+    return this.client.post<MaestroRunRequest, MaestroRunResponse>(MAESTRO_PATH, {
+      body,
+    } as RequestOptions<MaestroRunRequest>) as Promise<MaestroRunResponse>;
+  }
+
+  async get(runId: string): Promise<MaestroRunResponse> {
+    return this.client.get<string, MaestroRunResponse>(
+      `${MAESTRO_PATH}/${runId}`,
+    ) as Promise<MaestroRunResponse>;
+  }
+
+  private async poll({
+    runId,
+    timeout,
+    interval,
+  }: {
+    runId: string;
+    timeout: number;
+    interval: number;
+  }): Promise<MaestroRunResponse> {
+    const startTime = Date.now();
+
+    while (Date.now() - startTime < timeout) {
+      const response = await this.get(runId);
+      if (response.status === 'completed') {
+        return response;
+      }
+
+      await new Promise((resolve) => setTimeout(resolve, interval));
+    }
+
+    throw new Error(`Maestro run ${runId} timed out after ${timeout}ms`);
+  }
+
+  async create_and_poll(
+    body: MaestroRunRequest,
+    options?: MaestroRunRequestOptions,
+  ): Promise<MaestroRunResponse> {
+    const response = await this.create(body);
+    return this.poll({
+      runId: response.id,
+      timeout: options?.timeout ?? DEFAULT_TIMEOUT,
+      interval: options?.interval ?? DEFAULT_INTERVAL,
+    });
+  }
+}

src/types/index.ts

@@ -53,3 +53,5 @@ export {
   type UpdateFileRequest,
   type FilePathOrFileObject,
 } from './files';
+
+export { type MaestroRunRequest, type MaestroRunResponse, type MaestroRunRequestOptions } from './maestro';