Skip to content

Commit

Permalink
enhancement(frontend): added useVaultSecret hook and handled cloud-de…
Browse files Browse the repository at this point in the history
…v logic
  • Loading branch information
bekossy committed Dec 17, 2024
1 parent 2efc221 commit 5f4a7dc
Show file tree
Hide file tree
Showing 5 changed files with 137 additions and 62 deletions.
4 changes: 2 additions & 2 deletions agenta-web/src/components/pages/app-management/index.tsx
Original file line number Diff line number Diff line change
Expand Up @@ -8,7 +8,7 @@ import {createUseStyles} from "react-jss"
import {useAppsData} from "@/contexts/app.context"
import {useProfileData} from "@/contexts/profile.context"
import {usePostHogAg} from "@/hooks/usePostHogAg"
import {LlmProvider, getAllProviderLlmKeys} from "@/lib/helpers/llmProviders"
import {type LlmProvider} from "@/lib/helpers/llmProviders"
import {dynamicComponent, dynamicContext} from "@/lib/helpers/dynamic"
import dayjs from "dayjs"
import {useAppTheme} from "@/components/Layout/ThemeContextProvider"
Expand Down Expand Up @@ -120,7 +120,7 @@ const AppManagement: React.FC = () => {
setStatusModalOpen(true)

// attempt to create and start the template, notify user of the progress
const apiKeys = isDemo() ? secrets : getAllProviderLlmKeys()
const apiKeys = secrets
await createAndStartTemplate({
appName: newApp,
templateId: template_id,
Expand Down
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
import {useAppId} from "@/hooks/useAppId"
import {JSSTheme, LLMRunRateLimit, testset, Variant} from "@/lib/Types"
import {evaluatorConfigsAtom, evaluatorsAtom} from "@/lib/atoms/evaluation"
import {apiKeyObject, isDemo, redirectIfNoLLMKeys} from "@/lib/helpers/utils"
import {apiKeyObject, redirectIfNoLLMKeys} from "@/lib/helpers/utils"
import {fetchSingleProfile, fetchVariants} from "@/services/api"
import {createEvalutaiton} from "@/services/evaluations/api"
import {fetchTestsets} from "@/services/testsets/api"
Expand All @@ -15,7 +15,6 @@ import SelectVariantSection from "./SelectVariantSection"
import SelectEvaluatorSection from "./SelectEvaluatorSection"
import {dynamicComponent} from "@/lib/helpers/dynamic"
import {useVaultSecret} from "@/hooks/useVaultSecret"
import {getAllProviderLlmKeys} from "@/lib/helpers/llmProviders"

const AdvancedSettingsPopover: any = dynamicComponent(
"pages/evaluations/NewEvaluation/AdvancedSettingsPopover",
Expand Down Expand Up @@ -165,7 +164,7 @@ const NewEvaluationModal: React.FC<Props> = ({onSuccess, ...props}) => {
variant_ids: selectedVariantIds,
evaluators_configs: selectedEvalConfigs,
rate_limit: rateLimitValues,
lm_providers_keys: apiKeyObject(isDemo() ? secrets : getAllProviderLlmKeys()),
lm_providers_keys: apiKeyObject(secrets),
correct_answer_column: correctAnswerColumn,
})
.then(onSuccess)
Expand Down
40 changes: 15 additions & 25 deletions agenta-web/src/components/pages/settings/Secrets/Secrets.tsx
Original file line number Diff line number Diff line change
@@ -1,11 +1,5 @@
import {useVaultSecret} from "@/hooks/useVaultSecret"
import {
getLlmProviderKey,
saveLlmProviderKey,
removeSingleLlmProviderKey,
getAllProviderLlmKeys,
LlmProvider,
} from "@/lib/helpers/llmProviders"
import {getLlmProviderKey, type LlmProvider} from "@/lib/helpers/llmProviders"
import {isDemo} from "@/lib/helpers/utils"
import {Button, Input, Space, Typography, message} from "antd"
import {useEffect, useState} from "react"
Expand All @@ -14,9 +8,7 @@ const {Title, Text} = Typography

export default function Secrets() {
const {secrets, handleModifyVaultSecret, handleDeleteVaultSecret} = useVaultSecret()
const [llmProviderKeys, setLlmProviderKeys] = useState<LlmProvider[]>(
isDemo() ? [] : getAllProviderLlmKeys(),
)
const [llmProviderKeys, setLlmProviderKeys] = useState<LlmProvider[]>([])
const [messageAPI, contextHolder] = message.useMessage()

useEffect(() => {
Expand Down Expand Up @@ -61,16 +53,13 @@ export default function Secrets() {
type="primary"
disabled={key === getLlmProviderKey(title) || !key}
onClick={async () => {
if (isDemo()) {
await handleModifyVaultSecret({
name,
title,
key,
id: secretId,
})
} else {
saveLlmProviderKey(title, key)
}
await handleModifyVaultSecret({
name,
title,
key,
id: secretId,
})

messageAPI.success("The secret is saved")
}}
>
Expand All @@ -79,11 +68,12 @@ export default function Secrets() {
<Button
disabled={!Boolean(key)}
onClick={async () => {
if (isDemo() && secretId) {
await handleDeleteVaultSecret(secretId)
} else {
removeSingleLlmProviderKey(title)
}
await handleDeleteVaultSecret({
name,
id: secretId,
title,
key,
})

const newLlmProviderKeys = [...llmProviderKeys]
newLlmProviderKeys[i].key = ""
Expand Down
118 changes: 118 additions & 0 deletions agenta-web/src/hooks/useVaultSecret.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,118 @@
import {useEffect, useState} from "react"
import {
getAllProviderLlmKeys,
llmAvailableProviders,
LlmProvider,
removeSingleLlmProviderKey,
saveLlmProviderKey,
} from "@/lib/helpers/llmProviders"
import {isDemo} from "@/lib/helpers/utils"

export const useVaultSecret = () => {
const [secrets, setSecrets] = useState<LlmProvider[]>(llmAvailableProviders)

const getVaultSecrets = async () => {
try {
if (isDemo()) {
const {fetchVaultSecret} = await import("@/services/vault/api")
const data = await fetchVaultSecret()

setSecrets((prevSecret) => {
return prevSecret.map((secret) => {
const match = data.find((item: LlmProvider) => item.name === secret.name)
if (match) {
return {
...secret,
key: match.key,
id: match.id,
}
} else {
return secret
}
})
})
} else {
setSecrets(getAllProviderLlmKeys())
}
} catch (error) {
console.error(error)
}
}

useEffect(() => {
getVaultSecrets()
}, [])

const handleModifyVaultSecret = async (provider: LlmProvider) => {
try {
if (isDemo()) {
const {createVaultSecret, updateVaultSecret} = await import("@/services/vault/api")
const {SecretDTOKind, SecretDTOProvider} = await import("@/lib/types_ee")

const envNameMap: Record<string, any> = {
OPENAI_API_KEY: SecretDTOProvider.OPENAI,
COHERE_API_KEY: SecretDTOProvider.COHERE,
ANYSCALE_API_KEY: SecretDTOProvider.ANYSCALE,
DEEPINFRA_API_KEY: SecretDTOProvider.DEEPINFRA,
ALEPHALPHA_API_KEY: SecretDTOProvider.ALEPHALPHA,
GROQ_API_KEY: SecretDTOProvider.GROQ,
MISTRAL_API_KEY: SecretDTOProvider.MISTRALAI,
ANTHROPIC_API_KEY: SecretDTOProvider.ANTHROPIC,
PERPLEXITYAI_API_KEY: SecretDTOProvider.PERPLEXITYAI,
TOGETHERAI_API_KEY: SecretDTOProvider.TOGETHERAI,
OPENROUTER_API_KEY: SecretDTOProvider.OPENROUTER,
GEMINI_API_KEY: SecretDTOProvider.GEMINI,
}

const payload = {
header: {
name: provider.title,
description: "",
},
secret: {
kind: SecretDTOKind.PROVIDER_KEY,
data: {
provider: envNameMap[provider.name],
key: provider.key,
},
},
}

const findSecret = secrets.find((s) => s.name === provider.name)

if (findSecret && provider.id) {
await updateVaultSecret({secret_id: provider.id, payload})
} else {
await createVaultSecret({payload})
}

await getVaultSecrets()
} else {
saveLlmProviderKey(provider.title, provider.key)
}
} catch (error) {
console.error(error)
}
}

const handleDeleteVaultSecret = async (provider: LlmProvider) => {
try {
if (isDemo() && provider.id) {
const {deleteVaultSecret} = await import("@/services/vault/api")

await deleteVaultSecret({secret_id: provider.id})
await getVaultSecrets()
} else {
removeSingleLlmProviderKey(provider.title)
}
} catch (error) {
console.error(error)
}
}

return {
secrets,
handleModifyVaultSecret,
handleDeleteVaultSecret,
}
}
32 changes: 0 additions & 32 deletions agenta-web/src/lib/helpers/llmProviders.ts
Original file line number Diff line number Diff line change
@@ -1,5 +1,4 @@
import cloneDeep from "lodash/cloneDeep"
import {SecretDTOProvider, VaultSecretDTO} from "../types_ee"

const llmAvailableProvidersToken = "llmAvailableProvidersToken"

Expand Down Expand Up @@ -59,37 +58,6 @@ export const saveLlmProviderKey = (providerName: string, keyValue: string) => {
export const getLlmProviderKey = (providerName: string) =>
getAllProviderLlmKeys().find((item: LlmProvider) => item.title === providerName)?.key

export const transformSecret = (secrets: VaultSecretDTO[]) => {
return secrets.reduce((acc, curr) => {
const name = curr.header?.name
const {key, provider} = curr.secret.data

const envNameMap: Record<string, string> = {
[SecretDTOProvider.OPENAI]: "OPENAI_API_KEY",
[SecretDTOProvider.COHERE]: "COHERE_API_KEY",
[SecretDTOProvider.ANYSCALE]: "ANYSCALE_API_KEY",
[SecretDTOProvider.DEEPINFRA]: "DEEPINFRA_API_KEY",
[SecretDTOProvider.ALEPHALPHA]: "ALEPHALPHA_API_KEY",
[SecretDTOProvider.GROQ]: "GROQ_API_KEY",
[SecretDTOProvider.MISTRALAI]: "MISTRAL_API_KEY",
[SecretDTOProvider.ANTHROPIC]: "ANTHROPIC_API_KEY",
[SecretDTOProvider.PERPLEXITYAI]: "PERPLEXITYAI_API_KEY",
[SecretDTOProvider.TOGETHERAI]: "TOGETHERAI_API_KEY",
[SecretDTOProvider.OPENROUTER]: "OPENROUTER_API_KEY",
[SecretDTOProvider.GEMINI]: "GEMINI_API_KEY",
}

acc.push({
title: name || "",
key: key,
name: envNameMap[provider] || "",
id: curr.id,
})

return acc
}, [] as LlmProvider[])
}

export const getAllProviderLlmKeys = () => {
const providers = cloneDeep(llmAvailableProviders)
try {
Expand Down

0 comments on commit 5f4a7dc

Please sign in to comment.