zen: tpm routing
This commit is contained in:
@@ -1,10 +1,10 @@
|
|||||||
import { and, Database, eq, 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 { ModelRateLimitTable } from "@opencode-ai/console-core/schema/ip.sql.js"
|
import { ModelTpmLimitTable } from "@opencode-ai/console-core/schema/ip.sql.js"
|
||||||
import { UsageInfo } from "./provider/provider"
|
import { UsageInfo } from "./provider/provider"
|
||||||
|
|
||||||
export function createModelTpmLimiter(providers: { id: string; model: string; tpmLimit?: number }[]) {
|
export function createModelTpmLimiter(providers: { id: string; model: string; tpmLimit?: number }[]) {
|
||||||
const keys = providers.filter((p) => p.tpmLimit).map((p) => `${p.id}/${p.model}`)
|
const ids = providers.filter((p) => p.tpmLimit).map((p) => `${p.id}/${p.model}`)
|
||||||
if (keys.length === 0) return
|
if (ids.length === 0) return
|
||||||
|
|
||||||
const yyyyMMddHHmm = new Date(Date.now())
|
const yyyyMMddHHmm = new Date(Date.now())
|
||||||
.toISOString()
|
.toISOString()
|
||||||
@@ -16,30 +16,39 @@ export function createModelTpmLimiter(providers: { id: string; model: string; tp
|
|||||||
const data = await Database.use((tx) =>
|
const data = await Database.use((tx) =>
|
||||||
tx
|
tx
|
||||||
.select()
|
.select()
|
||||||
.from(ModelRateLimitTable)
|
.from(ModelTpmLimitTable)
|
||||||
.where(and(inArray(ModelRateLimitTable.key, keys), eq(ModelRateLimitTable.interval, yyyyMMddHHmm))),
|
.where(
|
||||||
|
inArray(
|
||||||
|
ModelTpmLimitTable.id,
|
||||||
|
ids.map((id) => formatId(id, yyyyMMddHHmm)),
|
||||||
|
),
|
||||||
|
),
|
||||||
)
|
)
|
||||||
|
|
||||||
// convert to map of model to count
|
// convert to map of model to count
|
||||||
return data.reduce(
|
return data.reduce(
|
||||||
(acc, curr) => {
|
(acc, curr) => {
|
||||||
acc[curr.key] = curr.count
|
acc[curr.id] = curr.count
|
||||||
return acc
|
return acc
|
||||||
},
|
},
|
||||||
{} as Record<string, number>,
|
{} as Record<string, number>,
|
||||||
)
|
)
|
||||||
},
|
},
|
||||||
track: async (id: string, model: string, usageInfo: UsageInfo) => {
|
track: async (provider: string, model: string, usageInfo: UsageInfo) => {
|
||||||
const key = `${id}/${model}`
|
const id = `${provider}/${model}`
|
||||||
if (!keys.includes(key)) return
|
if (!ids.includes(id)) return
|
||||||
const usage = usageInfo.inputTokens
|
const usage = usageInfo.inputTokens
|
||||||
if (usage <= 0) return
|
if (usage <= 0) return
|
||||||
await Database.use((tx) =>
|
await Database.use((tx) =>
|
||||||
tx
|
tx
|
||||||
.insert(ModelRateLimitTable)
|
.insert(ModelTpmLimitTable)
|
||||||
.values({ key, interval: yyyyMMddHHmm, count: usage })
|
.values({ id: formatId(id, yyyyMMddHHmm), count: usage })
|
||||||
.onDuplicateKeyUpdate({ set: { count: sql`${ModelRateLimitTable.count} + ${usage}` } }),
|
.onDuplicateKeyUpdate({ set: { count: sql`${ModelTpmLimitTable.count} + ${usage}` } }),
|
||||||
)
|
)
|
||||||
},
|
},
|
||||||
}
|
}
|
||||||
|
|
||||||
|
function formatId(id: string, yyyyMMddHHmm: string) {
|
||||||
|
return `${id.substring(0, 200)}/${yyyyMMddHHmm}`
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -0,0 +1,6 @@
|
|||||||
|
CREATE TABLE `model_tpm_limit` (
|
||||||
|
`id` varchar(255) NOT NULL,
|
||||||
|
`interval` int NOT NULL,
|
||||||
|
`count` int NOT NULL,
|
||||||
|
CONSTRAINT PRIMARY KEY(`id`,`interval`)
|
||||||
|
);
|
||||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,3 @@
|
|||||||
|
ALTER TABLE `model_tpm_limit` DROP PRIMARY KEY;--> statement-breakpoint
|
||||||
|
ALTER TABLE `model_tpm_limit` ADD PRIMARY KEY (`id`);--> statement-breakpoint
|
||||||
|
ALTER TABLE `model_tpm_limit` DROP COLUMN `interval`;
|
||||||
+2710
File diff suppressed because it is too large
Load Diff
@@ -31,6 +31,15 @@ export const KeyRateLimitTable = mysqlTable(
|
|||||||
(table) => [primaryKey({ columns: [table.key, table.interval] })],
|
(table) => [primaryKey({ columns: [table.key, table.interval] })],
|
||||||
)
|
)
|
||||||
|
|
||||||
|
export const ModelTpmLimitTable = mysqlTable(
|
||||||
|
"model_tpm_limit",
|
||||||
|
{
|
||||||
|
id: varchar("id", { length: 255 }).notNull(),
|
||||||
|
count: int("count").notNull(),
|
||||||
|
},
|
||||||
|
(table) => [primaryKey({ columns: [table.id] })],
|
||||||
|
)
|
||||||
|
|
||||||
export const ModelRateLimitTable = mysqlTable(
|
export const ModelRateLimitTable = mysqlTable(
|
||||||
"model_rate_limit",
|
"model_rate_limit",
|
||||||
{
|
{
|
||||||
|
|||||||
Reference in New Issue
Block a user