From ee3a96c7f0744625750d164765e98dd12579e14c Mon Sep 17 00:00:00 2001 From: TroyceGowdy Date: Wed, 9 Oct 2024 17:15:30 -0400 Subject: [PATCH 1/6] feat: Added 'ai.chat.allowedModels' feature to restrict chat models usage - AIController.ts now rejects requests if the requested model is not in the allowed models list - Added error handling to provide feedback when disallowed model is requested --- src/aux-records/AIController.ts | 39 ++++++++++++++++++++ src/aux-records/SubscriptionConfiguration.ts | 8 ++++ 2 files changed, 47 insertions(+) diff --git a/src/aux-records/AIController.ts b/src/aux-records/AIController.ts index b64dd906e..36ca4329f 100644 --- a/src/aux-records/AIController.ts +++ b/src/aux-records/AIController.ts @@ -450,6 +450,26 @@ export class AIController { maxTokens, }); + if (allowedFeatures.ai.chat.allowedModels) { + const allowedModels = allowedFeatures.ai.chat.allowedModels; + if ( + !allowedModels || + allowedModels.length === 0 || + allowedModels.includes(model) + ) { + return { + success: true, + choices: result.choices, + }; + } + return { + success: false, + errorCode: 'not_authorized', + errorMessage: + 'The subscription does not permit the given model for AI Chat features.', + }; + } + if (result.totalTokens > 0) { await this._metrics.recordChatMetrics({ userId: request.userId, @@ -598,6 +618,25 @@ export class AIController { }; } + if (allowedFeatures.ai.chat.allowedModels) { + const allowedModels = allowedFeatures.ai.chat.allowedModels; + if ( + !allowedModels || + allowedModels.length === 0 || + allowedModels.includes(model) + ) { + return { + success: true, + }; + } + return { + success: false, + errorCode: 'not_authorized', + errorMessage: + 'The subscription does not permit the given model for AI Chat features.', + }; + } + let maxTokens: number = undefined; if (allowedFeatures.ai.chat.maxTokensPerPeriod) { maxTokens = diff --git a/src/aux-records/SubscriptionConfiguration.ts b/src/aux-records/SubscriptionConfiguration.ts index 36507956a..f4ee16819 100644 --- a/src/aux-records/SubscriptionConfiguration.ts +++ b/src/aux-records/SubscriptionConfiguration.ts @@ -141,6 +141,12 @@ export const subscriptionFeaturesSchema = z.object({ .int() .positive() .optional(), + allowedModels: z + .array(z.string()) + .describe( + 'The list of model IDs that are allowed for the subscription. If omitted, then all models are allowed.' + ) + .optional(), }), images: z.object({ allowed: z @@ -938,6 +944,8 @@ export interface AIChatFeaturesConfiguration { * If not specified, then there is no limit. */ maxTokensPerPeriod?: number; + + allowedModels?: string[]; } export interface AIImageFeaturesConfiguration { From 4ef7daad6ad469d961606ed61a50fe533973b41f Mon Sep 17 00:00:00 2001 From: TroyceGowdy Date: Wed, 9 Oct 2024 17:20:11 -0400 Subject: [PATCH 2/6] chore: Added tests for 'ai.chat.allowedModels' --- src/aux-records/AIController.spec.ts | 119 +++++++++++++++++++++++++++ 1 file changed, 119 insertions(+) diff --git a/src/aux-records/AIController.spec.ts b/src/aux-records/AIController.spec.ts index 67aa8471e..1c9b563a6 100644 --- a/src/aux-records/AIController.spec.ts +++ b/src/aux-records/AIController.spec.ts @@ -369,6 +369,125 @@ describe('AIController', () => { expect(chatInterface.chat).not.toBeCalled(); }); + it('should return success when allowedModels includes the model', () => { + const allowedFeatures = { + ai: { + chat: { + allowedModels: ['modelA', 'modelB'], + }, + }, + }; + const model = 'modelA'; + const result = { choices: ['choice1', 'choice2'] }; + + const response = (function () { + if (allowedFeatures.ai.chat.allowedModels) { + const allowedModels = allowedFeatures.ai.chat.allowedModels; + if ( + !allowedModels || + allowedModels.length === 0 || + allowedModels.includes(model) + ) { + return { + success: true, + choices: result.choices, + }; + } + return { + success: false, + errorCode: 'not_authorized', + errorMessage: + 'The subscription does not permit the given model for AI Chat features.', + }; + } + })(); + + expect(response).toEqual({ + success: true, + choices: ['choice1', 'choice2'], + }); + }); + + it('should return not_authorized when allowedModels does not include the model', () => { + const allowedFeatures = { + ai: { + chat: { + allowedModels: ['modelA', 'modelB'], + }, + }, + }; + const model = 'modelC'; + const result = { choices: ['choice1', 'choice2'] }; + + const response = (function () { + if (allowedFeatures.ai.chat.allowedModels) { + const allowedModels = allowedFeatures.ai.chat.allowedModels; + if ( + !allowedModels || + allowedModels.length === 0 || + allowedModels.includes(model) + ) { + return { + success: true, + choices: result.choices, + }; + } + return { + success: false, + errorCode: 'not_authorized', + errorMessage: + 'The subscription does not permit the given model for AI Chat features.', + }; + } + })(); + + expect(response).toEqual({ + success: false, + errorCode: 'not_authorized', + errorMessage: + 'The subscription does not permit the given model for AI Chat features.', + }); + }); + + // it('should return success when allowedModels includes the model', () => { + // const allowedFeatures = { + // ai: { + // chat: { + // allowedModels: ['model1', 'model2'] + // } + // } + // }; + // const isStringInArray = (arr: Object, str: string): boolean => { + // return arr.hasOwnProperty(str); + // }; + // const model = 'model3'; + // const result = isStringInArray(allowedFeatures, model); + // expect(result).toEqual({ + // success: true, + // choice: result.choices + // }); + // }); + + // it('should return not_authorized when allowedModels does not include the model', () => { + // const allowedFeatures = { + // ai: { + // chat: { + // allowedModels: ['model1', 'model2'] + // } + // } + // }; + // const isStringInArray = (arr: Object, str: string): boolean => { + // return arr.hasOwnProperty(str); + // }; + // const model = 'model3'; + // const result = isStringInArray(allowedFeatures, model); + // expect(result).toEqual({ + // success: false, + // errorCode: 'not_authorized', + // errorMessage: 'The subscription does not permit the given model for AI Chat features.' + // }); + // }); + it('should return an not_logged_in result if the given a null userId', async () => { const result = await controller.chat({ model: 'test-model1', From 5b706951e695b859e750d53deb478bdc65720381 Mon Sep 17 00:00:00 2001 From: TroyceGowdy Date: Wed, 9 Oct 2024 17:38:47 -0400 Subject: [PATCH 3/6] chore: update CHANGELOG.md --- CHANGELOG.md | 1 + src/aux-records/AIController.spec.ts | 39 ---------------------------- 2 files changed, 1 insertion(+), 39 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index f6a4abd53..5a17988f0 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -36,6 +36,7 @@ - Webhooks do not automatically install abCore. They only use the bots that are stored in the target. - Webhooks always act like static insts. - This means that any changes made to bots in the webhook are erased after the webhook finishes. +- Added `ai.chat.allowedModels` feature to enforce model usage limits, restricting access to specific models based on configuration. ### :bug: Bug Fixes diff --git a/src/aux-records/AIController.spec.ts b/src/aux-records/AIController.spec.ts index 1c9b563a6..2852ae5b5 100644 --- a/src/aux-records/AIController.spec.ts +++ b/src/aux-records/AIController.spec.ts @@ -449,45 +449,6 @@ describe('AIController', () => { }); }); - // it('should return success when allowedModels includes the model', () => { - // const allowedFeatures = { - // ai: { - // chat: { - // allowedModels: ['model1', 'model2'] - // } - // } - // }; - // const isStringInArray = (arr: Object, str: string): boolean => { - // return arr.hasOwnProperty(str); - // }; - // const model = 'model3'; - // const result = isStringInArray(allowedFeatures, model); - // expect(result).toEqual({ - // success: true, - // choice: result.choices - // }); - // }); - - // it('should return not_authorized when allowedModels does not include the model', () => { - // const allowedFeatures = { - // ai: { - // chat: { - // allowedModels: ['model1', 'model2'] - // } - // } - // }; - // const isStringInArray = (arr: Object, str: string): boolean => { - // return arr.hasOwnProperty(str); - // }; - // const model = 'model3'; - // const result = isStringInArray(allowedFeatures, model); - // expect(result).toEqual({ - // success: false, - // errorCode: 'not_authorized', - // errorMessage: 'The subscription does not permit the given model for AI Chat features.' - // }); - // }); - it('should return an not_logged_in result if the given a null userId', async () => { const result = await controller.chat({ model: 'test-model1', From 2b150fc562630d18f17e10b2717fe6b30beb1470 Mon Sep 17 00:00:00 2001 From: TroyceGowdy Date: Thu, 10 Oct 2024 13:26:48 -0400 Subject: [PATCH 4/6] chore: Update CHANGELOG.md chore: Added description to AIChatFeaturesConfiguration allowedModels interface --- CHANGELOG.md | 1 + src/aux-records/SubscriptionConfiguration.ts | 4 ++++ 2 files changed, 5 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index 696573c41..6e172b685 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -37,6 +37,7 @@ - Webhooks always act like static insts. - This means that any changes made to bots in the webhook are erased after the webhook finishes. - Added the ability to request consent again so that a parent can adjust the privacy features for their child. +- Added `ai.chat.allowedModels` feature to enforce model usage limits, restricting access to specific models based on configuration. ### :bug: Bug Fixes diff --git a/src/aux-records/SubscriptionConfiguration.ts b/src/aux-records/SubscriptionConfiguration.ts index f4ee16819..ec3c15982 100644 --- a/src/aux-records/SubscriptionConfiguration.ts +++ b/src/aux-records/SubscriptionConfiguration.ts @@ -945,6 +945,10 @@ export interface AIChatFeaturesConfiguration { */ maxTokensPerPeriod?: number; + /** + * The list of model IDs that are allowed for the subscription. + * If omitted, then all models are allowed. + */ allowedModels?: string[]; } From 0fdce041f1dbcc7804506b71d6e4f9b4eecc7ffc Mon Sep 17 00:00:00 2001 From: TroyceGowdy Date: Wed, 16 Oct 2024 13:02:39 -0400 Subject: [PATCH 5/6] fix: Refactored to properly use AIController.ts instead of using duplicate code. - Updated implementation to use AIController.ts - Removed redundant code that was duplicating AIController.ts functionality --- src/aux-records/AIController.spec.ts | 258 ++++++++++++++++++++------- 1 file changed, 194 insertions(+), 64 deletions(-) diff --git a/src/aux-records/AIController.spec.ts b/src/aux-records/AIController.spec.ts index 2852ae5b5..d4f0c9a66 100644 --- a/src/aux-records/AIController.spec.ts +++ b/src/aux-records/AIController.spec.ts @@ -369,84 +369,109 @@ describe('AIController', () => { expect(chatInterface.chat).not.toBeCalled(); }); - it('should return success when allowedModels includes the model', () => { - const allowedFeatures = { - ai: { - chat: { - allowedModels: ['modelA', 'modelB'], + it('should return success when allowedModels includes the model', async () => { + chatInterface.chat.mockReturnValueOnce( + Promise.resolve({ + choices: [ + { + role: 'user', + content: 'test', + finishReason: 'stop', + }, + ], + totalTokens: 1, + }) + ); + + const result = await controller.chat({ + model: 'test-model1', + messages: [ + { + role: 'user', + content: 'test', }, - }, - }; - const model = 'modelA'; - const result = { choices: ['choice1', 'choice2'] }; - - const response = (function () { - if (allowedFeatures.ai.chat.allowedModels) { - const allowedModels = allowedFeatures.ai.chat.allowedModels; - if ( - !allowedModels || - allowedModels.length === 0 || - allowedModels.includes(model) - ) { - return { - success: true, - choices: result.choices, - }; - } - return { - success: false, - errorCode: 'not_authorized', - errorMessage: - 'The subscription does not permit the given model for AI Chat features.', - }; - } - })(); + ], + temperature: 0.5, + userId, + userSubscriptionTier, + }); - expect(response).toEqual({ + expect(result).toEqual({ success: true, - choices: ['choice1', 'choice2'], + choices: [ + { + role: 'user', + content: 'test', + finishReason: 'stop', + }, + ], + }); + expect(chatInterface.chat).toBeCalledWith({ + model: 'test-model1', + messages: [ + { + role: 'user', + content: 'test', + }, + ], + temperature: 0.5, + userId: 'test-user', }); }); - it('should return not_authorized when allowedModels does not include the model', () => { - const allowedFeatures = { - ai: { - chat: { - allowedModels: ['modelA', 'modelB'], + it('should return not_authorized when allowedModels does not include the model', async () => { + controller = new AIController({ + chat: { + interfaces: { + provider1: chatInterface, + }, + options: { + defaultModel: 'default-model', + defaultModelProvider: 'provider1', + allowedChatModels: [ + { + provider: 'provider1', + model: 'modelA', + }, + { + provider: 'provider1', + model: 'modelB', + }, + ], + allowedChatSubscriptionTiers: ['test-tier'], }, }, - }; - const model = 'modelC'; - const result = { choices: ['choice1', 'choice2'] }; - - const response = (function () { - if (allowedFeatures.ai.chat.allowedModels) { - const allowedModels = allowedFeatures.ai.chat.allowedModels; - if ( - !allowedModels || - allowedModels.length === 0 || - allowedModels.includes(model) - ) { - return { - success: true, - choices: result.choices, - }; - } - return { - success: false, - errorCode: 'not_authorized', - errorMessage: - 'The subscription does not permit the given model for AI Chat features.', - }; - } - })(); + generateSkybox: null, + images: null, + metrics: store, + config: store, + hume: null, + sloyd: null, + policies: null, + policyController: policies, + records: store, + }); + + const result = await controller.chat({ + model: 'modelC', + messages: [ + { + role: 'user', + content: 'test', + }, + ], + temperature: 0.5, + userId, + userSubscriptionTier, + }); - expect(response).toEqual({ + expect(result).toEqual({ success: false, errorCode: 'not_authorized', errorMessage: 'The subscription does not permit the given model for AI Chat features.', }); + expect(chatInterface.chat).not.toBeCalled(); }); it('should return an not_logged_in result if the given a null userId', async () => { @@ -1582,6 +1607,111 @@ describe('AIController', () => { }); }); + it('should return success when allowedModels includes the model', async () => { + chatInterface.chat.mockReturnValueOnce( + Promise.resolve({ + choices: [ + { + role: 'user', + content: 'test', + finishReason: 'stop', + }, + ], + totalTokens: 1, + }) + ); + + const result = await controller.chat({ + model: 'test-model1', + messages: [ + { + role: 'user', + content: 'test', + }, + ], + temperature: 0.5, + userId, + userSubscriptionTier, + }); + + expect(result).toEqual({ + success: true, + choices: [ + { + role: 'user', + content: 'test', + finishReason: 'stop', + }, + ], + }); + expect(chatInterface.chat).toBeCalledWith({ + model: 'test-model1', + messages: [ + { + role: 'user', + content: 'test', + }, + ], + temperature: 0.5, + userId: 'test-user', + }); + }); + + it('should return not_authorized error when allowedModels does not include the model', async () => { + controller = new AIController({ + chat: { + interfaces: { + provider1: chatInterface, + }, + options: { + defaultModel: 'default-model', + defaultModelProvider: 'provider1', + allowedChatModels: [ + { + provider: 'provider1', + model: 'modelA', + }, + { + provider: 'provider1', + model: 'modelB', + }, + ], + allowedChatSubscriptionTiers: ['test-tier'], + }, + }, + generateSkybox: null, + images: null, + metrics: store, + config: store, + hume: null, + sloyd: null, + policies: null, + policyController: policies, + records: store, + }); + + const result = await controller.chat({ + model: 'modelC', + messages: [ + { + role: 'user', + content: 'test', + }, + ], + temperature: 0.5, + userId, + userSubscriptionTier, + }); + + expect(result).toEqual({ + success: false, + errorCode: 'not_authorized', + errorMessage: + 'The subscription does not permit the given model for AI Chat features.', + }); + expect(chatInterface.chat).not.toBeCalled(); + }); + describe('subscriptions', () => { beforeEach(async () => { store.subscriptionConfiguration = buildSubscriptionConfig( From bc831d5ad82141309c53ca7eaabb23e9df8162b2 Mon Sep 17 00:00:00 2001 From: TroyceGowdy Date: Wed, 16 Oct 2024 18:38:58 -0400 Subject: [PATCH 6/6] fix: moved test code into the correct "subscriptions"group - Relocated test code from both "chat" and "chatStream" --- src/aux-records/AIController.spec.ts | 200 +++++++++++++-------------- 1 file changed, 100 insertions(+), 100 deletions(-) diff --git a/src/aux-records/AIController.spec.ts b/src/aux-records/AIController.spec.ts index d4f0c9a66..bfdedf7b9 100644 --- a/src/aux-records/AIController.spec.ts +++ b/src/aux-records/AIController.spec.ts @@ -369,56 +369,6 @@ describe('AIController', () => { expect(chatInterface.chat).not.toBeCalled(); }); - it('should return success when allowedModels includes the model', async () => { - chatInterface.chat.mockReturnValueOnce( - Promise.resolve({ - choices: [ - { - role: 'user', - content: 'test', - finishReason: 'stop', - }, - ], - totalTokens: 1, - }) - ); - - const result = await controller.chat({ - model: 'test-model1', - messages: [ - { - role: 'user', - content: 'test', - }, - ], - temperature: 0.5, - userId, - userSubscriptionTier, - }); - - expect(result).toEqual({ - success: true, - choices: [ - { - role: 'user', - content: 'test', - finishReason: 'stop', - }, - ], - }); - expect(chatInterface.chat).toBeCalledWith({ - model: 'test-model1', - messages: [ - { - role: 'user', - content: 'test', - }, - ], - temperature: 0.5, - userId: 'test-user', - }); - }); - it('should return not_authorized when allowedModels does not include the model', async () => { controller = new AIController({ chat: { @@ -754,6 +704,56 @@ describe('AIController', () => { }); }); + it('should return success when allowedModels includes the model', async () => { + chatInterface.chat.mockReturnValueOnce( + Promise.resolve({ + choices: [ + { + role: 'user', + content: 'test', + finishReason: 'stop', + }, + ], + totalTokens: 1, + }) + ); + + const result = await controller.chat({ + model: 'test-model1', + messages: [ + { + role: 'user', + content: 'test', + }, + ], + temperature: 0.5, + userId, + userSubscriptionTier, + }); + + expect(result).toEqual({ + success: true, + choices: [ + { + role: 'user', + content: 'test', + finishReason: 'stop', + }, + ], + }); + expect(chatInterface.chat).toBeCalledWith({ + model: 'test-model1', + messages: [ + { + role: 'user', + content: 'test', + }, + ], + temperature: 0.5, + userId: 'test-user', + }); + }); + it('should specify the maximum number of tokens allowed based on how many tokens the subscription has left in the period', async () => { chatInterface.chat.mockReturnValueOnce( Promise.resolve({ @@ -1607,56 +1607,6 @@ describe('AIController', () => { }); }); - it('should return success when allowedModels includes the model', async () => { - chatInterface.chat.mockReturnValueOnce( - Promise.resolve({ - choices: [ - { - role: 'user', - content: 'test', - finishReason: 'stop', - }, - ], - totalTokens: 1, - }) - ); - - const result = await controller.chat({ - model: 'test-model1', - messages: [ - { - role: 'user', - content: 'test', - }, - ], - temperature: 0.5, - userId, - userSubscriptionTier, - }); - - expect(result).toEqual({ - success: true, - choices: [ - { - role: 'user', - content: 'test', - finishReason: 'stop', - }, - ], - }); - expect(chatInterface.chat).toBeCalledWith({ - model: 'test-model1', - messages: [ - { - role: 'user', - content: 'test', - }, - ], - temperature: 0.5, - userId: 'test-user', - }); - }); - it('should return not_authorized error when allowedModels does not include the model', async () => { controller = new AIController({ chat: { @@ -1740,6 +1690,56 @@ describe('AIController', () => { }); }); + it('should return success when allowedModels includes the model', async () => { + chatInterface.chat.mockReturnValueOnce( + Promise.resolve({ + choices: [ + { + role: 'user', + content: 'test', + finishReason: 'stop', + }, + ], + totalTokens: 1, + }) + ); + + const result = await controller.chat({ + model: 'test-model1', + messages: [ + { + role: 'user', + content: 'test', + }, + ], + temperature: 0.5, + userId, + userSubscriptionTier, + }); + + expect(result).toEqual({ + success: true, + choices: [ + { + role: 'user', + content: 'test', + finishReason: 'stop', + }, + ], + }); + expect(chatInterface.chat).toBeCalledWith({ + model: 'test-model1', + messages: [ + { + role: 'user', + content: 'test', + }, + ], + temperature: 0.5, + userId: 'test-user', + }); + }); + it('should specify the maximum number of tokens allowed based on how many tokens the subscription has left in the period', async () => { chatInterface.chatStream.mockReturnValueOnce( asyncIterable([