file-batches.js 5.7 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127
  1. "use strict";
  2. // File generated from our OpenAPI spec by Stainless. See CONTRIBUTING.md for details.
  3. Object.defineProperty(exports, "__esModule", { value: true });
  4. exports.VectorStoreFilesPage = exports.FileBatches = void 0;
  5. const resource_1 = require("../../resource.js");
  6. const core_1 = require("../../core.js");
  7. const core_2 = require("../../core.js");
  8. const Util_1 = require("../../lib/Util.js");
  9. const files_1 = require("./files.js");
  10. Object.defineProperty(exports, "VectorStoreFilesPage", { enumerable: true, get: function () { return files_1.VectorStoreFilesPage; } });
  11. class FileBatches extends resource_1.APIResource {
  12. /**
  13. * Create a vector store file batch.
  14. */
  15. create(vectorStoreId, body, options) {
  16. return this._client.post(`/vector_stores/${vectorStoreId}/file_batches`, {
  17. body,
  18. ...options,
  19. headers: { 'OpenAI-Beta': 'assistants=v2', ...options?.headers },
  20. });
  21. }
  22. /**
  23. * Retrieves a vector store file batch.
  24. */
  25. retrieve(vectorStoreId, batchId, options) {
  26. return this._client.get(`/vector_stores/${vectorStoreId}/file_batches/${batchId}`, {
  27. ...options,
  28. headers: { 'OpenAI-Beta': 'assistants=v2', ...options?.headers },
  29. });
  30. }
  31. /**
  32. * Cancel a vector store file batch. This attempts to cancel the processing of
  33. * files in this batch as soon as possible.
  34. */
  35. cancel(vectorStoreId, batchId, options) {
  36. return this._client.post(`/vector_stores/${vectorStoreId}/file_batches/${batchId}/cancel`, {
  37. ...options,
  38. headers: { 'OpenAI-Beta': 'assistants=v2', ...options?.headers },
  39. });
  40. }
  41. /**
  42. * Create a vector store batch and poll until all files have been processed.
  43. */
  44. async createAndPoll(vectorStoreId, body, options) {
  45. const batch = await this.create(vectorStoreId, body);
  46. return await this.poll(vectorStoreId, batch.id, options);
  47. }
  48. listFiles(vectorStoreId, batchId, query = {}, options) {
  49. if ((0, core_1.isRequestOptions)(query)) {
  50. return this.listFiles(vectorStoreId, batchId, {}, query);
  51. }
  52. return this._client.getAPIList(`/vector_stores/${vectorStoreId}/file_batches/${batchId}/files`, files_1.VectorStoreFilesPage, { query, ...options, headers: { 'OpenAI-Beta': 'assistants=v2', ...options?.headers } });
  53. }
  54. /**
  55. * Wait for the given file batch to be processed.
  56. *
  57. * Note: this will return even if one of the files failed to process, you need to
  58. * check batch.file_counts.failed_count to handle this case.
  59. */
  60. async poll(vectorStoreId, batchId, options) {
  61. const headers = { ...options?.headers, 'X-Stainless-Poll-Helper': 'true' };
  62. if (options?.pollIntervalMs) {
  63. headers['X-Stainless-Custom-Poll-Interval'] = options.pollIntervalMs.toString();
  64. }
  65. while (true) {
  66. const { data: batch, response } = await this.retrieve(vectorStoreId, batchId, {
  67. ...options,
  68. headers,
  69. }).withResponse();
  70. switch (batch.status) {
  71. case 'in_progress':
  72. let sleepInterval = 5000;
  73. if (options?.pollIntervalMs) {
  74. sleepInterval = options.pollIntervalMs;
  75. }
  76. else {
  77. const headerInterval = response.headers.get('openai-poll-after-ms');
  78. if (headerInterval) {
  79. const headerIntervalMs = parseInt(headerInterval);
  80. if (!isNaN(headerIntervalMs)) {
  81. sleepInterval = headerIntervalMs;
  82. }
  83. }
  84. }
  85. await (0, core_2.sleep)(sleepInterval);
  86. break;
  87. case 'failed':
  88. case 'cancelled':
  89. case 'completed':
  90. return batch;
  91. }
  92. }
  93. }
  94. /**
  95. * Uploads the given files concurrently and then creates a vector store file batch.
  96. *
  97. * The concurrency limit is configurable using the `maxConcurrency` parameter.
  98. */
  99. async uploadAndPoll(vectorStoreId, { files, fileIds = [] }, options) {
  100. if (files == null || files.length == 0) {
  101. throw new Error(`No \`files\` provided to process. If you've already uploaded files you should use \`.createAndPoll()\` instead`);
  102. }
  103. const configuredConcurrency = options?.maxConcurrency ?? 5;
  104. // We cap the number of workers at the number of files (so we don't start any unnecessary workers)
  105. const concurrencyLimit = Math.min(configuredConcurrency, files.length);
  106. const client = this._client;
  107. const fileIterator = files.values();
  108. const allFileIds = [...fileIds];
  109. // This code is based on this design. The libraries don't accommodate our environment limits.
  110. // https://stackoverflow.com/questions/40639432/what-is-the-best-way-to-limit-concurrency-when-using-es6s-promise-all
  111. async function processFiles(iterator) {
  112. for (let item of iterator) {
  113. const fileObj = await client.files.create({ file: item, purpose: 'assistants' }, options);
  114. allFileIds.push(fileObj.id);
  115. }
  116. }
  117. // Start workers to process results
  118. const workers = Array(concurrencyLimit).fill(fileIterator).map(processFiles);
  119. // Wait for all processing to complete.
  120. await (0, Util_1.allSettledWithThrow)(workers);
  121. return await this.createAndPoll(vectorStoreId, {
  122. file_ids: allFileIds,
  123. });
  124. }
  125. }
  126. exports.FileBatches = FileBatches;
  127. //# sourceMappingURL=file-batches.js.map