fixes blank AWS responses due to reqs sometimes using wrong handler

This commit is contained in:
nai-degen 2023-10-23 22:23:06 -05:00
parent 8bd2f749c1
commit 52ec2ec265
4 changed files with 11 additions and 10 deletions

View File

@ -34,10 +34,6 @@ export const addKey: ProxyRequestMiddleware = (proxyReq, req) => {
throw new Error("You must specify a model with your request.");
}
// TODO: use separate middleware to deal with stream flags
req.isStreaming = req.body.stream === true || req.body.stream === "true";
req.body.stream = req.isStreaming;
if (req.inboundApi === req.outboundApi) {
assignedKey = keyPool.get(req.body.model);
} else {

View File

@ -14,6 +14,12 @@ export const createOnProxyReqHandler = ({
pipeline,
}: RewriterOptions): ProxyReqCallback => {
return (proxyReq, req, res, options) => {
// The streaming flag must be set before any other middleware runs, because
// it may influence which other middleware a particular API pipeline wants
// to run.
req.isStreaming = req.body.stream === true || req.body.stream === "true";
req.body.stream = req.isStreaming;
try {
for (const validator of beforeRewrite) {
validator(proxyReq, req, res, options);

View File

@ -3,7 +3,6 @@ import { Request, Response } from "express";
import * as http from "http";
import util from "util";
import zlib from "zlib";
import { logger } from "../../../logger";
import { enqueue, trackWaitTime } from "../../queue";
import { HttpError } from "../../../shared/errors";
import { keyPool } from "../../../shared/key-management";
@ -188,7 +187,7 @@ export const decodeResponseBody: RawResponseBodyHandler = async (
body = await decoder(body);
} else {
const errorMessage = `Proxy received response with unsupported content-encoding: ${contentEncoding}`;
logger.warn({ contentEncoding, key: req.key?.hash }, errorMessage);
req.log.warn({ contentEncoding, key: req.key?.hash }, errorMessage);
writeErrorResponse(req, res, 500, {
error: errorMessage,
contentEncoding,
@ -205,7 +204,7 @@ export const decodeResponseBody: RawResponseBodyHandler = async (
return resolve(body.toString());
} catch (error: any) {
const errorMessage = `Proxy received response with invalid JSON: ${error.message}`;
logger.warn({ error: error.stack, key: req.key?.hash }, errorMessage);
req.log.warn({ error: error.stack, key: req.key?.hash }, errorMessage);
writeErrorResponse(req, res, 500, { error: errorMessage });
return reject(errorMessage);
}
@ -251,7 +250,7 @@ const handleUpstreamErrors: ProxyResHandlerWithBody = async (
// Likely Bad Gateway or Gateway Timeout from upstream's reverse proxy
const hash = req.key?.hash;
const statusMessage = proxyRes.statusMessage || "Unknown error";
logger.warn({ statusCode, statusMessage, key: hash }, parseError.message);
req.log.warn({ statusCode, statusMessage, key: hash }, parseError.message);
const errorObject = {
statusCode,
@ -268,7 +267,7 @@ const handleUpstreamErrors: ProxyResHandlerWithBody = async (
errorPayload.error?.type ||
getAwsErrorType(proxyRes.headers["x-amzn-errortype"]);
logger.warn(
req.log.warn(
{ statusCode, type: errorType, errorPayload, key: req.key?.hash },
`Received error response from upstream. (${proxyRes.statusMessage})`
);

View File

@ -182,7 +182,7 @@ export class AwsKeyChecker extends KeyCheckerBase<AwsBedrockKey> {
}
this.log.debug(
{ key: key.hash, errorType, data, status },
{ key: key.hash, errorType, data, status, model },
"Liveness test complete."
);
}