|
|
@@ -1,4 +1,4 @@
|
|
|
-import { and, Database, inArray, sql } from "@opencode-ai/console-core/drizzle/index.js"
|
|
|
+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"
|
|
|
|
|
|
@@ -6,16 +6,12 @@ export function createModelTpmLimiter(providers: { id: string; model: string; tp
|
|
|
const ids = providers.filter((p) => p.tpmLimit).map((p) => `${p.id}/${p.model}`)
|
|
|
if (ids.length === 0) return
|
|
|
|
|
|
- const toInterval = (date: Date) =>
|
|
|
- parseInt(
|
|
|
- date
|
|
|
- .toISOString()
|
|
|
- .replace(/[^0-9]/g, "")
|
|
|
- .substring(0, 12),
|
|
|
- )
|
|
|
- const now = Date.now()
|
|
|
- const currInterval = toInterval(new Date(now))
|
|
|
- const prevInterval = toInterval(new Date(now - 60_000))
|
|
|
+ const yyyyMMddHHmm = parseInt(
|
|
|
+ new Date(Date.now())
|
|
|
+ .toISOString()
|
|
|
+ .replace(/[^0-9]/g, "")
|
|
|
+ .substring(0, 12),
|
|
|
+ )
|
|
|
|
|
|
return {
|
|
|
check: async () => {
|
|
|
@@ -23,18 +19,13 @@ export function createModelTpmLimiter(providers: { id: string; model: string; tp
|
|
|
tx
|
|
|
.select()
|
|
|
.from(ModelTpmRateLimitTable)
|
|
|
- .where(
|
|
|
- and(
|
|
|
- inArray(ModelTpmRateLimitTable.id, ids),
|
|
|
- inArray(ModelTpmRateLimitTable.interval, [currInterval, prevInterval]),
|
|
|
- ),
|
|
|
- ),
|
|
|
+ .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] = Math.max(acc[curr.id] ?? 0, curr.count)
|
|
|
+ acc[curr.id] = curr.count
|
|
|
return acc
|
|
|
},
|
|
|
{} as Record<string, number>,
|
|
|
@@ -48,7 +39,7 @@ export function createModelTpmLimiter(providers: { id: string; model: string; tp
|
|
|
await Database.use((tx) =>
|
|
|
tx
|
|
|
.insert(ModelTpmRateLimitTable)
|
|
|
- .values({ id, interval: currInterval, count: usage })
|
|
|
+ .values({ id, interval: yyyyMMddHHmm, count: usage })
|
|
|
.onDuplicateKeyUpdate({ set: { count: sql`${ModelTpmRateLimitTable.count} + ${usage}` } }),
|
|
|
)
|
|
|
},
|