Skip to content

Commit 6891c04

Browse files
yuhengshsyuhengshs
authored
(feat): Add granular user agent for Ai (#13835)
2 parents 3ee8932 + 865847e commit 6891c04

File tree

7 files changed

+131
-19
lines changed

7 files changed

+131
-19
lines changed

‎packages/api-graphql/__tests__/GraphQLAPI.test.ts

Lines changed: 56 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,8 @@ import * as typedQueries from './fixtures/with-types/queries';
66
import * as typedSubscriptions from './fixtures/with-types/subscriptions';
77
import { expectGet } from './utils/expects';
88
import { InternalGraphQLAPIClass } from '../src/internals/InternalGraphQLAPI';
9+
import { GraphQLAuthMode } from '@aws-amplify/core/internals/utils';
10+
import { INTERNAL_USER_AGENT_OVERRIDE } from '@aws-amplify/data-schema/runtime';
911

1012
import {
1113
__amplify,
@@ -1614,4 +1616,58 @@ describe('API test', () => {
16141616
const subscribeOptions = spyon_appsync_realtime.mock.calls[0][0];
16151617
expect(subscribeOptions).toBe(resolvedUrl);
16161618
});
1619+
test('graphql method handles INTERNAL_USER_AGENT_OVERRIDE correctly', async () => {
1620+
Amplify.configure({
1621+
API: {
1622+
GraphQL: {
1623+
defaultAuthMode: 'apiKey',
1624+
apiKey: 'FAKE-KEY',
1625+
endpoint: 'https://localhost/graphql',
1626+
region: 'local-host-h4x',
1627+
},
1628+
},
1629+
});
1630+
1631+
const mockPost = jest.fn().mockResolvedValue({
1632+
body: {
1633+
json: () => ({ data: { test: 'result' } }),
1634+
},
1635+
});
1636+
(raw.GraphQLAPI as any)._api.post = mockPost;
1637+
1638+
const graphqlOptions = {
1639+
query: 'query TestQuery { test }',
1640+
variables: { id: 'some-id' },
1641+
authMode: 'apiKey' as GraphQLAuthMode,
1642+
[INTERNAL_USER_AGENT_OVERRIDE]: {
1643+
category: 'CustomCategory',
1644+
action: 'CustomAction',
1645+
},
1646+
};
1647+
1648+
await client.graphql(graphqlOptions);
1649+
1650+
// Check if the INTERNAL_USER_AGENT_OVERRIDE was properly handled
1651+
expect(mockPost).toHaveBeenCalledWith(
1652+
expect.anything(),
1653+
expect.objectContaining({
1654+
options: expect.objectContaining({
1655+
headers: expect.objectContaining({
1656+
'x-amz-user-agent': expect.stringContaining(
1657+
'CustomCategory/CustomAction',
1658+
),
1659+
}),
1660+
}),
1661+
}),
1662+
);
1663+
// Ensure the INTERNAL_USER_AGENT_OVERRIDE was not passed along in the options
1664+
expect(mockPost).not.toHaveBeenCalledWith(
1665+
expect.anything(),
1666+
expect.objectContaining({
1667+
options: expect.objectContaining({
1668+
[INTERNAL_USER_AGENT_OVERRIDE]: expect.anything(),
1669+
}),
1670+
}),
1671+
);
1672+
});
16171673
});

‎packages/api-graphql/package.json

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -84,9 +84,12 @@
8484
"server"
8585
],
8686
"dependencies": {
87+
"@aws-amplify/api-rest": "4.0.49",
88+
"@aws-amplify/core": "6.4.2",
89+
"@aws-amplify/data-schema": "^1.7.0",
8790
"@aws-amplify/api-rest": "4.0.50",
8891
"@aws-amplify/core": "6.4.3",
89-
"@aws-amplify/data-schema": "^1.5.0",
92+
"@aws-amplify/data-schema": "^1.7.0",
9093
"@aws-sdk/types": "3.387.0",
9194
"graphql": "15.8.0",
9295
"rxjs": "^7.8.1",

‎packages/api-graphql/src/GraphQLAPI.ts

Lines changed: 34 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,13 +1,28 @@
11
// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
22
// SPDX-License-Identifier: Apache-2.0
33
import { AmplifyClassV6 } from '@aws-amplify/core';
4-
import { ApiAction, Category } from '@aws-amplify/core/internals/utils';
5-
import { CustomHeaders } from '@aws-amplify/data-schema/runtime';
4+
import {
5+
ApiAction,
6+
Category,
7+
CustomUserAgentDetails,
8+
} from '@aws-amplify/core/internals/utils';
9+
import {
10+
CustomHeaders,
11+
INTERNAL_USER_AGENT_OVERRIDE,
12+
} from '@aws-amplify/data-schema/runtime';
613
import { Observable } from 'rxjs';
714

815
import { GraphQLOptions, GraphQLResult } from './types';
916
import { InternalGraphQLAPIClass } from './internals/InternalGraphQLAPI';
1017

18+
function isGraphQLOptionsWithOverride(
19+
options: GraphQLOptions,
20+
): options is GraphQLOptions & {
21+
[INTERNAL_USER_AGENT_OVERRIDE]: CustomUserAgentDetails;
22+
} {
23+
return INTERNAL_USER_AGENT_OVERRIDE in options;
24+
}
25+
1126
export const graphqlOperation = (
1227
query: any,
1328
variables = {},
@@ -38,9 +53,25 @@ export class GraphQLAPIClass extends InternalGraphQLAPIClass {
3853
options: GraphQLOptions,
3954
additionalHeaders?: CustomHeaders,
4055
): Observable<GraphQLResult<T>> | Promise<GraphQLResult<T>> {
41-
return super.graphql(amplify, options, additionalHeaders, {
56+
const userAgentDetails: CustomUserAgentDetails = {
4257
category: Category.API,
4358
action: ApiAction.GraphQl,
59+
};
60+
61+
if (isGraphQLOptionsWithOverride(options)) {
62+
const {
63+
[INTERNAL_USER_AGENT_OVERRIDE]: internalUserAgentOverride,
64+
...cleanOptions
65+
} = options;
66+
67+
return super.graphql(amplify, cleanOptions, additionalHeaders, {
68+
...userAgentDetails,
69+
...internalUserAgentOverride,
70+
});
71+
}
72+
73+
return super.graphql(amplify, options, additionalHeaders, {
74+
...userAgentDetails,
4475
});
4576
}
4677

‎packages/aws-amplify/package.json

Lines changed: 10 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -293,7 +293,7 @@
293293
"name": "[Analytics] record (Pinpoint)",
294294
"path": "./dist/esm/analytics/index.mjs",
295295
"import": "{ record }",
296-
"limit": "17.25 kB"
296+
"limit": "17.35 kB"
297297
},
298298
{
299299
"name": "[Analytics] record (Kinesis)",
@@ -317,7 +317,7 @@
317317
"name": "[Analytics] identifyUser (Pinpoint)",
318318
"path": "./dist/esm/analytics/index.mjs",
319319
"import": "{ identifyUser }",
320-
"limit": "15.75 kB"
320+
"limit": "15.85 kB"
321321
},
322322
{
323323
"name": "[Analytics] enable",
@@ -335,7 +335,7 @@
335335
"name": "[API] generateClient (AppSync)",
336336
"path": "./dist/esm/api/index.mjs",
337337
"import": "{ generateClient }",
338-
"limit": "43.1 kB"
338+
"limit": "43.40 kB"
339339
},
340340
{
341341
"name": "[API] REST API handlers",
@@ -461,43 +461,43 @@
461461
"name": "[Storage] copy (S3)",
462462
"path": "./dist/esm/storage/index.mjs",
463463
"import": "{ copy }",
464-
"limit": "14.86 kB"
464+
"limit": "14.96 kB"
465465
},
466466
{
467467
"name": "[Storage] downloadData (S3)",
468468
"path": "./dist/esm/storage/index.mjs",
469469
"import": "{ downloadData }",
470-
"limit": "15.45 kB"
470+
"limit": "15.55 kB"
471471
},
472472
{
473473
"name": "[Storage] getProperties (S3)",
474474
"path": "./dist/esm/storage/index.mjs",
475475
"import": "{ getProperties }",
476-
"limit": "14.71 kB"
476+
"limit": "14.81 kB"
477477
},
478478
{
479479
"name": "[Storage] getUrl (S3)",
480480
"path": "./dist/esm/storage/index.mjs",
481481
"import": "{ getUrl }",
482-
"limit": "15.95 kB"
482+
"limit": "16.05 kB"
483483
},
484484
{
485485
"name": "[Storage] list (S3)",
486486
"path": "./dist/esm/storage/index.mjs",
487487
"import": "{ list }",
488-
"limit": "15.31 kB"
488+
"limit": "15.41 kB"
489489
},
490490
{
491491
"name": "[Storage] remove (S3)",
492492
"path": "./dist/esm/storage/index.mjs",
493493
"import": "{ remove }",
494-
"limit": "14.57 kB"
494+
"limit": "14.67 kB"
495495
},
496496
{
497497
"name": "[Storage] uploadData (S3)",
498498
"path": "./dist/esm/storage/index.mjs",
499499
"import": "{ uploadData }",
500-
"limit": "19.95 kB"
500+
"limit": "20.05 kB"
501501
}
502502
]
503503
}

‎packages/core/src/Platform/types.ts

Lines changed: 22 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,7 @@ export enum Framework {
2626
}
2727

2828
export enum Category {
29+
AI = 'ai',
2930
API = 'api',
3031
Auth = 'auth',
3132
Analytics = 'analytics',
@@ -39,6 +40,17 @@ export enum Category {
3940
Storage = 'storage',
4041
}
4142

43+
export enum AiAction {
44+
CreateConversation = '1',
45+
GetConversation = '2',
46+
ListConversations = '3',
47+
DeleteConversation = '4',
48+
SendMessage = '5',
49+
ListMessages = '6',
50+
OnMessage = '7',
51+
Generation = '8',
52+
}
53+
4254
export enum AnalyticsAction {
4355
Record = '1',
4456
IdentifyUser = '2',
@@ -123,6 +135,7 @@ export enum StorageAction {
123135
}
124136

125137
interface ActionMap {
138+
[Category.AI]: AiAction;
126139
[Category.Auth]: AuthAction;
127140
[Category.API]: ApiAction;
128141
[Category.Analytics]: AnalyticsAction;
@@ -148,6 +161,7 @@ interface CustomUserAgentDetailsBase {
148161

149162
export type CustomUserAgentDetails =
150163
| (CustomUserAgentDetailsBase & { category?: never; action?: never })
164+
| UserAgentDetailsWithCategory<Category.AI>
151165
| UserAgentDetailsWithCategory<Category.API>
152166
| UserAgentDetailsWithCategory<Category.Auth>
153167
| UserAgentDetailsWithCategory<Category.Analytics>
@@ -180,6 +194,12 @@ export interface StorageUserAgentInput {
180194
additionalDetails: AdditionalDetails;
181195
}
182196

197+
export interface AiUserAgentInput {
198+
category: Category.AI;
199+
apis: AiAction[];
200+
additionalDetails: AdditionalDetails;
201+
}
202+
183203
export interface AuthUserAgentInput {
184204
category: Category.Auth;
185205
apis: AuthAction[];
@@ -202,4 +222,5 @@ export type SetCustomUserAgentInput =
202222
| StorageUserAgentInput
203223
| AuthUserAgentInput
204224
| InAppMessagingUserAgentInput
205-
| GeoUserAgentInput;
225+
| GeoUserAgentInput
226+
| AiUserAgentInput;

‎packages/core/src/libraryUtils.ts

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -78,6 +78,7 @@ export {
7878
getAmplifyUserAgent,
7979
} from './Platform';
8080
export {
81+
AiAction,
8182
ApiAction,
8283
AuthAction,
8384
AnalyticsAction,

‎yarn.lock

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

0 commit comments

Comments
 (0)