1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
|
import { and, Database, eq, inArray, sql } from "@opencode-ai/console-core/drizzle/index.js"
import { ModelTpmRateLimitTable } from "@opencode-ai/console-core/schema/ip.sql.js"
import { UsageInfo } from "./provider/provider"
export function createModelTpmLimiter(providers: { id: string; model: string; tpmLimit?: number }[]) {
const ids = providers.filter((p) => p.tpmLimit).map((p) => `${p.id}/${p.model}`)
if (ids.length === 0) return
const yyyyMMddHHmm = parseInt(
new Date(Date.now())
.toISOString()
.replace(/[^0-9]/g, "")
.substring(0, 12),
)
return {
check: async () => {
const data = await Database.use((tx) =>
tx
.select()
.from(ModelTpmRateLimitTable)
.where(and(inArray(ModelTpmRateLimitTable.id, ids), eq(ModelTpmRateLimitTable.interval, yyyyMMddHHmm))),
)
// convert to map of model to count
return data.reduce(
(acc, curr) => {
acc[curr.id] = curr.count
return acc
},
{} as Record<string, number>,
)
},
track: async (provider: string, model: string, usageInfo: UsageInfo) => {
const id = `${provider}/${model}`
if (!ids.includes(id)) return
const usage = usageInfo.inputTokens
if (usage <= 0) return
await Database.use((tx) =>
tx
.insert(ModelTpmRateLimitTable)
.values({ id, interval: yyyyMMddHHmm, count: usage })
.onDuplicateKeyUpdate({ set: { count: sql`${ModelTpmRateLimitTable.count} + ${usage}` } }),
)
},
}
}
|
