v1 answer refactor (onyx-dot-app#4721)

* v1 answer refactor

* fix tests

* good catch, tests

* more cleanup

Author: evan-onyx

backend/onyx/agents/agent_search/dc_search_analysis/nodes/a1_search_objects.py

@@ -34,16 +34,14 @@ def search_objects(
     """
 
     graph_config = cast(GraphConfig, config["metadata"]["config"])
-    question = graph_config.inputs.search_request.query
+    question = graph_config.inputs.prompt_builder.raw_user_query
     search_tool = graph_config.tooling.search_tool
 
-    if search_tool is None or graph_config.inputs.search_request.persona is None:
+    if search_tool is None or graph_config.inputs.persona is None:
         raise ValueError("Search tool and persona must be provided for DivCon search")
 
     try:
-        instructions = graph_config.inputs.search_request.persona.prompts[
-            0
-        ].system_prompt
+        instructions = graph_config.inputs.persona.prompts[0].system_prompt
 
         agent_1_instructions = extract_section(
             instructions, "Agent Step 1:", "Agent Step 2:"

backend/onyx/agents/agent_search/dc_search_analysis/nodes/a2_research_object_source.py

@@ -35,18 +35,15 @@ def research_object_source(
     datetime.now()
 
     graph_config = cast(GraphConfig, config["metadata"]["config"])
-    graph_config.inputs.search_request.query
     search_tool = graph_config.tooling.search_tool
-    question = graph_config.inputs.search_request.query
+    question = graph_config.inputs.prompt_builder.raw_user_query
     object, document_source = state.object_source_combination
 
-    if search_tool is None or graph_config.inputs.search_request.persona is None:
+    if search_tool is None or graph_config.inputs.persona is None:
         raise ValueError("Search tool and persona must be provided for DivCon search")
 
     try:
-        instructions = graph_config.inputs.search_request.persona.prompts[
-            0
-        ].system_prompt
+        instructions = graph_config.inputs.persona.prompts[0].system_prompt
 
         agent_2_instructions = extract_section(
             instructions, "Agent Step 2:", "Agent Step 3:"

backend/onyx/agents/agent_search/dc_search_analysis/nodes/a3_structure_research_by_object.py

@@ -1,6 +1,4 @@
 from collections import defaultdict
-from datetime import datetime
-from typing import cast
 from typing import Dict
 from typing import List
 
@@ -11,7 +9,6 @@
 from onyx.agents.agent_search.dc_search_analysis.states import (
     ObjectResearchInformationUpdate,
 )
-from onyx.agents.agent_search.models import GraphConfig
 from onyx.agents.agent_search.shared_graph_utils.utils import write_custom_event
 from onyx.chat.models import AgentAnswerPiece
 from onyx.utils.logger import setup_logger
@@ -25,10 +22,6 @@ def structure_research_by_object(
     """
     LangGraph node to start the agentic search process.
     """
-    datetime.now()
-
-    graph_config = cast(GraphConfig, config["metadata"]["config"])
-    graph_config.inputs.search_request.query
 
     write_custom_event(
         "initial_agent_answer",

backend/onyx/agents/agent_search/dc_search_analysis/nodes/a4_consolidate_object_research.py

@@ -27,14 +27,13 @@ def consolidate_object_research(
     LangGraph node to start the agentic search process.
     """
     graph_config = cast(GraphConfig, config["metadata"]["config"])
-    graph_config.inputs.search_request.query
     search_tool = graph_config.tooling.search_tool
-    question = graph_config.inputs.search_request.query
+    question = graph_config.inputs.prompt_builder.raw_user_query
 
-    if search_tool is None or graph_config.inputs.search_request.persona is None:
+    if search_tool is None or graph_config.inputs.persona is None:
         raise ValueError("Search tool and persona must be provided for DivCon search")
 
-    instructions = graph_config.inputs.search_request.persona.prompts[0].system_prompt
+    instructions = graph_config.inputs.persona.prompts[0].system_prompt
 
     agent_4_instructions = extract_section(
         instructions, "Agent Step 4:", "Agent Step 5:"

backend/onyx/agents/agent_search/dc_search_analysis/nodes/a5_consolidate_research.py

@@ -28,10 +28,8 @@ def consolidate_research(
     """
     LangGraph node to start the agentic search process.
     """
-    datetime.now()
 
     graph_config = cast(GraphConfig, config["metadata"]["config"])
-    graph_config.inputs.search_request.query
 
     search_tool = graph_config.tooling.search_tool
 
@@ -46,11 +44,11 @@ def consolidate_research(
         writer,
     )
 
-    if search_tool is None or graph_config.inputs.search_request.persona is None:
+    if search_tool is None or graph_config.inputs.persona is None:
         raise ValueError("Search tool and persona must be provided for DivCon search")
 
     # Populate prompt
-    instructions = graph_config.inputs.search_request.persona.prompts[0].system_prompt
+    instructions = graph_config.inputs.persona.prompts[0].system_prompt
 
     try:
         agent_5_instructions = extract_section(

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

@@ -85,7 +85,7 @@ def generate_sub_answer(
     context_docs = dedup_sort_inference_section_list(context_docs)
 
     persona_contextualized_prompt = get_persona_agent_prompt_expressions(
-        graph_config.inputs.search_request.persona
+        graph_config.inputs.persona
     ).contextualized_prompt
 
     if len(context_docs) == 0:
@@ -106,7 +106,7 @@ def generate_sub_answer(
         fast_llm = graph_config.tooling.fast_llm
         msg = build_sub_question_answer_prompt(
             question=question,
-            original_question=graph_config.inputs.search_request.query,
+            original_question=graph_config.inputs.prompt_builder.raw_user_query,
             docs=context_docs,
             persona_specification=persona_contextualized_prompt,
             config=fast_llm.config,

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

@@ -105,7 +105,7 @@ def generate_initial_answer(
     node_start_time = datetime.now()
 
     graph_config = cast(GraphConfig, config["metadata"]["config"])
-    question = graph_config.inputs.search_request.query
+    question = graph_config.inputs.prompt_builder.raw_user_query
     prompt_enrichment_components = get_prompt_enrichment_components(graph_config)
 
     # get all documents cited in sub-questions

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

@@ -75,7 +75,7 @@ def decompose_orig_question(
     node_start_time = datetime.now()
 
     graph_config = cast(GraphConfig, config["metadata"]["config"])
-    question = graph_config.inputs.search_request.query
+    question = graph_config.inputs.prompt_builder.raw_user_query
     perform_initial_search_decomposition = (
         graph_config.behavior.perform_initial_search_decomposition
     )

backend/onyx/agents/agent_search/deep_search/initial/retrieve_orig_question_docs/nodes/format_orig_question_search_input.py

@@ -21,7 +21,7 @@ def format_orig_question_search_input(
     logger.debug("generate_raw_search_data")
     graph_config = cast(GraphConfig, config["metadata"]["config"])
     return ExpandedRetrievalInput(
-        question=graph_config.inputs.search_request.query,
+        question=graph_config.inputs.prompt_builder.raw_user_query,
         base_search=True,
         sub_question_id=None,  # This graph is always and only used for the original question
         log_messages=[],

backend/onyx/agents/agent_search/deep_search/main/edges.py

@@ -30,18 +30,18 @@ def route_initial_tool_choice(
     LangGraph edge to route to agent search.
     """
     agent_config = cast(GraphConfig, config["metadata"]["config"])
-    if state.tool_choice is not None:
-        if (
-            agent_config.behavior.use_agentic_search
-            and agent_config.tooling.search_tool is not None
-            and state.tool_choice.tool.name == agent_config.tooling.search_tool.name
-        ):
-            return "start_agent_search"
-        else:
-            return "call_tool"
-    else:
+    if state.tool_choice is None:
         return "logging_node"
 
+    if (
+        agent_config.behavior.use_agentic_search
+        and agent_config.tooling.search_tool is not None
+        and state.tool_choice.tool.name == agent_config.tooling.search_tool.name
+    ):
+        return "start_agent_search"
+    else:
+        return "call_tool"
+
 
 def parallelize_initial_sub_question_answering(
     state: MainState,