Add latency metrics (#4472)

* k

* update

* Update chat_backend.py

nit

---------

Co-authored-by: evan-danswer <evan@danswer.ai>

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,

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())
 
 

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")