Token Rate Limiting

WIP

Cleanup 🧹

Remove existing rate limiting logic

Cleanup 🧼

Undo nit

Cleanup 🧽

Move db constants (avoids circular import)

WIP

WIP

Cleanup

Lint

Resolve alembic conflict

Fix mypy

Add backfill to migration

Update comment

Make unauthenticated users still adhere to global limits

Use Depends

Remove enum from table

Update migration error handling + deletion

Address verbal feedback, cleanup urls, minor nits

web/next.config.js

@@ -48,6 +48,11 @@ const nextConfig = {
               source: "/admin/performance/custom-analytics",
               destination: "/ee/admin/performance/custom-analytics",
             },
+            // token rate limits
+            {
+              source: "/admin/token-rate-limits",
+              destination: "/ee/admin/token-rate-limits",
+            },
           ]
         : [];
 

web/src/app/admin/models/llm/page.tsx

@@ -1,176 +1,9 @@
 "use client";
 
-import { Form, Formik } from "formik";
-import { useEffect, useState } from "react";
 import { AdminPageTitle } from "@/components/admin/Title";
-import {
-  BooleanFormField,
-  SectionHeader,
-  TextFormField,
-} from "@/components/admin/connectors/Field";
-import { Popup } from "@/components/admin/connectors/Popup";
-import { Button, Divider, Text } from "@tremor/react";
 import { FiCpu } from "react-icons/fi";
 import { LLMConfiguration } from "./LLMConfiguration";
 
-const LLMOptions = () => {
-  const [popup, setPopup] = useState<{
-    message: string;
-    type: "success" | "error";
-  } | null>(null);
-
-  const [tokenBudgetGloballyEnabled, setTokenBudgetGloballyEnabled] =
-    useState(false);
-  const [initialValues, setInitialValues] = useState({
-    enable_token_budget: false,
-    token_budget: "",
-    token_budget_time_period: "",
-  });
-
-  const fetchConfig = async () => {
-    const response = await fetch("/api/manage/admin/token-budget-settings");
-    if (response.ok) {
-      const config = await response.json();
-      // Assuming the config object directly matches the structure needed for initialValues
-      setInitialValues({
-        enable_token_budget: config.enable_token_budget || false,
-        token_budget: config.token_budget || "",
-        token_budget_time_period: config.token_budget_time_period || "",
-      });
-      setTokenBudgetGloballyEnabled(true);
-    } else {
-      // Handle error or provide fallback values
-      setPopup({
-        message: "Failed to load current LLM options.",
-        type: "error",
-      });
-    }
-  };
-
-  // Fetch current config when the component mounts
-  useEffect(() => {
-    fetchConfig();
-  }, []);
-
-  if (!tokenBudgetGloballyEnabled) {
-    return null;
-  }
-
-  return (
-    <>
-      {popup && <Popup message={popup.message} type={popup.type} />}
-      <Formik
-        enableReinitialize={true}
-        initialValues={initialValues}
-        onSubmit={async (values) => {
-          const response = await fetch(
-            "/api/manage/admin/token-budget-settings",
-            {
-              method: "PUT",
-              headers: {
-                "Content-Type": "application/json",
-              },
-              body: JSON.stringify(values),
-            }
-          );
-          if (response.ok) {
-            setPopup({
-              message: "Updated LLM Options",
-              type: "success",
-            });
-            await fetchConfig();
-          } else {
-            const body = await response.json();
-            if (body.detail) {
-              setPopup({ message: body.detail, type: "error" });
-            } else {
-              setPopup({
-                message: "Unable to update LLM options.",
-                type: "error",
-              });
-            }
-            setTimeout(() => {
-              setPopup(null);
-            }, 4000);
-          }
-        }}
-      >
-        {({ isSubmitting, values, setFieldValue }) => {
-          return (
-            <Form>
-              <Divider />
-              <>
-                <SectionHeader>Token Budget</SectionHeader>
-                <Text>
-                  Set a maximum token use per time period. If the token budget
-                  is exceeded, Danswer will not be able to respond to queries
-                  until the next time period.
-                </Text>
-                <br />
-                <BooleanFormField
-                  name="enable_token_budget"
-                  label="Enable Token Budget"
-                  subtext="If enabled, Danswer will be limited to the token budget specified below."
-                  onChange={(e) => {
-                    setFieldValue("enable_token_budget", e.target.checked);
-                  }}
-                />
-                {values.enable_token_budget && (
-                  <>
-                    <TextFormField
-                      name="token_budget"
-                      label="Token Budget"
-                      subtext={
-                        <div>
-                          How many tokens (in thousands) can be used per time
-                          period? If unspecified, no limit will be set.
-                        </div>
-                      }
-                      onChange={(e) => {
-                        const value = e.target.value;
-                        // Allow only integer values
-                        if (value === "" || /^[0-9]+$/.test(value)) {
-                          setFieldValue("token_budget", value);
-                        }
-                      }}
-                    />
-                    <TextFormField
-                      name="token_budget_time_period"
-                      label="Token Budget Time Period (hours)"
-                      subtext={
-                        <div>
-                          Specify the length of the time period, in hours, over
-                          which the token budget will be applied.
-                        </div>
-                      }
-                      onChange={(e) => {
-                        const value = e.target.value;
-                        // Allow only integer values
-                        if (value === "" || /^[0-9]+$/.test(value)) {
-                          setFieldValue("token_budget_time_period", value);
-                        }
-                      }}
-                    />
-                  </>
-                )}
-              </>
-              <div className="flex">
-                <Button
-                  className="w-64 mx-auto"
-                  type="submit"
-                  disabled={isSubmitting}
-                >
-                  Submit
-                </Button>
-              </div>
-            </Form>
-          );
-        }}
-      </Formik>
-    </>
-  );
-};
-
 const Page = () => {
   return (
     <div className="mx-auto container">
@@ -181,7 +14,6 @@ const Page = () => {
 
       <LLMConfiguration />
 
-      <LLMOptions />
     </div>
   );
 };

web/src/app/admin/token-rate-limits/CreateRateLimitModal.tsx

@@ -0,0 +1,175 @@
+"use client";
+
+import * as Yup from "yup";
+import { Button } from "@tremor/react";
+import { useEffect, useState } from "react";
+import { Modal } from "@/components/Modal";
+import { Form, Formik } from "formik";
+import {
+  SelectorFormField,
+  TextFormField,
+} from "@/components/admin/connectors/Field";
+import { UserGroup } from "@/lib/types";
+import { Scope } from "./types";
+import { PopupSpec } from "@/components/admin/connectors/Popup";
+
+interface CreateRateLimitModalProps {
+  isOpen: boolean;
+  setIsOpen: (isOpen: boolean) => void;
+  onSubmit: (
+    target_scope: Scope,
+    period_hours: number,
+    token_budget: number,
+    group_id: number
+  ) => void;
+  setPopup: (popupSpec: PopupSpec | null) => void;
+  forSpecificScope?: Scope;
+  forSpecificUserGroup?: number;
+}
+
+export const CreateRateLimitModal = ({
+  isOpen,
+  setIsOpen,
+  onSubmit,
+  setPopup,
+  forSpecificScope,
+  forSpecificUserGroup,
+}: CreateRateLimitModalProps) => {
+  const [modalUserGroups, setModalUserGroups] = useState([]);
+  const [shouldFetchUserGroups, setShouldFetchUserGroups] = useState(
+    forSpecificScope === Scope.USER_GROUP
+  );
+
+  useEffect(() => {
+    const fetchData = async () => {
+      try {
+        const response = await fetch("/api/manage/admin/user-group");
+        const data = await response.json();
+        const options = data.map((userGroup: UserGroup) => ({
+          name: userGroup.name,
+          value: userGroup.id,
+        }));
+        setModalUserGroups(options);
+        setShouldFetchUserGroups(false);
+      } catch (error) {
+        setPopup({
+          type: "error",
+          message: `Failed to fetch user groups: ${error}`,
+        });
+      }
+    };
+
+    if (shouldFetchUserGroups) {
+      fetchData();
+    }
+  }, [shouldFetchUserGroups, setPopup]);
+
+  if (!isOpen) {
+    return null;
+  }
+
+  return (
+    <Modal
+      title={"Create a Token Rate Limit"}
+      onOutsideClick={() => setIsOpen(false)}
+      width="w-2/6"
+    >
+      <Formik
+        initialValues={{
+          enabled: true,
+          period_hours: "",
+          token_budget: "",
+          target_scope: forSpecificScope || Scope.GLOBAL,
+          user_group_id: forSpecificUserGroup,
+        }}
+        validationSchema={Yup.object().shape({
+          period_hours: Yup.number()
+            .required("Time Window is a required field")
+            .min(1, "Time Window must be at least 1 hour"),
+          token_budget: Yup.number()
+            .required("Token Budget is a required field")
+            .min(1, "Token Budget must be at least 1"),
+          target_scope: Yup.string().required(
+            "Target Scope is a required field"
+          ),
+          user_group_id: Yup.string().test(
+            "user_group_id",
+            "User Group is a required field",
+            (value, context) => {
+              return (
+                context.parent.target_scope !== "user_group" ||
+                (context.parent.target_scope === "user_group" &&
+                  value !== undefined)
+              );
+            }
+          ),
+        })}
+        onSubmit={async (values, formikHelpers) => {
+          formikHelpers.setSubmitting(true);
+          onSubmit(
+            values.target_scope,
+            Number(values.period_hours),
+            Number(values.token_budget),
+            Number(values.user_group_id)
+          );
+          return formikHelpers.setSubmitting(false);
+        }}
+      >
+        {({ isSubmitting, values, setFieldValue }) => (
+          <Form>
+            {!forSpecificScope && (
+              <SelectorFormField
+                name="target_scope"
+                label="Target Scope"
+                options={[
+                  { name: "Global", value: Scope.GLOBAL },
+                  { name: "User", value: Scope.USER },
+                  { name: "User Group", value: Scope.USER_GROUP },
+                ]}
+                includeDefault={false}
+                onSelect={(selected) => {
+                  setFieldValue("target_scope", selected)
+                  if (selected === Scope.USER_GROUP) {
+                    setShouldFetchUserGroups(true);
+                  }
+                }}
+              />
+            )}
+            {forSpecificUserGroup === undefined &&
+              values.target_scope === Scope.USER_GROUP && (
+                <SelectorFormField
+                  name="user_group_id"
+                  label="User Group"
+                  options={modalUserGroups}
+                  includeDefault={false}
+                />
+              )}
+            <TextFormField
+              name="period_hours"
+              label="Time Window (Hours)"
+              type="number"
+              placeholder=""
+            />
+            <TextFormField
+              name="token_budget"
+              label="Token Budget (Thousands)"
+              type="number"
+              placeholder=""
+            />
+            <div className="flex">
+              <Button
+                type="submit"
+                size="xs"
+                color="green"
+                disabled={isSubmitting}
+                className="mx-auto w-64"
+              >
+                Create!
+              </Button>
+            </div>
+          </Form>
+        )}
+      </Formik>
+    </Modal>
+  );
+};

web/src/app/admin/token-rate-limits/TokenRateLimitTables.tsx

@@ -0,0 +1,169 @@
+"use client";
+
+import {
+  Table,
+  TableHead,
+  TableRow,
+  TableHeaderCell,
+  TableBody,
+  TableCell,
+  Title,
+  Text,
+} from "@tremor/react";
+import { DeleteButton } from "@/components/DeleteButton";
+import { deleteTokenRateLimit, updateTokenRateLimit } from "./lib";
+import { ThreeDotsLoader } from "@/components/Loading";
+import { TokenRateLimitDisplay } from "./types";
+import { errorHandlingFetcher } from "@/lib/fetcher";
+import useSWR, { mutate } from "swr";
+import { CustomCheckbox } from "@/components/CustomCheckbox";
+
+type TokenRateLimitTableArgs = {
+  tokenRateLimits: TokenRateLimitDisplay[];
+  title?: string;
+  description?: string;
+  fetchUrl: string;
+  hideHeading?: boolean;
+};
+
+export const TokenRateLimitTable = ({
+  tokenRateLimits,
+  title,
+  description,
+  fetchUrl,
+  hideHeading,
+}: TokenRateLimitTableArgs) => {
+  const shouldRenderGroupName = () =>
+    tokenRateLimits.length > 0 && tokenRateLimits[0].group_name !== undefined;
+
+  const handleEnabledChange = (id: number) => {
+    const tokenRateLimit = tokenRateLimits.find(
+      (tokenRateLimit) => tokenRateLimit.token_id === id
+    );
+
+    if (!tokenRateLimit) {
+      return;
+    }
+
+    updateTokenRateLimit(id, {
+      token_budget: tokenRateLimit.token_budget,
+      period_hours: tokenRateLimit.period_hours,
+      enabled: !tokenRateLimit.enabled,
+    }).then(() => {
+      mutate(fetchUrl);
+    });
+  };
+
+  const handleDelete = (id: number) =>
+    deleteTokenRateLimit(id).then(() => {
+      mutate(fetchUrl);
+    });
+
+  if (tokenRateLimits.length === 0) {
+    return (
+      <div>
+        {!hideHeading && title && <Title>{title}</Title>}
+        {!hideHeading && description && (
+          <Text className="my-2">{description}</Text>
+        )}
+        <Text className={`${!hideHeading && "my-8"}`}>
+          No token rate limits set!
+        </Text>
+      </div>
+    );
+  }
+
+  return (
+    <div>
+      {!hideHeading && title && <Title>{title}</Title>}
+      {!hideHeading && description && (
+        <Text className="my-2">{description}</Text>
+      )}
+      <Table className={`overflow-visible ${!hideHeading && "my-8"}`}>
+        <TableHead>
+          <TableRow>
+            <TableHeaderCell>Enabled</TableHeaderCell>
+            {shouldRenderGroupName() && (
+              <TableHeaderCell>Group Name</TableHeaderCell>
+            )}
+            <TableHeaderCell>Time Window (Hours)</TableHeaderCell>
+            <TableHeaderCell>Token Budget (Thousands)</TableHeaderCell>
+            <TableHeaderCell>Delete</TableHeaderCell>
+          </TableRow>
+        </TableHead>
+        <TableBody>
+          {tokenRateLimits.map((tokenRateLimit) => {
+            return (
+              <TableRow key={tokenRateLimit.token_id}>
+                <TableCell>
+                  <div
+                    onClick={() => handleEnabledChange(tokenRateLimit.token_id)}
+                    className="px-1 py-0.5 hover:bg-hover-light rounded flex cursor-pointer select-none w-24 flex"
+                  >
+                    <div className="mx-auto flex">
+                      <CustomCheckbox checked={tokenRateLimit.enabled} />
+                      <p className="ml-2">
+                        {tokenRateLimit.enabled ? "Enabled" : "Disabled"}
+                      </p>
+                    </div>
+                  </div>
+                </TableCell>
+                {shouldRenderGroupName() && (
+                  <TableCell className="font-bold text-emphasis">
+                    {tokenRateLimit.group_name}
+                  </TableCell>
+                )}
+                <TableCell>{tokenRateLimit.period_hours}</TableCell>
+                <TableCell>{tokenRateLimit.token_budget}</TableCell>
+                <TableCell>
+                  <DeleteButton
+                    onClick={() => handleDelete(tokenRateLimit.token_id)}
+                  />
+                </TableCell>
+              </TableRow>
+            );
+          })}
+        </TableBody>
+      </Table>
+    </div>
+  );
+};
+
+export const GenericTokenRateLimitTable = ({
+  fetchUrl,
+  title,
+  description,
+  hideHeading,
+  responseMapper,
+}: {
+  fetchUrl: string;
+  title?: string;
+  description?: string;
+  hideHeading?: boolean;
+  responseMapper?: (data: any) => TokenRateLimitDisplay[];
+}) => {
+  const { data, isLoading, error } = useSWR(fetchUrl, errorHandlingFetcher);
+
+  if (isLoading) {
+    return <ThreeDotsLoader />;
+  }
+
+  if (!isLoading && error) {
+    return <Text>Failed to load token rate limits</Text>;
+  }
+
+  let processedData = data;
+  if (responseMapper) {
+    processedData = responseMapper(data);
+  }
+
+  return (
+    <TokenRateLimitTable
+      tokenRateLimits={processedData}
+      fetchUrl={fetchUrl}
+      title={title}
+      description={description}
+      hideHeading={hideHeading}
+    />
+  );
+};

web/src/app/admin/token-rate-limits/lib.ts

@@ -0,0 +1,64 @@
+import { TokenRateLimitArgs } from "./types";
+
+const API_PREFIX = "/api/admin/token-rate-limits";
+
+// Global Token Limits
+export const insertGlobalTokenRateLimit = async (
+  tokenRateLimit: TokenRateLimitArgs
+) => {
+  return await fetch(`${API_PREFIX}/global`, {
+    method: "POST",
+    headers: {
+      "Content-Type": "application/json",
+    },
+    body: JSON.stringify(tokenRateLimit),
+  });
+};
+
+// User Token Limits
+export const insertUserTokenRateLimit = async (
+  tokenRateLimit: TokenRateLimitArgs
+) => {
+  return await fetch(`${API_PREFIX}/users`, {
+    method: "POST",
+    headers: {
+      "Content-Type": "application/json",
+    },
+    body: JSON.stringify(tokenRateLimit),
+  });
+};
+
+// User Group Token Limits (EE Only)
+export const insertGroupTokenRateLimit = async (
+  tokenRateLimit: TokenRateLimitArgs,
+  group_id: number
+) => {
+  return await fetch(`${API_PREFIX}/user-group/${group_id}`, {
+    method: "POST",
+    headers: {
+      "Content-Type": "application/json",
+    },
+    body: JSON.stringify(tokenRateLimit),
+  });
+};
+
+// Common Endpoints
+
+export const deleteTokenRateLimit = async (token_rate_limit_id: number) => {
+  return await fetch(`${API_PREFIX}/rate-limit/${token_rate_limit_id}`, {
+    method: "DELETE",
+  });
+};
+
+export const updateTokenRateLimit = async (
+  token_rate_limit_id: number,
+  tokenRateLimit: TokenRateLimitArgs
+) => {
+  return await fetch(`${API_PREFIX}/rate-limit/${token_rate_limit_id}`, {
+    method: "PUT",
+    headers: {
+      "Content-Type": "application/json",
+    },
+    body: JSON.stringify(tokenRateLimit),
+  });
+};

web/src/app/admin/token-rate-limits/page.tsx

@@ -0,0 +1,223 @@
+"use client";
+
+import { AdminPageTitle } from "@/components/admin/Title";
+import {
+  Button,
+  Tab,
+  TabGroup,
+  TabList,
+  TabPanel,
+  TabPanels,
+  Text,
+} from "@tremor/react";
+import { useState } from "react";
+import { FiGlobe, FiShield, FiUser, FiUsers } from "react-icons/fi";
+import {
+  insertGlobalTokenRateLimit,
+  insertGroupTokenRateLimit,
+  insertUserTokenRateLimit,
+} from "./lib";
+import { Scope, TokenRateLimit } from "./types";
+import { GenericTokenRateLimitTable } from "./TokenRateLimitTables";
+import { mutate } from "swr";
+import { usePopup } from "@/components/admin/connectors/Popup";
+import { CreateRateLimitModal } from "./CreateRateLimitModal";
+import { EE_ENABLED } from "@/lib/constants";
+
+const BASE_URL = "/api/admin/token-rate-limits";
+const GLOBAL_TOKEN_FETCH_URL = `${BASE_URL}/global`;
+const USER_TOKEN_FETCH_URL = `${BASE_URL}/users`;
+const USER_GROUP_FETCH_URL = `${BASE_URL}/user-groups`;
+
+const GLOBAL_DESCRIPTION =
+  "Global rate limits apply to all users, user groups, and API keys. When the global \
+  rate limit is reached, no more tokens can be spent.";
+const USER_DESCRIPTION =
+  "User rate limits apply to individual users. When a user reaches a limit, they will \
+  be temporarily blocked from spending tokens.";
+const USER_GROUP_DESCRIPTION =
+  "User group rate limits apply to all users in a group. When a group reaches a limit, \
+  all users in the group will be temporarily blocked from spending tokens, regardless \
+  of their individual limits. If a user is in multiple groups, the most lenient limit \
+  will apply.";
+
+const handleCreateTokenRateLimit = async (
+  target_scope: Scope,
+  period_hours: number,
+  token_budget: number,
+  group_id: number = -1
+) => {
+  const tokenRateLimitArgs = {
+    enabled: true,
+    token_budget: token_budget,
+    period_hours: period_hours,
+  };
+
+  if (target_scope === Scope.GLOBAL) {
+    return await insertGlobalTokenRateLimit(tokenRateLimitArgs);
+  } else if (target_scope === Scope.USER) {
+    return await insertUserTokenRateLimit(tokenRateLimitArgs);
+  } else if (target_scope === Scope.USER_GROUP) {
+    return await insertGroupTokenRateLimit(tokenRateLimitArgs, group_id);
+  } else {
+    throw new Error(`Invalid target_scope: ${target_scope}`);
+  }
+};
+
+function Main() {
+  const [tabIndex, setTabIndex] = useState(0);
+  const [modalIsOpen, setModalIsOpen] = useState(false);
+  const { popup, setPopup } = usePopup();
+
+  const updateTable = (target_scope: Scope) => {
+    if (target_scope === Scope.GLOBAL) {
+      mutate(GLOBAL_TOKEN_FETCH_URL);
+      setTabIndex(0);
+    } else if (target_scope === Scope.USER) {
+      mutate(USER_TOKEN_FETCH_URL);
+      setTabIndex(1);
+    } else if (target_scope === Scope.USER_GROUP) {
+      mutate(USER_GROUP_FETCH_URL);
+      setTabIndex(2);
+    }
+  };
+
+  const handleSubmit = (
+    target_scope: Scope,
+    period_hours: number,
+    token_budget: number,
+    group_id: number = -1
+  ) => {
+    handleCreateTokenRateLimit(
+      target_scope,
+      period_hours,
+      token_budget,
+      group_id
+    )
+      .then(() => {
+        setModalIsOpen(false);
+        setPopup({ type: "success", message: "Token rate limit created!" });
+        updateTable(target_scope);
+      })
+      .catch((error) => {
+        setPopup({ type: "error", message: error.message });
+      });
+  };
+
+  return (
+    <div>
+      {popup}
+
+      <Text className="mb-2">
+        Token rate limits enable you control how many tokens can be spent in a
+        given time period. With token rate limits, you can:
+      </Text>
+
+      <ul className="list-disc mt-2 ml-4 mb-2">
+        <li>
+          <Text>
+            Set a global rate limit to control your organization&apos;s overall
+            token spend.
+          </Text>
+        </li>
+        {EE_ENABLED && (
+          <>
+            <li>
+              <Text>
+                Set rate limits for users to ensure that no single user can
+                spend too many tokens.
+              </Text>
+            </li>
+            <li>
+              <Text>
+                Set rate limits for user groups to control token spend for your
+                teams.
+              </Text>
+            </li>
+          </>
+        )}
+        <li>
+          <Text>Enable and disable rate limits on the fly.</Text>
+        </li>
+      </ul>
+
+      <Button
+        color="green"
+        size="xs"
+        className="mt-3"
+        onClick={() => setModalIsOpen(true)}
+      >
+        Create a Token Rate Limit
+      </Button>
+
+      {EE_ENABLED && (
+        <TabGroup className="mt-6" index={tabIndex} onIndexChange={setTabIndex}>
+          <TabList variant="line">
+            <Tab icon={FiGlobe}>Global</Tab>
+            <Tab icon={FiUser}>User</Tab>
+            <Tab icon={FiUsers}>User Groups</Tab>
+          </TabList>
+          <TabPanels className="mt-6">
+            <TabPanel>
+              <GenericTokenRateLimitTable
+                fetchUrl={GLOBAL_TOKEN_FETCH_URL}
+                title={"Global Token Rate Limits"}
+                description={GLOBAL_DESCRIPTION}
+              />
+            </TabPanel>
+            <TabPanel>
+              <GenericTokenRateLimitTable
+                fetchUrl={USER_TOKEN_FETCH_URL}
+                title={"User Token Rate Limits"}
+                description={USER_DESCRIPTION}
+              />
+            </TabPanel>
+            <TabPanel>
+              <GenericTokenRateLimitTable
+                fetchUrl={USER_GROUP_FETCH_URL}
+                title={"User Group Token Rate Limits"}
+                description={USER_GROUP_DESCRIPTION}
+                responseMapper={(data: Record<string, TokenRateLimit[]>) =>
+                  Object.entries(data).flatMap(([group_name, elements]) =>
+                    elements.map((element) => ({
+                      ...element,
+                      group_name,
+                    }))
+                  )
+                }
+              />
+            </TabPanel>
+          </TabPanels>
+        </TabGroup>
+      )}
+
+      {!EE_ENABLED && (
+        <div className="mt-6">
+          <GenericTokenRateLimitTable
+            fetchUrl={GLOBAL_TOKEN_FETCH_URL}
+            title={"Global Token Rate Limits"}
+            description={GLOBAL_DESCRIPTION}
+          />
+        </div>
+      )}
+
+      <CreateRateLimitModal
+        isOpen={modalIsOpen}
+        setIsOpen={() => setModalIsOpen(false)}
+        setPopup={setPopup}
+        onSubmit={handleSubmit}
+        forSpecificScope={EE_ENABLED ? undefined : Scope.GLOBAL}
+      />
+    </div>
+  );
+}
+
+export default function Page() {
+  return (
+    <div className="mx-auto container">
+      <AdminPageTitle title="Token Rate Limits" icon={<FiShield size={32} />} />
+
+      <Main />
+    </div>
+  );
+}

web/src/app/admin/token-rate-limits/types.ts

@@ -0,0 +1,22 @@
+export enum Scope {
+  USER = "user",
+  USER_GROUP = "user_group",
+  GLOBAL = "global",
+}
+
+export interface TokenRateLimitArgs {
+  enabled: boolean;
+  token_budget: number;
+  period_hours: number;
+}
+
+export interface TokenRateLimit {
+  token_id: number;
+  enabled: boolean;
+  token_budget: number;
+  period_hours: number;
+}
+
+export interface TokenRateLimitDisplay extends TokenRateLimit {
+  group_name?: string;
+}

web/src/app/ee/admin/groups/[groupId]/AddTokenRateLimitForm.tsx

@@ -0,0 +1,60 @@
+import { PopupSpec } from "@/components/admin/connectors/Popup";
+import { CreateRateLimitModal } from "../../../../admin/token-rate-limits/CreateRateLimitModal";
+import { Scope } from "../../../../admin/token-rate-limits/types";
+import { insertGroupTokenRateLimit } from "../../../../admin/token-rate-limits/lib";
+import { mutate } from "swr";
+
+interface AddMemberFormProps {
+  isOpen: boolean;
+  setIsOpen: (isOpen: boolean) => void;
+  setPopup: (popupSpec: PopupSpec | null) => void;
+  userGroupId: number;
+}
+
+const handleCreateGroupTokenRateLimit = async (
+  period_hours: number,
+  token_budget: number,
+  group_id: number = -1
+) => {
+  const tokenRateLimitArgs = {
+    enabled: true,
+    token_budget: token_budget,
+    period_hours: period_hours,
+  };
+  return await insertGroupTokenRateLimit(tokenRateLimitArgs, group_id);
+};
+
+export const AddTokenRateLimitForm: React.FC<AddMemberFormProps> = ({
+  isOpen,
+  setIsOpen,
+  setPopup,
+  userGroupId,
+}) => {
+  const handleSubmit = (
+    _: Scope,
+    period_hours: number,
+    token_budget: number,
+    group_id: number = -1
+  ) => {
+    handleCreateGroupTokenRateLimit(period_hours, token_budget, group_id)
+      .then(() => {
+        setIsOpen(false);
+        setPopup({ type: "success", message: "Token rate limit created!" });
+        mutate(`/api/admin/token-rate-limits/user-group/${userGroupId}`);
+      })
+      .catch((error) => {
+        setPopup({ type: "error", message: error.message });
+      });
+  };
+
+  return (
+    <CreateRateLimitModal
+      isOpen={isOpen}
+      setIsOpen={setIsOpen}
+      onSubmit={handleSubmit}
+      setPopup={setPopup}
+      forSpecificScope={Scope.USER_GROUP}
+      forSpecificUserGroup={userGroupId}
+    />
+  );
+};

web/src/app/ee/admin/groups/[groupId]/GroupDisplay.tsx

@@ -22,6 +22,8 @@ import {
 import { DeleteButton } from "@/components/DeleteButton";
 import { Bubble } from "@/components/Bubble";
 import { BookmarkIcon, RobotIcon } from "@/components/icons/icons";
+import { AddTokenRateLimitForm } from "./AddTokenRateLimitForm";
+import { GenericTokenRateLimitTable } from "@/app/admin/token-rate-limits/TokenRateLimitTables";
 
 interface GroupDisplayProps {
   users: User[];
@@ -39,6 +41,7 @@ export const GroupDisplay = ({
   const { popup, setPopup } = usePopup();
   const [addMemberFormVisible, setAddMemberFormVisible] = useState(false);
   const [addConnectorFormVisible, setAddConnectorFormVisible] = useState(false);
+  const [addRateLimitFormVisible, setAddRateLimitFormVisible] = useState(false);
 
   return (
     <div>
@@ -301,6 +304,31 @@ export const GroupDisplay = ({
           </>
         )}
       </div>
+
+      <Divider />
+
+      <h2 className="text-xl font-bold mt-8 mb-2">Token Rate Limits</h2>
+
+      <AddTokenRateLimitForm
+        isOpen={addRateLimitFormVisible}
+        setIsOpen={setAddRateLimitFormVisible}
+        setPopup={setPopup}
+        userGroupId={userGroup.id}
+      />
+
+      <GenericTokenRateLimitTable
+        fetchUrl={`/api/admin/token-rate-limits/user-group/${userGroup.id}`}
+        hideHeading
+      />
+
+      <Button
+        color="green"
+        size="xs"
+        className="mt-3"
+        onClick={() => setAddRateLimitFormVisible(true)}
+      >
+        Create a Token Rate Limit
+      </Button>
     </div>
   );
 };

web/src/components/admin/Layout.tsx

@@ -28,6 +28,7 @@ import {
   FiImage,
   FiPackage,
   FiSettings,
+  FiShield,
   FiSlack,
   FiTool,
 } from "react-icons/fi";
@@ -215,6 +216,15 @@ export async function Layout({ children }: { children: React.ReactNode }) {
                         },
                       ]
                     : []),
+                  {
+                    name: (
+                      <div className="flex">
+                        <FiShield size={18} />
+                        <div className="ml-1">Token Rate Limits</div>
+                      </div>
+                    ),
+                    link: "/admin/token-rate-limits",
+                  },
                 ],
               },
               ...(EE_ENABLED