wip: zen

packages/console/app/src/routes/zen/util/handler.ts

@@ -81,12 +81,13 @@ export async function handler(
     const isTrial = await trialLimiter?.isTrial()
     const rateLimiter = createRateLimiter(modelInfo.rateLimit, ip)
     await rateLimiter?.check()
-    const stickyTracker = createStickyTracker(modelInfo.stickyProvider ?? false, sessionId)
+    const stickyTracker = createStickyTracker(modelInfo.stickyProvider, sessionId)
     const stickyProvider = await stickyTracker?.get()
     const authInfo = await authenticate(modelInfo)
 
     const retriableRequest = async (retry: RetryOptions = { excludeProviders: [], retryCount: 0 }) => {
       const providerInfo = selectProvider(
+        model,
         zenData,
         authInfo,
         modelInfo,
@@ -101,7 +102,7 @@ export async function handler(
       logger.metric({ provider: providerInfo.id })
 
       const startTimestamp = Date.now()
-      const reqUrl = providerInfo.modifyUrl(providerInfo.api, providerInfo.model, isStream)
+      const reqUrl = providerInfo.modifyUrl(providerInfo.api, isStream)
       const reqBody = JSON.stringify(
         providerInfo.modifyBody({
           ...createBodyConverter(opts.format, providerInfo.format)(body),
@@ -135,7 +136,7 @@ export async function handler(
         // ie. openai 404 error: Item with id 'msg_0ead8b004a3b165d0069436a6b6834819896da85b63b196a3f' not found.
         res.status !== 404 &&
         // ie. cannot change codex model providers mid-session
-        !modelInfo.stickyProvider &&
+        modelInfo.stickyProvider !== "strict" &&
         modelInfo.fallbackProvider &&
         providerInfo.id !== modelInfo.fallbackProvider
       ) {
@@ -194,17 +195,19 @@ export async function handler(
     // Handle streaming response
     const streamConverter = createStreamPartConverter(providerInfo.format, opts.format)
     const usageParser = providerInfo.createUsageParser()
+    const binaryDecoder = providerInfo.createBinaryStreamDecoder()
     const stream = new ReadableStream({
       start(c) {
         const reader = res.body?.getReader()
         const decoder = new TextDecoder()
         const encoder = new TextEncoder()
+
         let buffer = ""
         let responseLength = 0
 
         function pump(): Promise<void> {
           return (
-            reader?.read().then(async ({ done, value }) => {
+            reader?.read().then(async ({ done, value: rawValue }) => {
               if (done) {
                 logger.metric({
                   response_length: responseLength,
@@ -230,6 +233,10 @@ export async function handler(
                   "timestamp.first_byte": now,
                 })
               }
+
+              const value = binaryDecoder ? binaryDecoder(rawValue) : rawValue
+              if (!value) return
+
               responseLength += value.length
               buffer += decoder.decode(value, { stream: true })
               dataDumper?.provideStream(buffer)
@@ -331,6 +338,7 @@ export async function handler(
   }
 
   function selectProvider(
+    reqModel: string,
     zenData: ZenData,
     authInfo: AuthInfo,
     modelInfo: ModelInfo,
@@ -339,7 +347,7 @@ export async function handler(
     retry: RetryOptions,
     stickyProvider: string | undefined,
   ) {
-    const provider = (() => {
+    const modelProvider = (() => {
       if (authInfo?.provider?.credentials) {
         return modelInfo.providers.find((provider) => provider.id === modelInfo.byokProvider)
       }
@@ -372,18 +380,19 @@ export async function handler(
       return providers[index || 0]
     })()
 
-    if (!provider) throw new ModelError("No provider available")
-    if (!(provider.id in zenData.providers)) throw new ModelError(`Provider ${provider.id} not supported`)
+    if (!modelProvider) throw new ModelError("No provider available")
+    if (!(modelProvider.id in zenData.providers)) throw new ModelError(`Provider ${modelProvider.id} not supported`)
 
     return {
-      ...provider,
-      ...zenData.providers[provider.id],
+      ...modelProvider,
+      ...zenData.providers[modelProvider.id],
       ...(() => {
-        const format = zenData.providers[provider.id].format
-        if (format === "anthropic") return anthropicHelper
-        if (format === "google") return googleHelper
-        if (format === "openai") return openaiHelper
-        return oaCompatHelper
+        const format = zenData.providers[modelProvider.id].format
+        const providerModel = modelProvider.model
+        if (format === "anthropic") return anthropicHelper({ reqModel, providerModel })
+        if (format === "google") return googleHelper({ reqModel, providerModel })
+        if (format === "openai") return openaiHelper({ reqModel, providerModel })
+        return oaCompatHelper({ reqModel, providerModel })
       })(),
     }
   }