(feat): Add granular user agent for Ai

packages/api-graphql/src/GraphQLAPI.ts

@@ -1,7 +1,11 @@
 // Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
 // SPDX-License-Identifier: Apache-2.0
 import { AmplifyClassV6 } from '@aws-amplify/core';
-import { ApiAction, Category } from '@aws-amplify/core/internals/utils';
+import {
+	ApiAction,
+	Category,
+	INTERNAL_USER_AGENT_OVERRIDE,
+} from '@aws-amplify/core/internals/utils';
 import { CustomHeaders } from '@aws-amplify/data-schema/runtime';
 import { Observable } from 'rxjs';
 
@@ -38,9 +42,16 @@ export class GraphQLAPIClass extends InternalGraphQLAPIClass {
 		options: GraphQLOptions,
 		additionalHeaders?: CustomHeaders,
 	): Observable<GraphQLResult<T>> | Promise<GraphQLResult<T>> {
-		return super.graphql(amplify, options, additionalHeaders, {
+		const internalUserAgentOverride = (options as any)[
+			INTERNAL_USER_AGENT_OVERRIDE
+		];
+		const { [INTERNAL_USER_AGENT_OVERRIDE]: _, ...cleanOptions } =
+			options as any;
+
+		return super.graphql(amplify, cleanOptions, additionalHeaders, {
 			category: Category.API,
 			action: ApiAction.GraphQl,
+			...(internalUserAgentOverride || {}),
 		});
 	}
 

packages/core/src/Platform/types.ts

@@ -26,6 +26,7 @@ export enum Framework {
 }
 
 export enum Category {
+	AI = 'ai',
 	API = 'api',
 	Auth = 'auth',
 	Analytics = 'analytics',
@@ -39,6 +40,16 @@ export enum Category {
 	Storage = 'storage',
 }
 
+export enum AiAction {
+	CreateConversation = '1',
+	GetConversation = '2',
+	ListConversations = '3',
+	SendMessage = '4',
+	ListMessages = '5',
+	OnMessage = '6',
+	Generation = '7',
+}
+
 export enum AnalyticsAction {
 	Record = '1',
 	IdentifyUser = '2',
@@ -123,6 +134,7 @@ export enum StorageAction {
 }
 
 interface ActionMap {
+	[Category.AI]: AiAction;
 	[Category.Auth]: AuthAction;
 	[Category.API]: ApiAction;
 	[Category.Analytics]: AnalyticsAction;
@@ -148,6 +160,7 @@ interface CustomUserAgentDetailsBase {
 
 export type CustomUserAgentDetails =
 	| (CustomUserAgentDetailsBase & { category?: never; action?: never })
+	| UserAgentDetailsWithCategory<Category.AI>
 	| UserAgentDetailsWithCategory<Category.API>
 	| UserAgentDetailsWithCategory<Category.Auth>
 	| UserAgentDetailsWithCategory<Category.Analytics>
@@ -180,6 +193,12 @@ export interface StorageUserAgentInput {
 	additionalDetails: AdditionalDetails;
 }
 
+export interface AiUserAgentInput {
+	category: Category.AI;
+	apis: AiAction[];
+	additionalDetails: AdditionalDetails;
+}
+
 export interface AuthUserAgentInput {
 	category: Category.Auth;
 	apis: AuthAction[];
@@ -202,4 +221,5 @@ export type SetCustomUserAgentInput =
 	| StorageUserAgentInput
 	| AuthUserAgentInput
 	| InAppMessagingUserAgentInput
-	| GeoUserAgentInput;
+	| GeoUserAgentInput
+	| AiUserAgentInput;

packages/core/src/constants.ts

@@ -9,3 +9,8 @@ export const USER_AGENT_HEADER = 'x-amz-user-agent';
 // Error exception code constants
 export const NO_HUBCALLBACK_PROVIDED_EXCEPTION =
 	'NoHubcallbackProvidedException';
+
+// User Agents Override Symbol
+export const INTERNAL_USER_AGENT_OVERRIDE = Symbol(
+	'INTERNAL_USER_AGENT_OVERRIDE',
+);

packages/core/src/libraryUtils.ts

@@ -78,6 +78,7 @@ export {
 	getAmplifyUserAgent,
 } from './Platform';
 export {
+	AiAction,
 	ApiAction,
 	AuthAction,
 	AnalyticsAction,
@@ -99,6 +100,7 @@ export {
 	GeoUserAgentInput,
 } from './Platform/types';
 export { setCustomUserAgent } from './Platform/customUserAgent';
+export { INTERNAL_USER_AGENT_OVERRIDE } from './constants';
 
 // Error handling
 export {