You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

71 lines
2.0 KiB

11 months ago
  1. class VoyageAiEmbedder {
  2. constructor() {
  3. if (!process.env.VOYAGEAI_API_KEY)
  4. throw new Error("No Voyage AI API key was set.");
  5. const {
  6. VoyageEmbeddings,
  7. } = require("@langchain/community/embeddings/voyage");
  8. this.model = process.env.EMBEDDING_MODEL_PREF || "voyage-3-lite";
  9. this.voyage = new VoyageEmbeddings({
  10. apiKey: process.env.VOYAGEAI_API_KEY,
  11. modelName: this.model,
  12. // Voyage AI's limit per request is 128 https://docs.voyageai.com/docs/rate-limits#use-larger-batches
  13. batchSize: 128,
  14. });
  15. this.embeddingMaxChunkLength = this.#getMaxEmbeddingLength();
  16. }
  17. // https://docs.voyageai.com/docs/embeddings
  18. #getMaxEmbeddingLength() {
  19. switch (this.model) {
  20. case "voyage-finance-2":
  21. case "voyage-multilingual-2":
  22. case "voyage-3":
  23. case "voyage-3-lite":
  24. case "voyage-3-large":
  25. case "voyage-code-3":
  26. return 32_000;
  27. case "voyage-large-2-instruct":
  28. case "voyage-law-2":
  29. case "voyage-code-2":
  30. case "voyage-large-2":
  31. return 16_000;
  32. case "voyage-2":
  33. return 4_000;
  34. default:
  35. return 4_000;
  36. }
  37. }
  38. async embedTextInput(textInput) {
  39. const result = await this.voyage.embedDocuments(
  40. Array.isArray(textInput) ? textInput : [textInput]
  41. );
  42. // If given an array return the native Array[Array] format since that should be the outcome.
  43. // But if given a single string, we need to flatten it so that we have a 1D array.
  44. return (Array.isArray(textInput) ? result : result.flat()) || [];
  45. }
  46. async embedChunks(textChunks = []) {
  47. try {
  48. const embeddings = await this.voyage.embedDocuments(textChunks);
  49. return embeddings;
  50. } catch (error) {
  51. console.error("Voyage AI Failed to embed:", error);
  52. if (
  53. error.message.includes(
  54. "Cannot read properties of undefined (reading '0')"
  55. )
  56. )
  57. throw new Error("Voyage AI failed to embed: Rate limit reached");
  58. throw error;
  59. }
  60. }
  61. }
  62. module.exports = {
  63. VoyageAiEmbedder,
  64. };