basant307/AI_Governance_Project

1import { describe, it, expect } from 'vitest';2import {3  normalize,4  tokenLimit,5  knownTokenLimit,6  escalatedOutputTokenLimit,7  DEFAULT_TOKEN_LIMIT,8  DEFAULT_OUTPUT_TOKEN_LIMIT,9  ESCALATED_MAX_TOKENS,10} from './tokenLimits.js';11 12describe('normalize', () => {13  it('should lowercase and trim the model string', () => {14    expect(normalize('  GEMINI-1.5-PRO  ')).toBe('gemini-1.5-pro');15  });16 17  it('should strip provider prefixes', () => {18    expect(normalize('google/gemini-1.5-pro')).toBe('gemini-1.5-pro');19    expect(normalize('anthropic/claude-3.5-sonnet')).toBe('claude-3.5-sonnet');20  });21 22  it('should handle pipe and colon separators', () => {23    expect(normalize('qwen|qwen2.5:qwen2.5-1m')).toBe('qwen2.5-1m');24  });25 26  it('should collapse whitespace to a single hyphen', () => {27    expect(normalize('claude 3.5 sonnet')).toBe('claude-3.5-sonnet');28  });29 30  it('should remove date and version suffixes', () => {31    expect(normalize('gemini-1.5-pro-20250219')).toBe('gemini-1.5-pro');32    expect(normalize('gpt-4o-mini-v1')).toBe('gpt-4o-mini');33    expect(normalize('claude-3.7-sonnet-20240715')).toBe('claude-3.7-sonnet');34    expect(normalize('gpt-4.1-latest')).toBe('gpt-4.1');35    expect(normalize('gemini-2.0-flash-preview-20250520')).toBe(36      'gemini-2.0-flash',37    );38  });39 40  it('should remove quantization and numeric suffixes', () => {41    expect(normalize('qwen3-coder-7b-4bit')).toBe('qwen3-coder-7b');42    expect(normalize('llama-4-scout-int8')).toBe('llama-4-scout');43    expect(normalize('mistral-large-2-bf16')).toBe('mistral-large-2');44    expect(normalize('deepseek-v3.1-q4')).toBe('deepseek-v3.1');45    expect(normalize('qwen2.5-quantized')).toBe('qwen2.5');46  });47 48  it('should handle a combination of normalization rules', () => {49    expect(normalize('  Google/GEMINI-2.5-PRO:gemini-2.5-pro-20250605  ')).toBe(50      'gemini-2.5-pro',51    );52  });53 54  it('should handle empty or null input', () => {55    expect(normalize('')).toBe('');56    expect(normalize(undefined as unknown as string)).toBe('');57    expect(normalize(null as unknown as string)).toBe('');58  });59 60  it('should remove preview suffixes', () => {61    expect(normalize('gemini-2.0-flash-preview')).toBe('gemini-2.0-flash');62  });63 64  it('should not remove "-latest" from specific Qwen model names', () => {65    expect(normalize('qwen-plus-latest')).toBe('qwen-plus-latest');66    expect(normalize('qwen-flash-latest')).toBe('qwen-flash-latest');67    expect(normalize('qwen-vl-max-latest')).toBe('qwen-vl-max-latest');68  });69 70  it('should preserve date suffixes for Kimi K2 models', () => {71    expect(normalize('kimi-k2-0905-preview')).toBe('kimi-k2-0905');72    expect(normalize('kimi-k2-0711-preview')).toBe('kimi-k2-0711');73    expect(normalize('kimi-k2-turbo-preview')).toBe('kimi-k2-turbo');74  });75 76  it('should remove date like suffixes', () => {77    expect(normalize('deepseek-r1-0528')).toBe('deepseek-r1');78  });79 80  it('should remove literal "-latest" "-exp" suffixes', () => {81    expect(normalize('gpt-4.1-latest')).toBe('gpt-4.1');82    expect(normalize('deepseek-v3.2-exp')).toBe('deepseek-v3.2');83  });84 85  it('should remove suffix version numbers with "v" prefix', () => {86    expect(normalize('model-test-v1.1')).toBe('model-test');87    expect(normalize('model-v1.1')).toBe('model');88  });89 90  it('should remove suffix version numbers w/o "v" prefix only if they are preceded by another dash', () => {91    expect(normalize('model-test-1.1')).toBe('model-test');92    expect(normalize('gpt-4.1')).toBe('gpt-4.1');93  });94});95 96describe('tokenLimit', () => {97  it('uses 200K as the global default context window', () => {98    expect(DEFAULT_TOKEN_LIMIT).toBe(200_000);99  });100 101  describe('Google Gemini', () => {102    it('should return 1M for Gemini 3.x (latest)', () => {103      expect(tokenLimit('gemini-3-pro-preview')).toBe(1000000);104      expect(tokenLimit('gemini-3-flash-preview')).toBe(1000000);105      expect(tokenLimit('gemini-3.1-pro-preview')).toBe(1000000);106    });107 108    it('should return 1M for legacy Gemini (fallback)', () => {109      expect(tokenLimit('gemini-2.5-pro')).toBe(1000000);110      expect(tokenLimit('gemini-2.5-flash')).toBe(1000000);111      expect(tokenLimit('gemini-2.0-flash')).toBe(1000000);112      expect(tokenLimit('gemini-1.5-pro')).toBe(1000000);113      expect(tokenLimit('gemini-1.5-flash')).toBe(1000000);114    });115  });116 117  describe('OpenAI', () => {118    it('should return 272K for GPT-5.x (latest)', () => {119      expect(tokenLimit('gpt-5')).toBe(272000);120      expect(tokenLimit('gpt-5-mini')).toBe(272000);121      expect(tokenLimit('gpt-5.2')).toBe(272000);122      expect(tokenLimit('gpt-5.2-pro')).toBe(272000);123    });124 125    it('should return 128K for legacy GPT (fallback)', () => {126      expect(tokenLimit('gpt-4o')).toBe(131072);127      expect(tokenLimit('gpt-4o-mini')).toBe(131072);128      expect(tokenLimit('gpt-4.1')).toBe(131072);129      expect(tokenLimit('gpt-4')).toBe(131072);130    });131 132    it('should return 200K for o-series', () => {133      expect(tokenLimit('o3')).toBe(200000);134      expect(tokenLimit('o3-mini')).toBe(200000);135      expect(tokenLimit('o4-mini')).toBe(200000);136    });137  });138 139  describe('Anthropic Claude', () => {140    it('should return 200K for all Claude models', () => {141      expect(tokenLimit('claude-opus-4-6')).toBe(200000);142      expect(tokenLimit('claude-sonnet-4-6')).toBe(200000);143      expect(tokenLimit('claude-sonnet-4')).toBe(200000);144      expect(tokenLimit('claude-opus-4')).toBe(200000);145      expect(tokenLimit('claude-3.5-sonnet')).toBe(200000);146      expect(tokenLimit('claude-3.7-sonnet')).toBe(200000);147    });148  });149 150  describe('Alibaba Qwen', () => {151    it('should return 1M for commercial Qwen3 models', () => {152      expect(tokenLimit('qwen3-coder-plus')).toBe(1000000);153      expect(tokenLimit('qwen3-coder-plus-20250601')).toBe(1000000);154      expect(tokenLimit('qwen3-coder-flash')).toBe(1000000);155      expect(tokenLimit('qwen3.5-plus')).toBe(1000000);156      expect(tokenLimit('coder-model')).toBe(1000000);157    });158 159    it('should return 256K for Qwen3 non-commercial models', () => {160      expect(tokenLimit('qwen3-max')).toBe(262144);161      expect(tokenLimit('qwen3-max-2026-01-23')).toBe(262144);162      expect(tokenLimit('qwen3-vl-plus')).toBe(262144);163      expect(tokenLimit('qwen3-coder-7b')).toBe(262144);164      expect(tokenLimit('qwen3-coder-next')).toBe(262144);165    });166 167    it('should return 1M for studio latest models', () => {168      expect(tokenLimit('qwen-plus-latest')).toBe(1000000);169      expect(tokenLimit('qwen-flash-latest')).toBe(1000000);170    });171 172    it('should return 256K for Qwen fallback', () => {173      expect(tokenLimit('qwen-plus')).toBe(262144);174      expect(tokenLimit('qwen-turbo')).toBe(262144);175      expect(tokenLimit('qwen2.5')).toBe(262144);176      expect(tokenLimit('qwen-vl-max-latest')).toBe(262144);177    });178  });179 180  describe('DeepSeek', () => {181    it('should return 1M for DeepSeek V4 models', () => {182      expect(tokenLimit('deepseek-v4-flash')).toBe(1000000);183      expect(tokenLimit('deepseek-v4-pro')).toBe(1000000);184    });185 186    it('should return 128K for DeepSeek models', () => {187      expect(tokenLimit('deepseek-r1')).toBe(131072);188      expect(tokenLimit('deepseek-v3')).toBe(131072);189      expect(tokenLimit('deepseek-chat')).toBe(131072);190    });191  });192 193  describe('Zhipu GLM', () => {194    it('should default GLM-5.2+ and GLM-6.x onward to 1M (forward default)', () => {195      expect(tokenLimit('glm-5.2')).toBe(1000000);196      expect(tokenLimit('GLM-5.2')).toBe(1000000);197      expect(tokenLimit('glm-5.3')).toBe(1000000);198      expect(tokenLimit('glm-6')).toBe(1000000);199      expect(tokenLimit('glm-6.5')).toBe(1000000);200      expect(tokenLimit('glm-10')).toBe(1000000); // two-digit major201    });202 203    it('should strip third-party deploy prefixes before matching', () => {204      expect(tokenLimit('zai/GLM-5.2')).toBe(1000000);205      expect(tokenLimit('pai/glm-5.3')).toBe(1000000);206      expect(tokenLimit('pai/glm-5.1')).toBe(202752);207    });208 209    it('should pin GLM-5 / 5.1 and GLM-4.x to 200K', () => {210      expect(tokenLimit('glm-5')).toBe(202752);211      expect(tokenLimit('glm-5.0')).toBe(202752);212      expect(tokenLimit('glm-5.1')).toBe(202752);213      expect(tokenLimit('glm-4.7')).toBe(202752);214    });215 216    it('should keep non-numeric GLM names on the conservative fallback', () => {217      expect(tokenLimit('glm-z1')).toBe(202752);218    });219 220    it('should return 200K for legacy GLM (fallback)', () => {221      expect(tokenLimit('glm-4.5')).toBe(202752);222      expect(tokenLimit('glm-4.5v')).toBe(202752);223      expect(tokenLimit('glm-4.5-air')).toBe(202752);224    });225  });226 227  describe('MiniMax', () => {228    it('should return 1M for MiniMax-M3', () => {229      expect(tokenLimit('MiniMax-M3')).toBe(1000000);230    });231 232    it('should return 196608 for MiniMax-M2.5 (latest)', () => {233      expect(tokenLimit('MiniMax-M2.5')).toBe(196608);234    });235 236    it('should return 200K for MiniMax fallback', () => {237      expect(tokenLimit('MiniMax-M2.1')).toBe(200000);238    });239  });240 241  describe('Moonshot Kimi', () => {242    it('should return 256K for Kimi models', () => {243      expect(tokenLimit('kimi-k2.5')).toBe(262144);244      expect(tokenLimit('kimi-k2-0905')).toBe(262144);245      expect(tokenLimit('kimi-k2-turbo')).toBe(262144);246    });247  });248 249  describe('Other models', () => {250    it('should return correct limits for other known models', () => {251      expect(tokenLimit('seed-oss')).toBe(524288);252    });253 254    it('should return the default token limit for unknown models', () => {255      expect(tokenLimit('llama-4-scout')).toBe(DEFAULT_TOKEN_LIMIT);256    });257  });258 259  it('should return the default token limit for an unknown model', () => {260    expect(tokenLimit('unknown-model-v1.0')).toBe(DEFAULT_TOKEN_LIMIT);261    expect(tokenLimit('mistral-large-2')).toBe(DEFAULT_TOKEN_LIMIT);262  });263 264  it('should return the correct limit for a complex model string', () => {265    expect(tokenLimit('  a/b/c|GPT-4o:gpt-4o-2024-05-13-q4  ')).toBe(131072);266  });267 268  it('should handle case-insensitive model names', () => {269    expect(tokenLimit('GPT-4O')).toBe(131072);270    expect(tokenLimit('CLAUDE-3.5-SONNET')).toBe(200000);271  });272});273 274describe('knownTokenLimit', () => {275  it('returns a limit for known input models', () => {276    expect(knownTokenLimit('qwen3-max')).toBe(262144);277    expect(knownTokenLimit('gpt-5')).toBe(272000);278  });279 280  it('returns a limit for known output models', () => {281    expect(knownTokenLimit('qwen3-max', 'output')).toBe(32768);282  });283 284  it('returns undefined for unknown models instead of the default fallback', () => {285    expect(knownTokenLimit('unknown-model-v1.0')).toBeUndefined();286  });287});288 289describe('tokenLimit with output type', () => {290  describe('latest models output limits', () => {291    it('should return correct output limits for GPT-5.x', () => {292      expect(tokenLimit('gpt-5.2', 'output')).toBe(131072);293      expect(tokenLimit('gpt-5-mini', 'output')).toBe(131072);294    });295 296    it('should return correct output limits for Gemini 3.x', () => {297      expect(tokenLimit('gemini-3-pro-preview', 'output')).toBe(65536);298      expect(tokenLimit('gemini-3-flash-preview', 'output')).toBe(65536);299    });300 301    it('should return correct output limits for Claude 4.6', () => {302      expect(tokenLimit('claude-opus-4-6', 'output')).toBe(131072);303      expect(tokenLimit('claude-sonnet-4-6', 'output')).toBe(65536);304    });305  });306 307  describe('legacy model output fallbacks', () => {308    it('should return fallback output limits for legacy GPT', () => {309      expect(tokenLimit('gpt-4o', 'output')).toBe(16384);310    });311 312    it('should return fallback output limits for legacy Gemini', () => {313      expect(tokenLimit('gemini-2.5-pro', 'output')).toBe(8192);314    });315 316    it('should return fallback output limits for legacy Claude', () => {317      expect(tokenLimit('claude-sonnet-4', 'output')).toBe(65536);318      expect(tokenLimit('claude-opus-4', 'output')).toBe(65536);319    });320  });321 322  describe('Qwen output limits', () => {323    it('should return correct output limits for Qwen models', () => {324      expect(tokenLimit('qwen3.5-plus', 'output')).toBe(65536);325      expect(tokenLimit('qwen3.6-plus', 'output')).toBe(65536);326      expect(tokenLimit('coder-model', 'output')).toBe(65536);327      // Models without specific output limits fall back to Qwen default (32K)328      expect(tokenLimit('qwen3-max', 'output')).toBe(32768);329      expect(tokenLimit('qwen3-max-2026-01-23', 'output')).toBe(32768);330    });331  });332 333  describe('other output limits', () => {334    it('should return correct output limits for DeepSeek', () => {335      expect(tokenLimit('deepseek-v4-flash', 'output')).toBe(384000);336      expect(tokenLimit('deepseek-v4-pro', 'output')).toBe(384000);337      expect(tokenLimit('deepseek-reasoner', 'output')).toBe(65536);338      expect(tokenLimit('deepseek-r1', 'output')).toBe(65536);339      expect(tokenLimit('deepseek-r1-0528', 'output')).toBe(65536);340      expect(tokenLimit('deepseek-chat', 'output')).toBe(8192);341    });342 343    it('should return correct output limits for GLM', () => {344      expect(tokenLimit('glm-5.2', 'output')).toBe(131072);345      expect(tokenLimit('GLM-5.2', 'output')).toBe(131072);346      expect(tokenLimit('glm-5.1', 'output')).toBe(131072);347      expect(tokenLimit('glm-5', 'output')).toBe(131072);348      expect(tokenLimit('glm-5-turbo', 'output')).toBe(131072);349      expect(tokenLimit('glm-4.7', 'output')).toBe(16384);350    });351 352    it('should return correct output limits for MiniMax', () => {353      expect(tokenLimit('MiniMax-M2.5', 'output')).toBe(65536);354    });355 356    it('should return correct output limits for Kimi', () => {357      expect(tokenLimit('kimi-k2.5', 'output')).toBe(32768);358    });359  });360 361  describe('default output limits', () => {362    it('should return the default output limit for unknown models', () => {363      expect(tokenLimit('unknown-model', 'output')).toBe(364        DEFAULT_OUTPUT_TOKEN_LIMIT,365      );366    });367  });368 369  describe('input vs output comparison', () => {370    it('should return different limits for input vs output', () => {371      expect(tokenLimit('qwen3-max', 'input')).toBe(262144);372      expect(tokenLimit('qwen3-max', 'output')).toBe(32768);373    });374 375    it('should default to input type when no type is specified', () => {376      expect(tokenLimit('qwen3-coder-plus')).toBe(1000000);377      expect(tokenLimit('unknown-model')).toBe(DEFAULT_TOKEN_LIMIT);378    });379  });380 381  describe('normalization with output limits', () => {382    it('should handle normalized model names for output limits', () => {383      expect(tokenLimit('QWEN3-MAX', 'output')).toBe(32768);384      expect(tokenLimit('qwen3-max-20250601', 'output')).toBe(32768);385    });386  });387});388 389describe('escalatedOutputTokenLimit', () => {390  it('uses the escalated floor for unknown models with a large window', () => {391    expect(escalatedOutputTokenLimit('unknown-model', 200_000)).toBe(392      ESCALATED_MAX_TOKENS,393    );394  });395 396  it('caps the reservation at half a small custom window (issue #6144)', () => {397    // A 64K local model must not have the flat 64K escalation floor398    // reserved — that would leave only 1,536 tokens of input budget.399    expect(escalatedOutputTokenLimit('qwen3coder-64k', 65_536)).toBe(32_768);400  });401 402  it('uses the model output limit when larger than the floor and within the cap', () => {403    // claude-sonnet-4-6 declares a 65,536 output limit; half of 200K is 100K.404    expect(escalatedOutputTokenLimit('claude-sonnet-4-6', 200_000)).toBe(405      65_536,406    );407  });408 409  it('falls back to DEFAULT_TOKEN_LIMIT for the cap when the window is unset', () => {410    expect(escalatedOutputTokenLimit('unknown-model')).toBe(411      Math.min(ESCALATED_MAX_TOKENS, Math.floor(DEFAULT_TOKEN_LIMIT / 2)),412    );413    expect(escalatedOutputTokenLimit('unknown-model', 0)).toBe(414      Math.min(ESCALATED_MAX_TOKENS, Math.floor(DEFAULT_TOKEN_LIMIT / 2)),415    );416  });417});418