feat: Introduce Request Timeout Configuration

PiperOrigin-RevId: 611274545

Author: yyyu-google

src/functions/count_tokens.ts

@@ -15,7 +15,11 @@
  * limitations under the License.
  */
 
-import {CountTokensRequest, CountTokensResponse} from '../types/content';
+import {
+  CountTokensRequest,
+  CountTokensResponse,
+  RequestOptions,
+} from '../types/content';
 import {GoogleGenerativeAIError} from '../types/errors';
 import {
   throwErrorIfNotOK,
@@ -34,7 +38,8 @@ export async function countTokens(
   publisherModelEndpoint: string,
   token: Promise<any>,
   request: CountTokensRequest,
-  apiEndpoint?: string
+  apiEndpoint?: string,
+  requestOptions?: RequestOptions
 ): Promise<CountTokensResponse> {
   const response = await postRequest({
     region: location,
@@ -44,6 +49,7 @@ export async function countTokens(
     token: await token,
     data: request,
     apiEndpoint: apiEndpoint,
+    requestOptions: requestOptions,
   }).catch(e => {
     throw new GoogleGenerativeAIError('exception posting request', e);
   });

src/functions/generate_content.ts

@@ -25,6 +25,7 @@ import {
   GenerateContentRequest,
   GenerateContentResult,
   GenerationConfig,
+  RequestOptions,
   SafetySetting,
   StreamGenerateContentResult,
 } from '../types/content';
@@ -51,7 +52,8 @@ export async function generateContent(
   request: GenerateContentRequest | string,
   apiEndpoint?: string,
   generation_config?: GenerationConfig,
-  safety_settings?: SafetySetting[]
+  safety_settings?: SafetySetting[],
+  requestOptions?: RequestOptions
 ): Promise<GenerateContentResult> {
   request = formatContentRequest(request, generation_config, safety_settings);
 
@@ -78,6 +80,7 @@ export async function generateContent(
     token: await token,
     data: generateContentRequest,
     apiEndpoint: apiEndpoint,
+    requestOptions: requestOptions,
     apiVersion: apiVersion,
   }).catch(e => {
     throw new GoogleGenerativeAIError('exception posting request', e);
@@ -103,7 +106,8 @@ export async function generateContentStream(
   request: GenerateContentRequest | string,
   apiEndpoint?: string,
   generation_config?: GenerationConfig,
-  safety_settings?: SafetySetting[]
+  safety_settings?: SafetySetting[],
+  requestOptions?: RequestOptions
 ): Promise<StreamGenerateContentResult> {
   request = formatContentRequest(request, generation_config, safety_settings);
   validateGenerateContentRequest(request);
@@ -129,6 +133,7 @@ export async function generateContentStream(
     token: await token,
     data: generateContentRequest,
     apiEndpoint: apiEndpoint,
+    requestOptions: requestOptions,
     apiVersion: apiVersion,
   }).catch(e => {
     throw new GoogleGenerativeAIError('exception posting request', e);

src/functions/post_request.ts

@@ -17,7 +17,11 @@
 
 const API_BASE_PATH = 'aiplatform.googleapis.com';
 
-import {GenerateContentRequest, CountTokensRequest} from '../types/content';
+import {
+  GenerateContentRequest,
+  CountTokensRequest,
+  RequestOptions,
+} from '../types/content';
 import * as constants from '../util/constants';
 
 /**
@@ -32,6 +36,7 @@ export async function postRequest({
   token,
   data,
   apiEndpoint,
+  requestOptions,
   apiVersion = 'v1',
 }: {
   region: string;
@@ -41,6 +46,7 @@ export async function postRequest({
   token: string;
   data: GenerateContentRequest | CountTokensRequest;
   apiEndpoint?: string;
+  requestOptions?: RequestOptions;
   apiVersion?: string;
 }): Promise<Response | undefined> {
   const vertexBaseEndpoint = apiEndpoint ?? `${region}-${API_BASE_PATH}`;
@@ -53,6 +59,7 @@ export async function postRequest({
   }
 
   return fetch(vertexEndpoint, {
+    ...getFetchOptions(requestOptions),
     method: 'POST',
     headers: {
       Authorization: `Bearer ${token}`,
@@ -62,3 +69,19 @@ export async function postRequest({
     body: JSON.stringify(data),
   });
 }
+
+function getFetchOptions(requestOptions?: RequestOptions): RequestInit {
+  const fetchOptions = {} as RequestInit;
+  if (
+    !requestOptions ||
+    requestOptions.timeoutMillis === undefined ||
+    requestOptions.timeoutMillis < 0
+  ) {
+    return fetchOptions;
+  }
+  const abortController = new AbortController();
+  const signal = abortController.signal;
+  setTimeout(() => abortController.abort, requestOptions.timeoutMillis);
+  fetchOptions.signal = signal;
+  return fetchOptions;
+}