Skip to content

Commit 1c91466

Browse files
VisargDVisargD
authored
Merge branch 'main' into chore/test-workflow
2 parents 1b37ca7 + 3460aca commit 1c91466

File tree

18 files changed

+1151
-1257
lines changed

18 files changed

+1151
-1257
lines changed

‎package-lock.json‎

Lines changed: 2 additions & 2 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎package.json‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
{
22
"name": "@portkey-ai/gateway",
3-
"version": "1.7.7",
3+
"version": "1.8.0",
44
"description": "A fast AI gateway by Portkey",
55
"repository": {
66
"type": "git",

‎src/globals.ts‎

Lines changed: 49 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -139,3 +139,52 @@ export const MULTIPART_FORM_DATA_ENDPOINTS: endpointStrings[] = [
139139
'createTranscription',
140140
'createTranslation',
141141
];
142+
143+
export const fileExtensionMimeTypeMap = {
144+
mp4: 'video/mp4',
145+
jpeg: 'image/jpeg',
146+
jpg: 'image/jpeg',
147+
png: 'image/png',
148+
bmp: 'image/bmp',
149+
tiff: 'image/tiff',
150+
webp: 'image/webp',
151+
pdf: 'application/pdf',
152+
csv: 'text/csv',
153+
doc: 'application/msword',
154+
docx: 'application/vnd.openxmlformats-officedocument.wordprocessingml.document',
155+
xls: 'application/vnd.ms-excel',
156+
xlsx: 'application/vnd.openxmlformats-officedocument.spreadsheetml.sheet',
157+
html: 'text/html',
158+
md: 'text/markdown',
159+
mp3: 'audio/mp3',
160+
wav: 'audio/wav',
161+
txt: 'text/plain',
162+
mov: 'video/mov',
163+
mpeg: 'video/mpeg',
164+
mpg: 'video/mpg',
165+
avi: 'video/avi',
166+
wmv: 'video/wmv',
167+
mpegps: 'video/mpegps',
168+
flv: 'video/flv',
169+
};
170+
171+
export const imagesMimeTypes = [
172+
fileExtensionMimeTypeMap.jpeg,
173+
fileExtensionMimeTypeMap.jpg,
174+
fileExtensionMimeTypeMap.png,
175+
fileExtensionMimeTypeMap.bmp,
176+
fileExtensionMimeTypeMap.tiff,
177+
fileExtensionMimeTypeMap.webp,
178+
];
179+
180+
export const documentMimeTypes = [
181+
fileExtensionMimeTypeMap.pdf,
182+
fileExtensionMimeTypeMap.csv,
183+
fileExtensionMimeTypeMap.doc,
184+
fileExtensionMimeTypeMap.docx,
185+
fileExtensionMimeTypeMap.xls,
186+
fileExtensionMimeTypeMap.xlsx,
187+
fileExtensionMimeTypeMap.html,
188+
fileExtensionMimeTypeMap.md,
189+
fileExtensionMimeTypeMap.txt,
190+
];

‎src/handlers/handlerUtils.ts‎

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -269,6 +269,7 @@ export async function tryPostProxy(
269269
: (providerOption.urlToFetch as string);
270270

271271
const headers = await apiConfig.headers({
272+
c,
272273
providerOptions: providerOption,
273274
fn,
274275
transformedRequestBody: params,
@@ -520,6 +521,7 @@ export async function tryPost(
520521
const url = `${baseUrl}${endpoint}`;
521522

522523
const headers = await apiConfig.headers({
524+
c,
523525
providerOptions: providerOption,
524526
fn,
525527
transformedRequestBody,
@@ -1011,6 +1013,13 @@ export function constructConfigFromRequestHeaders(
10111013
resourceName: requestHeaders[`x-${POWERED_BY}-azure-resource-name`],
10121014
deploymentId: requestHeaders[`x-${POWERED_BY}-azure-deployment-id`],
10131015
apiVersion: requestHeaders[`x-${POWERED_BY}-azure-api-version`],
1016+
azureAuthMode: requestHeaders[`x-${POWERED_BY}-azure-auth-mode`],
1017+
azureManagedClientId:
1018+
requestHeaders[`x-${POWERED_BY}-azure-managed-client-id`],
1019+
azureEntraClientId: requestHeaders[`x-${POWERED_BY}-azure-entra-client-id`],
1020+
azureEntraClientSecret:
1021+
requestHeaders[`x-${POWERED_BY}-azure-entra-client-secret`],
1022+
azureEntraTenantId: requestHeaders[`x-${POWERED_BY}-azure-entra-tenant-id`],
10141023
azureModelName: requestHeaders[`x-${POWERED_BY}-azure-model-name`],
10151024
};
10161025

@@ -1037,6 +1046,9 @@ export function constructConfigFromRequestHeaders(
10371046
awsSecretAccessKey: requestHeaders[`x-${POWERED_BY}-aws-secret-access-key`],
10381047
awsSessionToken: requestHeaders[`x-${POWERED_BY}-aws-session-token`],
10391048
awsRegion: requestHeaders[`x-${POWERED_BY}-aws-region`],
1049+
awsRoleArn: requestHeaders[`x-${POWERED_BY}-aws-role-arn`],
1050+
awsAuthType: requestHeaders[`x-${POWERED_BY}-aws-auth-type`],
1051+
awsExternalId: requestHeaders[`x-${POWERED_BY}-aws-external-id`],
10401052
};
10411053

10421054
const workersAiConfig = {

‎src/handlers/streamHandler.ts‎

Lines changed: 6 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -33,7 +33,10 @@ function getPayloadFromAWSChunk(chunk: Uint8Array): string {
3333

3434
const payloadLength = chunkLength - headersEnd - 4; // Subtracting 4 for the message crc
3535
const payload = chunk.slice(headersEnd, headersEnd + payloadLength);
36-
return decoder.decode(payload);
36+
const decodedJson = JSON.parse(decoder.decode(payload));
37+
return decodedJson.bytes
38+
? Buffer.from(decodedJson.bytes, 'base64').toString()
39+
: JSON.stringify(decodedJson);
3740
}
3841

3942
function concatenateUint8Arrays(a: Uint8Array, b: Uint8Array): Uint8Array {
@@ -60,10 +63,7 @@ export async function* readAWSStream(
6063
const data = buffer.subarray(0, expectedLength);
6164
buffer = buffer.subarray(expectedLength);
6265
expectedLength = readUInt32BE(buffer, 0);
63-
const payload = Buffer.from(
64-
JSON.parse(getPayloadFromAWSChunk(data)).bytes,
65-
'base64'
66-
).toString();
66+
const payload = getPayloadFromAWSChunk(data);
6767
if (transformFunction) {
6868
const transformedChunk = transformFunction(
6969
payload,
@@ -96,11 +96,7 @@ export async function* readAWSStream(
9696
buffer = buffer.subarray(expectedLength);
9797

9898
expectedLength = readUInt32BE(buffer, 0);
99-
const payload = Buffer.from(
100-
JSON.parse(getPayloadFromAWSChunk(data)).bytes,
101-
'base64'
102-
).toString();
103-
99+
const payload = getPayloadFromAWSChunk(data);
104100
if (transformFunction) {
105101
const transformedChunk = transformFunction(
106102
payload,

‎src/providers/anthropic/api.ts‎

Lines changed: 9 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -2,14 +2,20 @@ import { ProviderAPIConfig } from '../types';
22

33
const AnthropicAPIConfig: ProviderAPIConfig = {
44
getBaseURL: () => 'https://api.anthropic.com/v1',
5-
headers: ({ providerOptions, fn }) => {
5+
headers: ({ providerOptions, fn, gatewayRequestBody }) => {
66
const headers: Record<string, string> = {
77
'X-API-Key': `${providerOptions.apiKey}`,
88
};
99

10+
// Accept anthropic_beta and anthropic_version in body to support enviroments which cannot send it in headers.
1011
const betaHeader =
11-
providerOptions?.['anthropicBeta'] ?? 'messages-2023-12-15';
12-
const version = providerOptions?.['anthropicVersion'] ?? '2023-06-01';
12+
providerOptions?.['anthropicBeta'] ??
13+
gatewayRequestBody?.['anthropic_beta'] ??
14+
'messages-2023-12-15';
15+
const version =
16+
providerOptions?.['anthropicVersion'] ??
17+
gatewayRequestBody?.['anthropic_version'] ??
18+
'2023-06-01';
1319

1420
if (fn === 'chatComplete') {
1521
headers['anthropic-beta'] = betaHeader;

‎src/providers/azure-openai/api.ts‎

Lines changed: 35 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,13 +1,46 @@
11
import { ProviderAPIConfig } from '../types';
2+
import {
3+
getAccessTokenFromEntraId,
4+
getAzureManagedIdentityToken,
5+
} from './utils';
26

37
const AzureOpenAIAPIConfig: ProviderAPIConfig = {
48
getBaseURL: ({ providerOptions }) => {
59
const { resourceName, deploymentId } = providerOptions;
610
return `https://${resourceName}.openai.azure.com/openai/deployments/${deploymentId}`;
711
},
8-
headers: ({ providerOptions, fn }) => {
12+
headers: async ({ providerOptions, fn }) => {
13+
const { apiKey, azureAuthMode } = providerOptions;
14+
15+
if (azureAuthMode === 'entra') {
16+
const { azureEntraTenantId, azureEntraClientId, azureEntraClientSecret } =
17+
providerOptions;
18+
if (azureEntraTenantId && azureEntraClientId && azureEntraClientSecret) {
19+
const scope = 'https://cognitiveservices.azure.com/.default';
20+
const accessToken = await getAccessTokenFromEntraId(
21+
azureEntraTenantId,
22+
azureEntraClientId,
23+
azureEntraClientSecret,
24+
scope
25+
);
26+
return {
27+
Authorization: `Bearer ${accessToken}`,
28+
};
29+
}
30+
}
31+
if (azureAuthMode === 'managed') {
32+
const { azureManagedClientId } = providerOptions;
33+
const resource = 'https://cognitiveservices.azure.com/';
34+
const accessToken = await getAzureManagedIdentityToken(
35+
resource,
36+
azureManagedClientId
37+
);
38+
return {
39+
Authorization: `Bearer ${accessToken}`,
40+
};
41+
}
942
const headersObj: Record<string, string> = {
10-
'api-key': `${providerOptions.apiKey}`,
43+
'api-key': `${apiKey}`,
1144
};
1245
if (fn === 'createTranscription' || fn === 'createTranslation')
1346
headersObj['Content-Type'] = 'multipart/form-data';

‎src/providers/azure-openai/utils.ts‎

Lines changed: 60 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,60 @@
1+
export async function getAccessTokenFromEntraId(
2+
tenantId: string,
3+
clientId: string,
4+
clientSecret: string,
5+
scope = 'https://cognitiveservices.azure.com/.default'
6+
) {
7+
try {
8+
const url = `https://login.microsoftonline.com/${tenantId}/oauth2/v2.0/token`;
9+
const params = new URLSearchParams({
10+
client_id: clientId,
11+
client_secret: clientSecret,
12+
scope: scope,
13+
grant_type: 'client_credentials',
14+
});
15+
16+
const response = await fetch(url, {
17+
method: 'POST',
18+
headers: {
19+
'Content-Type': 'application/x-www-form-urlencoded',
20+
},
21+
body: params,
22+
});
23+
24+
if (!response.ok) {
25+
const errorMessage = await response.text();
26+
console.log({ message: `Error from Entra ${errorMessage}` });
27+
return undefined;
28+
}
29+
const data: { access_token: string } = await response.json();
30+
return data.access_token;
31+
} catch (error) {
32+
console.log(error);
33+
}
34+
}
35+
36+
export async function getAzureManagedIdentityToken(
37+
resource: string,
38+
clientId?: string
39+
) {
40+
try {
41+
const response = await fetch(
42+
`http://169.254.169.254/metadata/identity/oauth2/token?api-version=2018-02-01&resource=${encodeURIComponent(resource)}${clientId ? `&client_id=${encodeURIComponent(clientId)}` : ''}`,
43+
{
44+
method: 'GET',
45+
headers: {
46+
Metadata: 'true',
47+
},
48+
}
49+
);
50+
if (!response.ok) {
51+
const errorMessage = await response.text();
52+
console.log({ message: `Error from Managed ${errorMessage}` });
53+
return undefined;
54+
}
55+
const data: { access_token: string } = await response.json();
56+
return data.access_token;
57+
} catch (error) {
58+
console.log({ error });
59+
}
60+
}

‎src/providers/bedrock/api.ts‎

Lines changed: 50 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,14 @@
1+
import { env } from 'hono/adapter';
2+
import { GatewayError } from '../../errors/GatewayError';
13
import { ProviderAPIConfig } from '../types';
2-
import { generateAWSHeaders } from './utils';
4+
import { bedrockInvokeModels } from './constants';
5+
import { generateAWSHeaders, getAssumedRoleCredentials } from './utils';
36

47
const BedrockAPIConfig: ProviderAPIConfig = {
58
getBaseURL: ({ providerOptions }) =>
69
`https://bedrock-runtime.${providerOptions.awsRegion || 'us-east-1'}.amazonaws.com`,
710
headers: async ({
11+
c,
812
providerOptions,
913
transformedRequestBody,
1014
transformedRequestUrl,
@@ -13,6 +17,41 @@ const BedrockAPIConfig: ProviderAPIConfig = {
1317
'content-type': 'application/json',
1418
};
1519

20+
if (providerOptions.awsAuthType === 'assumedRole') {
21+
try {
22+
// Assume the role in the source account
23+
const sourceRoleCredentials = await getAssumedRoleCredentials(
24+
c,
25+
env(c).AWS_ASSUME_ROLE_SOURCE_ARN, // Role ARN in the source account
26+
env(c).AWS_ASSUME_ROLE_SOURCE_EXTERNAL_ID || '', // External ID for source role (if needed)
27+
providerOptions.awsRegion || ''
28+
);
29+
30+
if (!sourceRoleCredentials) {
31+
throw new Error('Server Error while assuming internal role');
32+
}
33+
34+
// Assume role in destination account using temporary creds obtained in first step
35+
const { accessKeyId, secretAccessKey, sessionToken } =
36+
(await getAssumedRoleCredentials(
37+
c,
38+
providerOptions.awsRoleArn || '',
39+
providerOptions.awsExternalId || '',
40+
providerOptions.awsRegion || '',
41+
{
42+
accessKeyId: sourceRoleCredentials.accessKeyId,
43+
secretAccessKey: sourceRoleCredentials.secretAccessKey,
44+
sessionToken: sourceRoleCredentials.sessionToken,
45+
}
46+
)) || {};
47+
providerOptions.awsAccessKeyId = accessKeyId;
48+
providerOptions.awsSecretAccessKey = secretAccessKey;
49+
providerOptions.awsSessionToken = sessionToken;
50+
} catch (e) {
51+
throw new GatewayError('Error while assuming bedrock role');
52+
}
53+
}
54+
1655
return generateAWSHeaders(
1756
transformedRequestBody,
1857
headers,
@@ -27,12 +66,20 @@ const BedrockAPIConfig: ProviderAPIConfig = {
2766
},
2867
getEndpoint: ({ fn, gatewayRequestBody }) => {
2968
const { model, stream } = gatewayRequestBody;
69+
if (!model) throw new GatewayError('Model is required');
3070
let mappedFn = fn;
3171
if (stream) {
3272
mappedFn = `stream-${fn}`;
3373
}
34-
const endpoint = `/model/${model}/invoke`;
35-
const streamEndpoint = `/model/${model}/invoke-with-response-stream`;
74+
let endpoint = `/model/${model}/invoke`;
75+
let streamEndpoint = `/model/${model}/invoke-with-response-stream`;
76+
if (
77+
(mappedFn === 'chatComplete' || mappedFn === 'stream-chatComplete') &&
78+
!bedrockInvokeModels.includes(model)
79+
) {
80+
endpoint = `/model/${model}/converse`;
81+
streamEndpoint = `/model/${model}/converse-stream`;
82+
}
3683
switch (mappedFn) {
3784
case 'chatComplete': {
3885
return endpoint;

0 commit comments

Comments
 (0)