Embedding header providers (#59)

* couple tiny fixes

* added embedding headers providers (no integration + docs + testing)

* added embedding headers providers (no docs + testing)

* patched a few implemenation bugs

* added embedding headers providers (no testing)

* added embedding headers providers unit tests

* headers integration testing

Author: toptobes

scripts/list-embedding-providers.sh

@@ -1,18 +1,9 @@
+# Errors on using unbound variables
+set -u
+
 # Properly sources the .env file to bring env variables into scope
 eval "$(tr -d '\r' < .env)"
 
-# Sanity check
-if [ -z "${APPLICATION_URI}" ]; then
-  echo "Error: APPLICATION_URI is not set."
-  exit 1
-fi
-
-# Sanity check
-if [ -z "${APPLICATION_TOKEN}" ]; then
-  echo "Error: APPLICATION_TOKEN is not set."
-  exit 1
-fi
-
 # jq strips away the irrelevant fields we don't care about
 curl -sL "${APPLICATION_URI}/api/json/v1" \
 --header "Token: ${APPLICATION_TOKEN}" \

src/api/clients/data-api-http-client.ts

@@ -15,9 +15,8 @@
 
 import {
   DEFAULT_DATA_API_AUTH_HEADER,
-  DEFAULT_EMBEDDING_API_KEY_HEADER,
   DEFAULT_NAMESPACE,
-  DEFAULT_TIMEOUT,
+  DEFAULT_TIMEOUT, HeaderProvider,
   hrTimeMs,
   HttpClient,
   HTTPClientOptions,
@@ -37,7 +36,8 @@ import {
   AdminCommandSucceededEvent,
   AdminSpawnOptions,
 } from '@/src/devops';
-import { TokenProvider } from '@/src/common';
+import { nullish, TokenProvider } from '@/src/common';
+import { EmbeddingHeadersProvider } from '@/src/data-api/embedding-providers';
 
 /**
  * @internal
@@ -55,11 +55,6 @@ interface ExecuteCommandOptions {
   collection?: string,
 }
 
-interface DataAPIHttpClientOptions extends HTTPClientOptions {
-  namespace: string | undefined,
-  emissionStrategy: EmissionStrategy,
-}
-
 type EmissionStrategy = (emitter: TypedEmitter<DataAPIClientEvents>) => {
   emitCommandStarted(info: DataAPIRequestInfo): void,
   emitCommandFailed(info: DataAPIRequestInfo, error: Error, started: number): void,
@@ -98,6 +93,13 @@ const adaptInfo4Devops = (info: DataAPIRequestInfo) => (<const>{
   path: info.url,
 });
 
+interface DataAPIHttpClientOpts extends HTTPClientOptions {
+  namespace: string | undefined,
+  emissionStrategy: EmissionStrategy,
+  embeddingHeaders: EmbeddingHeadersProvider,
+  tokenProvider: TokenProvider,
+}
+
 /**
  * @internal
  */
@@ -106,44 +108,52 @@ export class DataAPIHttpClient extends HttpClient {
   public namespace?: string;
   public maxTimeMS: number;
   public emissionStrategy: ReturnType<EmissionStrategy>
-  readonly #props: DataAPIHttpClientOptions;
+  readonly #props: DataAPIHttpClientOpts;
 
-  constructor(props: DataAPIHttpClientOptions, embeddingApiKey?: string) {
-    super(props, mkHeaders(embeddingApiKey));
+  constructor(props: DataAPIHttpClientOpts) {
+    super(props, [mkAuthHeaderProvider(props.tokenProvider), props.embeddingHeaders.getHeaders.bind(props.embeddingHeaders)]);
     this.namespace = 'namespace' in props ? props.namespace : DEFAULT_NAMESPACE;
     this.#props = props;
     this.maxTimeMS = this.fetchCtx.maxTimeMS ?? DEFAULT_TIMEOUT;
     this.emissionStrategy = props.emissionStrategy(props.emitter);
   }
 
   public forCollection(namespace: string, collection: string, opts: CollectionSpawnOptions | undefined): DataAPIHttpClient {
-    const clone = new DataAPIHttpClient(this.#props, opts?.embeddingApiKey);
+    const clone = new DataAPIHttpClient({
+      ...this.#props,
+      embeddingHeaders: EmbeddingHeadersProvider.parseHeaders(opts?.embeddingApiKey),
+      namespace: namespace,
+    });
+
     clone.collection = collection;
-    clone.namespace = namespace;
     clone.maxTimeMS = opts?.defaultMaxTimeMS ?? this.maxTimeMS;
+
     return clone;
   }
 
   public forDbAdmin(opts: AdminSpawnOptions | undefined): DataAPIHttpClient {
     const clone = new DataAPIHttpClient({
       ...this.#props,
+      tokenProvider: opts?.adminToken ? TokenProvider.parseToken(opts?.adminToken) : this.#props.tokenProvider,
       monitorCommands: opts?.monitorCommands || this.#props.monitorCommands,
-      applicationToken: TokenProvider.parseToken(opts?.adminToken || this.#props.applicationToken),
       baseUrl: opts?.endpointUrl || this.#props.baseUrl,
       baseApiPath: opts?.endpointUrl ? '' : this.#props.baseApiPath,
     });
+
     clone.emissionStrategy = EmissionStrategy.Admin(this.emitter);
     clone.collection = undefined;
     clone.namespace = undefined;
+
     return clone;
   }
 
-  public timeoutManager(timeoutMs: number | undefined) {
-    return this._mkTimeoutManager(timeoutMs);
+  public timeoutManager(timeout: number | undefined) {
+    timeout ??= this.maxTimeMS;
+    return new TimeoutManager(timeout, () => new DataAPITimeoutError(timeout));
   }
 
   public async executeCommand(command: Record<string, any>, options: TimeoutOptions & ExecuteCommandOptions | undefined) {
-    const timeoutManager = options?.timeoutManager ?? this._mkTimeoutManager(options?.maxTimeMS);
+    const timeoutManager = options?.timeoutManager ?? this.timeoutManager(options?.maxTimeMS);
 
     return await this._requestDataAPI({
       url: this.baseUrl,
@@ -218,11 +228,6 @@ export class DataAPIHttpClient extends HttpClient {
       throw e;
     }
   }
-
-  private _mkTimeoutManager(timeout: number | undefined) {
-    timeout ??= this.maxTimeMS;
-    return new TimeoutManager(timeout, () => new DataAPITimeoutError(timeout));
-  }
 }
 
 const mkFauxErroredResponse = (message: string): RawDataAPIResponse => {
@@ -269,15 +274,14 @@ export function reviver(_: string, value: any): any {
   return value;
 }
 
-function mkHeaders(embeddingApiKey: string | undefined) {
-  if (embeddingApiKey) {
-    return (token: string | undefined) => ({
-      [DEFAULT_EMBEDDING_API_KEY_HEADER]: embeddingApiKey,
-      [DEFAULT_DATA_API_AUTH_HEADER]: token,
-    });
-  } else {
-    return (token: string | undefined) => ({
-      [DEFAULT_DATA_API_AUTH_HEADER]: token,
-    });
-  }
+const mkAuthHeaderProvider = (tp: TokenProvider): HeaderProvider => () => {
+  const token = tp.getToken();
+
+  return (token instanceof Promise)
+    ? token.then(mkAuthHeader)
+    : mkAuthHeader(token);
 }
+
+const mkAuthHeader = (token: string | nullish): Record<string, string> => (token)
+  ? { [DEFAULT_DATA_API_AUTH_HEADER]: token }
+  : {};

src/api/clients/devops-api-http-client.ts

@@ -24,7 +24,8 @@ import {
   AdminCommandStartedEvent,
   AdminCommandSucceededEvent,
 } from '@/src/devops';
-import { HTTPClientOptions, HttpMethodStrings } from '@/src/api/clients/types';
+import { HeaderProvider, HTTPClientOptions, HttpMethodStrings } from '@/src/api/clients/types';
+import { nullish, TokenProvider } from '@/src/common';
 
 /**
  * @internal
@@ -50,19 +51,23 @@ interface DevopsAPIResponse {
   status: number,
 }
 
+interface DevOpsAPIHttpClientOpts extends HTTPClientOptions {
+  tokenProvider: TokenProvider,
+}
+
 /**
  * @internal
  */
 export class DevOpsAPIHttpClient extends HttpClient {
-  constructor(opts: HTTPClientOptions) {
-    super(opts, mkHeaders);
+  constructor(opts: DevOpsAPIHttpClientOpts) {
+    super(opts, [mkAuthHeaderProvider(opts.tokenProvider)]);
   }
 
   public async request(req: DevOpsAPIRequestInfo, options: TimeoutOptions | undefined, started: number = 0): Promise<DevopsAPIResponse> {
     const isLongRunning = started !== 0;
 
     try {
-      const timeoutManager = options?.timeoutManager ?? this._mkTimeoutManager(options?.maxTimeMS);
+      const timeoutManager = options?.timeoutManager ?? this._timeoutManager(options?.maxTimeMS);
       const url = this.baseUrl + req.path;
 
       if (this.monitorCommands && !isLongRunning) {
@@ -109,7 +114,7 @@ export class DevOpsAPIHttpClient extends HttpClient {
   }
 
   public async requestLongRunning(req: DevOpsAPIRequestInfo, info: LongRunningRequestInfo): Promise<DevopsAPIResponse> {
-    const timeoutManager = this._mkTimeoutManager(info.options?.maxTimeMS);
+    const timeoutManager = this._timeoutManager(info.options?.maxTimeMS);
     const isLongRunning = info.options?.blocking !== false;
 
     if (this.monitorCommands) {
@@ -132,6 +137,11 @@ export class DevOpsAPIHttpClient extends HttpClient {
     return resp;
   }
 
+  private _timeoutManager(timeout: number | undefined) {
+    timeout ??= this.fetchCtx.maxTimeMS ?? (12 * 60 * 1000);
+    return new TimeoutManager(timeout, (info) => new DevOpsAPITimeoutError(info.url, timeout));
+  }
+
   private async _awaitStatus(id: string, req: DevOpsAPIRequestInfo, info: LongRunningRequestInfo, timeoutManager: TimeoutManager, started: number): Promise<void> {
     if (info.options?.blocking === false) {
       return;
@@ -180,15 +190,16 @@ export class DevOpsAPIHttpClient extends HttpClient {
       });
     }
   }
-
-  private _mkTimeoutManager(timeout: number | undefined) {
-    timeout ??= this.fetchCtx.maxTimeMS ?? (12 * 60 * 1000);
-    return new TimeoutManager(timeout, (info) => new DevOpsAPITimeoutError(info.url, timeout));
-  }
 }
 
-function mkHeaders(token: string | undefined) {
-  return (token)
-    ? { [DEFAULT_DEVOPS_API_AUTH_HEADER]: `Bearer ${token}` }
-    : {}
+const mkAuthHeaderProvider = (tp: TokenProvider): HeaderProvider => () => {
+  const token = tp.getToken();
+
+  return (token instanceof Promise)
+    ? token.then(mkAuthHeader)
+    : mkAuthHeader(token);
 }
+
+const mkAuthHeader = (token: string | nullish): Record<string, string> => (token)
+  ? { [DEFAULT_DEVOPS_API_AUTH_HEADER]: `Bearer ${token}` }
+  : {};

src/api/clients/http-client.ts

@@ -16,8 +16,7 @@ import { CLIENT_USER_AGENT, RAGSTACK_REQUESTED_WITH } from '@/src/api/constants'
 import { Caller, DataAPIClientEvents } from '@/src/client';
 import TypedEmitter from 'typed-emitter';
 import { FetchCtx, FetcherResponseInfo } from '@/src/api/fetch/types';
-import { HTTPClientOptions, HTTPRequestInfo, MkReqHeaders } from '@/src/api/clients/types';
-import { TokenProvider } from '@/src/common';
+import { HeaderProvider, HTTPClientOptions, HTTPRequestInfo } from '@/src/api/clients/types';
 
 /**
  * @internal
@@ -27,11 +26,9 @@ export abstract class HttpClient {
   readonly emitter: TypedEmitter<DataAPIClientEvents>;
   readonly monitorCommands: boolean;
   readonly fetchCtx: FetchCtx;
-  readonly applicationToken: TokenProvider;
   readonly baseHeaders: Record<string, any>;
 
-  protected constructor(options: HTTPClientOptions, readonly mkReqHeaders: MkReqHeaders) {
-    this.applicationToken = options.applicationToken;
+  protected constructor(options: HTTPClientOptions, readonly headerProviders: HeaderProvider[]) {
     this.baseUrl = options.baseUrl;
     this.emitter = options.emitter;
     this.monitorCommands = options.monitorCommands;
@@ -63,9 +60,17 @@ export abstract class HttpClient {
       ? `${info.url}?${new URLSearchParams(params).toString()}`
       : info.url;
 
-    const token = await this.applicationToken.getToken();
-    const reqHeaders = this.mkReqHeaders(token ?? undefined);
-    Object.assign(reqHeaders, this.baseHeaders);
+    const reqHeaders = { ...this.baseHeaders };
+
+    for (const provider of this.headerProviders) {
+      const maybePromise = provider();
+
+      const newHeaders = ('then' in maybePromise)
+        ? await maybePromise
+        : maybePromise;
+
+      Object.assign(reqHeaders, newHeaders);
+    }
 
     return await this.fetchCtx.ctx.fetch({
       url: url,

src/api/clients/types.ts

@@ -16,15 +16,13 @@ import type TypedEmitter from 'typed-emitter';
 import type { DataAPICommandEvents } from '@/src/data-api';
 import type { FetchCtx, HttpMethods } from '@/src/api';
 import type { TimeoutManager } from '@/src/api/timeout-managers';
-import { TokenProvider } from '@/src/common';
 
 /**
  * @internal
  */
 export interface HTTPClientOptions {
   baseUrl: string,
   baseApiPath?: string | null,
-  applicationToken: TokenProvider,
   emitter: TypedEmitter<DataAPICommandEvents>,
   monitorCommands: boolean,
   fetchCtx: FetchCtx,
@@ -34,7 +32,7 @@ export interface HTTPClientOptions {
 /**
  * @internal
  */
-export type MkReqHeaders = (token: string | undefined) => Record<string, any>;
+export type HeaderProvider = (() => Promise<Record<string, string>> | Record<string, string>);
 
 /**
  * @internal

src/common/token-providers/static-token-provider.ts

@@ -53,7 +53,7 @@ export class StaticTokenProvider extends TokenProvider {
    *
    * @returns the string the token provider was instantiated with.
    */
-  override getToken(): Promise<string | nullish> {
-    return Promise.resolve(this.#token);
+  override getToken(): string | nullish {
+    return this.#token;
   }
 }

src/common/token-providers/token-provider.ts

@@ -43,7 +43,7 @@ export abstract class TokenProvider {
    * The function which provides the token. It may do any I/O as it wishes to obtain/refresh the token, as it's called
    * every time the token is required for use, whether it be for the Data API, or the DevOps API.
    */
-  abstract getToken(): Promise<string | nullish>;
+  abstract getToken(): string | nullish | Promise<string | nullish>;
 
   /**
    * Turns a string token into a {@link StaticTokenProvider} if necessary. Throws an error if

src/common/token-providers/userpass-token-providers.ts

@@ -48,8 +48,8 @@ export class UsernamePasswordTokenProvider extends TokenProvider {
    *
    * @returns the token in the format `cassandra:[username_b64]:[password_b64]`
    */
-  override getToken(): Promise<string> {
-    return Promise.resolve(this.#token);
+  override getToken(): string {
+    return this.#token;
   }
 
   private _encodeB64(input: string) {