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_ID = os.environ.get("OAUTH_GOOGLE_DRIVE_CLIENT_ID", ""
 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 get_tenant_id_for_email
 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.llm.llm_provider_options import OPENAI_PROVIDER_NAME
 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 typing import cast
 
 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="anonymous@onyx.app",
+        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.app_configs import WEB_DOMAIN
 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.models import User
 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 @@ class UserManager(UUIDIDMixin, BaseUserManager[User, uuid.UUID]):
         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 @@ class UserManager(UUIDIDMixin, BaseUserManager[User, uuid.UUID]):
             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.background.indexing.tracer import OnyxTracer
 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_heartbeat import IndexingHeartbeatInterface
 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 CHAT_TARGET_CHUNK_PERCENTAGE
 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 reserve_message_id
 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_implementations.search.search_tool import (
 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 @@ ID_SEPARATOR = ":;:"
 DEFAULT_BOOST = 0
 SESSION_KEY = "session"
 
+NO_AUTH_USER_ID = "__no_auth_user__"
+NO_AUTH_USER_EMAIL = "anonymous@onyx.app"
+
 # 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()
 

backend/onyx/db/milestone.py

@@ -0,0 +1,99 @@
+from sqlalchemy import select
+from sqlalchemy.exc import IntegrityError
+from sqlalchemy.orm import Session
+from sqlalchemy.orm.attributes import flag_modified
+
+from onyx.configs.constants import MilestoneRecordType
+from onyx.db.models import Milestone
+from onyx.db.models import User
+
+
+USER_ASSISTANT_PREFIX = "user_assistants_used_"
+MULTI_ASSISTANT_USED = "multi_assistant_used"
+
+
+def create_milestone(
+    user: User | None,
+    event_type: MilestoneRecordType,
+    db_session: Session,
+) -> Milestone:
+    milestone = Milestone(
+        event_type=event_type,
+        user_id=user.id if user else None,
+    )
+    db_session.add(milestone)
+    db_session.commit()
+
+    return milestone
+
+
+def create_milestone_if_not_exists(
+    user: User | None, event_type: MilestoneRecordType, db_session: Session
+) -> tuple[Milestone, bool]:
+    # Check if it exists
+    milestone = db_session.execute(
+        select(Milestone).where(Milestone.event_type == event_type)
+    ).scalar_one_or_none()
+
+    if milestone is not None:
+        return milestone, False
+
+    # If it doesn't exist, try to create it.
+    try:
+        milestone = create_milestone(user, event_type, db_session)
+        return milestone, True
+    except IntegrityError:
+        # Another thread or process inserted it in the meantime
+        db_session.rollback()
+        # Fetch again to return the existing record
+        milestone = db_session.execute(
+            select(Milestone).where(Milestone.event_type == event_type)
+        ).scalar_one()  # Now should exist
+        return milestone, False
+
+
+def update_user_assistant_milestone(
+    milestone: Milestone,
+    user_id: str | None,
+    assistant_id: int,
+    db_session: Session,
+) -> None:
+    event_tracker = milestone.event_tracker
+    if event_tracker is None:
+        milestone.event_tracker = event_tracker = {}
+
+    if event_tracker.get(MULTI_ASSISTANT_USED):
+        # No need to keep tracking and populating if the milestone has already been hit
+        return
+
+    user_key = f"{USER_ASSISTANT_PREFIX}{user_id}"
+
+    if event_tracker.get(user_key) is None:
+        event_tracker[user_key] = [assistant_id]
+    elif assistant_id not in event_tracker[user_key]:
+        event_tracker[user_key].append(assistant_id)
+
+    flag_modified(milestone, "event_tracker")
+    db_session.commit()
+
+
+def check_multi_assistant_milestone(
+    milestone: Milestone,
+    db_session: Session,
+) -> tuple[bool, bool]:
+    """Returns if the milestone was hit and if it was just hit for the first time"""
+    event_tracker = milestone.event_tracker
+    if event_tracker is None:
+        return False, False
+
+    if event_tracker.get(MULTI_ASSISTANT_USED):
+        return True, False
+
+    for key, value in event_tracker.items():
+        if key.startswith(USER_ASSISTANT_PREFIX) and len(value) > 1:
+            event_tracker[MULTI_ASSISTANT_USED] = True
+            flag_modified(milestone, "event_tracker")
+            db_session.commit()
+            return True, True
+
+    return False, False

backend/onyx/db/models.py

@@ -37,7 +37,7 @@ from sqlalchemy.types import TypeDecorator
 
 from onyx.auth.schemas import UserRole
 from onyx.configs.chat_configs import NUM_POSTPROCESSED_RESULTS
-from onyx.configs.constants import DEFAULT_BOOST
+from onyx.configs.constants import DEFAULT_BOOST, MilestoneRecordType
 from onyx.configs.constants import DocumentSource
 from onyx.configs.constants import FileOrigin
 from onyx.configs.constants import MessageType
@@ -1534,6 +1534,32 @@ class SlackBot(Base):
     )
 
 
+class Milestone(Base):
+    # This table is used to track significant events for a deployment towards finding value
+    # The table is currently not used for features but it may be used in the future to inform
+    # users about the product features and encourage usage/exploration.
+    __tablename__ = "milestone"
+
+    id: Mapped[UUID] = mapped_column(
+        PGUUID(as_uuid=True), primary_key=True, default=uuid4
+    )
+    user_id: Mapped[UUID | None] = mapped_column(
+        ForeignKey("user.id", ondelete="CASCADE"), nullable=True
+    )
+    event_type: Mapped[MilestoneRecordType] = mapped_column(String)
+    # Need to track counts and specific ids of certain events to know if the Milestone has been reached
+    event_tracker: Mapped[dict | None] = mapped_column(
+        postgresql.JSONB(), nullable=True
+    )
+    time_created: Mapped[datetime.datetime] = mapped_column(
+        DateTime(timezone=True), server_default=func.now()
+    )
+
+    user: Mapped[User | None] = relationship("User")
+
+    __table_args__ = (UniqueConstraint("event_type", name="uq_milestone_event_type"),)
+
+
 class TaskQueueState(Base):
     # Currently refers to Celery Tasks
     __tablename__ = "task_queue_jobs"

backend/onyx/server/documents/connector.py

@@ -21,6 +21,7 @@ from onyx.background.celery.versioned_apps.primary import app as primary_app
 from onyx.configs.app_configs import ENABLED_CONNECTOR_TYPES
 from onyx.configs.constants import DocumentSource
 from onyx.configs.constants import FileOrigin
+from onyx.configs.constants import MilestoneRecordType
 from onyx.configs.constants import OnyxCeleryPriority
 from onyx.configs.constants import OnyxCeleryTask
 from onyx.connectors.google_utils.google_auth import (
@@ -110,6 +111,7 @@ from onyx.server.documents.models import ObjectCreationIdResponse
 from onyx.server.documents.models import RunConnectorRequest
 from onyx.server.models import StatusResponse
 from onyx.utils.logger import setup_logger
+from onyx.utils.telemetry import create_milestone_and_report
 from onyx.utils.variable_functionality import fetch_ee_implementation_or_noop
 
 logger = setup_logger()
@@ -639,6 +641,15 @@ def get_connector_indexing_status(
             )
         )
 
+    # Visiting admin page brings the user to the current connectors page which calls this endpoint
+    create_milestone_and_report(
+        user=user,
+        distinct_id=user.email if user else tenant_id or "N/A",
+        event_type=MilestoneRecordType.VISITED_ADMIN_PAGE,
+        properties=None,
+        db_session=db_session,
+    )
+
     return indexing_statuses
 
 
@@ -663,6 +674,7 @@ def create_connector_from_model(
     connector_data: ConnectorUpdateRequest,
     user: User = Depends(current_curator_or_admin_user),
     db_session: Session = Depends(get_session),
+    tenant_id: str = Depends(get_current_tenant_id),
 ) -> ObjectCreationIdResponse:
     try:
         _validate_connector_allowed(connector_data.source)
@@ -677,10 +689,20 @@ def create_connector_from_model(
             object_is_perm_sync=connector_data.access_type == AccessType.SYNC,
         )
         connector_base = connector_data.to_connector_base()
-        return create_connector(
+        connector_response = create_connector(
             db_session=db_session,
             connector_data=connector_base,
         )
+
+        create_milestone_and_report(
+            user=user,
+            distinct_id=user.email if user else tenant_id or "N/A",
+            event_type=MilestoneRecordType.CREATED_CONNECTOR,
+            properties=None,
+            db_session=db_session,
+        )
+
+        return connector_response
     except ValueError as e:
         logger.error(f"Error creating connector: {e}")
         raise HTTPException(status_code=400, detail=str(e))
@@ -691,6 +713,7 @@ def create_connector_with_mock_credential(
     connector_data: ConnectorUpdateRequest,
     user: User = Depends(current_curator_or_admin_user),
     db_session: Session = Depends(get_session),
+    tenant_id: str = Depends(get_current_tenant_id),
 ) -> StatusResponse:
     fetch_ee_implementation_or_noop(
         "onyx.db.user_group", "validate_user_creation_permissions", None
@@ -728,6 +751,15 @@ def create_connector_with_mock_credential(
             cc_pair_name=connector_data.name,
             groups=connector_data.groups,
         )
+
+        create_milestone_and_report(
+            user=user,
+            distinct_id=user.email if user else tenant_id or "N/A",
+            event_type=MilestoneRecordType.CREATED_CONNECTOR,
+            properties=None,
+            db_session=db_session,
+        )
+
         return response
 
     except ValueError as e:

backend/onyx/server/features/persona/api.py

@@ -15,7 +15,9 @@ from onyx.auth.users import current_limited_user
 from onyx.auth.users import current_user
 from onyx.chat.prompt_builder.utils import build_dummy_prompt
 from onyx.configs.constants import FileOrigin
+from onyx.configs.constants import MilestoneRecordType
 from onyx.configs.constants import NotificationType
+from onyx.db.engine import get_current_tenant_id
 from onyx.db.engine import get_session
 from onyx.db.models import User
 from onyx.db.notification import create_notification
@@ -44,6 +46,7 @@ from onyx.server.features.persona.models import PromptTemplateResponse
 from onyx.server.models import DisplayPriorityRequest
 from onyx.tools.utils import is_image_generation_available
 from onyx.utils.logger import setup_logger
+from onyx.utils.telemetry import create_milestone_and_report
 
 
 logger = setup_logger()
@@ -167,14 +170,25 @@ def create_persona(
     create_persona_request: CreatePersonaRequest,
     user: User | None = Depends(current_user),
     db_session: Session = Depends(get_session),
+    tenant_id: str | None = Depends(get_current_tenant_id),
 ) -> PersonaSnapshot:
-    return create_update_persona(
+    persona_snapshot = create_update_persona(
         persona_id=None,
         create_persona_request=create_persona_request,
         user=user,
         db_session=db_session,
     )
 
+    create_milestone_and_report(
+        user=user,
+        distinct_id=tenant_id or "N/A",
+        event_type=MilestoneRecordType.CREATED_ASSISTANT,
+        properties=None,
+        db_session=db_session,
+    )
+
+    return persona_snapshot
+
 
 # NOTE: This endpoint cannot update persona configuration options that
 # are core to the persona, such as its display priority and

backend/onyx/server/manage/slack_bot.py

@@ -4,7 +4,9 @@ from fastapi import HTTPException
 from sqlalchemy.orm import Session
 
 from onyx.auth.users import current_admin_user
+from onyx.configs.constants import MilestoneRecordType
 from onyx.db.constants import SLACK_BOT_PERSONA_PREFIX
+from onyx.db.engine import get_current_tenant_id
 from onyx.db.engine import get_session
 from onyx.db.models import ChannelConfig
 from onyx.db.models import User
@@ -25,6 +27,7 @@ from onyx.server.manage.models import SlackBot
 from onyx.server.manage.models import SlackBotCreationRequest
 from onyx.server.manage.models import SlackChannelConfig
 from onyx.server.manage.models import SlackChannelConfigCreationRequest
+from onyx.utils.telemetry import create_milestone_and_report
 
 
 router = APIRouter(prefix="/manage")
@@ -217,6 +220,7 @@ def create_bot(
     slack_bot_creation_request: SlackBotCreationRequest,
     db_session: Session = Depends(get_session),
     _: User | None = Depends(current_admin_user),
+    tenant_id: str | None = Depends(get_current_tenant_id),
 ) -> SlackBot:
     slack_bot_model = insert_slack_bot(
         db_session=db_session,
@@ -225,6 +229,15 @@ def create_bot(
         bot_token=slack_bot_creation_request.bot_token,
         app_token=slack_bot_creation_request.app_token,
     )
+
+    create_milestone_and_report(
+        user=None,
+        distinct_id=tenant_id or "N/A",
+        event_type=MilestoneRecordType.CREATED_ONYX_BOT,
+        properties=None,
+        db_session=db_session,
+    )
+
     return SlackBot.from_model(slack_bot_model)
 
 

backend/onyx/server/query_and_chat/chat_backend.py

@@ -30,6 +30,7 @@ from onyx.chat.prompt_builder.citations_prompt import (
 from onyx.configs.app_configs import WEB_DOMAIN
 from onyx.configs.constants import FileOrigin
 from onyx.configs.constants import MessageType
+from onyx.configs.constants import MilestoneRecordType
 from onyx.configs.model_configs import LITELLM_PASS_THROUGH_HEADERS
 from onyx.db.chat import add_chats_to_session_from_slack_thread
 from onyx.db.chat import create_chat_session
@@ -44,7 +45,9 @@ from onyx.db.chat import get_or_create_root_message
 from onyx.db.chat import set_as_latest_chat_message
 from onyx.db.chat import translate_db_message_to_chat_message_detail
 from onyx.db.chat import update_chat_session
+from onyx.db.engine import get_current_tenant_id
 from onyx.db.engine import get_session
+from onyx.db.engine import get_session_with_tenant
 from onyx.db.feedback import create_chat_message_feedback
 from onyx.db.feedback import create_doc_retrieval_feedback
 from onyx.db.models import User
@@ -81,6 +84,7 @@ from onyx.server.query_and_chat.models import UpdateChatSessionThreadRequest
 from onyx.server.query_and_chat.token_limit import check_token_rate_limits
 from onyx.utils.headers import get_custom_tool_additional_request_headers
 from onyx.utils.logger import setup_logger
+from onyx.utils.telemetry import create_milestone_and_report
 
 
 logger = setup_logger()
@@ -315,8 +319,9 @@ def handle_new_chat_message(
     chat_message_req: CreateChatMessageRequest,
     request: Request,
     user: User | None = Depends(current_limited_user),
-    _: None = Depends(check_token_rate_limits),
+    _rate_limit_check: None = Depends(check_token_rate_limits),
     is_connected_func: Callable[[], bool] = Depends(is_connected),
+    tenant_id: str = Depends(get_current_tenant_id),
 ) -> StreamingResponse:
     """
     This endpoint is both used for all the following purposes:
@@ -347,6 +352,15 @@ def handle_new_chat_message(
     ):
         raise HTTPException(status_code=400, detail="Empty chat message is invalid")
 
+    with get_session_with_tenant(tenant_id) as db_session:
+        create_milestone_and_report(
+            user=user,
+            distinct_id=user.email if user else tenant_id or "N/A",
+            event_type=MilestoneRecordType.RAN_QUERY,
+            properties=None,
+            db_session=db_session,
+        )
+
     def stream_generator() -> Generator[str, None, None]:
         try:
             for packet in stream_chat_message(

backend/onyx/utils/telemetry.py

@@ -10,10 +10,17 @@ from onyx.configs.app_configs import DISABLE_TELEMETRY
 from onyx.configs.app_configs import ENTERPRISE_EDITION_ENABLED
 from onyx.configs.constants import KV_CUSTOMER_UUID_KEY
 from onyx.configs.constants import KV_INSTANCE_DOMAIN_KEY
+from onyx.configs.constants import MilestoneRecordType
 from onyx.db.engine import get_sqlalchemy_engine
+from onyx.db.milestone import create_milestone_if_not_exists
 from onyx.db.models import User
 from onyx.key_value_store.factory import get_kv_store
 from onyx.key_value_store.interface import KvKeyNotFoundError
+from onyx.utils.variable_functionality import (
+    fetch_versioned_implementation_with_fallback,
+)
+from onyx.utils.variable_functionality import noop_fallback
+from shared_configs.configs import MULTI_TENANT
 
 _DANSWER_TELEMETRY_ENDPOINT = "https://telemetry.onyx.app/anonymous_telemetry"
 _CACHED_UUID: str | None = None
@@ -103,3 +110,37 @@ def optional_telemetry(
     except Exception:
         # Should never interfere with normal functions of Onyx
         pass
+
+
+def mt_cloud_telemetry(
+    distinct_id: str,
+    event: MilestoneRecordType,
+    properties: dict | None = None,
+) -> None:
+    if not MULTI_TENANT:
+        return
+
+    # MIT version should not need to include any Posthog code
+    # This is only for Onyx MT Cloud, this code should also never be hit, no reason for any orgs to
+    # be running the Multi Tenant version of Onyx.
+    fetch_versioned_implementation_with_fallback(
+        module="onyx.utils.telemetry",
+        attribute="event_telemetry",
+        fallback=noop_fallback,
+    )(distinct_id, event, properties)
+
+
+def create_milestone_and_report(
+    user: User | None,
+    distinct_id: str,
+    event_type: MilestoneRecordType,
+    properties: dict | None,
+    db_session: Session,
+) -> None:
+    _, is_new = create_milestone_if_not_exists(user, event_type, db_session)
+    if is_new:
+        mt_cloud_telemetry(
+            distinct_id=distinct_id,
+            event=event_type,
+            properties=properties,
+        )

backend/requirements/dev.txt

@@ -9,6 +9,7 @@ mypy-extensions==1.0.0
 mypy==1.8.0
 pandas-stubs==2.2.3.241009
 pandas==2.2.3
+posthog==3.7.4
 pre-commit==3.2.2
 pytest-asyncio==0.22.0
 pytest==7.4.4

backend/requirements/ee.txt

@@ -1,2 +1,3 @@
+cohere==5.6.1
+posthog==3.7.4
 python3-saml==1.15.0
-cohere==5.6.1

backend/scripts/restart_containers.sh

@@ -48,4 +48,7 @@ sleep 1
 echo "Running Alembic migration..."
 alembic upgrade head
 
+# Run the following instead of the above if using MT cloud
+# alembic -n schema_private upgrade head
+
 echo "Containers restarted and migration completed."