Export Cohere type

[Rolilink]

I have been working on a CohereEmbeddings class for https://hwchase17.github.io/langchainjs since I need it but have been having issues importing the Cohere type, I had to copy and paste it into my code.

Can types be exported in a clearer way?

Code:

import cohere from "cohere-ai";
import { cohereResponse, embedRequest, embedResponse } from "cohere-ai/dist/models";
import { backOff } from "exponential-backoff";

export const chunkArray = <T>(arr: T[], chunkSize: number) =>
  arr.reduce((chunks, elem, index) => {
    const chunkIndex = Math.floor(index / chunkSize);
    const chunk = chunks[chunkIndex] || [];
    // eslint-disable-next-line no-param-reassign
    chunks[chunkIndex] = chunk.concat([elem]);
    return chunks;
  }, [] as T[][]);

export abstract class Embeddings {
  abstract embedDocuments(documents: string[]): Promise<number[][]>;

  abstract embedQuery(document: string): Promise<number[]>;
}

// Look for a way to export types
interface CohereService {
    init(key: string): void;
    embed(config: embedRequest): Promise<cohereResponse<embedResponse>>;
}

interface ModelParams {
  modelName: string;
    verbose?: boolean;
    batchSize?: number;
    maxRetries?: number;
    cohereApiKey?: string;
}

export class CohereEmbeddings extends Embeddings implements ModelParams {
  modelName = "multilingual-22-12";

  private apiKey: string;
    batchSize = 20;
  maxRetries = 6;
  private client: CohereService = cohere; // fix this types issue

  constructor(
    fields?: Partial<ModelParams> & {
      verbose?: boolean;
      batchSize?: number;
      maxRetries?: number;
      cohereApiKey?: string;
    }
  ) {
    super();

    const apiKey = fields?.cohereApiKey ?? process.env.OPENAI_API_KEY;
    if (!apiKey) {
      throw new Error("OpenAI API key not found");
    }

    this.modelName = fields?.modelName ?? this.modelName;
    this.batchSize = fields?.batchSize ?? this.batchSize;
    this.apiKey = apiKey;
    this.maxRetries = fields?.maxRetries ?? this.maxRetries;

        // initialize embeddings client
        this.client.init(this.apiKey);
  }

  async embedDocuments(texts: string[]): Promise<number[][]> {
    const subPrompts = chunkArray(texts, this.batchSize);

    const embeddings: number[][] = [];

    for (let i = 0; i < subPrompts.length; i += 1) {
      const input = subPrompts[i];
      const { body } = await this.embeddingWithRetry({
        model: this.modelName,
        texts: input,
      });
      for (let j = 0; j < input.length; j += 1) {
        embeddings.push(body.embeddings[j]);
      }
    }

    return embeddings;
  }

  async embedQuery(text: string): Promise<number[]> {
    const { body } = await this.embeddingWithRetry({
      model: this.modelName,
      texts: [text],
    });
    return body.embeddings[0];
  }

  private async embeddingWithRetry(request: embedRequest) {
    const makeCompletionRequest = () => this.client.embed(request);
    return backOff<cohereResponse<embedResponse>>(makeCompletionRequest, {
      startingDelay: 4,
      maxDelay: 10,
      numOfAttempts: this.maxRetries,
    });
  }
}

[suenalaba]

Hi @Rolilink I was facing a similar issue as you and I opened PR #88 to address the issue you mentioned, could you take a look?

[dsinghvi]

@suenalaba @Rolilink the latest SDK now exports a Cohere namespace and when you do Cohere. you will see all the relevant SDK types: