From 8c42f463c3b855618fb57cd20b661b59c0c6d894 Mon Sep 17 00:00:00 2001 From: pablonyx Date: Mon, 7 Apr 2025 13:28:39 -0700 Subject: [PATCH 1/3] k --- backend/onyx/auth/users.py | 4 ++++ backend/onyx/chat/process_message.py | 6 +----- backend/onyx/server/query_and_chat/chat_backend.py | 5 +++++ 3 files changed, 10 insertions(+), 5 deletions(-) diff --git a/backend/onyx/auth/users.py b/backend/onyx/auth/users.py index f852947200..aff5da9d46 100644 --- a/backend/onyx/auth/users.py +++ b/backend/onyx/auth/users.py @@ -105,6 +105,7 @@ 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 @@ async def validate_password(self, password: str, _: schemas.UC | models.UP) -> N ) return + @log_function_time(print_only=True) async def oauth_callback( self, oauth_name: str, @@ -609,6 +611,7 @@ async def on_after_request_verify( 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 @@ async def authorize( 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 eea541aca7..7b52c271b3 100644 --- a/backend/onyx/chat/process_message.py +++ b/backend/onyx/chat/process_message.py @@ -914,6 +914,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, @@ -1113,9 +1114,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 = [] @@ -1151,8 +1149,6 @@ def stream_chat_message_objects( if f_id in file_id_to_user_file ] - time.time() - ordering_start - yield UserKnowledgeFilePacket( user_files=[ FileDescriptor( diff --git a/backend/onyx/server/query_and_chat/chat_backend.py b/backend/onyx/server/query_and_chat/chat_backend.py index 55ac428e59..f5155218bf 100644 --- a/backend/onyx/server/query_and_chat/chat_backend.py +++ b/backend/onyx/server/query_and_chat/chat_backend.py @@ -418,6 +418,7 @@ def handle_new_chat_message( """ tenant_id = get_current_tenant_id() logger.debug(f"Received new chat message: {chat_message_req.message}") + start_time = time.time() if ( not chat_message_req.message @@ -448,6 +449,10 @@ def stream_generator() -> Generator[str, None, None]: ), is_connected=is_connected_func, ): + if isinstance(packet, dict) and "top_documents" in packet: + document_retrieval_latency = time.time() - start_time + logger.debug(f"First doc time: {document_retrieval_latency}") + yield json.dumps(packet) if isinstance(packet, dict) else packet except Exception as e: From 80c2c0349622d2d0ef360fbec8d977140c9885ba Mon Sep 17 00:00:00 2001 From: pablonyx Date: Tue, 8 Apr 2025 10:17:44 -0700 Subject: [PATCH 2/3] update --- backend/onyx/chat/process_message.py | 7 +++++++ backend/onyx/server/query_and_chat/chat_backend.py | 8 ++------ 2 files changed, 9 insertions(+), 6 deletions(-) diff --git a/backend/onyx/chat/process_message.py b/backend/onyx/chat/process_message.py index 7b52c271b3..f8c92a9cf3 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 @@ -1436,6 +1437,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 +1448,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 f5155218bf..75026752e5 100644 --- a/backend/onyx/server/query_and_chat/chat_backend.py +++ b/backend/onyx/server/query_and_chat/chat_backend.py @@ -418,7 +418,7 @@ def handle_new_chat_message( """ tenant_id = get_current_tenant_id() logger.debug(f"Received new chat message: {chat_message_req.message}") - start_time = time.time() + time.time() if ( not chat_message_req.message @@ -449,11 +449,7 @@ def stream_generator() -> Generator[str, None, None]: ), is_connected=is_connected_func, ): - if isinstance(packet, dict) and "top_documents" in packet: - document_retrieval_latency = time.time() - start_time - logger.debug(f"First doc time: {document_retrieval_latency}") - - yield json.dumps(packet) if isinstance(packet, dict) else packet + yield packet except Exception as e: logger.exception("Error in chat message streaming") From 6389b1f2d33500a803b70c3d826573c97372db54 Mon Sep 17 00:00:00 2001 From: evan-danswer Date: Tue, 8 Apr 2025 13:37:29 -0700 Subject: [PATCH 3/3] Update chat_backend.py nit --- backend/onyx/server/query_and_chat/chat_backend.py | 1 - 1 file changed, 1 deletion(-) diff --git a/backend/onyx/server/query_and_chat/chat_backend.py b/backend/onyx/server/query_and_chat/chat_backend.py index 75026752e5..998d295152 100644 --- a/backend/onyx/server/query_and_chat/chat_backend.py +++ b/backend/onyx/server/query_and_chat/chat_backend.py @@ -418,7 +418,6 @@ def handle_new_chat_message( """ tenant_id = get_current_tenant_id() logger.debug(f"Received new chat message: {chat_message_req.message}") - time.time() if ( not chat_message_req.message