From c7dd6c6e1d3b237133c698e4147cf7f5395e3ed5 Mon Sep 17 00:00:00 2001 From: Yuhong Sun Date: Sat, 14 Dec 2024 13:45:22 -0800 Subject: [PATCH] k --- backend/onyx/auth/noauth_user.py | 6 +- backend/onyx/chat/process_message.py | 32 +++++++++++ backend/onyx/configs/constants.py | 3 + backend/onyx/db/milestone.py | 57 +++++++++++++++++++ backend/onyx/db/models.py | 2 +- backend/onyx/server/features/persona/api.py | 17 +++++- backend/onyx/server/manage/slack_bot.py | 14 +++++ .../server/query_and_chat/chat_backend.py | 2 +- 8 files changed, 128 insertions(+), 5 deletions(-) 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/chat/process_message.py b/backend/onyx/chat/process_message.py index 87dbb500e84..61350b04c2a 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,32 @@ 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, + tenant_id=tenant_id, + 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 48dbbf0f7e2..11ff69840e8 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" diff --git a/backend/onyx/db/milestone.py b/backend/onyx/db/milestone.py index 5982a04782a..a21fbdeb3cf 100644 --- a/backend/onyx/db/milestone.py +++ b/backend/onyx/db/milestone.py @@ -6,6 +6,10 @@ from onyx.db.models import User +USER_ASSISTANT_PREFIX = "user_assistants_used_" +MULTI_ASSISTANT_USED = "multi_assistant_used" + + def create_milestone( user: User | None, tenant_id: str | None, @@ -50,3 +54,56 @@ def create_milestone_if_not_exists( return create_milestone(user, tenant_id, event_type, db_session), True return milestones[0], False + + +def update_milestone_event_tracker( + milestone: Milestone, + event_update: dict, + db_session: Session, +) -> None: + if milestone.event_tracker is None: + milestone.event_tracker = event_update + else: + milestone.event_tracker.update(event_update) + db_session.commit() + + +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 = {} + + user_key = f"{USER_ASSISTANT_PREFIX}{user_id}" + + if event_tracker.get(user_key) is None: + event_tracker[user_key] = [assistant_id] + else: + event_tracker[user_key].append(assistant_id) + + 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 + db_session.commit() + return True, True + + return False, False diff --git a/backend/onyx/db/models.py b/backend/onyx/db/models.py index 2bf436df855..ce1efb957c2 100644 --- a/backend/onyx/db/models.py +++ b/backend/onyx/db/models.py @@ -1549,7 +1549,7 @@ class Milestone(Base): ) 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[list[int] | None] = mapped_column( + event_tracker: Mapped[dict | None] = mapped_column( postgresql.JSONB(), nullable=True ) time_created: Mapped[datetime.datetime] = mapped_column( diff --git a/backend/onyx/server/features/persona/api.py b/backend/onyx/server/features/persona/api.py index 497bbf0aacf..876aaaa533b 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,26 @@ 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, + tenant_id=tenant_id, + 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..d13f779e475 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,16 @@ 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, + tenant_id=tenant_id, + 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 acb0f5ea978..9a8d0027ac8 100644 --- a/backend/onyx/server/query_and_chat/chat_backend.py +++ b/backend/onyx/server/query_and_chat/chat_backend.py @@ -319,7 +319,7 @@ 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: