1- import { and , Database , inArray , sql } from "@opencode-ai/console-core/drizzle/index.js"
1+ import { and , Database , eq , inArray , sql } from "@opencode-ai/console-core/drizzle/index.js"
22import { ModelTpmRateLimitTable } from "@opencode-ai/console-core/schema/ip.sql.js"
33import { UsageInfo } from "./provider/provider"
44
55export function createModelTpmLimiter ( providers : { id : string ; model : string ; tpmLimit ?: number } [ ] ) {
66 const ids = providers . filter ( ( p ) => p . tpmLimit ) . map ( ( p ) => `${ p . id } /${ p . model } ` )
77 if ( ids . length === 0 ) return
88
9- const toInterval = ( date : Date ) =>
10- parseInt (
11- date
12- . toISOString ( )
13- . replace ( / [ ^ 0 - 9 ] / g, "" )
14- . substring ( 0 , 12 ) ,
15- )
16- const now = Date . now ( )
17- const currInterval = toInterval ( new Date ( now ) )
18- const prevInterval = toInterval ( new Date ( now - 60_000 ) )
9+ const yyyyMMddHHmm = parseInt (
10+ new Date ( Date . now ( ) )
11+ . toISOString ( )
12+ . replace ( / [ ^ 0 - 9 ] / g, "" )
13+ . substring ( 0 , 12 ) ,
14+ )
1915
2016 return {
2117 check : async ( ) => {
2218 const data = await Database . use ( ( tx ) =>
2319 tx
2420 . select ( )
2521 . from ( ModelTpmRateLimitTable )
26- . where (
27- and (
28- inArray ( ModelTpmRateLimitTable . id , ids ) ,
29- inArray ( ModelTpmRateLimitTable . interval , [ currInterval , prevInterval ] ) ,
30- ) ,
31- ) ,
22+ . where ( and ( inArray ( ModelTpmRateLimitTable . id , ids ) , eq ( ModelTpmRateLimitTable . interval , yyyyMMddHHmm ) ) ) ,
3223 )
3324
3425 // convert to map of model to count
3526 return data . reduce (
3627 ( acc , curr ) => {
37- acc [ curr . id ] = Math . max ( acc [ curr . id ] ?? 0 , curr . count )
28+ acc [ curr . id ] = curr . count
3829 return acc
3930 } ,
4031 { } as Record < string , number > ,
@@ -48,7 +39,7 @@ export function createModelTpmLimiter(providers: { id: string; model: string; tp
4839 await Database . use ( ( tx ) =>
4940 tx
5041 . insert ( ModelTpmRateLimitTable )
51- . values ( { id, interval : currInterval , count : usage } )
42+ . values ( { id, interval : yyyyMMddHHmm , count : usage } )
5243 . onDuplicateKeyUpdate ( { set : { count : sql `${ ModelTpmRateLimitTable . count } + ${ usage } ` } } ) ,
5344 )
5445 } ,
0 commit comments