mirror of
https://github.com/anomalyco/opencode.git
synced 2026-05-05 07:52:03 +00:00
zen: routing logic
This commit is contained in:
parent
cb425ac927
commit
3707e4a49c
7 changed files with 2723 additions and 6 deletions
|
|
@ -45,6 +45,7 @@ import { LiteData } from "@opencode-ai/console-core/lite.js"
|
|||
import { Resource } from "@opencode-ai/console-resource"
|
||||
import { i18n, type Key } from "~/i18n"
|
||||
import { localeFromRequest } from "~/lib/language"
|
||||
import { createModelTpmLimiter } from "./modelTpmLimiter"
|
||||
|
||||
type ZenData = Awaited<ReturnType<typeof ZenData.list>>
|
||||
type RetryOptions = {
|
||||
|
|
@ -121,6 +122,8 @@ export async function handler(
|
|||
const authInfo = await authenticate(modelInfo, zenApiKey)
|
||||
const billingSource = validateBilling(authInfo, modelInfo)
|
||||
logger.metric({ source: billingSource })
|
||||
const modelTpmLimiter = createModelTpmLimiter(modelInfo.providers)
|
||||
const modelTpmLimits = await modelTpmLimiter?.check()
|
||||
|
||||
const retriableRequest = async (retry: RetryOptions = { excludeProviders: [], retryCount: 0 }) => {
|
||||
const providerInfo = selectProvider(
|
||||
|
|
@ -133,6 +136,7 @@ export async function handler(
|
|||
trialProviders,
|
||||
retry,
|
||||
stickyProvider,
|
||||
modelTpmLimits,
|
||||
)
|
||||
validateModelSettings(billingSource, authInfo)
|
||||
updateProviderKey(authInfo, providerInfo)
|
||||
|
|
@ -229,6 +233,7 @@ export async function handler(
|
|||
const usageInfo = providerInfo.normalizeUsage(json.usage)
|
||||
const costInfo = calculateCost(modelInfo, usageInfo)
|
||||
await trialLimiter?.track(usageInfo)
|
||||
await modelTpmLimiter?.track(providerInfo.id, providerInfo.model, usageInfo)
|
||||
await trackUsage(sessionId, billingSource, authInfo, modelInfo, providerInfo, usageInfo, costInfo)
|
||||
await reload(billingSource, authInfo, costInfo)
|
||||
json.cost = calculateOccurredCost(billingSource, costInfo)
|
||||
|
|
@ -278,6 +283,7 @@ export async function handler(
|
|||
const usageInfo = providerInfo.normalizeUsage(usage)
|
||||
const costInfo = calculateCost(modelInfo, usageInfo)
|
||||
await trialLimiter?.track(usageInfo)
|
||||
await modelTpmLimiter?.track(providerInfo.id, providerInfo.model, usageInfo)
|
||||
await trackUsage(sessionId, billingSource, authInfo, modelInfo, providerInfo, usageInfo, costInfo)
|
||||
await reload(billingSource, authInfo, costInfo)
|
||||
const cost = calculateOccurredCost(billingSource, costInfo)
|
||||
|
|
@ -433,12 +439,16 @@ export async function handler(
|
|||
trialProviders: string[] | undefined,
|
||||
retry: RetryOptions,
|
||||
stickyProvider: string | undefined,
|
||||
modelTpmLimits: Record<string, number> | undefined,
|
||||
) {
|
||||
const modelProvider = (() => {
|
||||
// Byok is top priority b/c if user set their own API key, we should use it
|
||||
// instead of using the sticky provider for the same session
|
||||
if (authInfo?.provider?.credentials) {
|
||||
return modelInfo.providers.find((provider) => provider.id === modelInfo.byokProvider)
|
||||
}
|
||||
|
||||
// Always use the same provider for the same session
|
||||
if (stickyProvider) {
|
||||
const provider = modelInfo.providers.find((provider) => provider.id === stickyProvider)
|
||||
if (provider) return provider
|
||||
|
|
@ -451,10 +461,20 @@ export async function handler(
|
|||
}
|
||||
|
||||
if (retry.retryCount !== MAX_FAILOVER_RETRIES) {
|
||||
const providers = modelInfo.providers
|
||||
const allProviders = modelInfo.providers
|
||||
.filter((provider) => !provider.disabled)
|
||||
.filter((provider) => provider.weight !== 0)
|
||||
.filter((provider) => !retry.excludeProviders.includes(provider.id))
|
||||
.flatMap((provider) => Array<typeof provider>(provider.weight ?? 1).fill(provider))
|
||||
.filter((provider) => {
|
||||
if (!provider.tpmLimit) return true
|
||||
const usage = modelTpmLimits?.[`${provider.id}/${provider.model}`] ?? 0
|
||||
return usage < provider.tpmLimit * 1_000_000
|
||||
})
|
||||
|
||||
const topPriority = Math.min(...allProviders.map((p) => p.priority))
|
||||
const providers = allProviders
|
||||
.filter((p) => p.priority <= topPriority)
|
||||
.flatMap((provider) => Array<typeof provider>(provider.weight).fill(provider))
|
||||
|
||||
// Use the last 4 characters of session ID to select a provider
|
||||
const identifier = sessionId.length ? sessionId : ip
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue