Skip to content

Commit

Permalink
synced changes
Browse files Browse the repository at this point in the history
  • Loading branch information
bekossy committed Dec 26, 2024
1 parent 2b57096 commit 5dc60f1
Show file tree
Hide file tree
Showing 7 changed files with 248 additions and 62 deletions.
6 changes: 4 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 "@/lib/helpers/analytics/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 All @@ -17,6 +17,7 @@ import GetStartedSection from "./components/GetStartedSection"
import ApplicationManagementSection from "./components/ApplicationManagementSection"
import ResultComponent from "@/components/ResultComponent/ResultComponent"
import {useProjectData} from "@/contexts/project.context"
import {useVaultSecret} from "@/hooks/useVaultSecret"

const CreateAppStatusModal: any = dynamicComponent(
"pages/app-management/modals/CreateAppStatusModal",
Expand Down Expand Up @@ -83,6 +84,7 @@ const AppManagement: React.FC = () => {
details: undefined,
appId: undefined,
})
const {secrets} = useVaultSecret()

const {project} = useProjectData()
const [useOrgData, setUseOrgData] = useState<Function>(() => () => "")
Expand Down Expand Up @@ -118,7 +120,7 @@ const AppManagement: React.FC = () => {
setStatusModalOpen(true)

// attempt to create and start the template, notify user of the progress
const apiKeys = getAllProviderLlmKeys()
const apiKeys = secrets
await createAndStartTemplate({
appName: newApp,
templateId: template_id,
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,7 @@ import SelectTestsetSection from "./SelectTestsetSection"
import SelectVariantSection from "./SelectVariantSection"
import SelectEvaluatorSection from "./SelectEvaluatorSection"
import {dynamicComponent} from "@/lib/helpers/dynamic"
import {useVaultSecret} from "@/hooks/useVaultSecret"

const AdvancedSettingsPopover: any = dynamicComponent(
"pages/evaluations/NewEvaluation/AdvancedSettingsPopover",
Expand Down Expand Up @@ -73,6 +74,7 @@ const NewEvaluationModal: React.FC<Props> = ({onSuccess, ...props}) => {
const [selectedTestsetId, setSelectedTestsetId] = useState("")
const [selectedVariantIds, setSelectedVariantIds] = useState<string[]>([])
const [selectedEvalConfigs, setSelectedEvalConfigs] = useState<string[]>([])
const {secrets} = useVaultSecret()

const [activePanel, setActivePanel] = useState<string | null>("testsetPanel")
const handlePanelChange = (key: string | string[]) => {
Expand Down Expand Up @@ -163,7 +165,7 @@ const NewEvaluationModal: React.FC<Props> = ({onSuccess, ...props}) => {
variant_ids: selectedVariantIds,
evaluators_configs: selectedEvalConfigs,
rate_limit: rateLimitValues,
lm_providers_keys: apiKeyObject(),
lm_providers_keys: apiKeyObject(secrets),
correct_answer_column: correctAnswerColumn,
})
.then(onSuccess)
Expand Down
128 changes: 78 additions & 50 deletions agenta-web/src/components/pages/settings/Secrets/Secrets.tsx
Original file line number Diff line number Diff line change
@@ -1,19 +1,25 @@
import {
getLlmProviderKey,
saveLlmProviderKey,
removeSingleLlmProviderKey,
getAllProviderLlmKeys,
LlmProvider,
} from "@/lib/helpers/llmProviders"
import {useVaultSecret} from "@/hooks/useVaultSecret"
import {type LlmProvider} from "@/lib/helpers/llmProviders"
import {Button, Input, Space, Typography, message} from "antd"
import {useState} from "react"
import {useEffect, useState} from "react"

const {Title, Text} = Typography

export default function Secrets() {
const [llmProviderKeys, setLlmProviderKeys] = useState(getAllProviderLlmKeys())
const {secrets, handleModifyVaultSecret, handleDeleteVaultSecret} = useVaultSecret()
const [llmProviderKeys, setLlmProviderKeys] = useState<LlmProvider[]>([])
const [loadingSecrets, setLoadingSecrets] = useState<Record<string, boolean>>({})
const [messageAPI, contextHolder] = message.useMessage()

useEffect(() => {
setLlmProviderKeys(secrets)
}, [secrets])

const setSecretLoading = (id: string | undefined, isLoading: boolean) => {
if (!id) return
setLoadingSecrets((prev) => ({...prev, [id]: isLoading}))
}

return (
<div data-cy="secrets">
{contextHolder}
Expand All @@ -30,47 +36,69 @@ export default function Secrets() {
<Title level={5}>Available Providers</Title>

<div>
{llmProviderKeys.map(({title, key}: LlmProvider, i: number) => (
<Space direction="horizontal" key={i} className="mb-2 ml-2">
<Input.Password
data-cy="openai-api-input"
value={key}
onChange={(e) => {
const newLlmProviderKeys = [...llmProviderKeys]
newLlmProviderKeys[i].key = e.target.value
setLlmProviderKeys(newLlmProviderKeys)
}}
addonBefore={`${title}`}
visibilityToggle={false}
className={"w-[420px]"}
/>
<Button
data-cy="openai-api-save"
type="primary"
disabled={key === getLlmProviderKey(title) || !key}
onClick={() => {
saveLlmProviderKey(title, key)
messageAPI.success("The secret is saved")
}}
>
Save
</Button>
<Button
disabled={!Boolean(key)}
onClick={() => {
removeSingleLlmProviderKey(title)

const newLlmProviderKeys = [...llmProviderKeys]
newLlmProviderKeys[i].key = ""
setLlmProviderKeys(newLlmProviderKeys)

messageAPI.warning("The secret is deleted")
}}
>
Delete
</Button>
</Space>
))}
{llmProviderKeys.map(
({name, title, key, id: secretId}: LlmProvider, i: number) => (
<Space direction="horizontal" key={i} className="mb-2 ml-2">
<Input.Password
data-cy="openai-api-input"
value={key}
onChange={(e) => {
const newLlmProviderKeys = [...llmProviderKeys]
newLlmProviderKeys[i].key = e.target.value
setLlmProviderKeys(newLlmProviderKeys)
}}
addonBefore={`${title}`}
visibilityToggle={false}
className={"w-[420px]"}
/>
<Button
data-cy="openai-api-save"
type="primary"
disabled={!key}
loading={loadingSecrets[secretId || ""] === true}
onClick={async () => {
try {
setSecretLoading(secretId, true)
await handleModifyVaultSecret({
name,
title,
key,
id: secretId,
})
messageAPI.success("The secret is saved")
} finally {
setSecretLoading(secretId, false)
}
}}
>
Save
</Button>
<Button
disabled={!Boolean(key)}
loading={loadingSecrets[secretId || ""] === true}
onClick={async () => {
try {
setSecretLoading(secretId, true)
await handleDeleteVaultSecret({
name,
id: secretId,
title,
key,
})
const newLlmProviderKeys = [...llmProviderKeys]
newLlmProviderKeys[i].key = ""
setLlmProviderKeys(newLlmProviderKeys)
messageAPI.warning("The secret is deleted")
} finally {
setSecretLoading(secretId, false)
}
}}
>
Delete
</Button>
</Space>
),
)}
</div>
</div>
</div>
Expand Down
151 changes: 151 additions & 0 deletions agenta-web/src/hooks/useVaultSecret.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,151 @@
import {useEffect, useRef, useState} from "react"
import {
getAllProviderLlmKeys,
llmAvailableProviders,
llmAvailableProvidersToken,
LlmProvider,
removeSingleLlmProviderKey,
saveLlmProviderKey,
} from "@/lib/helpers/llmProviders"
import {isDemo} from "@/lib/helpers/utils"
import {dynamicLib, dynamicService} from "@/lib/helpers/dynamic"

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

const getVaultSecrets = async () => {
try {
if (isDemo()) {
const {fetchVaultSecret} = await dynamicService("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 migrateProviderKeys = async () => {
try {
const localStorageProviders = localStorage.getItem(llmAvailableProvidersToken)

if (localStorageProviders) {
const providers = JSON.parse(localStorageProviders)

for (const provider of providers) {
if (provider.key) {
await handleModifyVaultSecret(provider as LlmProvider)
}
}

localStorage.setItem(`${llmAvailableProvidersToken}Backup`, localStorageProviders)

localStorage.removeItem(llmAvailableProvidersToken)
}
} catch (error) {
console.error(error)
}
}

useEffect(() => {
if (shouldRunMigration.current) {
shouldRunMigration.current = false
if (isDemo()) {
migrateProviderKeys()
}
}
}, [])

const handleModifyVaultSecret = async (provider: LlmProvider) => {
try {
if (isDemo()) {
const {updateVaultSecret, createVaultSecret} = await dynamicService("vault/api")
const {SecretDTOProvider, SecretDTOKind} = await dynamicLib("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 dynamicService("vault/api")
await deleteVaultSecret({secret_id: provider.id})
await getVaultSecrets()
} else {
removeSingleLlmProviderKey(provider.title)
}
} catch (error) {
console.error(error)
}
}

return {
secrets,
handleModifyVaultSecret,
handleDeleteVaultSecret,
}
}
8 changes: 8 additions & 0 deletions agenta-web/src/lib/helpers/dynamic.ts
Original file line number Diff line number Diff line change
Expand Up @@ -30,3 +30,11 @@ export async function dynamicService(path: string, fallback?: any) {
return fallback
}
}

export async function dynamicLib(path: string, fallback?: any) {
try {
return await import(`@/lib/${path}`)
} catch (error) {
return fallback
}
}
7 changes: 2 additions & 5 deletions agenta-web/src/lib/helpers/llmProviders.ts
Original file line number Diff line number Diff line change
@@ -1,12 +1,12 @@
import cloneDeep from "lodash/cloneDeep"
import {camelToSnake} from "./utils"

const llmAvailableProvidersToken = "llmAvailableProvidersToken"
export const llmAvailableProvidersToken = "llmAvailableProvidersToken"

export type LlmProvider = {
title: string
key: string
name: string
id?: string
}

export const llmAvailableProviders: LlmProvider[] = [
Expand Down Expand Up @@ -55,9 +55,6 @@ export const saveLlmProviderKey = (providerName: string, keyValue: string) => {
localStorage.setItem(llmAvailableProvidersToken, JSON.stringify(keys))
}

export const getLlmProviderKey = (providerName: string) =>
getAllProviderLlmKeys().find((item: LlmProvider) => item.title === providerName)?.key

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

0 comments on commit 5dc60f1

Please sign in to comment.