Skip to content

Commit 0845ead

Browse files
committed
update
1 parent 0d40243 commit 0845ead

File tree

2 files changed

+104
-34
lines changed

2 files changed

+104
-34
lines changed

lib/pricing.test.ts

Lines changed: 56 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -9,59 +9,87 @@ import {
99
} from "./pricing.ts";
1010

1111
describe("generateLookupCandidates", () => {
12-
describe("openrouter provider", () => {
12+
describe("openrouter provider (slash format)", () => {
1313
it("should generate candidates for openrouter/anthropic/model", () => {
1414
const candidates = generateLookupCandidates(
1515
"openrouter/anthropic/claude-sonnet-4",
1616
);
17-
expect(candidates).toEqual([
18-
"openrouter/anthropic/claude-sonnet-4",
19-
"anthropic/claude-sonnet-4",
20-
"claude-sonnet-4",
21-
]);
17+
expect(candidates).toContain("openrouter/anthropic/claude-sonnet-4");
18+
expect(candidates).toContain("anthropic/claude-sonnet-4");
19+
expect(candidates).toContain("claude-sonnet-4");
2220
});
2321

2422
it("should generate candidates for openrouter/openai/model", () => {
2523
const candidates = generateLookupCandidates("openrouter/openai/gpt-4o");
26-
expect(candidates).toEqual([
27-
"openrouter/openai/gpt-4o",
28-
"openai/gpt-4o",
29-
"gpt-4o",
30-
]);
24+
expect(candidates).toContain("openrouter/openai/gpt-4o");
25+
expect(candidates).toContain("openai/gpt-4o");
26+
expect(candidates).toContain("gpt-4o");
3127
});
3228

3329
it("should handle openrouter with simple model name", () => {
3430
const candidates = generateLookupCandidates("openrouter/some-model");
35-
expect(candidates).toEqual(["openrouter/some-model", "some-model"]);
31+
expect(candidates).toContain("openrouter/some-model");
32+
expect(candidates).toContain("some-model");
3633
});
3734
});
3835

39-
describe("anthropic provider", () => {
36+
describe("anthropic provider (slash format)", () => {
4037
it("should generate candidates for anthropic/model", () => {
4138
const candidates = generateLookupCandidates("anthropic/claude-haiku-4-5");
42-
expect(candidates).toEqual([
43-
"anthropic/claude-haiku-4-5",
44-
"claude-haiku-4-5",
45-
]);
39+
expect(candidates).toContain("anthropic/claude-haiku-4-5");
40+
expect(candidates).toContain("claude-haiku-4-5");
41+
});
42+
});
43+
44+
describe("Vercel AI Gateway format (colon separator)", () => {
45+
it("should generate candidates for anthropic:model", () => {
46+
const candidates = generateLookupCandidates("anthropic:claude-sonnet-4");
47+
expect(candidates).toContain("anthropic:claude-sonnet-4");
48+
expect(candidates).toContain("anthropic/claude-sonnet-4");
49+
expect(candidates).toContain("claude-sonnet-4");
50+
});
51+
52+
it("should generate candidates for openai:model", () => {
53+
const candidates = generateLookupCandidates("openai:gpt-4o");
54+
expect(candidates).toContain("openai:gpt-4o");
55+
expect(candidates).toContain("openai/gpt-4o");
56+
expect(candidates).toContain("gpt-4o");
57+
});
58+
59+
it("should generate candidates for google:model", () => {
60+
const candidates = generateLookupCandidates("google:gemini-2.0-flash");
61+
expect(candidates).toContain("google:gemini-2.0-flash");
62+
expect(candidates).toContain("google/gemini-2.0-flash");
63+
expect(candidates).toContain("gemini-2.0-flash");
64+
});
65+
66+
it("should generate candidates for x-ai:model", () => {
67+
const candidates = generateLookupCandidates("x-ai:grok-2");
68+
expect(candidates).toContain("x-ai:grok-2");
69+
expect(candidates).toContain("x-ai/grok-2");
70+
expect(candidates).toContain("grok-2");
4671
});
4772
});
4873

49-
describe("openai provider", () => {
74+
describe("openai provider (slash format)", () => {
5075
it("should generate candidates for openai/model", () => {
5176
const candidates = generateLookupCandidates("openai/gpt-4o");
52-
expect(candidates).toEqual(["openai/gpt-4o", "gpt-4o"]);
77+
expect(candidates).toContain("openai/gpt-4o");
78+
expect(candidates).toContain("gpt-4o");
5379
});
5480

5581
it("should handle openai/gpt-4o-mini", () => {
5682
const candidates = generateLookupCandidates("openai/gpt-4o-mini");
57-
expect(candidates).toEqual(["openai/gpt-4o-mini", "gpt-4o-mini"]);
83+
expect(candidates).toContain("openai/gpt-4o-mini");
84+
expect(candidates).toContain("gpt-4o-mini");
5885
});
5986
});
6087

6188
describe("lmstudio provider", () => {
6289
it("should generate candidates for lmstudio/model", () => {
6390
const candidates = generateLookupCandidates("lmstudio/llama-3-8b");
64-
expect(candidates).toEqual(["lmstudio/llama-3-8b", "llama-3-8b"]);
91+
expect(candidates).toContain("lmstudio/llama-3-8b");
92+
expect(candidates).toContain("llama-3-8b");
6593
});
6694
});
6795

@@ -70,10 +98,16 @@ describe("generateLookupCandidates", () => {
7098
const candidates = generateLookupCandidates("unknown/some-model");
7199
expect(candidates).toEqual(["unknown/some-model"]);
72100
});
101+
102+
it("should handle unknown provider with colon format", () => {
103+
const candidates = generateLookupCandidates("unknown:some-model");
104+
expect(candidates).toContain("unknown:some-model");
105+
expect(candidates).toContain("unknown/some-model");
106+
});
73107
});
74108

75109
describe("no provider prefix", () => {
76-
it("should return only the original string when no slash", () => {
110+
it("should return only the original string when no separator", () => {
77111
const candidates = generateLookupCandidates("claude-sonnet-4");
78112
expect(candidates).toEqual(["claude-sonnet-4"]);
79113
});

lib/pricing.ts

Lines changed: 48 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -61,6 +61,12 @@ const PROVIDER_CONFIGS: Record<string, ProviderConfig> = {
6161
anthropic: { strip: true, keepNested: false },
6262
openai: { strip: true, keepNested: false },
6363
lmstudio: { strip: true, keepNested: false },
64+
google: { strip: true, keepNested: false },
65+
meta: { strip: true, keepNested: false },
66+
mistral: { strip: true, keepNested: false },
67+
cohere: { strip: true, keepNested: false },
68+
"x-ai": { strip: true, keepNested: false },
69+
deepseek: { strip: true, keepNested: false },
6470
};
6571

6672
// Cache the loaded pricing data
@@ -91,35 +97,58 @@ function loadPricingData(): Record<string, unknown> {
9197
}
9298
}
9399

100+
/**
101+
* Normalize a model string by converting different separator formats
102+
* Vercel AI Gateway uses "provider:model" format
103+
* Old multi-provider setup used "provider/model" format
104+
* LiteLLM pricing data uses various formats
105+
*/
106+
function normalizeModelString(modelString: string): string {
107+
// Convert colon separator to slash for unified processing
108+
// e.g., "anthropic:claude-sonnet-4" -> "anthropic/claude-sonnet-4"
109+
return modelString.replace(":", "/");
110+
}
111+
94112
/**
95113
* Generate lookup candidates for a model string using provider configuration
96114
* Returns candidates in priority order (most specific first)
115+
*
116+
* Supports both formats:
117+
* - Vercel AI Gateway: "anthropic:claude-sonnet-4"
118+
* - Multi-provider: "anthropic/claude-sonnet-4"
97119
*/
98120
function generateLookupCandidates(modelString: string): string[] {
99121
const candidates: string[] = [];
122+
123+
// Normalize the model string (convert : to /)
124+
const normalizedString = normalizeModelString(modelString);
100125

101126
// Find matching provider config
102-
const slashIndex = modelString.indexOf("/");
127+
const slashIndex = normalizedString.indexOf("/");
103128
if (slashIndex === -1) {
104129
// No provider prefix, just use as-is
105-
return [modelString];
130+
return [modelString, normalizedString].filter((v, i, a) => a.indexOf(v) === i);
106131
}
107132

108-
const provider = modelString.slice(0, slashIndex);
133+
const provider = normalizedString.slice(0, slashIndex);
109134
const config = PROVIDER_CONFIGS[provider];
135+
const remainder = normalizedString.slice(slashIndex + 1);
110136

111-
if (!config) {
112-
// Unknown provider, try original string only
113-
return [modelString];
137+
// Always try the original string first
138+
candidates.push(modelString);
139+
140+
// Also try the normalized version (with / instead of :)
141+
if (normalizedString !== modelString) {
142+
candidates.push(normalizedString);
114143
}
115144

116-
const remainder = modelString.slice(slashIndex + 1);
117-
118-
// Always try the full string first (with provider prefix for the pricing file format)
119-
candidates.push(modelString);
145+
if (!config) {
146+
// Unknown provider, try original strings only
147+
return candidates.filter((v, i, a) => a.indexOf(v) === i);
148+
}
120149

121150
if (config.strip) {
122-
// Try without our provider prefix
151+
// Try without our provider prefix (just the model name)
123152
candidates.push(remainder);
124153
}
125154

@@ -131,7 +160,14 @@ function generateLookupCandidates(modelString: string): string[] {
131160
}
132161
}
133162

134-
return candidates;
163+
// Also try with common LiteLLM prefixes
164+
// LiteLLM often uses "provider/model" format
165+
if (!normalizedString.startsWith(provider + "/")) {
166+
candidates.push(`${provider}/${remainder}`);
167+
}
168+
169+
// Remove duplicates while preserving order
170+
return candidates.filter((v, i, a) => a.indexOf(v) === i);
135171
}
136172

137173
/**

0 commit comments

Comments
 (0)