diff --git a/packages/gen-ai-hub/src/utils/deployment-cache.test.ts b/packages/gen-ai-hub/src/utils/deployment-cache.test.ts index 81b23ee8..881c084f 100644 --- a/packages/gen-ai-hub/src/utils/deployment-cache.test.ts +++ b/packages/gen-ai-hub/src/utils/deployment-cache.test.ts @@ -28,42 +28,65 @@ describe('deployment cache', () => { }); }); - it('should cache all deployments independent of potentially given models', () => { - const opts = { - scenarioId: 'foundation-models', - model: { name: 'gpt-4o', version: 'latest' } - }; + describe('should cache all deployments independent of potentially given models', () => { + beforeEach(() => { + deploymentCache.setAll( + { + scenarioId: 'foundation-models', + model: { name: 'gpt-4o', version: 'latest' } + }, + [ + mockAiDeployment('deployment-id1', { + name: 'gpt-35-turbo', + version: 'latest' + }), + mockAiDeployment('deployment-id2', { + name: 'gpt-35-turbo', + version: '123' + }) + ] + ); + }); - deploymentCache.setAll(opts, [ - mockAiDeployment('deployment-id1', { - name: 'gpt-35-turbo', - version: 'latest' - }), - mockAiDeployment('deployment-id2', { - name: 'gpt-35-turbo', - version: '123' - }) - ]); + it('cache nothing for unlisted model names', () => { + expect( + deploymentCache.get({ + scenarioId: 'foundation-models', + model: { name: 'gpt-4o', version: 'latest' } + }) + ).toBeUndefined(); + }); - expect(deploymentCache.get(opts)).toBeUndefined(); - expect( - deploymentCache.get({ - ...opts, - model: { - name: 'gpt-35-turbo', - version: 'latest' - } - })?.id - ).toEqual('deployment-id1'); - expect( - deploymentCache.get({ - ...opts, - model: { - name: 'gpt-35-turbo', - version: '123' - } - })?.id - ).toEqual('deployment-id2'); + it('retrieve the deployment matching the model name and version', () => { + expect( + deploymentCache.get({ + scenarioId: 'foundation-models', + model: { + name: 'gpt-35-turbo', + version: '123' + } + })?.id + ).toEqual('deployment-id2'); + }); + + it('retrieve the first deployment matching the model name when version is missing', () => { + expect( + deploymentCache.get({ + scenarioId: 'foundation-models', + model: { + name: 'gpt-35-turbo' + } + })?.id + ).toEqual('deployment-id1'); + }); + + it('retrieve the deployment when model is missing', () => { + expect( + deploymentCache.get({ + scenarioId: 'foundation-models' + })?.id + ).toEqual('deployment-id1'); + }); }); it('should cache only the first deployments for equal models and versions', () => { diff --git a/packages/gen-ai-hub/src/utils/deployment-cache.ts b/packages/gen-ai-hub/src/utils/deployment-cache.ts index 010a0604..324ca97a 100644 --- a/packages/gen-ai-hub/src/utils/deployment-cache.ts +++ b/packages/gen-ai-hub/src/utils/deployment-cache.ts @@ -9,7 +9,7 @@ function getCacheKey({ model, resourceGroup = 'default' }: DeploymentResolutionOptions) { - return `${scenarioId}-${executableId}-${model?.name ?? ''}-${model?.version ?? ''}-${resourceGroup}`; + return `${scenarioId}:${executableId}:${model?.name ?? ''}:${model?.version ?? ''}:${resourceGroup}`; } interface Deployment { @@ -51,15 +51,25 @@ function createDeploymentCache(cache: Cache) { * @param deployments - Deployments to retrieve the IDs and models from. */ setAll: ( - opts: Omit, + opts: DeploymentResolutionOptions, deployments: AiDeployment[] ): void => { // go backwards to cache the first deployment ID for each model - [...deployments].reverse().forEach(deployment => { - cache.set(getCacheKey({ ...opts, model: extractModel(deployment) }), { - entry: transformDeploymentForCache(deployment) + [...deployments] + .reverse() + .map(deployment => transformDeploymentForCache(deployment)) + .flatMap(entry => [ + entry, + { id: entry.id }, + ...(entry.model + ? [{ id: entry.id, model: { name: entry.model.name } }] + : []) + ]) + .forEach(entry => { + cache.set(getCacheKey({ ...opts, model: entry.model }), { + entry + }); }); - }); }, clear: () => cache.clear() }; diff --git a/packages/gen-ai-hub/src/utils/deployment-resolver.test.ts b/packages/gen-ai-hub/src/utils/deployment-resolver.test.ts index c6c08da1..0730dd0c 100644 --- a/packages/gen-ai-hub/src/utils/deployment-resolver.test.ts +++ b/packages/gen-ai-hub/src/utils/deployment-resolver.test.ts @@ -29,7 +29,7 @@ describe('deployment resolver', () => { const id = await resolveDeploymentId({ scenarioId: 'foundation-models' }); - expect(id).toBe('1'); + expect(id).toEqual('1'); }); it('should return the first deployment with the correct model name', async () => { @@ -37,7 +37,7 @@ describe('deployment resolver', () => { scenarioId: 'foundation-models', model: { name: 'gpt-4o' } }); - expect(id).toBe('1'); + expect(id).toEqual('1'); }); it('should return the deployment with the correct model name and version', async () => { @@ -45,7 +45,7 @@ describe('deployment resolver', () => { scenarioId: 'foundation-models', model: { name: 'gpt-4o', version: '0613' } }); - expect(id).toBe('2'); + expect(id).toEqual('2'); }); it('should retrieve deployment from cache if available', async () => { @@ -55,8 +55,8 @@ describe('deployment resolver', () => { }; deploymentCache.set(opts, { id: '1' } as AiDeployment); const id = await resolveDeploymentId(opts); - expect(id).toBe('1'); - expect(nock.isDone()).toBe(false); + expect(id).toEqual('1'); + expect(nock.isDone()).toEqual(false); }); it('should throw in case no deployment with the given model name is found', async () => { @@ -131,7 +131,7 @@ describe('deployment resolver', () => { resourceGroup: 'otherId' }); - expect(id).toBe('5'); + expect(id).toEqual('5'); }); });