From 81e9880d9d17f5168ee3b0bc53b7eec03b902bec Mon Sep 17 00:00:00 2001 From: Alan Hagedorn <66842958+avhagedorn@users.noreply.github.com> Date: Sat, 13 Apr 2024 12:08:36 -0700 Subject: [PATCH] Add names to API Keys (#63) --- backend/ee/danswer/auth/api_key.py | 1 + backend/ee/danswer/db/api_key.py | 27 ++++- backend/ee/danswer/server/api_key/api.py | 17 ++- backend/ee/danswer/server/api_key/models.py | 5 + .../ee/admin/api-key/DanswerApiKeyForm.tsx | 108 ++++++++++++++++++ web/src/app/ee/admin/api-key/lib.ts | 37 ++++++ web/src/app/ee/admin/api-key/page.tsx | 102 +++++++++++------ web/src/app/ee/admin/api-key/types.ts | 11 ++ 8 files changed, 267 insertions(+), 41 deletions(-) create mode 100644 backend/ee/danswer/server/api_key/models.py create mode 100644 web/src/app/ee/admin/api-key/DanswerApiKeyForm.tsx create mode 100644 web/src/app/ee/admin/api-key/lib.ts create mode 100644 web/src/app/ee/admin/api-key/types.ts diff --git a/backend/ee/danswer/auth/api_key.py b/backend/ee/danswer/auth/api_key.py index d1ce3b016..9f9d2bac4 100644 --- a/backend/ee/danswer/auth/api_key.py +++ b/backend/ee/danswer/auth/api_key.py @@ -18,6 +18,7 @@ class ApiKeyDescriptor(BaseModel): api_key_id: int api_key_display: str api_key: str | None = None # only present on initial creation + api_key_name: str | None = None user_id: uuid.UUID diff --git a/backend/ee/danswer/db/api_key.py b/backend/ee/danswer/db/api_key.py index cc9c37924..8b734f6f7 100644 --- a/backend/ee/danswer/db/api_key.py +++ b/backend/ee/danswer/db/api_key.py @@ -12,6 +12,7 @@ from ee.danswer.auth.api_key import build_displayable_api_key from ee.danswer.auth.api_key import generate_api_key from ee.danswer.auth.api_key import hash_api_key from ee.danswer.db.constants import DANSWER_API_KEY_DUMMY_EMAIL_DOMAIN +from ee.danswer.server.api_key.models import APIKeyArgs _DANSWER_API_KEY = "danswer_api_key" @@ -28,6 +29,7 @@ def fetch_api_keys(db_session: Session) -> list[ApiKeyDescriptor]: ApiKeyDescriptor( api_key_id=api_key.id, api_key_display=api_key.api_key_display, + api_key_name=api_key.name, user_id=api_key.user_id, ) for api_key in api_keys @@ -44,7 +46,9 @@ def fetch_user_for_api_key(hashed_api_key: str, db_session: Session) -> User | N return db_session.scalar(select(User).where(User.id == api_key.user_id)) # type: ignore -def insert_api_key(db_session: Session, user_id: uuid.UUID | None) -> ApiKeyDescriptor: +def insert_api_key( + db_session: Session, api_key_args: APIKeyArgs, user_id: uuid.UUID | None +) -> ApiKeyDescriptor: std_password_helper = PasswordHelper() api_key = generate_api_key() api_key_user_id = uuid.uuid4() @@ -62,6 +66,7 @@ def insert_api_key(db_session: Session, user_id: uuid.UUID | None) -> ApiKeyDesc db_session.add(api_key_user_row) api_key_row = ApiKey( + name=api_key_args.name, hashed_api_key=hash_api_key(api_key), api_key_display=build_displayable_api_key(api_key), user_id=api_key_user_id, @@ -74,10 +79,29 @@ def insert_api_key(db_session: Session, user_id: uuid.UUID | None) -> ApiKeyDesc api_key_id=api_key_row.id, api_key_display=api_key_row.api_key_display, api_key=api_key, + api_key_name=api_key_args.name, user_id=api_key_user_id, ) +def update_api_key( + db_session: Session, api_key_id: int, api_key_args: APIKeyArgs +) -> ApiKeyDescriptor: + existing_api_key = db_session.scalar(select(ApiKey).where(ApiKey.id == api_key_id)) + if existing_api_key is None: + raise ValueError(f"API key with id {api_key_id} does not exist") + + existing_api_key.name = api_key_args.name + db_session.commit() + + return ApiKeyDescriptor( + api_key_id=existing_api_key.id, + api_key_display=existing_api_key.api_key_display, + api_key_name=api_key_args.name, + user_id=existing_api_key.user_id, + ) + + def regenerate_api_key(db_session: Session, api_key_id: int) -> ApiKeyDescriptor: """NOTE: currently, any admin can regenerate any API key.""" existing_api_key = db_session.scalar(select(ApiKey).where(ApiKey.id == api_key_id)) @@ -93,6 +117,7 @@ def regenerate_api_key(db_session: Session, api_key_id: int) -> ApiKeyDescriptor api_key_id=existing_api_key.id, api_key_display=existing_api_key.api_key_display, api_key=new_api_key, + api_key_name=existing_api_key.name, user_id=existing_api_key.user_id, ) diff --git a/backend/ee/danswer/server/api_key/api.py b/backend/ee/danswer/server/api_key/api.py index f23321280..d9a62d065 100644 --- a/backend/ee/danswer/server/api_key/api.py +++ b/backend/ee/danswer/server/api_key/api.py @@ -10,6 +10,8 @@ from ee.danswer.db.api_key import fetch_api_keys from ee.danswer.db.api_key import insert_api_key from ee.danswer.db.api_key import regenerate_api_key from ee.danswer.db.api_key import remove_api_key +from ee.danswer.db.api_key import update_api_key +from ee.danswer.server.api_key.models import APIKeyArgs router = APIRouter(prefix="/admin/api-key") @@ -24,13 +26,14 @@ def list_api_keys( @router.post("") def create_api_key( + api_key_args: APIKeyArgs, user: db_models.User | None = Depends(current_admin_user), db_session: Session = Depends(get_session), ) -> ApiKeyDescriptor: - return insert_api_key(db_session, user.id if user else None) + return insert_api_key(db_session, api_key_args, user.id if user else None) -@router.patch("/{api_key_id}") +@router.post("/{api_key_id}/regenerate") def regenerate_existing_api_key( api_key_id: int, _: db_models.User | None = Depends(current_admin_user), @@ -39,6 +42,16 @@ def regenerate_existing_api_key( return regenerate_api_key(db_session, api_key_id) +@router.patch("/{api_key_id}") +def update_existing_api_key( + api_key_id: int, + api_key_args: APIKeyArgs, + _: db_models.User | None = Depends(current_admin_user), + db_session: Session = Depends(get_session), +) -> ApiKeyDescriptor: + return update_api_key(db_session, api_key_id, api_key_args) + + @router.delete("/{api_key_id}") def delete_api_key( api_key_id: int, diff --git a/backend/ee/danswer/server/api_key/models.py b/backend/ee/danswer/server/api_key/models.py new file mode 100644 index 000000000..a26e27512 --- /dev/null +++ b/backend/ee/danswer/server/api_key/models.py @@ -0,0 +1,5 @@ +from pydantic import BaseModel + + +class APIKeyArgs(BaseModel): + name: str | None = None diff --git a/web/src/app/ee/admin/api-key/DanswerApiKeyForm.tsx b/web/src/app/ee/admin/api-key/DanswerApiKeyForm.tsx new file mode 100644 index 000000000..5fb162cff --- /dev/null +++ b/web/src/app/ee/admin/api-key/DanswerApiKeyForm.tsx @@ -0,0 +1,108 @@ +import { Form, Formik } from "formik"; +import { PopupSpec } from "@/components/admin/connectors/Popup"; +import { TextFormField } from "@/components/admin/connectors/Field"; +import { createApiKey, updateApiKey } from "./lib"; +import { Modal } from "@/components/Modal"; +import { XIcon } from "@/components/icons/icons"; +import { Button, Divider, Text } from "@tremor/react"; + +interface DanswerApiKeyFormProps { + onClose: () => void; + setPopup: (popupSpec: PopupSpec | null) => void; + onCreateApiKey: (apiKey: APIKey) => void; + apiKey?: APIKey; +} + +export const DanswerApiKeyForm = ({ + onClose, + setPopup, + onCreateApiKey, + apiKey, +}: DanswerApiKeyFormProps) => { + const isUpdate = apiKey !== undefined; + + return ( + +
+

+ {isUpdate ? "Update API Key" : "Create a new API Key"} +
+ +
+

+ + + + { + formikHelpers.setSubmitting(true); + let response; + if (isUpdate) { + response = await updateApiKey(apiKey.api_key_id, values); + } else { + response = await createApiKey(values); + } + formikHelpers.setSubmitting(false); + if (response.ok) { + setPopup({ + message: isUpdate + ? "Successfully updated API key!" + : "Successfully created API key!", + type: "success", + }); + if (!isUpdate) { + onCreateApiKey(await response.json()); + } + onClose(); + } else { + const responseJson = await response.json(); + const errorMsg = responseJson.detail || responseJson.message; + setPopup({ + message: isUpdate + ? `Error updating API key - ${errorMsg}` + : `Error creating API key - ${errorMsg}`, + type: "error", + }); + } + }} + > + {({ isSubmitting, values, setFieldValue }) => ( +
+ + Choose a memorable name for your API key. This is optional and + can be added or changed later! + + + + +
+ +
+ + )} +
+
+ + ); +}; diff --git a/web/src/app/ee/admin/api-key/lib.ts b/web/src/app/ee/admin/api-key/lib.ts new file mode 100644 index 000000000..641cfb074 --- /dev/null +++ b/web/src/app/ee/admin/api-key/lib.ts @@ -0,0 +1,37 @@ +export const createApiKey = async (apiKeyArgs: APIKeyArgs) => { + return fetch("/api/admin/api-key", { + method: "POST", + headers: { + "Content-Type": "application/json", + }, + body: JSON.stringify(apiKeyArgs), + }); +}; + +export const regenerateApiKey = async (apiKey: APIKey) => { + return fetch(`/api/admin/api-key/${apiKey.api_key_id}/regenerate`, { + method: "POST", + headers: { + "Content-Type": "application/json", + }, + }); +}; + +export const updateApiKey = async ( + apiKeyId: number, + apiKeyArgs: APIKeyArgs +) => { + return fetch(`/api/admin/api-key/${apiKeyId}`, { + method: "PATCH", + headers: { + "Content-Type": "application/json", + }, + body: JSON.stringify(apiKeyArgs), + }); +}; + +export const deleteApiKey = async (apiKeyId: number) => { + return fetch(`/api/admin/api-key/${apiKeyId}`, { + method: "DELETE", + }); +}; diff --git a/web/src/app/ee/admin/api-key/page.tsx b/web/src/app/ee/admin/api-key/page.tsx index f0ac7e1a7..04ce59ebe 100644 --- a/web/src/app/ee/admin/api-key/page.tsx +++ b/web/src/app/ee/admin/api-key/page.tsx @@ -21,16 +21,11 @@ import { usePopup } from "@/components/admin/connectors/Popup"; import { useState } from "react"; import { Table } from "@tremor/react"; import { DeleteButton } from "@/components/DeleteButton"; -import { FiCopy, FiRefreshCw, FiX } from "react-icons/fi"; +import { FiCopy, FiEdit, FiRefreshCw, FiX } from "react-icons/fi"; import { Modal } from "@/components/Modal"; import { Spinner } from "@/components/Spinner"; - -interface APIKey { - api_key_id: number; - api_key_display: string; - api_key: string | null; // only present on initial creation - user_id: string; -} +import { deleteApiKey, regenerateApiKey } from "./lib"; +import { DanswerApiKeyForm } from "./DanswerApiKeyForm"; const API_KEY_TEXT = ` API Keys allow you to access Danswer APIs programmatically. Click the button below to generate a new API Key. @@ -97,6 +92,13 @@ function Main() { const [fullApiKey, setFullApiKey] = useState(null); const [keyIsGenerating, setKeyIsGenerating] = useState(false); + const [showCreateUpdateForm, setShowCreateUpdateForm] = useState(false); + const [selectedApiKey, setSelectedApiKey] = useState(); + + const handleEdit = (apiKey: APIKey) => { + setSelectedApiKey(apiKey); + setShowCreateUpdateForm(true); + }; if (isLoading) { return ; @@ -116,24 +118,7 @@ function Main() { color="green" size="xs" className="mt-3" - onClick={async () => { - setKeyIsGenerating(true); - const response = await fetch("/api/admin/api-key", { - method: "POST", - }); - setKeyIsGenerating(false); - if (!response.ok) { - const errorMsg = await response.text(); - setPopup({ - type: "error", - message: `Failed to create API Key: ${errorMsg}`, - }); - return; - } - const newKey = (await response.json()) as APIKey; - setFullApiKey(newKey.api_key); - mutate("/api/admin/api-key"); - }} + onClick={() => setShowCreateUpdateForm(true)} > Create API Key @@ -145,12 +130,29 @@ function Main() { {popup} {API_KEY_TEXT} {newApiKeyButton} + + {showCreateUpdateForm && ( + { + setFullApiKey(apiKey.api_key); + }} + onClose={() => { + setShowCreateUpdateForm(false); + setSelectedApiKey(undefined); + mutate("/api/admin/api-key"); + }} + setPopup={setPopup} + apiKey={selectedApiKey} + /> + )} ); } return (
+ {popup} + {fullApiKey && ( + Name API Key Regenerate Delete @@ -177,6 +180,24 @@ function Main() { {apiKeys.map((apiKey) => ( + +
handleEdit(apiKey)} + > + + {apiKey.api_key_name || null} +
+ {apiKey.api_key_display} @@ -194,12 +215,7 @@ function Main() { text-sm`} onClick={async () => { setKeyIsGenerating(true); - const response = await fetch( - `/api/admin/api-key/${apiKey.api_key_id}`, - { - method: "PATCH", - } - ); + const response = await regenerateApiKey(apiKey); setKeyIsGenerating(false); if (!response.ok) { const errorMsg = await response.text(); @@ -221,12 +237,7 @@ function Main() { { - const response = await fetch( - `/api/admin/api-key/${apiKey.api_key_id}`, - { - method: "DELETE", - } - ); + const response = await deleteApiKey(apiKey.api_key_id); if (!response.ok) { const errorMsg = await response.text(); setPopup({ @@ -243,6 +254,21 @@ function Main() { ))} + + {showCreateUpdateForm && ( + { + setFullApiKey(apiKey.api_key); + }} + onClose={() => { + setShowCreateUpdateForm(false); + setSelectedApiKey(undefined); + mutate("/api/admin/api-key"); + }} + setPopup={setPopup} + apiKey={selectedApiKey} + /> + )}
); } diff --git a/web/src/app/ee/admin/api-key/types.ts b/web/src/app/ee/admin/api-key/types.ts new file mode 100644 index 000000000..01984c776 --- /dev/null +++ b/web/src/app/ee/admin/api-key/types.ts @@ -0,0 +1,11 @@ +interface APIKey { + api_key_id: number; + api_key_display: string; + api_key: string | null; + api_key_name: string | null; + user_id: string; +} + +interface APIKeyArgs { + name?: string; +}