StreamWriter w/o default values

backend/onyx/agents/agent_search/deep_search/initial/generate_individual_sub_answer/nodes/generate_sub_answer.py

@@ -39,7 +39,7 @@ logger = setup_logger()
 def generate_sub_answer(
     state: AnswerQuestionState,
     config: RunnableConfig,
-    writer: StreamWriter = lambda _: None,
+    writer: StreamWriter,
 ) -> SubQuestionAnswerGenerationUpdate:
     """
     LangGraph node to generate a sub-answer.

backend/onyx/agents/agent_search/deep_search/initial/generate_initial_answer/nodes/generate_initial_answer.py

@@ -61,7 +61,7 @@ from onyx.tools.tool_implementations.search.search_tool import yield_search_resp
 def generate_initial_answer(
     state: SubQuestionRetrievalState,
     config: RunnableConfig,
-    writer: StreamWriter = lambda _: None,
+    writer: StreamWriter,
 ) -> InitialAnswerUpdate:
     """
     LangGraph node to generate the initial answer, using the initial sub-questions/sub-answers and the

backend/onyx/agents/agent_search/deep_search/initial/generate_sub_answers/nodes/decompose_orig_question.py

@@ -46,7 +46,7 @@ logger = setup_logger()
 def decompose_orig_question(
     state: SubQuestionRetrievalState,
     config: RunnableConfig,
-    writer: StreamWriter = lambda _: None,
+    writer: StreamWriter,
 ) -> InitialQuestionDecompositionUpdate:
     """
     LangGraph node to decompose the original question into sub-questions.

backend/onyx/agents/agent_search/deep_search/shared/expanded_retrieval/nodes/expand_queries.py

@@ -29,7 +29,7 @@ from onyx.prompts.agent_search import (
 def expand_queries(
     state: ExpandedRetrievalInput,
     config: RunnableConfig,
-    writer: StreamWriter = lambda _: None,
+    writer: StreamWriter,
 ) -> QueryExpansionUpdate:
     """
     LangGraph node to expand a question into multiple search queries.

backend/onyx/agents/agent_search/deep_search/shared/expanded_retrieval/nodes/format_results.py

@@ -28,7 +28,7 @@ from onyx.tools.tool_implementations.search.search_tool import yield_search_resp
 def format_results(
     state: ExpandedRetrievalState,
     config: RunnableConfig,
-    writer: StreamWriter = lambda _: None,
+    writer: StreamWriter,
 ) -> ExpandedRetrievalUpdate:
     """
     LangGraph node that constructs the proper expanded retrieval format.

backend/onyx/agents/agent_search/orchestration/nodes/llm_tool_choice.py

@@ -26,9 +26,7 @@ logger = setup_logger()
 # from the state and config
 # TODO: fan-out to multiple tool call nodes? Make this configurable?
 def llm_tool_choice(
-    state: ToolChoiceState,
-    config: RunnableConfig,
-    writer: StreamWriter = lambda _: None,
+    state: ToolChoiceState, config: RunnableConfig, writer: StreamWriter
 ) -> ToolChoiceUpdate:
     """
     This node is responsible for calling the LLM to choose a tool. If no tool is chosen,

backend/onyx/agents/agent_search/orchestration/nodes/tool_call.py

@@ -27,7 +27,7 @@ def emit_packet(packet: AnswerPacket, writer: StreamWriter) -> None:
 def tool_call(
     state: ToolChoiceUpdate,
     config: RunnableConfig,
-    writer: StreamWriter = lambda _: None,
+    writer: StreamWriter,
 ) -> ToolCallUpdate:
     """Calls the tool specified in the state and updates the state with the result"""
 

backend/onyx/agents/agent_search/shared_graph_utils/models.py

@@ -12,20 +12,6 @@ from onyx.context.search.models import InferenceSection
 from onyx.tools.models import SearchQueryInfo
 
 
-# Pydantic models for structured outputs
-# class RewrittenQueries(BaseModel):
-#     rewritten_queries: list[str]
-
-
-# class BinaryDecision(BaseModel):
-#     decision: Literal["yes", "no"]
-
-
-# class BinaryDecisionWithReasoning(BaseModel):
-#     reasoning: str
-#     decision: Literal["yes", "no"]
-
-
 class RetrievalFitScoreMetrics(BaseModel):
     scores: dict[str, float]
     chunk_ids: list[str]

backend/onyx/server/query_and_chat/chat_backend.py

@@ -420,6 +420,10 @@ def handle_new_chat_message(
 
     def stream_generator() -> Generator[str, None, None]:
         try:
+            import time
+
+            start_time = time.time()
+            n = 0
             for packet in stream_chat_message(
                 new_msg_req=chat_message_req,
                 user=user,
@@ -431,6 +435,19 @@ def handle_new_chat_message(
                 ),
                 is_connected=is_connected_func,
             ):
+                if "top_documents" in packet:
+                    to_first_docs = time.time() - start_time
+                    print(f"Time to first docs: {to_first_docs}")
+                    print(packet)
+                elif "answer_piece" in packet:
+                    to_answer_piece = time.time() - start_time
+                    if n == 1:
+                        print(f"Time to answer piece: {to_answer_piece}")
+                        print(packet)
+                    n += 1
+
+                # time_since_start = time.time() - start_time
+                # print(f"Time since start: {time_since_start}")
                 yield json.dumps(packet) if isinstance(packet, dict) else packet
 
         except Exception as e: