Feature Addition: Introduce Request Timeout Configuration (#31)

google-gemini · Jan 31, 2024 · c64fca1 · c64fca1
1 parent 985e01a
commit c64fca1
Show file tree

Hide file tree

Showing 21 changed files with 190 additions and 26 deletions.
diff --git a/.changeset/rotten-queens-cheat.md b/.changeset/rotten-queens-cheat.md
@@ -0,0 +1,5 @@
+---
+"@google/generative-ai": minor
+---
+
+add request timeout configuration
diff --git a/docs/reference/generative-ai.chatsession._constructor_.md b/docs/reference/generative-ai.chatsession._constructor_.md
@@ -9,7 +9,7 @@ Constructs a new instance of the `ChatSession` class
 **Signature:**
 
 ```typescript
-constructor(apiKey: string, model: string, params?: StartChatParams);
+constructor(apiKey: string, model: string, params?: StartChatParams, requestOptions?: RequestOptions);
 ```
 
 ## Parameters
@@ -19,4 +19,5 @@ constructor(apiKey: string, model: string, params?: StartChatParams);
 |  apiKey | string |  |
 |  model | string |  |
 |  params | [StartChatParams](./generative-ai.startchatparams.md) | _(Optional)_ |
+|  requestOptions | [RequestOptions](./generative-ai.requestoptions.md) | _(Optional)_ |
 
diff --git a/docs/reference/generative-ai.chatsession.md b/docs/reference/generative-ai.chatsession.md
@@ -16,14 +16,15 @@ export declare class ChatSession
 
 |  Constructor | Modifiers | Description |
 |  --- | --- | --- |
-|  [(constructor)(apiKey, model, params)](./generative-ai.chatsession._constructor_.md) |  | Constructs a new instance of the <code>ChatSession</code> class |
+|  [(constructor)(apiKey, model, params, requestOptions)](./generative-ai.chatsession._constructor_.md) |  | Constructs a new instance of the <code>ChatSession</code> class |
 
 ## Properties
 
 |  Property | Modifiers | Type | Description |
 |  --- | --- | --- | --- |
 |  [model](./generative-ai.chatsession.model.md) |  | string |  |
 |  [params?](./generative-ai.chatsession.params.md) |  | [StartChatParams](./generative-ai.startchatparams.md) | _(Optional)_ |
+|  [requestOptions?](./generative-ai.chatsession.requestoptions.md) |  | [RequestOptions](./generative-ai.requestoptions.md) | _(Optional)_ |
 
 ## Methods
 

diff --git a/docs/reference/generative-ai.chatsession.requestoptions.md b/docs/reference/generative-ai.chatsession.requestoptions.md
@@ -0,0 +1,11 @@
+<!-- Do not edit this file. It is automatically generated by API Documenter. -->
+
+[Home](./index.md) &gt; [@google/generative-ai](./generative-ai.md) &gt; [ChatSession](./generative-ai.chatsession.md) &gt; [requestOptions](./generative-ai.chatsession.requestoptions.md)
+
+## ChatSession.requestOptions property
+
+**Signature:**
+
+```typescript
+requestOptions?: RequestOptions;
+```
diff --git a/docs/reference/generative-ai.generativemodel._constructor_.md b/docs/reference/generative-ai.generativemodel._constructor_.md
@@ -9,7 +9,7 @@ Constructs a new instance of the `GenerativeModel` class
 **Signature:**
 
 ```typescript
-constructor(apiKey: string, modelParams: ModelParams);
+constructor(apiKey: string, modelParams: ModelParams, requestOptions?: RequestOptions);
 ```
 
 ## Parameters
@@ -18,4 +18,5 @@ constructor(apiKey: string, modelParams: ModelParams);
 |  --- | --- | --- |
 |  apiKey | string |  |
 |  modelParams | [ModelParams](./generative-ai.modelparams.md) |  |
+|  requestOptions | [RequestOptions](./generative-ai.requestoptions.md) | _(Optional)_ |
 
diff --git a/docs/reference/generative-ai.generativemodel.md b/docs/reference/generative-ai.generativemodel.md
@@ -16,7 +16,7 @@ export declare class GenerativeModel
 
 |  Constructor | Modifiers | Description |
 |  --- | --- | --- |
-|  [(constructor)(apiKey, modelParams)](./generative-ai.generativemodel._constructor_.md) |  | Constructs a new instance of the <code>GenerativeModel</code> class |
+|  [(constructor)(apiKey, modelParams, requestOptions)](./generative-ai.generativemodel._constructor_.md) |  | Constructs a new instance of the <code>GenerativeModel</code> class |
 
 ## Properties
 
@@ -25,6 +25,7 @@ export declare class GenerativeModel
 |  [apiKey](./generative-ai.generativemodel.apikey.md) |  | string |  |
 |  [generationConfig](./generative-ai.generativemodel.generationconfig.md) |  | [GenerationConfig](./generative-ai.generationconfig.md) |  |
 |  [model](./generative-ai.generativemodel.model.md) |  | string |  |
+|  [requestOptions](./generative-ai.generativemodel.requestoptions.md) |  | [RequestOptions](./generative-ai.requestoptions.md) |  |
 |  [safetySettings](./generative-ai.generativemodel.safetysettings.md) |  | [SafetySetting](./generative-ai.safetysetting.md)<!-- -->\[\] |  |
 
 ## Methods

diff --git a/docs/reference/generative-ai.generativemodel.requestoptions.md b/docs/reference/generative-ai.generativemodel.requestoptions.md
@@ -0,0 +1,11 @@
+<!-- Do not edit this file. It is automatically generated by API Documenter. -->
+
+[Home](./index.md) &gt; [@google/generative-ai](./generative-ai.md) &gt; [GenerativeModel](./generative-ai.generativemodel.md) &gt; [requestOptions](./generative-ai.generativemodel.requestoptions.md)
+
+## GenerativeModel.requestOptions property
+
+**Signature:**
+
+```typescript
+requestOptions: RequestOptions;
+```
diff --git a/docs/reference/generative-ai.googlegenerativeai.getgenerativemodel.md b/docs/reference/generative-ai.googlegenerativeai.getgenerativemodel.md
@@ -9,14 +9,15 @@ Gets a [GenerativeModel](./generative-ai.generativemodel.md) instance for the pr
 **Signature:**
 
 ```typescript
-getGenerativeModel(modelParams: ModelParams): GenerativeModel;
+getGenerativeModel(modelParams: ModelParams, requestOptions?: RequestOptions): GenerativeModel;
 ```
 
 ## Parameters
 
 |  Parameter | Type | Description |
 |  --- | --- | --- |
 |  modelParams | [ModelParams](./generative-ai.modelparams.md) |  |
+|  requestOptions | [RequestOptions](./generative-ai.requestoptions.md) | _(Optional)_ |
 
 **Returns:**
 

diff --git a/docs/reference/generative-ai.googlegenerativeai.md b/docs/reference/generative-ai.googlegenerativeai.md
@@ -28,5 +28,5 @@ export declare class GoogleGenerativeAI
 
 |  Method | Modifiers | Description |
 |  --- | --- | --- |
-|  [getGenerativeModel(modelParams)](./generative-ai.googlegenerativeai.getgenerativemodel.md) |  | Gets a [GenerativeModel](./generative-ai.generativemodel.md) instance for the provided model name. |
+|  [getGenerativeModel(modelParams, requestOptions)](./generative-ai.googlegenerativeai.getgenerativemodel.md) |  | Gets a [GenerativeModel](./generative-ai.generativemodel.md) instance for the provided model name. |
 
diff --git a/docs/reference/generative-ai.md b/docs/reference/generative-ai.md
@@ -50,6 +50,7 @@
 |  [InputContent](./generative-ai.inputcontent.md) | Content that can be provided as history input to startChat(). |
 |  [ModelParams](./generative-ai.modelparams.md) | Params passed to [GoogleGenerativeAI.getGenerativeModel()](./generative-ai.googlegenerativeai.getgenerativemodel.md)<!-- -->. |
 |  [PromptFeedback](./generative-ai.promptfeedback.md) | If the prompt was blocked, this will be populated with <code>blockReason</code> and the relevant <code>safetyRatings</code>. |
+|  [RequestOptions](./generative-ai.requestoptions.md) | Params passed to [GoogleGenerativeAI.getGenerativeModel()](./generative-ai.googlegenerativeai.getgenerativemodel.md)<!-- -->. |
 |  [SafetyRating](./generative-ai.safetyrating.md) | A safety rating associated with a [GenerateContentCandidate](./generative-ai.generatecontentcandidate.md) |
 |  [SafetySetting](./generative-ai.safetysetting.md) | Safety setting that can be sent as part of request parameters. |
 |  [StartChatParams](./generative-ai.startchatparams.md) | Params for [GenerativeModel.startChat()](./generative-ai.generativemodel.startchat.md)<!-- -->. |

diff --git a/docs/reference/generative-ai.requestoptions.md b/docs/reference/generative-ai.requestoptions.md
@@ -0,0 +1,20 @@
+<!-- Do not edit this file. It is automatically generated by API Documenter. -->
+
+[Home](./index.md) &gt; [@google/generative-ai](./generative-ai.md) &gt; [RequestOptions](./generative-ai.requestoptions.md)
+
+## RequestOptions interface
+
+Params passed to [GoogleGenerativeAI.getGenerativeModel()](./generative-ai.googlegenerativeai.getgenerativemodel.md)<!-- -->.
+
+**Signature:**
+
+```typescript
+export interface RequestOptions 
+```
+
+## Properties
+
+|  Property | Modifiers | Type | Description |
+|  --- | --- | --- | --- |
+|  [timeout?](./generative-ai.requestoptions.timeout.md) |  | number | _(Optional)_ |
+
diff --git a/docs/reference/generative-ai.requestoptions.timeout.md b/docs/reference/generative-ai.requestoptions.timeout.md
@@ -0,0 +1,11 @@
+<!-- Do not edit this file. It is automatically generated by API Documenter. -->
+
+[Home](./index.md) &gt; [@google/generative-ai](./generative-ai.md) &gt; [RequestOptions](./generative-ai.requestoptions.md) &gt; [timeout](./generative-ai.requestoptions.timeout.md)
+
+## RequestOptions.timeout property
+
+**Signature:**
+
+```typescript
+timeout?: number;
+```
diff --git a/packages/main/src/gen-ai.ts b/packages/main/src/gen-ai.ts
@@ -16,7 +16,7 @@
  */
 
 import { GoogleGenerativeAIError } from "./errors";
-import { ModelParams } from "../types";
+import { ModelParams, RequestOptions } from "../types";
 import { GenerativeModel } from "./models/generative-model";
 
 export { ChatSession } from "./methods/chat-session";
@@ -32,13 +32,16 @@ export class GoogleGenerativeAI {
   /**
    * Gets a {@link GenerativeModel} instance for the provided model name.
    */
-  getGenerativeModel(modelParams: ModelParams): GenerativeModel {
+  getGenerativeModel(
+    modelParams: ModelParams,
+    requestOptions?: RequestOptions,
+  ): GenerativeModel {
     if (!modelParams.model) {
       throw new GoogleGenerativeAIError(
         `Must provide a model name. ` +
           `Example: genai.getGenerativeModel({ model: 'my-model-name' })`,
       );
     }
-    return new GenerativeModel(this.apiKey, modelParams);
+    return new GenerativeModel(this.apiKey, modelParams, requestOptions);
   }
 }
diff --git a/packages/main/src/methods/chat-session.ts b/packages/main/src/methods/chat-session.ts
@@ -21,6 +21,7 @@ import {
   GenerateContentResult,
   GenerateContentStreamResult,
   Part,
+  RequestOptions,
   StartChatParams,
 } from "../../types";
 import { formatNewContent } from "../requests/request-helpers";
@@ -47,6 +48,7 @@ export class ChatSession {
     apiKey: string,
     public model: string,
     public params?: StartChatParams,
+    public requestOptions?: RequestOptions,
   ) {
     this._apiKey = apiKey;
     if (params?.history) {
@@ -89,7 +91,12 @@ export class ChatSession {
     // Add onto the chain.
     this._sendPromise = this._sendPromise
       .then(() =>
-        generateContent(this._apiKey, this.model, generateContentRequest),
+        generateContent(
+          this._apiKey,
+          this.model,
+          generateContentRequest,
+          this.requestOptions,
+        ),
       )
       .then((result) => {
         if (
@@ -137,6 +144,7 @@ export class ChatSession {
       this._apiKey,
       this.model,
       generateContentRequest,
+      this.requestOptions,
     );
 
     // Add onto the chain.

diff --git a/packages/main/src/methods/count-tokens.ts b/packages/main/src/methods/count-tokens.ts
@@ -15,15 +15,24 @@
  * limitations under the License.
  */
 
-import { CountTokensRequest, CountTokensResponse } from "../../types";
+import {
+  CountTokensRequest,
+  CountTokensResponse,
+  RequestOptions,
+} from "../../types";
 import { RequestUrl, Task, makeRequest } from "../requests/request";
 
 export async function countTokens(
   apiKey: string,
   model: string,
   params: CountTokensRequest,
+  requestOptions?: RequestOptions,
 ): Promise<CountTokensResponse> {
   const url = new RequestUrl(model, Task.COUNT_TOKENS, apiKey, false);
-  const response = await makeRequest(url, JSON.stringify({ ...params, model }));
+  const response = await makeRequest(
+    url,
+    JSON.stringify({ ...params, model }),
+    requestOptions,
+  );
   return response.json();
 }
diff --git a/packages/main/src/methods/embed-content.ts b/packages/main/src/methods/embed-content.ts
@@ -20,23 +20,30 @@ import {
   BatchEmbedContentsResponse,
   EmbedContentRequest,
   EmbedContentResponse,
+  RequestOptions,
 } from "../../types";
 import { RequestUrl, Task, makeRequest } from "../requests/request";
 
 export async function embedContent(
   apiKey: string,
   model: string,
   params: EmbedContentRequest,
+  requestOptions?: RequestOptions,
 ): Promise<EmbedContentResponse> {
   const url = new RequestUrl(model, Task.EMBED_CONTENT, apiKey, false);
-  const response = await makeRequest(url, JSON.stringify(params));
+  const response = await makeRequest(
+    url,
+    JSON.stringify(params),
+    requestOptions,
+  );
   return response.json();
 }
 
 export async function batchEmbedContents(
   apiKey: string,
   model: string,
   params: BatchEmbedContentsRequest,
+  requestOptions?: RequestOptions,
 ): Promise<BatchEmbedContentsResponse> {
   const url = new RequestUrl(model, Task.BATCH_EMBED_CONTENTS, apiKey, false);
   const requestsWithModel: EmbedContentRequest[] = params.requests.map(
@@ -47,6 +54,7 @@ export async function batchEmbedContents(
   const response = await makeRequest(
     url,
     JSON.stringify({ requests: requestsWithModel }),
+    requestOptions,
   );
   return response.json();
 }
diff --git a/packages/main/src/methods/generate-content.ts b/packages/main/src/methods/generate-content.ts
@@ -20,6 +20,7 @@ import {
   GenerateContentResponse,
   GenerateContentResult,
   GenerateContentStreamResult,
+  RequestOptions,
 } from "../../types";
 import { RequestUrl, Task, makeRequest } from "../requests/request";
 import { addHelpers } from "../requests/response-helpers";
@@ -29,29 +30,39 @@ export async function generateContentStream(
   apiKey: string,
   model: string,
   params: GenerateContentRequest,
+  requestOptions?: RequestOptions,
 ): Promise<GenerateContentStreamResult> {
   const url = new RequestUrl(
     model,
     Task.STREAM_GENERATE_CONTENT,
     apiKey,
     /* stream */ true,
   );
-  const response = await makeRequest(url, JSON.stringify(params));
+  const response = await makeRequest(
+    url,
+    JSON.stringify(params),
+    requestOptions,
+  );
   return processStream(response);
 }
 
 export async function generateContent(
   apiKey: string,
   model: string,
   params: GenerateContentRequest,
+  requestOptions?: RequestOptions,
 ): Promise<GenerateContentResult> {
   const url = new RequestUrl(
     model,
     Task.GENERATE_CONTENT,
     apiKey,
     /* stream */ false,
   );
-  const response = await makeRequest(url, JSON.stringify(params));
+  const response = await makeRequest(
+    url,
+    JSON.stringify(params),
+    requestOptions,
+  );
   const responseJson: GenerateContentResponse = await response.json();
   const enhancedResponse = addHelpers(responseJson);
   return {