From 814f97c2c7eed51d85c6d6810f7fb6eacc378b01 Mon Sep 17 00:00:00 2001 From: Yuhong Sun Date: Sun, 15 Dec 2024 16:05:03 -0800 Subject: [PATCH] MT Cloud Monitoring (#3465) --- .../versions/91a0a4d62b14_milestone.py | 45 +++++++++ backend/ee/onyx/configs/app_configs.py | 6 ++ .../ee/onyx/server/tenants/provisioning.py | 14 +++ backend/ee/onyx/utils/telemetry.py | 14 +++ backend/onyx/auth/noauth_user.py | 6 +- backend/onyx/auth/users.py | 24 ++++- .../onyx/background/indexing/run_indexing.py | 11 +++ backend/onyx/chat/process_message.py | 31 ++++++ backend/onyx/configs/constants.py | 16 +++ backend/onyx/db/milestone.py | 99 +++++++++++++++++++ backend/onyx/db/models.py | 28 +++++- backend/onyx/server/documents/connector.py | 34 ++++++- backend/onyx/server/features/persona/api.py | 16 ++- backend/onyx/server/manage/slack_bot.py | 13 +++ .../server/query_and_chat/chat_backend.py | 16 ++- backend/onyx/utils/telemetry.py | 41 ++++++++ backend/requirements/dev.txt | 1 + backend/requirements/ee.txt | 3 +- backend/scripts/restart_containers.sh | 3 + 19 files changed, 413 insertions(+), 8 deletions(-) create mode 100644 backend/alembic/versions/91a0a4d62b14_milestone.py create mode 100644 backend/ee/onyx/utils/telemetry.py create mode 100644 backend/onyx/db/milestone.py diff --git a/backend/alembic/versions/91a0a4d62b14_milestone.py b/backend/alembic/versions/91a0a4d62b14_milestone.py new file mode 100644 index 00000000000..f42b676db9d --- /dev/null +++ b/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") diff --git a/backend/ee/onyx/configs/app_configs.py b/backend/ee/onyx/configs/app_configs.py index a23ab4cc96c..07688ff146b 100644 --- a/backend/ee/onyx/configs/app_configs.py +++ b/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" diff --git a/backend/ee/onyx/server/tenants/provisioning.py b/backend/ee/onyx/server/tenants/provisioning.py index 65c49e58da4..c28de5c1e06 100644 --- a/backend/ee/onyx/server/tenants/provisioning.py +++ b/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( diff --git a/backend/ee/onyx/utils/telemetry.py b/backend/ee/onyx/utils/telemetry.py new file mode 100644 index 00000000000..960ae46939a --- /dev/null +++ b/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) diff --git a/backend/onyx/auth/noauth_user.py b/backend/onyx/auth/noauth_user.py index e26032a9a99..ac1557030c3 100644 --- a/backend/onyx/auth/noauth_user.py +++ b/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="anonymous@onyx.app", + id=NO_AUTH_USER_ID, + email=NO_AUTH_USER_EMAIL, is_active=True, is_superuser=False, is_verified=True, diff --git a/backend/onyx/auth/users.py b/backend/onyx/auth/users.py index 3347a14cbf2..9aaec58233d 100644 --- a/backend/onyx/auth/users.py +++ b/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, diff --git a/backend/onyx/background/indexing/run_indexing.py b/backend/onyx/background/indexing/run_indexing.py index 942a06f8ea5..35351de2f9f 100644 --- a/backend/onyx/background/indexing/run_indexing.py +++ b/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" diff --git a/backend/onyx/chat/process_message.py b/backend/onyx/chat/process_message.py index 87dbb500e84..28d67d03231 100644 --- a/backend/onyx/chat/process_message.py +++ b/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 diff --git a/backend/onyx/configs/constants.py b/backend/onyx/configs/constants.py index 1725de71788..ab13f2c68bd 100644 --- a/backend/onyx/configs/constants.py +++ b/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 = "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() diff --git a/backend/onyx/db/milestone.py b/backend/onyx/db/milestone.py new file mode 100644 index 00000000000..4d38d367578 --- /dev/null +++ b/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 diff --git a/backend/onyx/db/models.py b/backend/onyx/db/models.py index 4dae6e7d760..353eba9c9da 100644 --- a/backend/onyx/db/models.py +++ b/backend/onyx/db/models.py @@ -37,7 +37,7 @@ 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" diff --git a/backend/onyx/server/documents/connector.py b/backend/onyx/server/documents/connector.py index 67bca86fb81..4fab885195a 100644 --- a/backend/onyx/server/documents/connector.py +++ b/backend/onyx/server/documents/connector.py @@ -21,6 +21,7 @@ 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 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: diff --git a/backend/onyx/server/features/persona/api.py b/backend/onyx/server/features/persona/api.py index 497bbf0aacf..9911bfc184c 100644 --- a/backend/onyx/server/features/persona/api.py +++ b/backend/onyx/server/features/persona/api.py @@ -15,7 +15,9 @@ 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.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 diff --git a/backend/onyx/server/manage/slack_bot.py b/backend/onyx/server/manage/slack_bot.py index e84529ef066..44c283bbb48 100644 --- a/backend/onyx/server/manage/slack_bot.py +++ b/backend/onyx/server/manage/slack_bot.py @@ -4,7 +4,9 @@ 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 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) diff --git a/backend/onyx/server/query_and_chat/chat_backend.py b/backend/onyx/server/query_and_chat/chat_backend.py index 4291004a51d..21499032980 100644 --- a/backend/onyx/server/query_and_chat/chat_backend.py +++ b/backend/onyx/server/query_and_chat/chat_backend.py @@ -30,6 +30,7 @@ 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 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.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( diff --git a/backend/onyx/utils/telemetry.py b/backend/onyx/utils/telemetry.py index 33727cd5c2e..23a239a5564 100644 --- a/backend/onyx/utils/telemetry.py +++ b/backend/onyx/utils/telemetry.py @@ -10,10 +10,17 @@ 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 telemetry_logic() -> None: 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, + ) diff --git a/backend/requirements/dev.txt b/backend/requirements/dev.txt index a89b8db674d..4b2b227f445 100644 --- a/backend/requirements/dev.txt +++ b/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 diff --git a/backend/requirements/ee.txt b/backend/requirements/ee.txt index 18dc3200bdc..e587869fd75 100644 --- a/backend/requirements/ee.txt +++ b/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 \ No newline at end of file diff --git a/backend/scripts/restart_containers.sh b/backend/scripts/restart_containers.sh index adf5929c0f4..98316cd5d89 100755 --- a/backend/scripts/restart_containers.sh +++ b/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."