feature: add abortSignal to chunk endpoint routes

devflowinc · Aug 23, 2024 · 806e04b · 806e04b
1 parent 0aeb786
commit 806e04b
Show file tree

Hide file tree

Showing 2 changed files with 182 additions and 84 deletions.
diff --git a/clients/ts-sdk/src/fetch-client.ts b/clients/ts-sdk/src/fetch-client.ts
@@ -109,7 +109,8 @@ export class TrieveFetchClient {
   >(
     path: P,
     method: EJECT extends false ? M : HttpMethod,
-    params?: EJECT extends false ? RequestParams<P, M> : URQ
+    params?: EJECT extends false ? RequestParams<P, M> : URQ,
+    signal?: AbortSignal
   ): Promise<EJECT extends false ? ResponseBody<P, M> : URE> {
     let requestBody: unknown;
 
@@ -162,6 +163,7 @@ export class TrieveFetchClient {
       method,
       headers: headers,
       body: requestBody ? JSON.stringify(requestBody) : undefined,
+      signal: signal,
     });
 
     if (!response.ok) {

diff --git a/clients/ts-sdk/src/functions/chunks/index.ts b/clients/ts-sdk/src/functions/chunks/index.ts
@@ -28,176 +28,272 @@ import { TrieveSDK } from "../../sdk";
 export async function search(
   /** @hidden */
   this: TrieveSDK,
-  props: SearchChunksReqPayload
+  props: SearchChunksReqPayload,
+  signal?: AbortSignal
 ) {
-  const searchResults = (await this.trieve.fetch("/api/chunk/search", "post", {
-    xApiVersion: "V2",
-    data: props,
-    datasetId: this.datasetId,
-  })) as SearchResponseBody;
+  const searchResults = (await this.trieve.fetch(
+    "/api/chunk/search",
+    "post",
+    {
+      xApiVersion: "V2",
+      data: props,
+      datasetId: this.datasetId,
+    },
+    signal
+  )) as SearchResponseBody;
 
   return searchResults;
 }
 export async function createChunk(
   /** @hidden */
   this: TrieveSDK,
-  props: CreateChunkReqPayloadEnum
+  props: CreateChunkReqPayloadEnum,
+  signal?: AbortSignal
 ) {
-  return this.trieve.fetch("/api/chunk", "post", {
-    data: props,
-    datasetId: this.datasetId,
-  });
+  return this.trieve.fetch(
+    "/api/chunk",
+    "post",
+    {
+      data: props,
+      datasetId: this.datasetId,
+    },
+    signal
+  );
 }
 export async function autocomplete(
   /** @hidden */
   this: TrieveSDK,
-  props: AutocompleteReqPayload
+  props: AutocompleteReqPayload,
+  signal?: AbortSignal
 ) {
-  return this.trieve.fetch("/api/chunk/autocomplete", "post", {
-    data: props,
-    datasetId: this.datasetId,
-  });
+  return this.trieve.fetch(
+    "/api/chunk/autocomplete",
+    "post",
+    {
+      data: props,
+      datasetId: this.datasetId,
+    },
+    signal
+  );
 }
 
 export async function getRecommendedChunks(
   /** @hidden */
   this: TrieveSDK,
-  props: RecommendChunksRequest
+  props: RecommendChunksRequest,
+  signal?: AbortSignal
 ) {
-  return this.trieve.fetch("/api/chunk/recommend", "post", {
-    data: props,
-    datasetId: this.datasetId,
-  });
+  return this.trieve.fetch(
+    "/api/chunk/recommend",
+    "post",
+    {
+      data: props,
+      datasetId: this.datasetId,
+    },
+    signal
+  );
 }
 
 export async function ragOnChunk(
   /** @hidden */
   this: TrieveSDK,
-  props: GenerateOffChunksReqPayload
+  props: GenerateOffChunksReqPayload,
+  signal?: AbortSignal
 ) {
-  return this.trieve.fetch("/api/chunk/generate", "post", {
-    data: props,
-    datasetId: this.datasetId,
-  });
+  return this.trieve.fetch(
+    "/api/chunk/generate",
+    "post",
+    {
+      data: props,
+      datasetId: this.datasetId,
+    },
+    signal
+  );
 }
 
 export async function suggestedQueries(
   /** @hidden */
   this: TrieveSDK,
-  props: SuggestedQueriesReqPayload
+  props: SuggestedQueriesReqPayload,
+  signal?: AbortSignal
 ) {
-  return this.trieve.fetch("/api/chunk/suggestions", "post", {
-    data: props,
-    datasetId: this.datasetId,
-  });
+  return this.trieve.fetch(
+    "/api/chunk/suggestions",
+    "post",
+    {
+      data: props,
+      datasetId: this.datasetId,
+    },
+    signal
+  );
 }
 
 export async function countChunksAboveThreshold(
   /** @hidden */
   this: TrieveSDK,
-  props: CountChunksReqPayload
+  props: CountChunksReqPayload,
+  signal?: AbortSignal
 ) {
-  return this.trieve.fetch("/api/chunk/count", "post", {
-    data: props,
-    datasetId: this.datasetId,
-  });
+  return this.trieve.fetch(
+    "/api/chunk/count",
+    "post",
+    {
+      data: props,
+      datasetId: this.datasetId,
+    },
+    signal
+  );
 }
 
 export async function scroll(
   /** @hidden */
   this: TrieveSDK,
-  props: ScrollChunksReqPayload
+  props: ScrollChunksReqPayload,
+  signal?: AbortSignal
 ) {
-  return this.trieve.fetch("/api/chunks/scroll", "post", {
-    data: props,
-    datasetId: this.datasetId,
-  });
+  return this.trieve.fetch(
+    "/api/chunks/scroll",
+    "post",
+    {
+      data: props,
+      datasetId: this.datasetId,
+    },
+    signal
+  );
 }
 export async function updateChunk(
   /** @hidden */
   this: TrieveSDK,
-  props: UpdateChunkReqPayload
+  props: UpdateChunkReqPayload,
+  signal?: AbortSignal
 ) {
-  return this.trieve.fetch("/api/chunk", "put", {
-    data: props,
-    datasetId: this.datasetId,
-  });
+  return this.trieve.fetch(
+    "/api/chunk",
+    "put",
+    {
+      data: props,
+      datasetId: this.datasetId,
+    },
+    signal
+  );
 }
 export async function updateChunkByTrackingId(
   /** @hidden */
   this: TrieveSDK,
-  props: UpdateChunkByTrackingIdData
+  props: UpdateChunkByTrackingIdData,
+  signal?: AbortSignal
 ) {
-  return this.trieve.fetch("/api/chunk/tracking_id/update", "put", {
-    data: props,
-    datasetId: this.datasetId,
-  });
+  return this.trieve.fetch(
+    "/api/chunk/tracking_id/update",
+    "put",
+    {
+      data: props,
+      datasetId: this.datasetId,
+    },
+    signal
+  );
 }
 
 export async function getChunkByTrackingId(
   /** @hidden */
   this: TrieveSDK,
-  props: Omit<GetChunkByTrackingIdData, "trDataset">
+  props: Omit<GetChunkByTrackingIdData, "trDataset">,
+  signal?: AbortSignal
 ) {
-  return this.trieve.fetch("/api/chunk/tracking_id/{tracking_id}", "get", {
-    trackingId: props.trackingId,
-    datasetId: this.datasetId,
-    xApiVersion: props.xApiVersion ?? "V2",
-  });
+  return this.trieve.fetch(
+    "/api/chunk/tracking_id/{tracking_id}",
+    "get",
+    {
+      trackingId: props.trackingId,
+      datasetId: this.datasetId,
+      xApiVersion: props.xApiVersion ?? "V2",
+    },
+    signal
+  );
 }
 
 export async function deleteChunkByTrackingId(
   /** @hidden */
   this: TrieveSDK,
-  props: Omit<DeleteChunkByTrackingIdData, "trDataset">
+  props: Omit<DeleteChunkByTrackingIdData, "trDataset">,
+  signal?: AbortSignal
 ) {
-  return this.trieve.fetch("/api/chunk/tracking_id/{tracking_id}", "delete", {
-    trackingId: props.trackingId,
-    datasetId: this.datasetId,
-  });
+  return this.trieve.fetch(
+    "/api/chunk/tracking_id/{tracking_id}",
+    "delete",
+    {
+      trackingId: props.trackingId,
+      datasetId: this.datasetId,
+    },
+    signal
+  );
 }
 
 export async function getChunkById(
   /** @hidden */
   this: TrieveSDK,
-  props: Omit<GetChunkByIdData, "trDataset">
+  props: Omit<GetChunkByIdData, "trDataset">,
+  signal?: AbortSignal
 ) {
-  return this.trieve.fetch("/api/chunk/{chunk_id}", "get", {
-    chunkId: props.chunkId,
-    xApiVersion: props.xApiVersion ?? "V2",
-    datasetId: this.datasetId,
-  });
+  return this.trieve.fetch(
+    "/api/chunk/{chunk_id}",
+    "get",
+    {
+      chunkId: props.chunkId,
+      xApiVersion: props.xApiVersion ?? "V2",
+      datasetId: this.datasetId,
+    },
+    signal
+  );
 }
 
 export async function deleteChunkById(
   /** @hidden */
   this: TrieveSDK,
-  props: DeleteChunkData
+  props: DeleteChunkData,
+  signal?: AbortSignal
 ) {
-  return this.trieve.fetch("/api/chunk/{chunk_id}", "delete", {
-    chunkId: props.chunkId,
-    datasetId: this.datasetId,
-  });
+  return this.trieve.fetch(
+    "/api/chunk/{chunk_id}",
+    "delete",
+    {
+      chunkId: props.chunkId,
+      datasetId: this.datasetId,
+    },
+    signal
+  );
 }
 
 export async function getChunksByIds(
   /** @hidden */
   this: TrieveSDK,
-  props: GetChunksData
+  props: GetChunksData,
+  signal?: AbortSignal
 ) {
-  return this.trieve.fetch("/api/chunks", "post", {
-    data: props,
-    datasetId: this.datasetId,
-  });
+  return this.trieve.fetch(
+    "/api/chunks",
+    "post",
+    {
+      data: props,
+      datasetId: this.datasetId,
+    },
+    signal
+  );
 }
 
 export async function getChunksByTrackingIds(
   /** @hidden */
   this: TrieveSDK,
-  props: GetTrackingChunksData
+  props: GetTrackingChunksData,
+  signal?: AbortSignal
 ) {
-  return this.trieve.fetch("/api/chunks/tracking", "post", {
-    data: props,
-    datasetId: this.datasetId,
-  });
+  return this.trieve.fetch(
+    "/api/chunks/tracking",
+    "post",
+    {
+      data: props,
+      datasetId: this.datasetId,
+    },
+    signal
+  );
 }