zen: routing logic

2 files changed, 21 insertions, 4 deletions

diff --git a/packages/console/core/src/model.ts b/packages/console/core/src/model.ts
index 3d614d303..6281382d6 100644
--- a/packages/console/core/src/model.ts
+++ b/packages/console/core/src/model.ts
@@ -34,6 +34,8 @@ export namespace ZenData {
       z.object({
         id: z.string(),
         model: z.string(),
+        priority: z.number().optional(),
+        tpmLimit: z.number().optional(),
         weight: z.number().optional(),
         disabled: z.boolean().optional(),
         storeModel: z.string().optional(),
@@ -123,10 +125,16 @@ export namespace ZenData {
       ),
       models: (() => {
         const normalize = (model: z.infer<typeof ModelSchema>) => {
-          const composite = model.providers.find((p) => compositeProviders[p.id].length > 1)
+          const providers = model.providers.map((p) => ({
+            ...p,
+            priority: p.priority ?? Infinity,
+            weight: p.weight ?? 1,
+          }))
+          const composite = providers.find((p) => compositeProviders[p.id].length > 1)
           if (!composite)
             return {
               trialProvider: model.trialProvider ? [model.trialProvider] : undefined,
+              providers,
             }
 
           const weightMulti = compositeProviders[composite.id].length
@@ -137,17 +145,16 @@ export namespace ZenData {
               if (model.trialProvider === composite.id) return compositeProviders[composite.id].map((p) => p.id)
               return [model.trialProvider]
             })(),
-            providers: model.providers.flatMap((p) =>
+            providers: providers.flatMap((p) =>
               p.id === composite.id
                 ? compositeProviders[p.id].map((sub) => ({
                     ...p,
                     id: sub.id,
-                    weight: p.weight ?? 1,
                   }))
                 : [
                     {
                       ...p,
-                      weight: (p.weight ?? 1) * weightMulti,
+                      weight: p.weight * weightMulti,
                     },
                   ],
             ),
diff --git a/packages/console/core/src/schema/ip.sql.ts b/packages/console/core/src/schema/ip.sql.ts
index a840a78c1..830842e64 100644
--- a/packages/console/core/src/schema/ip.sql.ts
+++ b/packages/console/core/src/schema/ip.sql.ts
@@ -30,3 +30,13 @@ export const KeyRateLimitTable = mysqlTable(
   },
   (table) => [primaryKey({ columns: [table.key, table.interval] })],
 )
+
+export const ModelRateLimitTable = mysqlTable(
+  "model_rate_limit",
+  {
+    key: varchar("key", { length: 255 }).notNull(),
+    interval: varchar("interval", { length: 40 }).notNull(),
+    count: int("count").notNull(),
+  },
+  (table) => [primaryKey({ columns: [table.key, table.interval] })],
+)