|
| 1 | +import { describe, expect, it } from "vitest"; |
| 2 | +import { getContextWindow, inferContextWindow, DEFAULT_CONTEXT_WINDOW } from "./model_context"; |
| 3 | + |
| 4 | +describe("getContextWindow", () => { |
| 5 | + it("returns user-configured contextWindow when provided", () => { |
| 6 | + expect(getContextWindow({ model: "gpt-4o", contextWindow: 50_000 })).toBe(50_000); |
| 7 | + }); |
| 8 | + |
| 9 | + it("matches GPT-4o prefix", () => { |
| 10 | + expect(getContextWindow({ model: "gpt-4o" })).toBe(128_000); |
| 11 | + expect(getContextWindow({ model: "gpt-4o-mini" })).toBe(128_000); |
| 12 | + }); |
| 13 | + |
| 14 | + it("matches GPT-4.1 prefix before GPT-4", () => { |
| 15 | + expect(getContextWindow({ model: "gpt-4.1-nano" })).toBe(1_047_576); |
| 16 | + }); |
| 17 | + |
| 18 | + it("matches GPT-4 Turbo before GPT-4 base", () => { |
| 19 | + expect(getContextWindow({ model: "gpt-4-turbo-preview" })).toBe(128_000); |
| 20 | + }); |
| 21 | + |
| 22 | + it("matches GPT-4 base", () => { |
| 23 | + expect(getContextWindow({ model: "gpt-4-0613" })).toBe(8_192); |
| 24 | + }); |
| 25 | + |
| 26 | + it("matches Claude models", () => { |
| 27 | + expect(getContextWindow({ model: "claude-sonnet-4-20250514" })).toBe(200_000); |
| 28 | + expect(getContextWindow({ model: "claude-3-haiku" })).toBe(200_000); |
| 29 | + }); |
| 30 | + |
| 31 | + it("matches Gemini models", () => { |
| 32 | + expect(getContextWindow({ model: "gemini-2.0-flash" })).toBe(1_048_576); |
| 33 | + }); |
| 34 | + |
| 35 | + it("matches DeepSeek models", () => { |
| 36 | + expect(getContextWindow({ model: "deepseek-chat" })).toBe(64_000); |
| 37 | + }); |
| 38 | + |
| 39 | + it("matches Qwen models", () => { |
| 40 | + expect(getContextWindow({ model: "qwen-max" })).toBe(131_072); |
| 41 | + }); |
| 42 | + |
| 43 | + it("matches Llama-4 before Llama base", () => { |
| 44 | + expect(getContextWindow({ model: "llama-4-maverick" })).toBe(1_048_576); |
| 45 | + expect(getContextWindow({ model: "llama-3.1-70b" })).toBe(131_072); |
| 46 | + }); |
| 47 | + |
| 48 | + it("is case-insensitive", () => { |
| 49 | + expect(getContextWindow({ model: "GPT-4O" })).toBe(128_000); |
| 50 | + expect(getContextWindow({ model: "Claude-Sonnet-4" })).toBe(200_000); |
| 51 | + }); |
| 52 | + |
| 53 | + it("returns default for unknown models", () => { |
| 54 | + expect(getContextWindow({ model: "my-custom-model" })).toBe(DEFAULT_CONTEXT_WINDOW); |
| 55 | + }); |
| 56 | +}); |
| 57 | + |
| 58 | +describe("inferContextWindow", () => { |
| 59 | + it("returns prefix-matched value", () => { |
| 60 | + expect(inferContextWindow("gpt-4o")).toBe(128_000); |
| 61 | + }); |
| 62 | + |
| 63 | + it("returns default for unknown models", () => { |
| 64 | + expect(inferContextWindow("unknown-model")).toBe(DEFAULT_CONTEXT_WINDOW); |
| 65 | + }); |
| 66 | +}); |
0 commit comments