diff --git a/packages/spacecat-shared-data-access/src/models/organization/organization.model.js b/packages/spacecat-shared-data-access/src/models/organization/organization.model.js index 62efc670c..7d565f009 100755 --- a/packages/spacecat-shared-data-access/src/models/organization/organization.model.js +++ b/packages/spacecat-shared-data-access/src/models/organization/organization.model.js @@ -24,7 +24,15 @@ class Organization extends BaseModel { static IMS_ORG_ID_REGEX = /[a-z0-9]{24}@AdobeOrg/i; - // add your custom methods or overrides here + static LLM_BACKEND_AZURE = 'azure'; + + static LLM_BACKEND_BEDROCK = 'bedrock'; + + static LLM_BACKENDS = ['azure', 'bedrock']; + + getLlmBackend() { + return this.record.llmBackend ?? Organization.LLM_BACKEND_AZURE; + } } export default Organization; diff --git a/packages/spacecat-shared-data-access/src/models/organization/organization.schema.js b/packages/spacecat-shared-data-access/src/models/organization/organization.schema.js index 01b8830a5..8d032cbf6 100644 --- a/packages/spacecat-shared-data-access/src/models/organization/organization.schema.js +++ b/packages/spacecat-shared-data-access/src/models/organization/organization.schema.js @@ -40,6 +40,10 @@ const schema = new SchemaBuilder(Organization, OrganizationCollection) type: 'string', validate: (value) => !value || Organization.IMS_ORG_ID_REGEX.test(value), }) + .addAttribute('llmBackend', { + type: Organization.LLM_BACKENDS, + validate: (value) => !value || Organization.LLM_BACKENDS.includes(value), + }) .addAttribute('fulfillableItems', { type: 'any', validate: (value) => !value || isNonEmptyObject(value), diff --git a/packages/spacecat-shared-data-access/test/unit/models/organization/organization.model.test.js b/packages/spacecat-shared-data-access/test/unit/models/organization/organization.model.test.js index ee7ba7b0b..2661ae04e 100755 --- a/packages/spacecat-shared-data-access/test/unit/models/organization/organization.model.test.js +++ b/packages/spacecat-shared-data-access/test/unit/models/organization/organization.model.test.js @@ -94,4 +94,35 @@ describe('OrganizationModel', () => { expect(instance.getFulfillableItems()).to.deep.equal(['item3', 'item4']); }); }); + + describe('llmBackend', () => { + it('defaults to azure', () => { + expect(instance.getLlmBackend()).to.equal('azure'); + }); + + it('sets llmBackend to bedrock', () => { + instance.setLlmBackend('bedrock'); + expect(instance.getLlmBackend()).to.equal('bedrock'); + }); + + it('sets llmBackend back to azure', () => { + instance.setLlmBackend('bedrock'); + instance.setLlmBackend('azure'); + expect(instance.getLlmBackend()).to.equal('azure'); + }); + }); + + describe('LLM_BACKENDS constants', () => { + it('defines LLM_BACKEND_AZURE', () => { + expect(Organization.LLM_BACKEND_AZURE).to.equal('azure'); + }); + + it('defines LLM_BACKEND_BEDROCK', () => { + expect(Organization.LLM_BACKEND_BEDROCK).to.equal('bedrock'); + }); + + it('defines LLM_BACKENDS list', () => { + expect(Organization.LLM_BACKENDS).to.deep.equal(['azure', 'bedrock']); + }); + }); });