Store API keys in env to enable preset keys (#1093)

* Add API keys to profile if they exist in env

* Revert unnecessary changes

* Added optional API Keys

* add disable to profile if environment variables are set

* add disables for Azure API Keys + profile

* add more modularity to openai organization user profile disable

* refactor retrieval/keys to use reuseable functions; update model selector to 'see' environment variables

* nits for build (types in types folder)

* handle ui for profile envs

---------

Co-authored-by: Kelvin Wong <[email protected]>
Co-authored-by: Kelvin Wong <[email protected]>
Co-authored-by: Mckay Wrigley <[email protected]>

Author: 3 people

.env.local.example

@@ -7,3 +7,20 @@ SUPABASE_SERVICE_ROLE_KEY=
 
 # Ollama
 NEXT_PUBLIC_OLLAMA_URL=http://localhost:11434
+
+# API Keys (Optional: Entering an API key here overrides the API keys globally for all users.)
+OPENAI_API_KEY=
+ANTHROPIC_API_KEY=
+GOOGLE_GEMINI_API_KEY=
+MISTRAL_API_KEY=
+PERPLEXITY_API_KEY=
+AZURE_OPENAI_API_KEY=
+
+NEXT_PUBLIC_OPENAI_ORGANIZATION_ID=
+
+# Azure API Information
+AZURE_OPENAI_API_KEY=
+NEXT_PUBLIC_AZURE_OPENAI_ENDPOINT=
+NEXT_PUBLIC_AZURE_GPT_35_TURBO_ID=
+NEXT_PUBLIC_AZURE_GPT_45_VISION_ID=
+NEXT_PUBLIC_AZURE_GPT_45_TURBO_ID=

.husky/pre-commit

@@ -1,4 +1,5 @@
 #!/usr/bin/env sh
+
 . "$(dirname -- "$0")/_/husky.sh"
 
-npm run lint:fix && npm run format:write && git add .
+npm run lint:fix && npm run format:write && git add .

app/api/chat/anthropic/route.ts

@@ -1,5 +1,5 @@
 import { CHAT_SETTING_LIMITS } from "@/lib/chat-setting-limits"
-import { checkApiKey, getServerProfile } from "@/lib/server-chat-helpers"
+import { checkApiKey, getServerProfile } from "@/lib/server/server-chat-helpers"
 import { ChatSettings } from "@/types"
 import Anthropic from "@anthropic-ai/sdk"
 import { AnthropicStream, StreamingTextResponse } from "ai"

app/api/chat/azure/route.ts

@@ -1,5 +1,5 @@
 import { CHAT_SETTING_LIMITS } from "@/lib/chat-setting-limits"
-import { checkApiKey, getServerProfile } from "@/lib/server-chat-helpers"
+import { checkApiKey, getServerProfile } from "@/lib/server/server-chat-helpers"
 import { ChatAPIPayload } from "@/types"
 import { OpenAIStream, StreamingTextResponse } from "ai"
 import OpenAI from "openai"

app/api/chat/google/route.ts

@@ -1,4 +1,4 @@
-import { checkApiKey, getServerProfile } from "@/lib/server-chat-helpers"
+import { checkApiKey, getServerProfile } from "@/lib/server/server-chat-helpers"
 import { ChatSettings } from "@/types"
 import { GoogleGenerativeAI } from "@google/generative-ai"
 

app/api/chat/mistral/route.ts

@@ -1,4 +1,4 @@
-import { checkApiKey, getServerProfile } from "@/lib/server-chat-helpers"
+import { checkApiKey, getServerProfile } from "@/lib/server/server-chat-helpers"
 import { ChatSettings } from "@/types"
 
 export const runtime = "edge"

app/api/chat/openai/route.ts

@@ -1,5 +1,5 @@
 import { CHAT_SETTING_LIMITS } from "@/lib/chat-setting-limits"
-import { checkApiKey, getServerProfile } from "@/lib/server-chat-helpers"
+import { checkApiKey, getServerProfile } from "@/lib/server/server-chat-helpers"
 import { ChatSettings } from "@/types"
 import { OpenAIStream, StreamingTextResponse } from "ai"
 import OpenAI from "openai"

app/api/chat/perplexity/route.ts

@@ -1,4 +1,4 @@
-import { checkApiKey, getServerProfile } from "@/lib/server-chat-helpers"
+import { checkApiKey, getServerProfile } from "@/lib/server/server-chat-helpers"
 import { ChatSettings } from "@/types"
 
 export const runtime = "edge"

app/api/command/route.ts

@@ -1,5 +1,5 @@
 import { CHAT_SETTING_LIMITS } from "@/lib/chat-setting-limits"
-import { checkApiKey, getServerProfile } from "@/lib/server-chat-helpers"
+import { checkApiKey, getServerProfile } from "@/lib/server/server-chat-helpers"
 import OpenAI from "openai"
 
 export const runtime = "edge"

app/api/retrieval/keys/route.ts

@@ -0,0 +1,23 @@
+import { isUsingEnvironmentKey } from "@/lib/envs"
+import { createResponse } from "@/lib/server/server-utils"
+import { KeyTypeT } from "@/types/key-type"
+import { VALID_KEYS } from "@/types/valid-keys"
+
+export async function POST(request: Request) {
+  const json = await request.json()
+  const { key } = json as {
+    key: string
+  }
+
+  if (!key) {
+    return createResponse({ error: "Key type is required" }, 400)
+  }
+
+  if (!(key in VALID_KEYS)) {
+    return createResponse({ error: "Invalid key type" }, 400)
+  }
+
+  const isUsing = isUsingEnvironmentKey(key as KeyTypeT)
+
+  return createResponse({ isUsing }, 200)
+}

app/api/retrieval/process/docx/route.ts

@@ -1,6 +1,6 @@
 import { generateLocalEmbedding } from "@/lib/generate-local-embedding"
 import { processDocX } from "@/lib/retrieval/processing"
-import { checkApiKey, getServerProfile } from "@/lib/server-chat-helpers"
+import { checkApiKey, getServerProfile } from "@/lib/server/server-chat-helpers"
 import { Database } from "@/supabase/types"
 import { FileItemChunk } from "@/types"
 import { createClient } from "@supabase/supabase-js"

app/api/retrieval/process/route.ts

@@ -6,7 +6,7 @@ import {
   processPdf,
   processTxt
 } from "@/lib/retrieval/processing"
-import { checkApiKey, getServerProfile } from "@/lib/server-chat-helpers"
+import { checkApiKey, getServerProfile } from "@/lib/server/server-chat-helpers"
 import { Database } from "@/supabase/types"
 import { FileItemChunk } from "@/types"
 import { createClient } from "@supabase/supabase-js"

app/api/retrieval/retrieve/route.ts

@@ -1,5 +1,5 @@
 import { generateLocalEmbedding } from "@/lib/generate-local-embedding"
-import { checkApiKey, getServerProfile } from "@/lib/server-chat-helpers"
+import { checkApiKey, getServerProfile } from "@/lib/server/server-chat-helpers"
 import { Database } from "@/supabase/types"
 import { createClient } from "@supabase/supabase-js"
 import OpenAI from "openai"

components/models/model-option.tsx

@@ -2,7 +2,7 @@ import { ChatbotUIContext } from "@/context/context"
 import { isModelLocked } from "@/lib/is-model-locked"
 import { LLM, LLMID } from "@/types"
 import { IconLock } from "@tabler/icons-react"
-import { FC, useContext } from "react"
+import { FC, useContext, useEffect, useState } from "react"
 import { WithTooltip } from "../ui/with-tooltip"
 import { ModelIcon } from "./model-icon"
 
@@ -13,10 +13,19 @@ interface ModelOptionProps {
 
 export const ModelOption: FC<ModelOptionProps> = ({ model, onSelect }) => {
   const { profile } = useContext(ChatbotUIContext)
+  const [isLocked, setIsLocked] = useState<Boolean>(true)
 
-  if (!profile) return null
+  useEffect(() => {
+    async function setup() {
+      if (!profile) return null
+      const locked = await isModelLocked(model.provider, profile)
+
+      setIsLocked(locked)
+    }
+    setup()
+  }, [model, profile])
 
-  const isLocked = isModelLocked(model.provider, profile)
+  if (!profile) return null
 
   const handleSelectModel = () => {
     if (isLocked) return

components/models/model-select.tsx

@@ -37,6 +37,8 @@ export const ModelSelect: FC<ModelSelectProps> = ({
   const [search, setSearch] = useState("")
   const [tab, setTab] = useState<"hosted" | "local">("hosted")
 
+  const [isLocked, setIsLocked] = useState<Boolean>(true)
+
   useEffect(() => {
     if (isOpen) {
       setTimeout(() => {
@@ -45,6 +47,17 @@ export const ModelSelect: FC<ModelSelectProps> = ({
     }
   }, [isOpen])
 
+  useEffect(() => {
+    const checkModelLock = async () => {
+      if (SELECTED_MODEL && profile) {
+        const locked = await isModelLocked(SELECTED_MODEL.provider, profile)
+        setIsLocked(locked)
+      }
+    }
+
+    checkModelLock()
+  }, [profile])
+
   const handleSelectModel = (modelId: LLMID) => {
     onSelectModel(modelId)
     setIsOpen(false)
@@ -71,7 +84,6 @@ export const ModelSelect: FC<ModelSelectProps> = ({
   if (!SELECTED_MODEL) return null
   if (!profile) return null
 
-  const isLocked = isModelLocked(SELECTED_MODEL.provider, profile)
   const usingLocalModels = availableLocalModels.length > 0
 
   return (