feat: 🛡️ add bearerToken support for model providers

microsoft · Dec 19, 2024 · 4c25fe2 · 4c25fe2
1 parent 19236d2
commit 4c25fe2
Show file tree

Hide file tree

Showing 4 changed files with 39 additions and 30 deletions.
diff --git a/docs/public/schemas/llms.json b/docs/public/schemas/llms.json
@@ -54,6 +54,10 @@
                     "prediction": {
                         "type": "boolean",
                         "description": "Indicates if 'prediction' message are supported"
+                    },
+                    "bearerToken": {
+                        "type": "boolean",
+                        "description": "Indicates if bearer token is supported"
                     }
                 },
                 "additionalProperties": false,

diff --git a/packages/core/src/constants.ts b/packages/core/src/constants.ts
@@ -239,6 +239,7 @@ export const MODEL_PROVIDERS = Object.freeze<
         topLogprobs?: boolean
         topP?: boolean
         prediction?: boolean
+        bearerToken?: boolean
     }[]
 >(CONFIGURATION_DATA.providers)
 export const MODEL_PRICINGS = Object.freeze<

diff --git a/packages/core/src/llms.json b/packages/core/src/llms.json
@@ -3,28 +3,33 @@
     "providers": [
         {
             "id": "openai",
-            "detail": "OpenAI (or compatible)"
+            "detail": "OpenAI (or compatible)",
+            "bearerToken": true
         },
         {
             "id": "github",
             "detail": "GitHub Models",
             "logprobs": false,
             "topLogprobs": false,
             "limitations": "Smaller context windows, and rate limiting",
-            "prediction": false
+            "prediction": false,
+            "bearerToken": true
         },
         {
             "id": "azure",
-            "detail": "Azure OpenAI deployment"
+            "detail": "Azure OpenAI deployment",
+            "bearerToken": false
         },
         {
             "id": "azure_serverless",
-            "detail": "Azure AI OpenAI (serverless deployments)"
+            "detail": "Azure AI OpenAI (serverless deployments)",
+            "bearerToken": false
         },
         {
             "id": "azure_serverless_models",
             "detail": "Azure AI Models (serverless deployments, not OpenAI)",
-            "prediction": false
+            "prediction": false,
+            "bearerToken": true
         },
         {
             "id": "anthropic",
@@ -39,7 +44,8 @@
             "seed": false,
             "tools": false,
             "openaiCompatibility": "https://ai.google.dev/gemini-api/docs/openai",
-            "prediction": false
+            "prediction": false,
+            "bearerToken": true
         },
         {
             "id": "huggingface",
@@ -49,7 +55,8 @@
         {
             "id": "mistal",
             "detail": "Mistral AI",
-            "prediction": false
+            "prediction": false,
+            "bearerToken": true
         },
         {
             "id": "transformers",
@@ -78,7 +85,8 @@
             "detail": "Alibaba models",
             "openaiCompatibility": "https://www.alibabacloud.com/help/en/model-studio/developer-reference/compatibility-of-openai-with-dashscope",
             "tools": false,
-            "prediction": false
+            "prediction": false,
+            "bearerToken": true
         },
         {
             "id": "llamafile",

diff --git a/packages/core/src/openai.ts b/packages/core/src/openai.ts
@@ -44,35 +44,31 @@ import { INITryParse } from "./ini"
 import { serializeChunkChoiceToLogProbs } from "./logprob"
 
 export function getConfigHeaders(cfg: LanguageModelConfiguration) {
+    const { provider } = parseModelIdentifier(cfg.model)
     let { token, type, base } = cfg
     if (type === "azure_serverless_models") {
         const keys = INITryParse(token)
         if (keys && Object.keys(keys).length > 1) token = keys[cfg.model]
     }
+    const features = MODEL_PROVIDERS.find(({ id }) => id === provider)
+    const useBearer = features?.bearerToken === true
     const isBearer = /^Bearer /i.test(cfg.token)
-    const res: Record<string, string> = {
-        // openai
-        Authorization: isBearer
-            ? token
-            : token &&
-                (type === "openai" ||
-                    type === "localai" ||
-                    type === "azure_serverless_models" ||
-                    base === OPENROUTER_API_CHAT_URL)
-              ? `Bearer ${token}`
-              : undefined,
-        // azure
-        "api-key":
-            token &&
-            !isBearer &&
-            (type === "azure" ||
-                type === "azure_serverless" ||
-                type === "alibaba")
-                ? token
-                : undefined,
+    const Authorization = isBearer
+        ? token
+        : token &&
+            (useBearer ||
+                type === "openai" ||
+                type === "localai" ||
+                type === "azure_serverless_models" ||
+                base === OPENROUTER_API_CHAT_URL)
+          ? `Bearer ${token}`
+          : undefined
+    const apiKey = Authorization ? undefined : token
+    const res: Record<string, string> = deleteUndefinedValues({
+        Authorization,
+        "api-key": apiKey,
         "User-Agent": TOOL_ID,
-    }
-    deleteUndefinedValues(res)
+    })
     return res
 }