fix: CR

Author: asafgardin

src/AI21.ts

@@ -3,46 +3,54 @@ import { AI21EnvConfig } from './EnvConfig';
 import { MissingAPIKeyError } from './errors';
 import { Chat } from './resources/chat';
 import { APIClient } from './APIClient';
+import { DefaultQuery, Headers } from './types';
+
+export type ClientOptions = {
+  baseURL?: string;
+  apiKey?: string;
+  maxRetries?: number;
+  timeout?: number;
+  via?: string | null;
+  defaultQuery?: DefaultQuery;
+  defaultHeaders?: Headers;
+  dangerouslyAllowBrowser?: boolean;
+};
 
-export interface AI21Options {
-  apiKey: string;
-  baseURL: string;
-  via: string | null;
-  timeout: number;
-  maxRetries: number;
-}
 export class AI21 extends APIClient {
+  protected options: ClientOptions;
   private apiKey: string;
-  private _options: AI21Options;
+  private via: string | null;
 
-  constructor(
-    apiKey: string = AI21EnvConfig.API_KEY,
-    baseURL: string = AI21EnvConfig.BASE_URL,
-    timeout: number = AI21EnvConfig.TIMEOUT_SECONDS,
-    maxRetries: number = AI21EnvConfig.MAX_RETRIES,
-    via: string | null = null,
-  ) {
-    const options: AI21Options = {
+  constructor({
+    apiKey = AI21EnvConfig.API_KEY,
+    baseURL = AI21EnvConfig.BASE_URL,
+    timeout = AI21EnvConfig.TIMEOUT_SECONDS,
+    maxRetries,
+    via,
+    ...opts
+  }: ClientOptions) {
+    const options: ClientOptions = {
       apiKey,
       baseURL,
-      via,
       timeout,
       maxRetries,
+      via,
+      ...opts,
     };
 
     super({
       baseURL,
       timeout,
       maxRetries,
-      options: {},
     });
 
     if (!apiKey) {
       throw new MissingAPIKeyError();
     }
 
     this.apiKey = apiKey;
-    this._options = options;
+    this.via = via ?? null;
+    this.options = options;
   }
 
   // eslint-disable-next-line @typescript-eslint/no-unused-vars
@@ -59,11 +67,15 @@ export class AI21 extends APIClient {
     };
   }
 
+  protected defaultQuery(): DefaultQuery | undefined {
+    return this.options.defaultQuery;
+  }
+
   protected override getUserAgent(): string {
     let userAgent = super.getUserAgent();
 
-    if (this._options.via) {
-      userAgent += ` via ${this._options.via}`;
+    if (this.via) {
+      userAgent += ` via ${this.via}`;
     }
     return userAgent;
   }

src/APIClient.ts

@@ -25,20 +25,7 @@ const validatePositiveInteger = (name: string, n: unknown): number => {
   return n;
 };
 
-type ClientOptions = {
-  apiKey?: string;
-  organization?: string | null;
-  project?: string | null;
-  baseURL?: string;
-  maxRetries?: number;
-  timeout?: number;
-  defaultQuery?: DefaultQuery;
-  defaultHeaders?: Headers;
-  dangerouslyAllowBrowser?: boolean;
-};
-
 export abstract class APIClient {
-  protected options: ClientOptions;
   protected baseURL: string;
   protected maxRetries: number;
   protected timeout: number;
@@ -47,18 +34,14 @@ export abstract class APIClient {
     baseURL,
     maxRetries = AI21EnvConfig.MAX_RETRIES,
     timeout = AI21EnvConfig.TIMEOUT_SECONDS,
-    options,
   }: {
     baseURL: string;
     maxRetries?: number | undefined;
     timeout: number | undefined;
-    options: ClientOptions;
   }) {
     this.baseURL = baseURL;
     this.maxRetries = validatePositiveInteger('maxRetries', maxRetries);
     this.timeout = validatePositiveInteger('timeout', timeout);
-
-    this.options = options;
   }
   get<Req, Rsp>(path: string, opts?: PromiseOrValue<RequestOptions<Req>>): Promise<Rsp> {
     return this.makeRequest('get', path, opts);
@@ -98,9 +81,7 @@ export abstract class APIClient {
     return {};
   }
 
-  protected defaultQuery(): DefaultQuery | undefined {
-    return this.options.defaultQuery;
-  }
+  protected abstract defaultQuery(): DefaultQuery | undefined;
 
   private makeRequest<Req, Rsp>(
     method: HTTPMethod,

src/example.ts

@@ -4,7 +4,7 @@ import { AI21 } from './ai21';
   This is a temporary example to test the API streaming/non-streaming functionality.
 */
 async function main() {
-  const client = new AI21('L40MQGXxfbtQVnCRqNTTKaojD8Snt7nQ');
+  const client = new AI21({ apiKey: process.env.AI21_API_KEY });
 
   try {
     const response = await client.chat.completions.create({