file-batches.mjs 5.6 KB

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