diff --git a/api/utils/tokens.js b/api/utils/tokens.js index b0d9098821..3b973c40cf 100644 --- a/api/utils/tokens.js +++ b/api/utils/tokens.js @@ -59,6 +59,8 @@ const openAIModels = { 'gpt-3.5-turbo-1106': 16375, // -10 from max 'gpt-3.5-turbo-0125': 16375, // -10 from max 'mistral-': 31990, // -10 from max + llama3: 8187, // -5 from max + 'llama-3': 8187, // -5 from max }; const cohereModels = { diff --git a/api/utils/tokens.spec.js b/api/utils/tokens.spec.js index c51610c917..228e298605 100644 --- a/api/utils/tokens.spec.js +++ b/api/utils/tokens.spec.js @@ -20,6 +20,18 @@ describe('getModelMaxTokens', () => { ); }); + test('should return correct tokens for LLama 3 models', () => { + expect(getModelMaxTokens('meta-llama/llama-3-8b')).toBe( + maxTokensMap[EModelEndpoint.openAI]['llama-3'], + ); + expect(getModelMaxTokens('meta-llama/llama-3-8b')).toBe( + maxTokensMap[EModelEndpoint.openAI]['llama3'], + ); + expect(getModelMaxTokens('llama-3-500b')).toBe(maxTokensMap[EModelEndpoint.openAI]['llama-3']); + expect(getModelMaxTokens('llama3-70b')).toBe(maxTokensMap[EModelEndpoint.openAI]['llama3']); + expect(getModelMaxTokens('llama3:latest')).toBe(maxTokensMap[EModelEndpoint.openAI]['llama3']); + }); + test('should return undefined for no match', () => { expect(getModelMaxTokens('unknown-model')).toBeUndefined(); });