file-batches.mjs 5.4 KB

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