|
15 | 15 | * limitations under the License. |
16 | 16 | */ |
17 | 17 |
|
18 | | -import { ErrorDetails, RequestOptions, VertexAIErrorCode } from '../types'; |
| 18 | +import { |
| 19 | + ErrorDetails, |
| 20 | + SingleRequestOptions, |
| 21 | + VertexAIErrorCode |
| 22 | +} from '../types'; |
19 | 23 | import { VertexAIError } from '../errors'; |
20 | 24 | import { ApiSettings } from '../types/internal'; |
21 | 25 | import { |
|
27 | 31 | } from '../constants'; |
28 | 32 | import { logger } from '../logger'; |
29 | 33 |
|
| 34 | +const TIMEOUT_EXPIRED_MESSAGE = 'Timeout has expired.'; |
| 35 | +const ABORT_ERROR_NAME = 'AbortError'; |
| 36 | + |
30 | 37 | export enum Task { |
31 | 38 | GENERATE_CONTENT = 'generateContent', |
32 | 39 | STREAM_GENERATE_CONTENT = 'streamGenerateContent', |
|
40 | 47 | public task: Task, |
41 | 48 | public apiSettings: ApiSettings, |
42 | 49 | public stream: boolean, |
43 | | - public requestOptions?: RequestOptions |
| 50 | + public requestOptions?: SingleRequestOptions |
44 | 51 | ) {} |
45 | 52 | toString(): string { |
46 | 53 | // TODO: allow user-set option if that feature becomes available |
|
115 | 122 | apiSettings: ApiSettings, |
116 | 123 | stream: boolean, |
117 | 124 | body: string, |
118 | | - requestOptions?: RequestOptions |
| 125 | + singleRequestOptions?: SingleRequestOptions |
119 | 126 | ): Promise<{ url: string; fetchOptions: RequestInit }> { |
120 | | - const url = new RequestUrl(model, task, apiSettings, stream, requestOptions); |
| 127 | + const url = new RequestUrl( |
| 128 | + model, |
| 129 | + task, |
| 130 | + apiSettings, |
| 131 | + stream, |
| 132 | + singleRequestOptions |
| 133 | + ); |
121 | 134 | return { |
122 | 135 | url: url.toString(), |
123 | 136 | fetchOptions: { |
|
134 | 147 | apiSettings: ApiSettings, |
135 | 148 | stream: boolean, |
136 | 149 | body: string, |
137 | | - requestOptions?: RequestOptions |
| 150 | + singleRequestOptions?: SingleRequestOptions |
138 | 151 | ): Promise<Response> { |
139 | | - const url = new RequestUrl(model, task, apiSettings, stream, requestOptions); |
| 152 | + const url = new RequestUrl( |
| 153 | + model, |
| 154 | + task, |
| 155 | + apiSettings, |
| 156 | + stream, |
| 157 | + singleRequestOptions |
| 158 | + ); |
140 | 159 | let response; |
141 | | - let fetchTimeoutId: string | number | NodeJS.Timeout | undefined; |
| 160 | + |
| 161 | + const externalSignal = singleRequestOptions?.signal; |
| 162 | + const timeoutMillis = |
| 163 | + singleRequestOptions?.timeout != null && singleRequestOptions.timeout >= 0 |
| 164 | + ? singleRequestOptions.timeout |
| 165 | + : DEFAULT_FETCH_TIMEOUT_MS; |
| 166 | + const internalAbortController = new AbortController(); |
| 167 | + const fetchTimeoutId = setTimeout(() => { |
| 168 | + internalAbortController.abort(TIMEOUT_EXPIRED_MESSAGE); |
| 169 | + logger.debug( |
| 170 | + `Aborting request to ${url} due to timeout (${timeoutMillis}ms)` |
| 171 | + ); |
| 172 | + }, timeoutMillis); |
| 173 | + |
| 174 | + if (externalSignal) { |
| 175 | + if (externalSignal.aborted) { |
| 176 | + clearTimeout(fetchTimeoutId); |
| 177 | + throw new DOMException( |
| 178 | + externalSignal.reason ?? 'Aborted externally before fetch', |
| 179 | + ABORT_ERROR_NAME |
| 180 | + ); |
| 181 | + } |
| 182 | + |
| 183 | + const externalAbortListener = () => { |
| 184 | + logger.debug(`Aborting request to ${url} due to external abort signal.`); |
| 185 | + internalAbortController.abort(externalSignal.reason); |
| 186 | + }; |
| 187 | + |
| 188 | + externalSignal.addEventListener('abort', externalAbortListener, { |
| 189 | + once: true |
| 190 | + }); |
| 191 | + } |
| 192 | + |
142 | 193 | try { |
143 | 194 | const request = await constructRequest( |
144 | 195 | model, |
145 | 196 | task, |
146 | 197 | apiSettings, |
147 | 198 | stream, |
148 | 199 | body, |
149 | | - requestOptions |
| 200 | + singleRequestOptions |
150 | 201 | ); |
151 | | - // Timeout is 180s by default |
152 | | - const timeoutMillis = |
153 | | - requestOptions?.timeout != null && requestOptions.timeout >= 0 |
154 | | - ? requestOptions.timeout |
155 | | - : DEFAULT_FETCH_TIMEOUT_MS; |
156 | | - const abortController = new AbortController(); |
157 | | - fetchTimeoutId = setTimeout(() => abortController.abort(), timeoutMillis); |
158 | | - request.fetchOptions.signal = abortController.signal; |
| 202 | + request.fetchOptions.signal = internalAbortController.signal; |
159 | 203 |
|
160 | 204 | response = await fetch(request.url, request.fetchOptions); |
161 | 205 | if (!response.ok) { |
|
0 commit comments