renamed directories, prompts, and small citation fix

backend/onyx/agents/agent_search/deep_search_a/initial/individual_sub_answer_generation/edges.py

@@ -3,10 +3,10 @@ from datetime import datetime
 
 from langgraph.types import Send
 
-from onyx.agents.agent_search.deep_search_a.initial__individual_sub_answer__subgraph.states import (
+from onyx.agents.agent_search.deep_search_a.initial.individual_sub_answer_generation.states import (
     AnswerQuestionInput,
 )
-from onyx.agents.agent_search.deep_search_a.util__expanded_retrieval__subgraph.states import (
+from onyx.agents.agent_search.deep_search_a.shared.expanded_retrieval.states import (
     ExpandedRetrievalInput,
 )
 from onyx.utils.logger import setup_logger

backend/onyx/agents/agent_search/deep_search_a/initial/individual_sub_answer_generation/graph_builder.py

@@ -2,31 +2,31 @@ from langgraph.graph import END
 from langgraph.graph import START
 from langgraph.graph import StateGraph
 
-from onyx.agents.agent_search.deep_search_a.initial__individual_sub_answer__subgraph.edges import (
+from onyx.agents.agent_search.deep_search_a.initial.individual_sub_answer_generation.edges import (
     send_to_expanded_retrieval,
 )
-from onyx.agents.agent_search.deep_search_a.initial__individual_sub_answer__subgraph.nodes.answer_check import (
+from onyx.agents.agent_search.deep_search_a.initial.individual_sub_answer_generation.nodes.answer_check import (
     answer_check,
 )
-from onyx.agents.agent_search.deep_search_a.initial__individual_sub_answer__subgraph.nodes.answer_generation import (
+from onyx.agents.agent_search.deep_search_a.initial.individual_sub_answer_generation.nodes.answer_generation import (
     answer_generation,
 )
-from onyx.agents.agent_search.deep_search_a.initial__individual_sub_answer__subgraph.nodes.format_answer import (
+from onyx.agents.agent_search.deep_search_a.initial.individual_sub_answer_generation.nodes.format_answer import (
     format_answer,
 )
-from onyx.agents.agent_search.deep_search_a.initial__individual_sub_answer__subgraph.nodes.ingest_retrieval import (
+from onyx.agents.agent_search.deep_search_a.initial.individual_sub_answer_generation.nodes.ingest_retrieval import (
     ingest_retrieval,
 )
-from onyx.agents.agent_search.deep_search_a.initial__individual_sub_answer__subgraph.states import (
+from onyx.agents.agent_search.deep_search_a.initial.individual_sub_answer_generation.states import (
     AnswerQuestionInput,
 )
-from onyx.agents.agent_search.deep_search_a.initial__individual_sub_answer__subgraph.states import (
+from onyx.agents.agent_search.deep_search_a.initial.individual_sub_answer_generation.states import (
     AnswerQuestionOutput,
 )
-from onyx.agents.agent_search.deep_search_a.initial__individual_sub_answer__subgraph.states import (
+from onyx.agents.agent_search.deep_search_a.initial.individual_sub_answer_generation.states import (
     AnswerQuestionState,
 )
-from onyx.agents.agent_search.deep_search_a.util__expanded_retrieval__subgraph.graph_builder import (
+from onyx.agents.agent_search.deep_search_a.shared.expanded_retrieval.graph_builder import (
     expanded_retrieval_graph_builder,
 )
 from onyx.agents.agent_search.shared_graph_utils.utils import get_test_config

backend/onyx/agents/agent_search/deep_search_a/initial/individual_sub_answer_generation/nodes/answer_check.py

@@ -5,10 +5,10 @@ from langchain_core.messages import HumanMessage
 from langchain_core.messages import merge_message_runs
 from langchain_core.runnables.config import RunnableConfig
 
-from onyx.agents.agent_search.deep_search_a.initial__individual_sub_answer__subgraph.states import (
+from onyx.agents.agent_search.deep_search_a.initial.individual_sub_answer_generation.states import (
     AnswerQuestionState,
 )
-from onyx.agents.agent_search.deep_search_a.initial__individual_sub_answer__subgraph.states import (
+from onyx.agents.agent_search.deep_search_a.initial.individual_sub_answer_generation.states import (
     QACheckUpdate,
 )
 from onyx.agents.agent_search.models import AgentSearchConfig

backend/onyx/agents/agent_search/deep_search_a/initial/individual_sub_answer_generation/nodes/answer_generation.py

@@ -6,10 +6,10 @@ from langchain_core.callbacks.manager import dispatch_custom_event
 from langchain_core.messages import merge_message_runs
 from langchain_core.runnables.config import RunnableConfig
 
-from onyx.agents.agent_search.deep_search_a.initial__individual_sub_answer__subgraph.states import (
+from onyx.agents.agent_search.deep_search_a.initial.individual_sub_answer_generation.states import (
     AnswerQuestionState,
 )
-from onyx.agents.agent_search.deep_search_a.initial__individual_sub_answer__subgraph.states import (
+from onyx.agents.agent_search.deep_search_a.initial.individual_sub_answer_generation.states import (
     QAGenerationUpdate,
 )
 from onyx.agents.agent_search.models import AgentSearchConfig
@@ -102,7 +102,9 @@ def answer_generation(
         )
 
     answer_citation_ids = get_answer_citation_ids(answer_str)
-    cited_docs = [context_docs[id] for id in answer_citation_ids]
+    cited_docs = [
+        context_docs[id] for id in answer_citation_ids if id < len(context_docs)
+    ]
 
     stop_event = StreamStopInfo(
         stop_reason=StreamStopReason.FINISHED,

backend/onyx/agents/agent_search/deep_search_a/initial/individual_sub_answer_generation/nodes/format_answer.py

@@ -1,7 +1,7 @@
-from onyx.agents.agent_search.deep_search_a.initial__individual_sub_answer__subgraph.states import (
+from onyx.agents.agent_search.deep_search_a.initial.individual_sub_answer_generation.states import (
     AnswerQuestionOutput,
 )
-from onyx.agents.agent_search.deep_search_a.initial__individual_sub_answer__subgraph.states import (
+from onyx.agents.agent_search.deep_search_a.initial.individual_sub_answer_generation.states import (
     AnswerQuestionState,
 )
 from onyx.agents.agent_search.shared_graph_utils.models import (

backend/onyx/agents/agent_search/deep_search_a/initial/individual_sub_answer_generation/nodes/ingest_retrieval.py

@@ -1,7 +1,7 @@
-from onyx.agents.agent_search.deep_search_a.initial__individual_sub_answer__subgraph.states import (
+from onyx.agents.agent_search.deep_search_a.initial.individual_sub_answer_generation.states import (
     RetrievalIngestionUpdate,
 )
-from onyx.agents.agent_search.deep_search_a.util__expanded_retrieval__subgraph.states import (
+from onyx.agents.agent_search.deep_search_a.shared.expanded_retrieval.states import (
     ExpandedRetrievalOutput,
 )
 from onyx.agents.agent_search.shared_graph_utils.models import AgentChunkStats

backend/onyx/agents/agent_search/deep_search_a/initial/initial_answer_generation/edges.py

@@ -3,13 +3,13 @@ from datetime import datetime
 
 from langgraph.types import Send
 
-from onyx.agents.agent_search.deep_search_a.initial__individual_sub_answer__subgraph.states import (
+from onyx.agents.agent_search.deep_search_a.initial.individual_sub_answer_generation.states import (
     AnswerQuestionInput,
 )
-from onyx.agents.agent_search.deep_search_a.initial__individual_sub_answer__subgraph.states import (
+from onyx.agents.agent_search.deep_search_a.initial.individual_sub_answer_generation.states import (
     AnswerQuestionOutput,
 )
-from onyx.agents.agent_search.deep_search_a.initial__retrieval_sub_answers__subgraph.states import (
+from onyx.agents.agent_search.deep_search_a.initial.initial_answer_generation.states import (
     SearchSQState,
 )
 from onyx.agents.agent_search.shared_graph_utils.utils import make_question_id

backend/onyx/agents/agent_search/deep_search_a/initial/initial_answer_generation/graph_builder.py

@@ -2,27 +2,27 @@ from langgraph.graph import END
 from langgraph.graph import START
 from langgraph.graph import StateGraph
 
-from onyx.agents.agent_search.deep_search_a.initial__consolidate_sub_answers__subgraph.graph_builder import (
-    initial_sq_subgraph_builder,
-)
-from onyx.agents.agent_search.deep_search_a.initial__retrieval__subgraph.graph_builder import (
-    base_raw_search_graph_builder,
-)
-from onyx.agents.agent_search.deep_search_a.initial__retrieval_sub_answers__subgraph.nodes.generate_initial_answer import (
+from onyx.agents.agent_search.deep_search_a.initial.initial_answer_generation.nodes.generate_initial_answer import (
     generate_initial_answer,
 )
-from onyx.agents.agent_search.deep_search_a.initial__retrieval_sub_answers__subgraph.nodes.initial_answer_quality_check import (
+from onyx.agents.agent_search.deep_search_a.initial.initial_answer_generation.nodes.initial_answer_quality_check import (
     initial_answer_quality_check,
 )
-from onyx.agents.agent_search.deep_search_a.initial__retrieval_sub_answers__subgraph.nodes.retrieval_consolidation import (
+from onyx.agents.agent_search.deep_search_a.initial.initial_answer_generation.nodes.retrieval_consolidation import (
     retrieval_consolidation,
 )
-from onyx.agents.agent_search.deep_search_a.initial__retrieval_sub_answers__subgraph.states import (
+from onyx.agents.agent_search.deep_search_a.initial.initial_answer_generation.states import (
     SearchSQInput,
 )
-from onyx.agents.agent_search.deep_search_a.initial__retrieval_sub_answers__subgraph.states import (
+from onyx.agents.agent_search.deep_search_a.initial.initial_answer_generation.states import (
     SearchSQState,
 )
+from onyx.agents.agent_search.deep_search_a.initial.retrieval.graph_builder import (
+    base_raw_search_graph_builder,
+)
+from onyx.agents.agent_search.deep_search_a.initial.sub_answer_consolidation.graph_builder import (
+    initial_sq_subgraph_builder,
+)
 from onyx.utils.logger import setup_logger
 
 logger = setup_logger()

backend/onyx/agents/agent_search/deep_search_a/initial/initial_answer_generation/nodes/generate_initial_answer.py

@@ -7,19 +7,19 @@ from langchain_core.messages import HumanMessage
 from langchain_core.messages import merge_content
 from langchain_core.runnables import RunnableConfig
 
-from onyx.agents.agent_search.deep_search_a.initial__retrieval_sub_answers__subgraph.states import (
+from onyx.agents.agent_search.deep_search_a.initial.initial_answer_generation.states import (
     SearchSQState,
 )
-from onyx.agents.agent_search.deep_search_a.main__graph.models import AgentBaseMetrics
-from onyx.agents.agent_search.deep_search_a.main__graph.operations import (
+from onyx.agents.agent_search.deep_search_a.main.models import AgentBaseMetrics
+from onyx.agents.agent_search.deep_search_a.main.operations import (
     calculate_initial_agent_stats,
 )
-from onyx.agents.agent_search.deep_search_a.main__graph.operations import get_query_info
-from onyx.agents.agent_search.deep_search_a.main__graph.operations import logger
-from onyx.agents.agent_search.deep_search_a.main__graph.operations import (
+from onyx.agents.agent_search.deep_search_a.main.operations import get_query_info
+from onyx.agents.agent_search.deep_search_a.main.operations import logger
+from onyx.agents.agent_search.deep_search_a.main.operations import (
     remove_document_citations,
 )
-from onyx.agents.agent_search.deep_search_a.main__graph.states import (
+from onyx.agents.agent_search.deep_search_a.main.states import (
     InitialAnswerUpdate,
 )
 from onyx.agents.agent_search.models import AgentSearchConfig
@@ -127,7 +127,7 @@ def generate_initial_answer(
         )
 
     else:
-        decomp_answer_results = state.decomp_answer_results
+        decomp_answer_results = state.sub_question_results
 
         good_qa_list: list[str] = []
 
@@ -220,7 +220,7 @@ def generate_initial_answer(
         answer = cast(str, response)
 
         initial_agent_stats = calculate_initial_agent_stats(
-            state.decomp_answer_results, state.original_question_retrieval_stats
+            state.sub_question_results, state.original_question_retrieval_stats
         )
 
         logger.debug(

backend/onyx/agents/agent_search/deep_search_a/initial/initial_answer_generation/nodes/ingest_initial_base_retrieval.py

@@ -1,10 +1,10 @@
 from datetime import datetime
 
-from onyx.agents.agent_search.deep_search_a.initial__retrieval__subgraph.states import (
+from onyx.agents.agent_search.deep_search_a.initial.retrieval.states import (
     BaseRawSearchOutput,
 )
-from onyx.agents.agent_search.deep_search_a.main__graph.operations import logger
-from onyx.agents.agent_search.deep_search_a.main__graph.states import (
+from onyx.agents.agent_search.deep_search_a.main.operations import logger
+from onyx.agents.agent_search.deep_search_a.main.states import (
     ExpandedRetrievalUpdate,
 )
 from onyx.agents.agent_search.shared_graph_utils.models import AgentChunkStats

backend/onyx/agents/agent_search/deep_search_a/initial/initial_answer_generation/nodes/initial_answer_quality_check.py

@@ -1,10 +1,10 @@
 from datetime import datetime
 
-from onyx.agents.agent_search.deep_search_a.initial__retrieval_sub_answers__subgraph.states import (
+from onyx.agents.agent_search.deep_search_a.initial.initial_answer_generation.states import (
     SearchSQState,
 )
-from onyx.agents.agent_search.deep_search_a.main__graph.operations import logger
-from onyx.agents.agent_search.deep_search_a.main__graph.states import (
+from onyx.agents.agent_search.deep_search_a.main.operations import logger
+from onyx.agents.agent_search.deep_search_a.main.states import (
     InitialAnswerQualityUpdate,
 )
 

backend/onyx/agents/agent_search/deep_search_a/initial/initial_answer_generation/nodes/retrieval_consolidation.py

@@ -1,9 +1,9 @@
 from datetime import datetime
 
-from onyx.agents.agent_search.deep_search_a.initial__retrieval_sub_answers__subgraph.states import (
+from onyx.agents.agent_search.deep_search_a.initial.initial_answer_generation.states import (
     SearchSQState,
 )
-from onyx.agents.agent_search.deep_search_a.main__graph.states import LoggerUpdate
+from onyx.agents.agent_search.deep_search_a.main.states import LoggerUpdate
 
 
 def retrieval_consolidation(

backend/onyx/agents/agent_search/deep_search_a/initial/initial_answer_generation/states.py

@@ -3,23 +3,23 @@ from typing import Annotated
 from typing import TypedDict
 
 from onyx.agents.agent_search.core_state import CoreState
-from onyx.agents.agent_search.deep_search_a.main__graph.states import BaseDecompUpdate
-from onyx.agents.agent_search.deep_search_a.main__graph.states import (
+from onyx.agents.agent_search.deep_search_a.main.states import BaseDecompUpdate
+from onyx.agents.agent_search.deep_search_a.main.states import (
     DecompAnswersUpdate,
 )
-from onyx.agents.agent_search.deep_search_a.main__graph.states import (
+from onyx.agents.agent_search.deep_search_a.main.states import (
     ExpandedRetrievalUpdate,
 )
-from onyx.agents.agent_search.deep_search_a.main__graph.states import (
+from onyx.agents.agent_search.deep_search_a.main.states import (
     ExploratorySearchUpdate,
 )
-from onyx.agents.agent_search.deep_search_a.main__graph.states import (
+from onyx.agents.agent_search.deep_search_a.main.states import (
     InitialAnswerQualityUpdate,
 )
-from onyx.agents.agent_search.deep_search_a.main__graph.states import (
+from onyx.agents.agent_search.deep_search_a.main.states import (
     InitialAnswerUpdate,
 )
-from onyx.agents.agent_search.deep_search_a.util__expanded_retrieval__subgraph.models import (
+from onyx.agents.agent_search.deep_search_a.shared.expanded_retrieval.models import (
     ExpandedRetrievalResult,
 )
 

backend/onyx/agents/agent_search/deep_search_a/initial/retrieval/graph_builder.py

@@ -2,25 +2,25 @@ from langgraph.graph import END
 from langgraph.graph import START
 from langgraph.graph import StateGraph
 
-from onyx.agents.agent_search.deep_search_a.initial__retrieval__subgraph.nodes.format_raw_search_results import (
+from onyx.agents.agent_search.deep_search_a.initial.retrieval.nodes.format_raw_search_results import (
     format_raw_search_results,
 )
-from onyx.agents.agent_search.deep_search_a.initial__retrieval__subgraph.nodes.generate_raw_search_data import (
+from onyx.agents.agent_search.deep_search_a.initial.retrieval.nodes.generate_raw_search_data import (
     generate_raw_search_data,
 )
-from onyx.agents.agent_search.deep_search_a.initial__retrieval__subgraph.nodes.ingest_initial_base_retrieval import (
+from onyx.agents.agent_search.deep_search_a.initial.retrieval.nodes.ingest_initial_base_retrieval import (
     ingest_initial_base_retrieval,
 )
-from onyx.agents.agent_search.deep_search_a.initial__retrieval__subgraph.states import (
+from onyx.agents.agent_search.deep_search_a.initial.retrieval.states import (
     BaseRawSearchInput,
 )
-from onyx.agents.agent_search.deep_search_a.initial__retrieval__subgraph.states import (
+from onyx.agents.agent_search.deep_search_a.initial.retrieval.states import (
     BaseRawSearchOutput,
 )
-from onyx.agents.agent_search.deep_search_a.initial__retrieval__subgraph.states import (
+from onyx.agents.agent_search.deep_search_a.initial.retrieval.states import (
     BaseRawSearchState,
 )
-from onyx.agents.agent_search.deep_search_a.util__expanded_retrieval__subgraph.graph_builder import (
+from onyx.agents.agent_search.deep_search_a.shared.expanded_retrieval.graph_builder import (
     expanded_retrieval_graph_builder,
 )
 

backend/onyx/agents/agent_search/deep_search_a/initial/retrieval/nodes/format_raw_search_results.py

@@ -1,7 +1,7 @@
-from onyx.agents.agent_search.deep_search_a.initial__retrieval__subgraph.states import (
+from onyx.agents.agent_search.deep_search_a.initial.retrieval.states import (
     BaseRawSearchOutput,
 )
-from onyx.agents.agent_search.deep_search_a.util__expanded_retrieval__subgraph.states import (
+from onyx.agents.agent_search.deep_search_a.shared.expanded_retrieval.states import (
     ExpandedRetrievalOutput,
 )
 from onyx.utils.logger import setup_logger

backend/onyx/agents/agent_search/deep_search_a/initial/retrieval/nodes/generate_raw_search_data.py

@@ -3,7 +3,7 @@ from typing import cast
 from langchain_core.runnables.config import RunnableConfig
 
 from onyx.agents.agent_search.core_state import CoreState
-from onyx.agents.agent_search.deep_search_a.util__expanded_retrieval__subgraph.states import (
+from onyx.agents.agent_search.deep_search_a.shared.expanded_retrieval.states import (
     ExpandedRetrievalInput,
 )
 from onyx.agents.agent_search.models import AgentSearchConfig

backend/onyx/agents/agent_search/deep_search_a/initial/retrieval/nodes/ingest_initial_base_retrieval.py

@@ -1,10 +1,10 @@
 from datetime import datetime
 
-from onyx.agents.agent_search.deep_search_a.initial__retrieval__subgraph.states import (
+from onyx.agents.agent_search.deep_search_a.initial.retrieval.states import (
     BaseRawSearchOutput,
 )
-from onyx.agents.agent_search.deep_search_a.main__graph.operations import logger
-from onyx.agents.agent_search.deep_search_a.main__graph.states import (
+from onyx.agents.agent_search.deep_search_a.main.operations import logger
+from onyx.agents.agent_search.deep_search_a.main.states import (
     ExpandedRetrievalUpdate,
 )
 from onyx.agents.agent_search.shared_graph_utils.models import AgentChunkStats

backend/onyx/agents/agent_search/deep_search_a/initial/retrieval/states.py

@@ -1,12 +1,12 @@
 from pydantic import BaseModel
 
-from onyx.agents.agent_search.deep_search_a.main__graph.states import (
+from onyx.agents.agent_search.deep_search_a.main.states import (
     ExpandedRetrievalUpdate,
 )
-from onyx.agents.agent_search.deep_search_a.util__expanded_retrieval__subgraph.models import (
+from onyx.agents.agent_search.deep_search_a.shared.expanded_retrieval.models import (
     ExpandedRetrievalResult,
 )
-from onyx.agents.agent_search.deep_search_a.util__expanded_retrieval__subgraph.states import (
+from onyx.agents.agent_search.deep_search_a.shared.expanded_retrieval.states import (
     ExpandedRetrievalInput,
 )
 

backend/onyx/agents/agent_search/deep_search_a/initial/sub_answer_consolidation/edges.py

@@ -3,13 +3,13 @@ from datetime import datetime
 
 from langgraph.types import Send
 
-from onyx.agents.agent_search.deep_search_a.initial__individual_sub_answer__subgraph.states import (
+from onyx.agents.agent_search.deep_search_a.initial.individual_sub_answer_generation.states import (
     AnswerQuestionInput,
 )
-from onyx.agents.agent_search.deep_search_a.initial__individual_sub_answer__subgraph.states import (
+from onyx.agents.agent_search.deep_search_a.initial.individual_sub_answer_generation.states import (
     AnswerQuestionOutput,
 )
-from onyx.agents.agent_search.deep_search_a.initial__retrieval_sub_answers__subgraph.states import (
+from onyx.agents.agent_search.deep_search_a.initial.initial_answer_generation.states import (
     SearchSQState,
 )
 from onyx.agents.agent_search.shared_graph_utils.utils import make_question_id

backend/onyx/agents/agent_search/deep_search_a/initial/sub_answer_consolidation/graph_builder.py

@@ -2,24 +2,24 @@ from langgraph.graph import END
 from langgraph.graph import START
 from langgraph.graph import StateGraph
 
-from onyx.agents.agent_search.deep_search_a.initial__consolidate_sub_answers__subgraph.edges import (
+from onyx.agents.agent_search.deep_search_a.initial.individual_sub_answer_generation.graph_builder import (
+    answer_query_graph_builder,
+)
+from onyx.agents.agent_search.deep_search_a.initial.sub_answer_consolidation.edges import (
     parallelize_initial_sub_question_answering,
 )
-from onyx.agents.agent_search.deep_search_a.initial__consolidate_sub_answers__subgraph.nodes.ingest_initial_sub_answers import (
+from onyx.agents.agent_search.deep_search_a.initial.sub_answer_consolidation.nodes.ingest_initial_sub_answers import (
     ingest_initial_sub_answers,
 )
-from onyx.agents.agent_search.deep_search_a.initial__consolidate_sub_answers__subgraph.nodes.initial_decomposition import (
+from onyx.agents.agent_search.deep_search_a.initial.sub_answer_consolidation.nodes.initial_decomposition import (
     initial_sub_question_creation,
 )
-from onyx.agents.agent_search.deep_search_a.initial__consolidate_sub_answers__subgraph.states import (
+from onyx.agents.agent_search.deep_search_a.initial.sub_answer_consolidation.states import (
     SQInput,
 )
-from onyx.agents.agent_search.deep_search_a.initial__consolidate_sub_answers__subgraph.states import (
+from onyx.agents.agent_search.deep_search_a.initial.sub_answer_consolidation.states import (
     SQState,
 )
-from onyx.agents.agent_search.deep_search_a.initial__individual_sub_answer__subgraph.graph_builder import (
-    answer_query_graph_builder,
-)
 from onyx.utils.logger import setup_logger
 
 logger = setup_logger()

backend/onyx/agents/agent_search/deep_search_a/initial/sub_answer_consolidation/nodes/ingest_initial_sub_answers.py

@@ -1,10 +1,10 @@
 from datetime import datetime
 
-from onyx.agents.agent_search.deep_search_a.initial__individual_sub_answer__subgraph.states import (
+from onyx.agents.agent_search.deep_search_a.initial.individual_sub_answer_generation.states import (
     AnswerQuestionOutput,
 )
-from onyx.agents.agent_search.deep_search_a.main__graph.operations import logger
-from onyx.agents.agent_search.deep_search_a.main__graph.states import (
+from onyx.agents.agent_search.deep_search_a.main.operations import logger
+from onyx.agents.agent_search.deep_search_a.main.states import (
     DecompAnswersUpdate,
 )
 from onyx.agents.agent_search.shared_graph_utils.operators import (
@@ -38,7 +38,7 @@ def ingest_initial_sub_answers(
         documents=dedup_inference_sections(documents, []),
         context_documents=dedup_inference_sections(context_documents, []),
         cited_docs=dedup_inference_sections(cited_docs, []),
-        decomp_answer_results=answer_results,
+        sub_question_results=answer_results,
         log_messages=[
             f"{now_start} -- Main - Ingest initial processed sub questions,  Time taken: {now_end - now_start}"
         ],

backend/onyx/agents/agent_search/deep_search_a/initial/sub_answer_consolidation/nodes/initial_decomposition.py

@@ -6,17 +6,17 @@ from langchain_core.messages import HumanMessage
 from langchain_core.messages import merge_content
 from langchain_core.runnables import RunnableConfig
 
-from onyx.agents.agent_search.deep_search_a.initial__retrieval_sub_answers__subgraph.states import (
+from onyx.agents.agent_search.deep_search_a.initial.initial_answer_generation.states import (
     SearchSQState,
 )
-from onyx.agents.agent_search.deep_search_a.main__graph.models import (
+from onyx.agents.agent_search.deep_search_a.main.models import (
     AgentRefinedMetrics,
 )
-from onyx.agents.agent_search.deep_search_a.main__graph.operations import (
+from onyx.agents.agent_search.deep_search_a.main.operations import (
     dispatch_subquestion,
 )
-from onyx.agents.agent_search.deep_search_a.main__graph.operations import logger
-from onyx.agents.agent_search.deep_search_a.main__graph.states import BaseDecompUpdate
+from onyx.agents.agent_search.deep_search_a.main.operations import logger
+from onyx.agents.agent_search.deep_search_a.main.states import BaseDecompUpdate
 from onyx.agents.agent_search.models import AgentSearchConfig
 from onyx.agents.agent_search.shared_graph_utils.agent_prompt_ops import (
     build_history_prompt,

backend/onyx/agents/agent_search/deep_search_a/initial/sub_answer_consolidation/states.py

@@ -1,11 +1,11 @@
 from typing import TypedDict
 
 from onyx.agents.agent_search.core_state import CoreState
-from onyx.agents.agent_search.deep_search_a.main__graph.states import BaseDecompUpdate
-from onyx.agents.agent_search.deep_search_a.main__graph.states import (
+from onyx.agents.agent_search.deep_search_a.main.states import BaseDecompUpdate
+from onyx.agents.agent_search.deep_search_a.main.states import (
     DecompAnswersUpdate,
 )
-from onyx.agents.agent_search.deep_search_a.main__graph.states import (
+from onyx.agents.agent_search.deep_search_a.main.states import (
     InitialAnswerUpdate,
 )
 

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

@@ -6,14 +6,14 @@ from typing import Literal
 from langchain_core.runnables import RunnableConfig
 from langgraph.types import Send
 
-from onyx.agents.agent_search.deep_search_a.initial__individual_sub_answer__subgraph.states import (
+from onyx.agents.agent_search.deep_search_a.initial.individual_sub_answer_generation.states import (
     AnswerQuestionInput,
 )
-from onyx.agents.agent_search.deep_search_a.initial__individual_sub_answer__subgraph.states import (
+from onyx.agents.agent_search.deep_search_a.initial.individual_sub_answer_generation.states import (
     AnswerQuestionOutput,
 )
-from onyx.agents.agent_search.deep_search_a.main__graph.states import MainState
-from onyx.agents.agent_search.deep_search_a.main__graph.states import (
+from onyx.agents.agent_search.deep_search_a.main.states import MainState
+from onyx.agents.agent_search.deep_search_a.main.states import (
     RequireRefinedAnswerUpdate,
 )
 from onyx.agents.agent_search.models import AgentSearchConfig

backend/onyx/agents/agent_search/deep_search_a/main/graph_builder.py

@@ -2,45 +2,45 @@ from langgraph.graph import END
 from langgraph.graph import START
 from langgraph.graph import StateGraph
 
-from onyx.agents.agent_search.deep_search_a.initial__retrieval_sub_answers__subgraph.graph_builder import (
+from onyx.agents.agent_search.deep_search_a.initial.initial_answer_generation.graph_builder import (
     initial_search_sq_subgraph_builder,
 )
-from onyx.agents.agent_search.deep_search_a.main__graph.edges import (
+from onyx.agents.agent_search.deep_search_a.main.edges import (
     continue_to_refined_answer_or_end,
 )
-from onyx.agents.agent_search.deep_search_a.main__graph.edges import (
+from onyx.agents.agent_search.deep_search_a.main.edges import (
     parallelize_refined_sub_question_answering,
 )
-from onyx.agents.agent_search.deep_search_a.main__graph.edges import (
+from onyx.agents.agent_search.deep_search_a.main.edges import (
     route_initial_tool_choice,
 )
-from onyx.agents.agent_search.deep_search_a.main__graph.nodes.agent_logging import (
+from onyx.agents.agent_search.deep_search_a.main.nodes.agent_logging import (
     agent_logging,
 )
-from onyx.agents.agent_search.deep_search_a.main__graph.nodes.agent_search_start import (
+from onyx.agents.agent_search.deep_search_a.main.nodes.agent_search_start import (
     agent_search_start,
 )
-from onyx.agents.agent_search.deep_search_a.main__graph.nodes.answer_comparison import (
+from onyx.agents.agent_search.deep_search_a.main.nodes.answer_comparison import (
     answer_comparison,
 )
-from onyx.agents.agent_search.deep_search_a.main__graph.nodes.entity_term_extraction_llm import (
+from onyx.agents.agent_search.deep_search_a.main.nodes.entity_term_extraction_llm import (
     entity_term_extraction_llm,
 )
-from onyx.agents.agent_search.deep_search_a.main__graph.nodes.generate_refined_answer import (
+from onyx.agents.agent_search.deep_search_a.main.nodes.generate_refined_answer import (
     generate_refined_answer,
 )
-from onyx.agents.agent_search.deep_search_a.main__graph.nodes.ingest_refined_answers import (
+from onyx.agents.agent_search.deep_search_a.main.nodes.ingest_refined_answers import (
     ingest_refined_answers,
 )
-from onyx.agents.agent_search.deep_search_a.main__graph.nodes.refined_answer_decision import (
+from onyx.agents.agent_search.deep_search_a.main.nodes.refined_answer_decision import (
     refined_answer_decision,
 )
-from onyx.agents.agent_search.deep_search_a.main__graph.nodes.refined_sub_question_creation import (
+from onyx.agents.agent_search.deep_search_a.main.nodes.refined_sub_question_creation import (
     refined_sub_question_creation,
 )
-from onyx.agents.agent_search.deep_search_a.main__graph.states import MainInput
-from onyx.agents.agent_search.deep_search_a.main__graph.states import MainState
-from onyx.agents.agent_search.deep_search_a.refinement__consolidate_sub_answers__subgraph.graph_builder import (
+from onyx.agents.agent_search.deep_search_a.main.states import MainInput
+from onyx.agents.agent_search.deep_search_a.main.states import MainState
+from onyx.agents.agent_search.deep_search_a.refininement.sub_answer_consolidation.graph_builder import (
     answer_refined_query_graph_builder,
 )
 from onyx.agents.agent_search.orchestration.nodes.basic_use_tool_response import (

backend/onyx/agents/agent_search/deep_search_a/main/nodes/agent_logging.py

@@ -3,13 +3,13 @@ from typing import cast
 
 from langchain_core.runnables import RunnableConfig
 
-from onyx.agents.agent_search.deep_search_a.main__graph.models import (
+from onyx.agents.agent_search.deep_search_a.main.models import (
     AgentAdditionalMetrics,
 )
-from onyx.agents.agent_search.deep_search_a.main__graph.models import AgentTimings
-from onyx.agents.agent_search.deep_search_a.main__graph.operations import logger
-from onyx.agents.agent_search.deep_search_a.main__graph.states import MainOutput
-from onyx.agents.agent_search.deep_search_a.main__graph.states import MainState
+from onyx.agents.agent_search.deep_search_a.main.models import AgentTimings
+from onyx.agents.agent_search.deep_search_a.main.operations import logger
+from onyx.agents.agent_search.deep_search_a.main.states import MainOutput
+from onyx.agents.agent_search.deep_search_a.main.states import MainState
 from onyx.agents.agent_search.models import AgentSearchConfig
 from onyx.agents.agent_search.shared_graph_utils.models import CombinedAgentMetrics
 from onyx.db.chat import log_agent_metrics
@@ -85,7 +85,7 @@ def agent_logging(state: MainState, config: RunnableConfig) -> MainOutput:
             db_session = agent_a_config.db_session
             chat_session_id = agent_a_config.chat_session_id
             primary_message_id = agent_a_config.message_id
-            sub_question_answer_results = state.decomp_answer_results
+            sub_question_answer_results = state.sub_question_results
 
             log_agent_sub_question_results(
                 db_session=db_session,

backend/onyx/agents/agent_search/deep_search_a/main/nodes/agent_path_decision.py

@@ -3,9 +3,9 @@ from typing import cast
 
 from langchain_core.runnables import RunnableConfig
 
-from onyx.agents.agent_search.deep_search_a.main__graph.operations import logger
-from onyx.agents.agent_search.deep_search_a.main__graph.states import MainState
-from onyx.agents.agent_search.deep_search_a.main__graph.states import RoutingDecision
+from onyx.agents.agent_search.deep_search_a.main.operations import logger
+from onyx.agents.agent_search.deep_search_a.main.states import MainState
+from onyx.agents.agent_search.deep_search_a.main.states import RoutingDecision
 from onyx.agents.agent_search.models import AgentSearchConfig
 
 

backend/onyx/agents/agent_search/deep_search_a/main/nodes/agent_path_routing.py

@@ -3,7 +3,7 @@ from typing import Literal
 
 from langgraph.types import Command
 
-from onyx.agents.agent_search.deep_search_a.main__graph.states import MainState
+from onyx.agents.agent_search.deep_search_a.main.states import MainState
 
 
 def agent_path_routing(

backend/onyx/agents/agent_search/deep_search_a/main/nodes/agent_search_start.py

@@ -3,11 +3,11 @@ from typing import cast
 
 from langchain_core.runnables import RunnableConfig
 
-from onyx.agents.agent_search.deep_search_a.main__graph.operations import logger
-from onyx.agents.agent_search.deep_search_a.main__graph.states import (
+from onyx.agents.agent_search.deep_search_a.main.operations import logger
+from onyx.agents.agent_search.deep_search_a.main.states import (
     ExploratorySearchUpdate,
 )
-from onyx.agents.agent_search.deep_search_a.main__graph.states import MainState
+from onyx.agents.agent_search.deep_search_a.main.states import MainState
 from onyx.agents.agent_search.models import AgentSearchConfig
 from onyx.agents.agent_search.shared_graph_utils.agent_prompt_ops import (
     build_history_prompt,

backend/onyx/agents/agent_search/deep_search_a/main/nodes/answer_comparison.py

@@ -5,9 +5,9 @@ from langchain_core.callbacks.manager import dispatch_custom_event
 from langchain_core.messages import HumanMessage
 from langchain_core.runnables import RunnableConfig
 
-from onyx.agents.agent_search.deep_search_a.main__graph.operations import logger
-from onyx.agents.agent_search.deep_search_a.main__graph.states import AnswerComparison
-from onyx.agents.agent_search.deep_search_a.main__graph.states import MainState
+from onyx.agents.agent_search.deep_search_a.main.operations import logger
+from onyx.agents.agent_search.deep_search_a.main.states import AnswerComparison
+from onyx.agents.agent_search.deep_search_a.main.states import MainState
 from onyx.agents.agent_search.models import AgentSearchConfig
 from onyx.agents.agent_search.shared_graph_utils.prompts import ANSWER_COMPARISON_PROMPT
 from onyx.chat.models import RefinedAnswerImprovement

backend/onyx/agents/agent_search/deep_search_a/main/nodes/direct_llm_handling.py

@@ -7,11 +7,11 @@ from langchain_core.messages import HumanMessage
 from langchain_core.messages import merge_content
 from langchain_core.runnables import RunnableConfig
 
-from onyx.agents.agent_search.deep_search_a.main__graph.operations import logger
-from onyx.agents.agent_search.deep_search_a.main__graph.states import (
+from onyx.agents.agent_search.deep_search_a.main.operations import logger
+from onyx.agents.agent_search.deep_search_a.main.states import (
     InitialAnswerUpdate,
 )
-from onyx.agents.agent_search.deep_search_a.main__graph.states import MainState
+from onyx.agents.agent_search.deep_search_a.main.states import MainState
 from onyx.agents.agent_search.models import AgentSearchConfig
 from onyx.agents.agent_search.shared_graph_utils.prompts import DIRECT_LLM_PROMPT
 from onyx.agents.agent_search.shared_graph_utils.utils import (

backend/onyx/agents/agent_search/deep_search_a/main/nodes/entity_term_extraction_llm.py

@@ -6,11 +6,11 @@ from typing import cast
 from langchain_core.messages import HumanMessage
 from langchain_core.runnables import RunnableConfig
 
-from onyx.agents.agent_search.deep_search_a.main__graph.operations import logger
-from onyx.agents.agent_search.deep_search_a.main__graph.states import (
+from onyx.agents.agent_search.deep_search_a.main.operations import logger
+from onyx.agents.agent_search.deep_search_a.main.states import (
     EntityTermExtractionUpdate,
 )
-from onyx.agents.agent_search.deep_search_a.main__graph.states import MainState
+from onyx.agents.agent_search.deep_search_a.main.states import MainState
 from onyx.agents.agent_search.models import AgentSearchConfig
 from onyx.agents.agent_search.shared_graph_utils.agent_prompt_ops import (
     trim_prompt_piece,

backend/onyx/agents/agent_search/deep_search_a/main/nodes/generate_initial_base_search_only_answer.py

@@ -4,11 +4,11 @@ from typing import cast
 from langchain_core.messages import HumanMessage
 from langchain_core.runnables import RunnableConfig
 
-from onyx.agents.agent_search.deep_search_a.main__graph.operations import logger
-from onyx.agents.agent_search.deep_search_a.main__graph.states import (
+from onyx.agents.agent_search.deep_search_a.main.operations import logger
+from onyx.agents.agent_search.deep_search_a.main.states import (
     InitialAnswerBASEUpdate,
 )
-from onyx.agents.agent_search.deep_search_a.main__graph.states import MainState
+from onyx.agents.agent_search.deep_search_a.main.states import MainState
 from onyx.agents.agent_search.models import AgentSearchConfig
 from onyx.agents.agent_search.shared_graph_utils.agent_prompt_ops import (
     trim_prompt_piece,

backend/onyx/agents/agent_search/deep_search_a/main/nodes/generate_refined_answer.py

@@ -7,16 +7,16 @@ from langchain_core.messages import HumanMessage
 from langchain_core.messages import merge_content
 from langchain_core.runnables import RunnableConfig
 
-from onyx.agents.agent_search.deep_search_a.main__graph.models import (
+from onyx.agents.agent_search.deep_search_a.main.models import (
     AgentRefinedMetrics,
 )
-from onyx.agents.agent_search.deep_search_a.main__graph.operations import get_query_info
-from onyx.agents.agent_search.deep_search_a.main__graph.operations import logger
-from onyx.agents.agent_search.deep_search_a.main__graph.operations import (
+from onyx.agents.agent_search.deep_search_a.main.operations import get_query_info
+from onyx.agents.agent_search.deep_search_a.main.operations import logger
+from onyx.agents.agent_search.deep_search_a.main.operations import (
     remove_document_citations,
 )
-from onyx.agents.agent_search.deep_search_a.main__graph.states import MainState
-from onyx.agents.agent_search.deep_search_a.main__graph.states import (
+from onyx.agents.agent_search.deep_search_a.main.states import MainState
+from onyx.agents.agent_search.deep_search_a.main.states import (
     RefinedAnswerUpdate,
 )
 from onyx.agents.agent_search.models import AgentSearchConfig
@@ -36,7 +36,7 @@ from onyx.agents.agent_search.shared_graph_utils.prompts import (
     REVISED_RAG_PROMPT_NO_SUB_QUESTIONS,
 )
 from onyx.agents.agent_search.shared_graph_utils.prompts import (
-    SUB_QUESTION_ANSWER_TEMPLATE,
+    SUB_QUESTION_ANSWER_TEMPLATE_REVISED,
 )
 from onyx.agents.agent_search.shared_graph_utils.prompts import UNKNOWN_ANSWER
 from onyx.agents.agent_search.shared_graph_utils.utils import (
@@ -119,7 +119,7 @@ def generate_refined_answer(
     else:
         revision_doc_effectiveness = 10.0
 
-    decomp_answer_results = state.decomp_answer_results
+    decomp_answer_results = state.sub_question_results
     # revised_answer_results = state.refined_decomp_answer_results
 
     answered_qa_list: list[str] = []
@@ -141,17 +141,20 @@ def generate_refined_answer(
             and len(decomp_answer_result.answer) > 0
             and decomp_answer_result.answer != UNKNOWN_ANSWER
         ):
+            if question_level == 0:
+                initial_good_sub_questions.append(decomp_answer_result.question)
+                sub_question_type = "initial"
+            else:
+                new_revised_good_sub_questions.append(decomp_answer_result.question)
+                sub_question_type = "refined"
             answered_qa_list.append(
-                SUB_QUESTION_ANSWER_TEMPLATE.format(
+                SUB_QUESTION_ANSWER_TEMPLATE_REVISED.format(
                     sub_question=decomp_answer_result.question,
                     sub_answer=decomp_answer_result.answer,
                     sub_question_nr=sub_question_nr,
+                    sub_question_type=sub_question_type,
                 )
             )
-            if question_level == 0:
-                initial_good_sub_questions.append(decomp_answer_result.question)
-            else:
-                new_revised_good_sub_questions.append(decomp_answer_result.question)
 
         sub_question_nr += 1
 

backend/onyx/agents/agent_search/deep_search_a/main/nodes/ingest_refined_answers.py

@@ -1,10 +1,10 @@
 from datetime import datetime
 
-from onyx.agents.agent_search.deep_search_a.initial__individual_sub_answer__subgraph.states import (
+from onyx.agents.agent_search.deep_search_a.initial.individual_sub_answer_generation.states import (
     AnswerQuestionOutput,
 )
-from onyx.agents.agent_search.deep_search_a.main__graph.operations import logger
-from onyx.agents.agent_search.deep_search_a.main__graph.states import (
+from onyx.agents.agent_search.deep_search_a.main.operations import logger
+from onyx.agents.agent_search.deep_search_a.main.states import (
     DecompAnswersUpdate,
 )
 from onyx.agents.agent_search.shared_graph_utils.operators import (
@@ -34,7 +34,7 @@ def ingest_refined_answers(
         # Deduping is done by the documents operator for the main graph
         # so we might not need to dedup here
         documents=dedup_inference_sections(documents, []),
-        decomp_answer_results=answer_results,
+        sub_question_results=answer_results,
         log_messages=[
             f"{now_start} -- Main - Ingest refined answers,  Time taken: {now_end - now_start}"
         ],

backend/onyx/agents/agent_search/deep_search_a/main/nodes/refined_answer_decision.py

@@ -3,9 +3,9 @@ from typing import cast
 
 from langchain_core.runnables import RunnableConfig
 
-from onyx.agents.agent_search.deep_search_a.main__graph.operations import logger
-from onyx.agents.agent_search.deep_search_a.main__graph.states import MainState
-from onyx.agents.agent_search.deep_search_a.main__graph.states import (
+from onyx.agents.agent_search.deep_search_a.main.operations import logger
+from onyx.agents.agent_search.deep_search_a.main.states import MainState
+from onyx.agents.agent_search.deep_search_a.main.states import (
     RequireRefinedAnswerUpdate,
 )
 from onyx.agents.agent_search.models import AgentSearchConfig

backend/onyx/agents/agent_search/deep_search_a/main/nodes/refined_sub_question_creation.py

@@ -6,17 +6,17 @@ from langchain_core.messages import HumanMessage
 from langchain_core.messages import merge_content
 from langchain_core.runnables import RunnableConfig
 
-from onyx.agents.agent_search.deep_search_a.main__graph.models import (
+from onyx.agents.agent_search.deep_search_a.main.models import (
     FollowUpSubQuestion,
 )
-from onyx.agents.agent_search.deep_search_a.main__graph.operations import (
+from onyx.agents.agent_search.deep_search_a.main.operations import (
     dispatch_subquestion,
 )
-from onyx.agents.agent_search.deep_search_a.main__graph.operations import logger
-from onyx.agents.agent_search.deep_search_a.main__graph.states import (
+from onyx.agents.agent_search.deep_search_a.main.operations import logger
+from onyx.agents.agent_search.deep_search_a.main.states import (
     FollowUpSubQuestionsUpdate,
 )
-from onyx.agents.agent_search.deep_search_a.main__graph.states import MainState
+from onyx.agents.agent_search.deep_search_a.main.states import MainState
 from onyx.agents.agent_search.models import AgentSearchConfig
 from onyx.agents.agent_search.shared_graph_utils.agent_prompt_ops import (
     build_history_prompt,
@@ -64,7 +64,7 @@ def refined_sub_question_creation(
         entity_retlation_term_extractions
     )
 
-    initial_question_answers = state.decomp_answer_results
+    initial_question_answers = state.sub_question_results
 
     addressed_question_list = [
         x.question for x in initial_question_answers if x.verified_high_quality

backend/onyx/agents/agent_search/deep_search_a/main/states.py

@@ -6,14 +6,14 @@ from typing import TypedDict
 from pydantic import BaseModel
 
 from onyx.agents.agent_search.core_state import CoreState
-from onyx.agents.agent_search.deep_search_a.main__graph.models import AgentBaseMetrics
-from onyx.agents.agent_search.deep_search_a.main__graph.models import (
+from onyx.agents.agent_search.deep_search_a.main.models import AgentBaseMetrics
+from onyx.agents.agent_search.deep_search_a.main.models import (
     AgentRefinedMetrics,
 )
-from onyx.agents.agent_search.deep_search_a.main__graph.models import (
+from onyx.agents.agent_search.deep_search_a.main.models import (
     FollowUpSubQuestion,
 )
-from onyx.agents.agent_search.deep_search_a.util__expanded_retrieval__subgraph.models import (
+from onyx.agents.agent_search.deep_search_a.shared.expanded_retrieval.models import (
     ExpandedRetrievalResult,
 )
 from onyx.agents.agent_search.orchestration.states import ToolCallUpdate
@@ -107,7 +107,7 @@ class DecompAnswersUpdate(LoggerUpdate):
     cited_docs: Annotated[
         list[InferenceSection], dedup_inference_sections
     ] = []  # cited docs from sub-answers are used for answer context
-    decomp_answer_results: Annotated[
+    sub_question_results: Annotated[
         list[QuestionAnswerResults], dedup_question_answer_results
     ] = []
 

backend/onyx/agents/agent_search/deep_search_a/refininement/sub_answer_consolidation/edges.py

@@ -3,10 +3,10 @@ from datetime import datetime
 
 from langgraph.types import Send
 
-from onyx.agents.agent_search.deep_search_a.initial__individual_sub_answer__subgraph.states import (
+from onyx.agents.agent_search.deep_search_a.initial.individual_sub_answer_generation.states import (
     AnswerQuestionInput,
 )
-from onyx.agents.agent_search.deep_search_a.util__expanded_retrieval__subgraph.states import (
+from onyx.agents.agent_search.deep_search_a.shared.expanded_retrieval.states import (
     ExpandedRetrievalInput,
 )
 from onyx.utils.logger import setup_logger

backend/onyx/agents/agent_search/deep_search_a/refininement/sub_answer_consolidation/graph_builder.py

@@ -2,31 +2,31 @@ from langgraph.graph import END
 from langgraph.graph import START
 from langgraph.graph import StateGraph
 
-from onyx.agents.agent_search.deep_search_a.initial__individual_sub_answer__subgraph.nodes.answer_check import (
+from onyx.agents.agent_search.deep_search_a.initial.individual_sub_answer_generation.nodes.answer_check import (
     answer_check,
 )
-from onyx.agents.agent_search.deep_search_a.initial__individual_sub_answer__subgraph.nodes.answer_generation import (
+from onyx.agents.agent_search.deep_search_a.initial.individual_sub_answer_generation.nodes.answer_generation import (
     answer_generation,
 )
-from onyx.agents.agent_search.deep_search_a.initial__individual_sub_answer__subgraph.nodes.format_answer import (
+from onyx.agents.agent_search.deep_search_a.initial.individual_sub_answer_generation.nodes.format_answer import (
     format_answer,
 )
-from onyx.agents.agent_search.deep_search_a.initial__individual_sub_answer__subgraph.nodes.ingest_retrieval import (
+from onyx.agents.agent_search.deep_search_a.initial.individual_sub_answer_generation.nodes.ingest_retrieval import (
     ingest_retrieval,
 )
-from onyx.agents.agent_search.deep_search_a.initial__individual_sub_answer__subgraph.states import (
+from onyx.agents.agent_search.deep_search_a.initial.individual_sub_answer_generation.states import (
     AnswerQuestionInput,
 )
-from onyx.agents.agent_search.deep_search_a.initial__individual_sub_answer__subgraph.states import (
+from onyx.agents.agent_search.deep_search_a.initial.individual_sub_answer_generation.states import (
     AnswerQuestionOutput,
 )
-from onyx.agents.agent_search.deep_search_a.initial__individual_sub_answer__subgraph.states import (
+from onyx.agents.agent_search.deep_search_a.initial.individual_sub_answer_generation.states import (
     AnswerQuestionState,
 )
-from onyx.agents.agent_search.deep_search_a.refinement__consolidate_sub_answers__subgraph.edges import (
+from onyx.agents.agent_search.deep_search_a.refininement.sub_answer_consolidation.edges import (
     send_to_expanded_refined_retrieval,
 )
-from onyx.agents.agent_search.deep_search_a.util__expanded_retrieval__subgraph.graph_builder import (
+from onyx.agents.agent_search.deep_search_a.shared.expanded_retrieval.graph_builder import (
     expanded_retrieval_graph_builder,
 )
 from onyx.utils.logger import setup_logger

backend/onyx/agents/agent_search/deep_search_a/shared/expanded_retrieval/edges.py

@@ -4,10 +4,10 @@ from typing import cast
 from langchain_core.runnables.config import RunnableConfig
 from langgraph.types import Send
 
-from onyx.agents.agent_search.deep_search_a.util__expanded_retrieval__subgraph.states import (
+from onyx.agents.agent_search.deep_search_a.shared.expanded_retrieval.states import (
     ExpandedRetrievalState,
 )
-from onyx.agents.agent_search.deep_search_a.util__expanded_retrieval__subgraph.states import (
+from onyx.agents.agent_search.deep_search_a.shared.expanded_retrieval.states import (
     RetrievalInput,
 )
 from onyx.agents.agent_search.models import AgentSearchConfig

backend/onyx/agents/agent_search/deep_search_a/shared/expanded_retrieval/graph_builder.py

@@ -2,37 +2,37 @@ from langgraph.graph import END
 from langgraph.graph import START
 from langgraph.graph import StateGraph
 
-from onyx.agents.agent_search.deep_search_a.util__expanded_retrieval__subgraph.edges import (
+from onyx.agents.agent_search.deep_search_a.shared.expanded_retrieval.edges import (
     parallel_retrieval_edge,
 )
-from onyx.agents.agent_search.deep_search_a.util__expanded_retrieval__subgraph.nodes.doc_reranking import (
+from onyx.agents.agent_search.deep_search_a.shared.expanded_retrieval.nodes.doc_reranking import (
     doc_reranking,
 )
-from onyx.agents.agent_search.deep_search_a.util__expanded_retrieval__subgraph.nodes.doc_retrieval import (
+from onyx.agents.agent_search.deep_search_a.shared.expanded_retrieval.nodes.doc_retrieval import (
     doc_retrieval,
 )
-from onyx.agents.agent_search.deep_search_a.util__expanded_retrieval__subgraph.nodes.doc_verification import (
+from onyx.agents.agent_search.deep_search_a.shared.expanded_retrieval.nodes.doc_verification import (
     doc_verification,
 )
-from onyx.agents.agent_search.deep_search_a.util__expanded_retrieval__subgraph.nodes.dummy import (
+from onyx.agents.agent_search.deep_search_a.shared.expanded_retrieval.nodes.dummy import (
     dummy,
 )
-from onyx.agents.agent_search.deep_search_a.util__expanded_retrieval__subgraph.nodes.expand_queries import (
+from onyx.agents.agent_search.deep_search_a.shared.expanded_retrieval.nodes.expand_queries import (
     expand_queries,
 )
-from onyx.agents.agent_search.deep_search_a.util__expanded_retrieval__subgraph.nodes.format_results import (
+from onyx.agents.agent_search.deep_search_a.shared.expanded_retrieval.nodes.format_results import (
     format_results,
 )
-from onyx.agents.agent_search.deep_search_a.util__expanded_retrieval__subgraph.nodes.verification_kickoff import (
+from onyx.agents.agent_search.deep_search_a.shared.expanded_retrieval.nodes.verification_kickoff import (
     verification_kickoff,
 )
-from onyx.agents.agent_search.deep_search_a.util__expanded_retrieval__subgraph.states import (
+from onyx.agents.agent_search.deep_search_a.shared.expanded_retrieval.states import (
     ExpandedRetrievalInput,
 )
-from onyx.agents.agent_search.deep_search_a.util__expanded_retrieval__subgraph.states import (
+from onyx.agents.agent_search.deep_search_a.shared.expanded_retrieval.states import (
     ExpandedRetrievalOutput,
 )
-from onyx.agents.agent_search.deep_search_a.util__expanded_retrieval__subgraph.states import (
+from onyx.agents.agent_search.deep_search_a.shared.expanded_retrieval.states import (
     ExpandedRetrievalState,
 )
 from onyx.agents.agent_search.shared_graph_utils.utils import get_test_config

backend/onyx/agents/agent_search/deep_search_a/shared/expanded_retrieval/nodes/doc_reranking.py

@@ -3,13 +3,13 @@ from typing import cast
 
 from langchain_core.runnables.config import RunnableConfig
 
-from onyx.agents.agent_search.deep_search_a.util__expanded_retrieval__subgraph.operations import (
+from onyx.agents.agent_search.deep_search_a.shared.expanded_retrieval.operations import (
     logger,
 )
-from onyx.agents.agent_search.deep_search_a.util__expanded_retrieval__subgraph.states import (
+from onyx.agents.agent_search.deep_search_a.shared.expanded_retrieval.states import (
     DocRerankingUpdate,
 )
-from onyx.agents.agent_search.deep_search_a.util__expanded_retrieval__subgraph.states import (
+from onyx.agents.agent_search.deep_search_a.shared.expanded_retrieval.states import (
     ExpandedRetrievalState,
 )
 from onyx.agents.agent_search.models import AgentSearchConfig

backend/onyx/agents/agent_search/deep_search_a/shared/expanded_retrieval/nodes/doc_retrieval.py

@@ -3,13 +3,13 @@ from typing import cast
 
 from langchain_core.runnables.config import RunnableConfig
 
-from onyx.agents.agent_search.deep_search_a.util__expanded_retrieval__subgraph.operations import (
+from onyx.agents.agent_search.deep_search_a.shared.expanded_retrieval.operations import (
     logger,
 )
-from onyx.agents.agent_search.deep_search_a.util__expanded_retrieval__subgraph.states import (
+from onyx.agents.agent_search.deep_search_a.shared.expanded_retrieval.states import (
     DocRetrievalUpdate,
 )
-from onyx.agents.agent_search.deep_search_a.util__expanded_retrieval__subgraph.states import (
+from onyx.agents.agent_search.deep_search_a.shared.expanded_retrieval.states import (
     RetrievalInput,
 )
 from onyx.agents.agent_search.models import AgentSearchConfig

backend/onyx/agents/agent_search/deep_search_a/shared/expanded_retrieval/nodes/doc_verification.py

@@ -3,10 +3,10 @@ from typing import cast
 from langchain_core.messages import HumanMessage
 from langchain_core.runnables.config import RunnableConfig
 
-from onyx.agents.agent_search.deep_search_a.util__expanded_retrieval__subgraph.states import (
+from onyx.agents.agent_search.deep_search_a.shared.expanded_retrieval.states import (
     DocVerificationInput,
 )
-from onyx.agents.agent_search.deep_search_a.util__expanded_retrieval__subgraph.states import (
+from onyx.agents.agent_search.deep_search_a.shared.expanded_retrieval.states import (
     DocVerificationUpdate,
 )
 from onyx.agents.agent_search.models import AgentSearchConfig

backend/onyx/agents/agent_search/deep_search_a/shared/expanded_retrieval/nodes/dummy.py

@@ -1,9 +1,9 @@
 from langchain_core.runnables.config import RunnableConfig
 
-from onyx.agents.agent_search.deep_search_a.util__expanded_retrieval__subgraph.states import (
+from onyx.agents.agent_search.deep_search_a.shared.expanded_retrieval.states import (
     ExpandedRetrievalState,
 )
-from onyx.agents.agent_search.deep_search_a.util__expanded_retrieval__subgraph.states import (
+from onyx.agents.agent_search.deep_search_a.shared.expanded_retrieval.states import (
     QueryExpansionUpdate,
 )
 

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

@@ -5,16 +5,16 @@ from langchain_core.messages import HumanMessage
 from langchain_core.messages import merge_message_runs
 from langchain_core.runnables.config import RunnableConfig
 
-from onyx.agents.agent_search.deep_search_a.util__expanded_retrieval__subgraph.operations import (
+from onyx.agents.agent_search.deep_search_a.shared.expanded_retrieval.operations import (
     dispatch_subquery,
 )
-from onyx.agents.agent_search.deep_search_a.util__expanded_retrieval__subgraph.operations import (
+from onyx.agents.agent_search.deep_search_a.shared.expanded_retrieval.operations import (
     logger,
 )
-from onyx.agents.agent_search.deep_search_a.util__expanded_retrieval__subgraph.states import (
+from onyx.agents.agent_search.deep_search_a.shared.expanded_retrieval.states import (
     ExpandedRetrievalInput,
 )
-from onyx.agents.agent_search.deep_search_a.util__expanded_retrieval__subgraph.states import (
+from onyx.agents.agent_search.deep_search_a.shared.expanded_retrieval.states import (
     QueryExpansionUpdate,
 )
 from onyx.agents.agent_search.models import AgentSearchConfig

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

@@ -3,16 +3,16 @@ from typing import cast
 from langchain_core.callbacks.manager import dispatch_custom_event
 from langchain_core.runnables.config import RunnableConfig
 
-from onyx.agents.agent_search.deep_search_a.util__expanded_retrieval__subgraph.models import (
+from onyx.agents.agent_search.deep_search_a.shared.expanded_retrieval.models import (
     ExpandedRetrievalResult,
 )
-from onyx.agents.agent_search.deep_search_a.util__expanded_retrieval__subgraph.operations import (
+from onyx.agents.agent_search.deep_search_a.shared.expanded_retrieval.operations import (
     calculate_sub_question_retrieval_stats,
 )
-from onyx.agents.agent_search.deep_search_a.util__expanded_retrieval__subgraph.states import (
+from onyx.agents.agent_search.deep_search_a.shared.expanded_retrieval.states import (
     ExpandedRetrievalState,
 )
-from onyx.agents.agent_search.deep_search_a.util__expanded_retrieval__subgraph.states import (
+from onyx.agents.agent_search.deep_search_a.shared.expanded_retrieval.states import (
     ExpandedRetrievalUpdate,
 )
 from onyx.agents.agent_search.models import AgentSearchConfig

backend/onyx/agents/agent_search/deep_search_a/shared/expanded_retrieval/nodes/verification_kickoff.py

@@ -4,10 +4,10 @@ from langchain_core.runnables.config import RunnableConfig
 from langgraph.types import Command
 from langgraph.types import Send
 
-from onyx.agents.agent_search.deep_search_a.util__expanded_retrieval__subgraph.states import (
+from onyx.agents.agent_search.deep_search_a.shared.expanded_retrieval.states import (
     DocVerificationInput,
 )
-from onyx.agents.agent_search.deep_search_a.util__expanded_retrieval__subgraph.states import (
+from onyx.agents.agent_search.deep_search_a.shared.expanded_retrieval.states import (
     ExpandedRetrievalState,
 )
 

backend/onyx/agents/agent_search/deep_search_a/shared/expanded_retrieval/states.py

@@ -4,7 +4,7 @@ from typing import Annotated
 from pydantic import BaseModel
 
 from onyx.agents.agent_search.core_state import SubgraphCoreState
-from onyx.agents.agent_search.deep_search_a.util__expanded_retrieval__subgraph.models import (
+from onyx.agents.agent_search.deep_search_a.shared.expanded_retrieval.models import (
     ExpandedRetrievalResult,
 )
 from onyx.agents.agent_search.shared_graph_utils.models import QueryResult

backend/onyx/agents/agent_search/run_graph.py

@@ -9,10 +9,10 @@ from langgraph.graph.state import CompiledStateGraph
 
 from onyx.agents.agent_search.basic.graph_builder import basic_graph_builder
 from onyx.agents.agent_search.basic.states import BasicInput
-from onyx.agents.agent_search.deep_search_a.main__graph.graph_builder import (
+from onyx.agents.agent_search.deep_search_a.main.graph_builder import (
     main_graph_builder as main_graph_builder_a,
 )
-from onyx.agents.agent_search.deep_search_a.main__graph.states import (
+from onyx.agents.agent_search.deep_search_a.main.states import (
     MainInput as MainInput_a,
 )
 from onyx.agents.agent_search.models import AgentSearchConfig

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

@@ -2,14 +2,14 @@ from typing import Literal
 
 from pydantic import BaseModel
 
-from onyx.agents.agent_search.deep_search_a.main__graph.models import (
+from onyx.agents.agent_search.deep_search_a.main.models import (
     AgentAdditionalMetrics,
 )
-from onyx.agents.agent_search.deep_search_a.main__graph.models import AgentBaseMetrics
-from onyx.agents.agent_search.deep_search_a.main__graph.models import (
+from onyx.agents.agent_search.deep_search_a.main.models import AgentBaseMetrics
+from onyx.agents.agent_search.deep_search_a.main.models import (
     AgentRefinedMetrics,
 )
-from onyx.agents.agent_search.deep_search_a.main__graph.models import AgentTimings
+from onyx.agents.agent_search.deep_search_a.main.models import AgentTimings
 from onyx.context.search.models import InferenceSection
 from onyx.tools.models import SearchQueryInfo