MT Cloud Monitoring (#3465)

backend/alembic/versions/91a0a4d62b14_milestone.py

@@ -0,0 +1,45 @@
+"""Milestone
+
+Revision ID: 91a0a4d62b14
+Revises: dab04867cd88
+Create Date: 2024-12-13 19:03:30.947551
+
+"""
+from alembic import op
+import sqlalchemy as sa
+import fastapi_users_db_sqlalchemy
+from sqlalchemy.dialects import postgresql
+
+# revision identifiers, used by Alembic.
+revision = "91a0a4d62b14"
+down_revision = "dab04867cd88"
+branch_labels = None
+depends_on = None
+
+
+def upgrade() -> None:
+    op.create_table(
+        "milestone",
+        sa.Column("id", sa.UUID(), nullable=False),
+        sa.Column("tenant_id", sa.String(), nullable=True),
+        sa.Column(
+            "user_id",
+            fastapi_users_db_sqlalchemy.generics.GUID(),
+            nullable=True,
+        ),
+        sa.Column("event_type", sa.String(), nullable=False),
+        sa.Column(
+            "time_created",
+            sa.DateTime(timezone=True),
+            server_default=sa.text("now()"),
+            nullable=False,
+        ),
+        sa.Column("event_tracker", postgresql.JSONB(), nullable=True),
+        sa.ForeignKeyConstraint(["user_id"], ["user.id"], ondelete="CASCADE"),
+        sa.PrimaryKeyConstraint("id"),
+        sa.UniqueConstraint("event_type", name="uq_milestone_event_type"),
+    )
+
+
+def downgrade() -> None:
+    op.drop_table("milestone")

backend/ee/onyx/configs/app_configs.py

@@ -47,3 +47,9 @@
 OAUTH_GOOGLE_DRIVE_CLIENT_SECRET = os.environ.get(
     "OAUTH_GOOGLE_DRIVE_CLIENT_SECRET", ""
 )
+
+
+# The posthog client does not accept empty API keys or hosts however it fails silently
+# when the capture is called. These defaults prevent Posthog issues from breaking the Onyx app
+POSTHOG_API_KEY = os.environ.get("POSTHOG_API_KEY") or "FooBar"
+POSTHOG_HOST = os.environ.get("POSTHOG_HOST") or "https://us.i.posthog.com"

backend/ee/onyx/server/tenants/provisioning.py

@@ -20,6 +20,7 @@
 from ee.onyx.server.tenants.user_mapping import user_owns_a_tenant
 from onyx.auth.users import exceptions
 from onyx.configs.app_configs import CONTROL_PLANE_API_BASE_URL
+from onyx.configs.constants import MilestoneRecordType
 from onyx.db.engine import get_session_with_tenant
 from onyx.db.engine import get_sqlalchemy_engine
 from onyx.db.llm import update_default_provider
@@ -35,12 +36,14 @@
 from onyx.server.manage.embedding.models import CloudEmbeddingProviderCreationRequest
 from onyx.server.manage.llm.models import LLMProviderUpsertRequest
 from onyx.setup import setup_onyx
+from onyx.utils.telemetry import create_milestone_and_report
 from shared_configs.configs import MULTI_TENANT
 from shared_configs.configs import POSTGRES_DEFAULT_SCHEMA
 from shared_configs.configs import TENANT_ID_PREFIX
 from shared_configs.contextvars import CURRENT_TENANT_ID_CONTEXTVAR
 from shared_configs.enums import EmbeddingProvider
 
+
 logger = logging.getLogger(__name__)
 
 
@@ -122,6 +125,17 @@ async def provision_tenant(tenant_id: str, email: str) -> None:
 
         add_users_to_tenant([email], tenant_id)
 
+        with get_session_with_tenant(tenant_id) as db_session:
+            create_milestone_and_report(
+                user=None,
+                distinct_id=tenant_id,
+                event_type=MilestoneRecordType.TENANT_CREATED,
+                properties={
+                    "email": email,
+                },
+                db_session=db_session,
+            )
+
     except Exception as e:
         logger.exception(f"Failed to create tenant {tenant_id}")
         raise HTTPException(

backend/ee/onyx/utils/telemetry.py

@@ -0,0 +1,14 @@
+from posthog import Posthog
+
+from ee.onyx.configs.app_configs import POSTHOG_API_KEY
+from ee.onyx.configs.app_configs import POSTHOG_HOST
+
+posthog = Posthog(project_api_key=POSTHOG_API_KEY, host=POSTHOG_HOST)
+
+
+def event_telemetry(
+    distinct_id: str,
+    event: str,
+    properties: dict | None = None,
+) -> None:
+    posthog.capture(distinct_id, event, properties)

backend/onyx/auth/noauth_user.py

@@ -4,6 +4,8 @@
 
 from onyx.auth.schemas import UserRole
 from onyx.configs.constants import KV_NO_AUTH_USER_PREFERENCES_KEY
+from onyx.configs.constants import NO_AUTH_USER_EMAIL
+from onyx.configs.constants import NO_AUTH_USER_ID
 from onyx.key_value_store.store import KeyValueStore
 from onyx.key_value_store.store import KvKeyNotFoundError
 from onyx.server.manage.models import UserInfo
@@ -30,8 +32,8 @@ def load_no_auth_user_preferences(store: KeyValueStore) -> UserPreferences:
 
 def fetch_no_auth_user(store: KeyValueStore) -> UserInfo:
     return UserInfo(
-        id="__no_auth_user__",
-        email="[email protected]",
+        id=NO_AUTH_USER_ID,
+        email=NO_AUTH_USER_EMAIL,
         is_active=True,
         is_superuser=False,
         is_verified=True,

backend/onyx/auth/users.py

@@ -72,6 +72,7 @@
 from onyx.configs.constants import AuthType
 from onyx.configs.constants import DANSWER_API_KEY_DUMMY_EMAIL_DOMAIN
 from onyx.configs.constants import DANSWER_API_KEY_PREFIX
+from onyx.configs.constants import MilestoneRecordType
 from onyx.configs.constants import UNNAMED_KEY_PLACEHOLDER
 from onyx.db.api_key import fetch_user_for_api_key
 from onyx.db.auth import get_access_token_db
@@ -88,6 +89,7 @@
 from onyx.db.users import get_user_by_email
 from onyx.server.utils import BasicAuthenticationError
 from onyx.utils.logger import setup_logger
+from onyx.utils.telemetry import create_milestone_and_report
 from onyx.utils.telemetry import optional_telemetry
 from onyx.utils.telemetry import RecordType
 from onyx.utils.variable_functionality import fetch_ee_implementation_or_noop
@@ -225,6 +227,7 @@ async def create(
         safe: bool = False,
         request: Optional[Request] = None,
     ) -> User:
+        user_count: int | None = None
         referral_source = None
         if request is not None:
             referral_source = request.cookies.get("referral_source", None)
@@ -278,7 +281,26 @@ async def create(
             finally:
                 CURRENT_TENANT_ID_CONTEXTVAR.reset(token)
 
-            return user
+        # Blocking but this should be very quick
+        with get_session_with_tenant(tenant_id) as db_session:
+            if not user_count:
+                create_milestone_and_report(
+                    user=user,
+                    distinct_id=user.email,
+                    event_type=MilestoneRecordType.USER_SIGNED_UP,
+                    properties=None,
+                    db_session=db_session,
+                )
+            else:
+                create_milestone_and_report(
+                    user=user,
+                    distinct_id=user.email,
+                    event_type=MilestoneRecordType.MULTIPLE_USERS,
+                    properties=None,
+                    db_session=db_session,
+                )
+
+        return user
 
     async def oauth_callback(
         self,

backend/onyx/background/indexing/run_indexing.py

@@ -11,6 +11,7 @@
 from onyx.configs.app_configs import INDEXING_SIZE_WARNING_THRESHOLD
 from onyx.configs.app_configs import INDEXING_TRACER_INTERVAL
 from onyx.configs.app_configs import POLL_CONNECTOR_OFFSET
+from onyx.configs.constants import MilestoneRecordType
 from onyx.connectors.connector_runner import ConnectorRunner
 from onyx.connectors.factory import instantiate_connector
 from onyx.connectors.models import IndexAttemptMetadata
@@ -34,6 +35,7 @@
 from onyx.indexing.indexing_pipeline import build_indexing_pipeline
 from onyx.utils.logger import setup_logger
 from onyx.utils.logger import TaskAttemptSingleton
+from onyx.utils.telemetry import create_milestone_and_report
 from onyx.utils.variable_functionality import global_version
 
 logger = setup_logger()
@@ -396,6 +398,15 @@ def _run_indexing(
 
     if index_attempt_md.num_exceptions == 0:
         mark_attempt_succeeded(index_attempt, db_session)
+
+        create_milestone_and_report(
+            user=None,
+            distinct_id=tenant_id or "N/A",
+            event_type=MilestoneRecordType.CONNECTOR_SUCCEEDED,
+            properties=None,
+            db_session=db_session,
+        )
+
         logger.info(
             f"Connector succeeded: "
             f"docs={document_count} chunks={chunk_count} elapsed={elapsed_time:.2f}s"

backend/onyx/chat/process_message.py

@@ -31,6 +31,8 @@
 from onyx.configs.chat_configs import DISABLE_LLM_CHOOSE_SEARCH
 from onyx.configs.chat_configs import MAX_CHUNKS_FED_TO_CHAT
 from onyx.configs.constants import MessageType
+from onyx.configs.constants import MilestoneRecordType
+from onyx.configs.constants import NO_AUTH_USER_ID
 from onyx.context.search.enums import OptionalSearchSetting
 from onyx.context.search.enums import QueryFlow
 from onyx.context.search.enums import SearchType
@@ -53,6 +55,9 @@
 from onyx.db.chat import translate_db_message_to_chat_message_detail
 from onyx.db.chat import translate_db_search_doc_to_server_search_doc
 from onyx.db.engine import get_session_context_manager
+from onyx.db.milestone import check_multi_assistant_milestone
+from onyx.db.milestone import create_milestone_if_not_exists
+from onyx.db.milestone import update_user_assistant_milestone
 from onyx.db.models import SearchDoc as DbSearchDoc
 from onyx.db.models import ToolCall
 from onyx.db.models import User
@@ -117,6 +122,7 @@
 from onyx.tools.tool_runner import ToolCallFinalResult
 from onyx.utils.logger import setup_logger
 from onyx.utils.long_term_log import LongTermLogger
+from onyx.utils.telemetry import mt_cloud_telemetry
 from onyx.utils.timing import log_function_time
 from onyx.utils.timing import log_generator_function_time
 from shared_configs.contextvars import CURRENT_TENANT_ID_CONTEXTVAR
@@ -356,6 +362,31 @@ def stream_chat_message_objects(
         if not persona:
             raise RuntimeError("No persona specified or found for chat session")
 
+        multi_assistant_milestone, _is_new = create_milestone_if_not_exists(
+            user=user,
+            event_type=MilestoneRecordType.MULTIPLE_ASSISTANTS,
+            db_session=db_session,
+        )
+
+        update_user_assistant_milestone(
+            milestone=multi_assistant_milestone,
+            user_id=str(user.id) if user else NO_AUTH_USER_ID,
+            assistant_id=persona.id,
+            db_session=db_session,
+        )
+
+        _, just_hit_multi_assistant_milestone = check_multi_assistant_milestone(
+            milestone=multi_assistant_milestone,
+            db_session=db_session,
+        )
+
+        if just_hit_multi_assistant_milestone:
+            mt_cloud_telemetry(
+                distinct_id=tenant_id,
+                event=MilestoneRecordType.MULTIPLE_ASSISTANTS,
+                properties=None,
+            )
+
         # If a prompt override is specified via the API, use that with highest priority
         # but for saving it, we are just mapping it to an existing prompt
         prompt_id = new_msg_req.prompt_id

backend/onyx/configs/constants.py

@@ -15,6 +15,9 @@
 DEFAULT_BOOST = 0
 SESSION_KEY = "session"
 
+NO_AUTH_USER_ID = "__no_auth_user__"
+NO_AUTH_USER_EMAIL = "[email protected]"
+
 # For chunking/processing chunks
 RETURN_SEPARATOR = "\n\r\n"
 SECTION_SEPARATOR = "\n\n"
@@ -210,6 +213,19 @@ class FileOrigin(str, Enum):
     OTHER = "other"
 
 
+class MilestoneRecordType(str, Enum):
+    TENANT_CREATED = "tenant_created"
+    USER_SIGNED_UP = "user_signed_up"
+    MULTIPLE_USERS = "multiple_users"
+    VISITED_ADMIN_PAGE = "visited_admin_page"
+    CREATED_CONNECTOR = "created_connector"
+    CONNECTOR_SUCCEEDED = "connector_succeeded"
+    RAN_QUERY = "ran_query"
+    MULTIPLE_ASSISTANTS = "multiple_assistants"
+    CREATED_ASSISTANT = "created_assistant"
+    CREATED_ONYX_BOT = "created_onyx_bot"
+
+
 class PostgresAdvisoryLocks(Enum):
     KOMBU_MESSAGE_CLEANUP_LOCK_ID = auto()