Do not count API keys as users (#3022)

* don't count api keys as users * typing
danswer-ai · Nov 1, 2024 · 6d543f3 · 6d543f3
1 parent ccdc09e
commit 6d543f3
Show file tree

Hide file tree

Showing 3 changed files with 15 additions and 6 deletions.
diff --git a/backend/danswer/db/auth.py b/backend/danswer/db/auth.py
@@ -22,6 +22,7 @@
 from danswer.utils.variable_functionality import (
     fetch_versioned_implementation_with_fallback,
 )
+from ee.danswer.db.api_key import get_api_key_email_pattern
 
 
 def get_default_admin_user_emails() -> list[str]:
@@ -35,12 +36,16 @@ def get_default_admin_user_emails() -> list[str]:
     return get_default_admin_user_emails_fn()
 
 
-def get_total_users(db_session: Session) -> int:
+def get_total_users_count(db_session: Session) -> int:
     """
     Returns the total number of users in the system.
     This is the sum of users and invited users.
     """
-    user_count = db_session.query(User).count()
+    user_count = (
+        db_session.query(User)
+        .filter(~User.email.endswith(get_api_key_email_pattern()))  # type: ignore
+        .count()
+    )
     invited_users = len(get_invited_users())
     return user_count + invited_users
 

diff --git a/backend/danswer/server/manage/users.py b/backend/danswer/server/manage/users.py
@@ -37,7 +37,7 @@
 from danswer.configs.app_configs import SESSION_EXPIRE_TIME_SECONDS
 from danswer.configs.app_configs import VALID_EMAIL_DOMAINS
 from danswer.configs.constants import AuthType
-from danswer.db.auth import get_total_users
+from danswer.db.auth import get_total_users_count
 from danswer.db.engine import CURRENT_TENANT_ID_CONTEXTVAR
 from danswer.db.engine import get_session
 from danswer.db.models import AccessToken
@@ -226,7 +226,7 @@ def bulk_invite_users(
     try:
         logger.info("Registering tenant users")
         register_tenant_users(
-            CURRENT_TENANT_ID_CONTEXTVAR.get(), get_total_users(db_session)
+            CURRENT_TENANT_ID_CONTEXTVAR.get(), get_total_users_count(db_session)
         )
         if ENABLE_EMAIL_INVITES:
             try:
@@ -262,7 +262,7 @@ def remove_invited_user(
     try:
         if MULTI_TENANT:
             register_tenant_users(
-                CURRENT_TENANT_ID_CONTEXTVAR.get(), get_total_users(db_session)
+                CURRENT_TENANT_ID_CONTEXTVAR.get(), get_total_users_count(db_session)
             )
     except Exception:
         logger.error(

diff --git a/backend/ee/danswer/db/api_key.py b/backend/ee/danswer/db/api_key.py
@@ -17,8 +17,12 @@
 from ee.danswer.server.api_key.models import APIKeyArgs
 
 
+def get_api_key_email_pattern() -> str:
+    return DANSWER_API_KEY_DUMMY_EMAIL_DOMAIN
+
+
 def is_api_key_email_address(email: str) -> bool:
-    return email.endswith(f"{DANSWER_API_KEY_DUMMY_EMAIL_DOMAIN}")
+    return email.endswith(get_api_key_email_pattern())
 
 
 def fetch_api_keys(db_session: Session) -> list[ApiKeyDescriptor]: