diff --git a/backend/onyx/auth/users.py b/backend/onyx/auth/users.py index f85294720..aff5da9d4 100644 --- a/backend/onyx/auth/users.py +++ b/backend/onyx/auth/users.py @@ -105,6 +105,7 @@ 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.timing import log_function_time from onyx.utils.url import add_url_params from onyx.utils.variable_functionality import fetch_ee_implementation_or_noop from onyx.utils.variable_functionality import fetch_versioned_implementation @@ -363,6 +364,7 @@ class UserManager(UUIDIDMixin, BaseUserManager[User, uuid.UUID]): ) return + @log_function_time(print_only=True) async def oauth_callback( self, oauth_name: str, @@ -609,6 +611,7 @@ class UserManager(UUIDIDMixin, BaseUserManager[User, uuid.UUID]): user.email, token, new_organization=user_count == 1 ) + @log_function_time(print_only=True) async def authenticate( self, credentials: OAuth2PasswordRequestForm ) -> Optional[User]: @@ -1235,6 +1238,7 @@ def get_oauth_router( return OAuth2AuthorizeResponse(authorization_url=authorization_url) + @log_function_time(print_only=True) @router.get( "/callback", name=callback_route_name, diff --git a/backend/onyx/chat/process_message.py b/backend/onyx/chat/process_message.py index f2aaeeb5c..2e80d44b4 100644 --- a/backend/onyx/chat/process_message.py +++ b/backend/onyx/chat/process_message.py @@ -1,3 +1,4 @@ +import time import traceback from collections import defaultdict from collections.abc import Callable @@ -914,6 +915,7 @@ def stream_chat_message_objects( retrieval_options.filters.user_folder_ids = user_folder_ids # Create override kwargs for the search tool + override_kwargs = SearchToolOverrideKwargs( force_no_rerank=search_for_ordering_only, # Skip reranking for ordering-only alternate_db_session=None, @@ -1109,9 +1111,6 @@ def stream_chat_message_objects( logger.info( f"ORDERING: Processing search results for ordering {len(user_files)} user files" ) - import time - - ordering_start = time.time() # Extract document order from search results doc_order = [] @@ -1147,8 +1146,6 @@ def stream_chat_message_objects( if f_id in file_id_to_user_file ] - time.time() - ordering_start - yield UserKnowledgeFilePacket( user_files=[ FileDescriptor( @@ -1436,6 +1433,7 @@ def stream_chat_message( custom_tool_additional_headers: dict[str, str] | None = None, is_connected: Callable[[], bool] | None = None, ) -> Iterator[str]: + start_time = time.time() with get_session_context_manager() as db_session: objects = stream_chat_message_objects( new_msg_req=new_msg_req, @@ -1446,6 +1444,11 @@ def stream_chat_message( is_connected=is_connected, ) for obj in objects: + # Check if this is a QADocsResponse with document results + if isinstance(obj, QADocsResponse): + document_retrieval_latency = time.time() - start_time + logger.debug(f"First doc time: {document_retrieval_latency}") + yield get_json_line(obj.model_dump()) diff --git a/backend/onyx/server/query_and_chat/chat_backend.py b/backend/onyx/server/query_and_chat/chat_backend.py index 55ac428e5..998d29515 100644 --- a/backend/onyx/server/query_and_chat/chat_backend.py +++ b/backend/onyx/server/query_and_chat/chat_backend.py @@ -448,7 +448,7 @@ def handle_new_chat_message( ), is_connected=is_connected_func, ): - yield json.dumps(packet) if isinstance(packet, dict) else packet + yield packet except Exception as e: logger.exception("Error in chat message streaming")