Add names to API Keys (#63)

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
 

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,
     )
 

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,

backend/ee/danswer/server/api_key/models.py

@@ -0,0 +1,5 @@
+from pydantic import BaseModel
+
+
+class APIKeyArgs(BaseModel):
+    name: str | None = None

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 (
+    <Modal onOutsideClick={onClose} width="w-2/6">
+      <div className="px-8 py-6 bg-background">
+        <h2 className="text-xl font-bold flex">
+          {isUpdate ? "Update API Key" : "Create a new API Key"}
+          <div
+            onClick={onClose}
+            className="ml-auto hover:bg-hover p-1.5 rounded"
+          >
+            <XIcon
+              size={20}
+              className="my-auto flex flex-shrink-0 cursor-pointer"
+            />
+          </div>
+        </h2>
+
+        <Divider />
+
+        <Formik
+          initialValues={{
+            name: apiKey?.api_key_name || "",
+          }}
+          onSubmit={async (values, formikHelpers) => {
+            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 }) => (
+            <Form>
+              <Text className="mb-4 text-lg">
+                Choose a memorable name for your API key. This is optional and
+                can be added or changed later!
+              </Text>
+
+              <TextFormField
+                name="name"
+                label="Name (optional):"
+                autoCompleteDisabled={true}
+              />
+
+              <div className="flex">
+                <Button
+                  type="submit"
+                  size="xs"
+                  color="green"
+                  disabled={isSubmitting}
+                  className="mx-auto w-64"
+                >
+                  {isUpdate ? "Update!" : "Create!"}
+                </Button>
+              </div>
+            </Form>
+          )}
+        </Formik>
+      </div>
+    </Modal>
+  );
+};

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",
+  });
+};

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";
-
+import { deleteApiKey, regenerateApiKey } from "./lib";
-interface APIKey {
+import { DanswerApiKeyForm } from "./DanswerApiKeyForm";
-  api_key_id: number;
-  api_key_display: string;
-  api_key: string | null; // only present on initial creation
-  user_id: string;
-}
 
 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<string | null>(null);
   const [keyIsGenerating, setKeyIsGenerating] = useState(false);
+  const [showCreateUpdateForm, setShowCreateUpdateForm] = useState(false);
+  const [selectedApiKey, setSelectedApiKey] = useState<APIKey | undefined>();
+
+  const handleEdit = (apiKey: APIKey) => {
+    setSelectedApiKey(apiKey);
+    setShowCreateUpdateForm(true);
+  };
 
   if (isLoading) {
     return <ThreeDotsLoader />;
@@ -116,24 +118,7 @@ function Main() {
       color="green"
       size="xs"
       className="mt-3"
-      onClick={async () => {
+      onClick={() => setShowCreateUpdateForm(true)}
-        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");
-      }}
     >
       Create API Key
     </Button>
@@ -145,12 +130,29 @@ function Main() {
         {popup}
         <Text>{API_KEY_TEXT}</Text>
         {newApiKeyButton}
+
+        {showCreateUpdateForm && (
+          <DanswerApiKeyForm
+            onCreateApiKey={(apiKey) => {
+              setFullApiKey(apiKey.api_key);
+            }}
+            onClose={() => {
+              setShowCreateUpdateForm(false);
+              setSelectedApiKey(undefined);
+              mutate("/api/admin/api-key");
+            }}
+            setPopup={setPopup}
+            apiKey={selectedApiKey}
+          />
+        )}
       </div>
     );
   }
 
   return (
     <div>
+      {popup}
+
       {fullApiKey && (
         <NewApiKeyModal
           apiKey={fullApiKey}
@@ -169,6 +171,7 @@ function Main() {
       <Table className="overflow-visible">
         <TableHead>
           <TableRow>
+            <TableHeaderCell>Name</TableHeaderCell>
             <TableHeaderCell>API Key</TableHeaderCell>
             <TableHeaderCell>Regenerate</TableHeaderCell>
             <TableHeaderCell>Delete</TableHeaderCell>
@@ -177,6 +180,24 @@ function Main() {
         <TableBody>
           {apiKeys.map((apiKey) => (
             <TableRow key={apiKey.api_key_id}>
+              <TableCell>
+                <div
+                  className={`
+                  my-auto 
+                  flex 
+                  mb-1 
+                  w-fit 
+                  hover:bg-hover cursor-pointer
+                  p-2 
+                  rounded-lg
+                  border-border
+                  text-sm`}
+                  onClick={() => handleEdit(apiKey)}
+                >
+                  <FiEdit className="my-auto mr-2" />
+                  {apiKey.api_key_name || <i>null</i>}
+                </div>
+              </TableCell>
               <TableCell className="max-w-64">
                 {apiKey.api_key_display}
               </TableCell>
@@ -194,12 +215,7 @@ function Main() {
                   text-sm`}
                   onClick={async () => {
                     setKeyIsGenerating(true);
-                    const response = await fetch(
+                    const response = await regenerateApiKey(apiKey);
-                      `/api/admin/api-key/${apiKey.api_key_id}`,
-                      {
-                        method: "PATCH",
-                      }
-                    );
                     setKeyIsGenerating(false);
                     if (!response.ok) {
                       const errorMsg = await response.text();
@@ -221,12 +237,7 @@ function Main() {
               <TableCell>
                 <DeleteButton
                   onClick={async () => {
-                    const response = await fetch(
+                    const response = await deleteApiKey(apiKey.api_key_id);
-                      `/api/admin/api-key/${apiKey.api_key_id}`,
-                      {
-                        method: "DELETE",
-                      }
-                    );
                     if (!response.ok) {
                       const errorMsg = await response.text();
                       setPopup({
@@ -243,6 +254,21 @@ function Main() {
           ))}
         </TableBody>
       </Table>
+
+      {showCreateUpdateForm && (
+        <DanswerApiKeyForm
+          onCreateApiKey={(apiKey) => {
+            setFullApiKey(apiKey.api_key);
+          }}
+          onClose={() => {
+            setShowCreateUpdateForm(false);
+            setSelectedApiKey(undefined);
+            mutate("/api/admin/api-key");
+          }}
+          setPopup={setPopup}
+          apiKey={selectedApiKey}
+        />
+      )}
     </div>
   );
 }

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;
+}