graph directory renamings

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

@@ -3,13 +3,13 @@ from datetime import datetime
 
 from langgraph.types import Send
 
-from onyx.agents.agent_search.deep_search_a.answer_initial_sub_question.states import (
+from onyx.agents.agent_search.deep_search_a.initial__individual_sub_answer__subgraph.states import (
     AnswerQuestionInput,
 )
-from onyx.agents.agent_search.deep_search_a.answer_initial_sub_question.states import (
+from onyx.agents.agent_search.deep_search_a.initial__individual_sub_answer__subgraph.states import (
     AnswerQuestionOutput,
 )
-from onyx.agents.agent_search.deep_search_a.initial_search_sq_subgraph.states import (
+from onyx.agents.agent_search.deep_search_a.initial__retrieval_sub_answers__subgraph.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__consolidate_sub_answers__subgraph/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.answer_initial_sub_question.graph_builder import (
-    answer_query_graph_builder,
-)
-from onyx.agents.agent_search.deep_search_a.initial_sub_question_answering.edges import (
+from onyx.agents.agent_search.deep_search_a.initial__consolidate_sub_answers__subgraph.edges import (
     parallelize_initial_sub_question_answering,
 )
-from onyx.agents.agent_search.deep_search_a.initial_sub_question_answering.nodes.ingest_initial_sub_question_answers import (
-    ingest_initial_sub_question_answers,
+from onyx.agents.agent_search.deep_search_a.initial__consolidate_sub_answers__subgraph.nodes.ingest_initial_sub_answers import (
+    ingest_initial_sub_answers,
 )
-from onyx.agents.agent_search.deep_search_a.initial_sub_question_answering.nodes.initial_sub_question_creation import (
+from onyx.agents.agent_search.deep_search_a.initial__consolidate_sub_answers__subgraph.nodes.initial_decomposition import (
     initial_sub_question_creation,
 )
-from onyx.agents.agent_search.deep_search_a.initial_sub_question_answering.states import (
+from onyx.agents.agent_search.deep_search_a.initial__consolidate_sub_answers__subgraph.states import (
     SQInput,
 )
-from onyx.agents.agent_search.deep_search_a.initial_sub_question_answering.states import (
+from onyx.agents.agent_search.deep_search_a.initial__consolidate_sub_answers__subgraph.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()
@@ -45,7 +45,7 @@ def initial_sq_subgraph_builder(test_mode: bool = False) -> StateGraph:
 
     graph.add_node(
         node="ingest_initial_sub_question_answers",
-        action=ingest_initial_sub_question_answers,
+        action=ingest_initial_sub_answers,
     )
 
     ### Add edges ###

backend/onyx/agents/agent_search/deep_search_a/initial__consolidate_sub_answers__subgraph/nodes/ingest_initial_sub_answers.py

@@ -1,16 +1,18 @@
 from datetime import datetime
 
-from onyx.agents.agent_search.deep_search_a.answer_initial_sub_question.states import (
+from onyx.agents.agent_search.deep_search_a.initial__individual_sub_answer__subgraph.states import (
     AnswerQuestionOutput,
 )
-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.deep_search_a.main__graph.operations import logger
+from onyx.agents.agent_search.deep_search_a.main__graph.states import (
+    DecompAnswersUpdate,
+)
 from onyx.agents.agent_search.shared_graph_utils.operators import (
     dedup_inference_sections,
 )
 
 
-def ingest_initial_sub_question_answers(
+def ingest_initial_sub_answers(
     state: AnswerQuestionOutput,
 ) -> DecompAnswersUpdate:
     now_start = datetime.now()

backend/onyx/agents/agent_search/deep_search_a/initial__consolidate_sub_answers__subgraph/nodes/initial_decomposition.py

@@ -6,13 +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_search_sq_subgraph.states import (
+from onyx.agents.agent_search.deep_search_a.initial__retrieval_sub_answers__subgraph.states import (
     SearchSQState,
 )
-from onyx.agents.agent_search.deep_search_a.main.models import AgentRefinedMetrics
-from onyx.agents.agent_search.deep_search_a.main.operations import dispatch_subquestion
-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.deep_search_a.main__graph.models import (
+    AgentRefinedMetrics,
+)
+from onyx.agents.agent_search.deep_search_a.main__graph.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.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__consolidate_sub_answers__subgraph/states.py

@@ -1,9 +1,13 @@
 from typing import TypedDict
 
 from onyx.agents.agent_search.core_state import CoreState
-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.states import InitialAnswerUpdate
+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 (
+    DecompAnswersUpdate,
+)
+from onyx.agents.agent_search.deep_search_a.main__graph.states import (
+    InitialAnswerUpdate,
+)
 
 ### States ###
 

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

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

backend/onyx/agents/agent_search/deep_search_a/initial__individual_sub_answer__subgraph/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.answer_initial_sub_question.edges import (
+from onyx.agents.agent_search.deep_search_a.initial__individual_sub_answer__subgraph.edges import (
     send_to_expanded_retrieval,
 )
-from onyx.agents.agent_search.deep_search_a.answer_initial_sub_question.nodes.answer_check import (
+from onyx.agents.agent_search.deep_search_a.initial__individual_sub_answer__subgraph.nodes.answer_check import (
     answer_check,
 )
-from onyx.agents.agent_search.deep_search_a.answer_initial_sub_question.nodes.answer_generation import (
+from onyx.agents.agent_search.deep_search_a.initial__individual_sub_answer__subgraph.nodes.answer_generation import (
     answer_generation,
 )
-from onyx.agents.agent_search.deep_search_a.answer_initial_sub_question.nodes.format_answer import (
+from onyx.agents.agent_search.deep_search_a.initial__individual_sub_answer__subgraph.nodes.format_answer import (
     format_answer,
 )
-from onyx.agents.agent_search.deep_search_a.answer_initial_sub_question.nodes.ingest_retrieval import (
+from onyx.agents.agent_search.deep_search_a.initial__individual_sub_answer__subgraph.nodes.ingest_retrieval import (
     ingest_retrieval,
 )
-from onyx.agents.agent_search.deep_search_a.answer_initial_sub_question.states import (
+from onyx.agents.agent_search.deep_search_a.initial__individual_sub_answer__subgraph.states import (
     AnswerQuestionInput,
 )
-from onyx.agents.agent_search.deep_search_a.answer_initial_sub_question.states import (
+from onyx.agents.agent_search.deep_search_a.initial__individual_sub_answer__subgraph.states import (
     AnswerQuestionOutput,
 )
-from onyx.agents.agent_search.deep_search_a.answer_initial_sub_question.states import (
+from onyx.agents.agent_search.deep_search_a.initial__individual_sub_answer__subgraph.states import (
     AnswerQuestionState,
 )
-from onyx.agents.agent_search.deep_search_a.expanded_retrieval.graph_builder import (
+from onyx.agents.agent_search.deep_search_a.util__expanded_retrieval__subgraph.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__subgraph/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.answer_initial_sub_question.states import (
+from onyx.agents.agent_search.deep_search_a.initial__individual_sub_answer__subgraph.states import (
     AnswerQuestionState,
 )
-from onyx.agents.agent_search.deep_search_a.answer_initial_sub_question.states import (
+from onyx.agents.agent_search.deep_search_a.initial__individual_sub_answer__subgraph.states import (
     QACheckUpdate,
 )
 from onyx.agents.agent_search.models import AgentSearchConfig

backend/onyx/agents/agent_search/deep_search_a/initial__individual_sub_answer__subgraph/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.answer_initial_sub_question.states import (
+from onyx.agents.agent_search.deep_search_a.initial__individual_sub_answer__subgraph.states import (
     AnswerQuestionState,
 )
-from onyx.agents.agent_search.deep_search_a.answer_initial_sub_question.states import (
+from onyx.agents.agent_search.deep_search_a.initial__individual_sub_answer__subgraph.states import (
     QAGenerationUpdate,
 )
 from onyx.agents.agent_search.models import AgentSearchConfig

backend/onyx/agents/agent_search/deep_search_a/initial__individual_sub_answer__subgraph/nodes/format_answer.py

@@ -1,7 +1,7 @@
-from onyx.agents.agent_search.deep_search_a.answer_initial_sub_question.states import (
+from onyx.agents.agent_search.deep_search_a.initial__individual_sub_answer__subgraph.states import (
     AnswerQuestionOutput,
 )
-from onyx.agents.agent_search.deep_search_a.answer_initial_sub_question.states import (
+from onyx.agents.agent_search.deep_search_a.initial__individual_sub_answer__subgraph.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__subgraph/nodes/ingest_retrieval.py

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

backend/onyx/agents/agent_search/deep_search_a/initial__retrieval__subgraph/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.base_raw_search.nodes.format_raw_search_results import (
+from onyx.agents.agent_search.deep_search_a.initial__retrieval__subgraph.nodes.format_raw_search_results import (
     format_raw_search_results,
 )
-from onyx.agents.agent_search.deep_search_a.base_raw_search.nodes.generate_raw_search_data import (
+from onyx.agents.agent_search.deep_search_a.initial__retrieval__subgraph.nodes.generate_raw_search_data import (
     generate_raw_search_data,
 )
-from onyx.agents.agent_search.deep_search_a.base_raw_search.nodes.ingest_initial_base_retrieval import (
+from onyx.agents.agent_search.deep_search_a.initial__retrieval__subgraph.nodes.ingest_initial_base_retrieval import (
     ingest_initial_base_retrieval,
 )
-from onyx.agents.agent_search.deep_search_a.base_raw_search.states import (
+from onyx.agents.agent_search.deep_search_a.initial__retrieval__subgraph.states import (
     BaseRawSearchInput,
 )
-from onyx.agents.agent_search.deep_search_a.base_raw_search.states import (
+from onyx.agents.agent_search.deep_search_a.initial__retrieval__subgraph.states import (
     BaseRawSearchOutput,
 )
-from onyx.agents.agent_search.deep_search_a.base_raw_search.states import (
+from onyx.agents.agent_search.deep_search_a.initial__retrieval__subgraph.states import (
     BaseRawSearchState,
 )
-from onyx.agents.agent_search.deep_search_a.expanded_retrieval.graph_builder import (
+from onyx.agents.agent_search.deep_search_a.util__expanded_retrieval__subgraph.graph_builder import (
     expanded_retrieval_graph_builder,
 )
 

backend/onyx/agents/agent_search/deep_search_a/initial__retrieval__subgraph/nodes/format_raw_search_results.py

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

backend/onyx/agents/agent_search/deep_search_a/initial__retrieval__subgraph/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.expanded_retrieval.states import (
+from onyx.agents.agent_search.deep_search_a.util__expanded_retrieval__subgraph.states import (
     ExpandedRetrievalInput,
 )
 from onyx.agents.agent_search.models import AgentSearchConfig

backend/onyx/agents/agent_search/deep_search_a/initial__retrieval__subgraph/nodes/ingest_initial_base_retrieval.py

@@ -1,10 +1,12 @@
 from datetime import datetime
 
-from onyx.agents.agent_search.deep_search_a.base_raw_search.states import (
+from onyx.agents.agent_search.deep_search_a.initial__retrieval__subgraph.states import (
     BaseRawSearchOutput,
 )
-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.deep_search_a.main__graph.operations import logger
+from onyx.agents.agent_search.deep_search_a.main__graph.states import (
+    ExpandedRetrievalUpdate,
+)
 from onyx.agents.agent_search.shared_graph_utils.models import AgentChunkStats
 
 

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

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

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

@@ -3,13 +3,13 @@ from datetime import datetime
 
 from langgraph.types import Send
 
-from onyx.agents.agent_search.deep_search_a.answer_initial_sub_question.states import (
+from onyx.agents.agent_search.deep_search_a.initial__individual_sub_answer__subgraph.states import (
     AnswerQuestionInput,
 )
-from onyx.agents.agent_search.deep_search_a.answer_initial_sub_question.states import (
+from onyx.agents.agent_search.deep_search_a.initial__individual_sub_answer__subgraph.states import (
     AnswerQuestionOutput,
 )
-from onyx.agents.agent_search.deep_search_a.initial_search_sq_subgraph.states import (
+from onyx.agents.agent_search.deep_search_a.initial__retrieval_sub_answers__subgraph.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__retrieval_sub_answers__subgraph/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.base_raw_search.graph_builder import (
+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_search_sq_subgraph.nodes.generate_initial_answer import (
+from onyx.agents.agent_search.deep_search_a.initial__retrieval_sub_answers__subgraph.nodes.generate_initial_answer import (
     generate_initial_answer,
 )
-from onyx.agents.agent_search.deep_search_a.initial_search_sq_subgraph.nodes.initial_answer_quality_check import (
+from onyx.agents.agent_search.deep_search_a.initial__retrieval_sub_answers__subgraph.nodes.initial_answer_quality_check import (
     initial_answer_quality_check,
 )
-from onyx.agents.agent_search.deep_search_a.initial_search_sq_subgraph.nodes.retrieval_consolidation import (
+from onyx.agents.agent_search.deep_search_a.initial__retrieval_sub_answers__subgraph.nodes.retrieval_consolidation import (
     retrieval_consolidation,
 )
-from onyx.agents.agent_search.deep_search_a.initial_search_sq_subgraph.states import (
+from onyx.agents.agent_search.deep_search_a.initial__retrieval_sub_answers__subgraph.states import (
     SearchSQInput,
 )
-from onyx.agents.agent_search.deep_search_a.initial_search_sq_subgraph.states import (
+from onyx.agents.agent_search.deep_search_a.initial__retrieval_sub_answers__subgraph.states import (
     SearchSQState,
 )
-from onyx.agents.agent_search.deep_search_a.initial_sub_question_answering.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__retrieval_sub_answers__subgraph/nodes/generate_initial_answer.py

@@ -7,19 +7,21 @@ 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_search_sq_subgraph.states import (
+from onyx.agents.agent_search.deep_search_a.initial__retrieval_sub_answers__subgraph.states import (
     SearchSQState,
 )
-from onyx.agents.agent_search.deep_search_a.main.models import AgentBaseMetrics
-from onyx.agents.agent_search.deep_search_a.main.operations import (
+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 (
     calculate_initial_agent_stats,
 )
-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 (
+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 (
     remove_document_citations,
 )
-from onyx.agents.agent_search.deep_search_a.main.states import InitialAnswerUpdate
+from onyx.agents.agent_search.deep_search_a.main__graph.states import (
+    InitialAnswerUpdate,
+)
 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__retrieval_sub_answers__subgraph/nodes/ingest_initial_base_retrieval.py

@@ -1,10 +1,12 @@
 from datetime import datetime
 
-from onyx.agents.agent_search.deep_search_a.base_raw_search.states import (
+from onyx.agents.agent_search.deep_search_a.initial__retrieval__subgraph.states import (
     BaseRawSearchOutput,
 )
-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.deep_search_a.main__graph.operations import logger
+from onyx.agents.agent_search.deep_search_a.main__graph.states import (
+    ExpandedRetrievalUpdate,
+)
 from onyx.agents.agent_search.shared_graph_utils.models import AgentChunkStats
 
 

backend/onyx/agents/agent_search/deep_search_a/initial__retrieval_sub_answers__subgraph/nodes/initial_answer_quality_check.py

@@ -1,10 +1,10 @@
 from datetime import datetime
 
-from onyx.agents.agent_search.deep_search_a.initial_search_sq_subgraph.states import (
+from onyx.agents.agent_search.deep_search_a.initial__retrieval_sub_answers__subgraph.states import (
     SearchSQState,
 )
-from onyx.agents.agent_search.deep_search_a.main.operations import logger
-from onyx.agents.agent_search.deep_search_a.main.states import (
+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 (
     InitialAnswerQualityUpdate,
 )
 

backend/onyx/agents/agent_search/deep_search_a/initial__retrieval_sub_answers__subgraph/nodes/retrieval_consolidation.py

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

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

@@ -3,17 +3,25 @@ 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.expanded_retrieval.models import (
-    ExpandedRetrievalResult,
+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 (
+    DecompAnswersUpdate,
 )
-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.states import ExpandedRetrievalUpdate
-from onyx.agents.agent_search.deep_search_a.main.states import ExploratorySearchUpdate
-from onyx.agents.agent_search.deep_search_a.main.states import (
+from onyx.agents.agent_search.deep_search_a.main__graph.states import (
+    ExpandedRetrievalUpdate,
+)
+from onyx.agents.agent_search.deep_search_a.main__graph.states import (
+    ExploratorySearchUpdate,
+)
+from onyx.agents.agent_search.deep_search_a.main__graph.states import (
     InitialAnswerQualityUpdate,
 )
-from onyx.agents.agent_search.deep_search_a.main.states import InitialAnswerUpdate
+from onyx.agents.agent_search.deep_search_a.main__graph.states import (
+    InitialAnswerUpdate,
+)
+from onyx.agents.agent_search.deep_search_a.util__expanded_retrieval__subgraph.models import (
+    ExpandedRetrievalResult,
+)
 
 ### States ###
 

backend/onyx/agents/agent_search/deep_search_a/main__graph/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.answer_initial_sub_question.states import (
+from onyx.agents.agent_search.deep_search_a.initial__individual_sub_answer__subgraph.states import (
     AnswerQuestionInput,
 )
-from onyx.agents.agent_search.deep_search_a.answer_initial_sub_question.states import (
+from onyx.agents.agent_search.deep_search_a.initial__individual_sub_answer__subgraph.states import (
     AnswerQuestionOutput,
 )
-from onyx.agents.agent_search.deep_search_a.main.states import MainState
-from onyx.agents.agent_search.deep_search_a.main.states import (
+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 (
     RequireRefinedAnswerUpdate,
 )
 from onyx.agents.agent_search.models import AgentSearchConfig

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

@@ -2,47 +2,47 @@ from langgraph.graph import END
 from langgraph.graph import START
 from langgraph.graph import StateGraph
 
-from onyx.agents.agent_search.deep_search_a.answer_refinement_sub_question.graph_builder import (
-    answer_refined_query_graph_builder,
-)
-from onyx.agents.agent_search.deep_search_a.initial_search_sq_subgraph.graph_builder import (
+from onyx.agents.agent_search.deep_search_a.initial__retrieval_sub_answers__subgraph.graph_builder import (
     initial_search_sq_subgraph_builder,
 )
-from onyx.agents.agent_search.deep_search_a.main.edges import (
+from onyx.agents.agent_search.deep_search_a.main__graph.edges import (
     continue_to_refined_answer_or_end,
 )
-from onyx.agents.agent_search.deep_search_a.main.edges import (
+from onyx.agents.agent_search.deep_search_a.main__graph.edges import (
     parallelize_refined_sub_question_answering,
 )
-from onyx.agents.agent_search.deep_search_a.main.edges import (
+from onyx.agents.agent_search.deep_search_a.main__graph.edges import (
     route_initial_tool_choice,
 )
-from onyx.agents.agent_search.deep_search_a.main.nodes.agent_logging import (
+from onyx.agents.agent_search.deep_search_a.main__graph.nodes.agent_logging import (
     agent_logging,
 )
-from onyx.agents.agent_search.deep_search_a.main.nodes.agent_search_start import (
+from onyx.agents.agent_search.deep_search_a.main__graph.nodes.agent_search_start import (
     agent_search_start,
 )
-from onyx.agents.agent_search.deep_search_a.main.nodes.answer_comparison import (
+from onyx.agents.agent_search.deep_search_a.main__graph.nodes.answer_comparison import (
     answer_comparison,
 )
-from onyx.agents.agent_search.deep_search_a.main.nodes.entity_term_extraction_llm import (
+from onyx.agents.agent_search.deep_search_a.main__graph.nodes.entity_term_extraction_llm import (
     entity_term_extraction_llm,
 )
-from onyx.agents.agent_search.deep_search_a.main.nodes.generate_refined_answer import (
+from onyx.agents.agent_search.deep_search_a.main__graph.nodes.generate_refined_answer import (
     generate_refined_answer,
 )
-from onyx.agents.agent_search.deep_search_a.main.nodes.ingest_refined_answers import (
+from onyx.agents.agent_search.deep_search_a.main__graph.nodes.ingest_refined_answers import (
     ingest_refined_answers,
 )
-from onyx.agents.agent_search.deep_search_a.main.nodes.refined_answer_decision import (
+from onyx.agents.agent_search.deep_search_a.main__graph.nodes.refined_answer_decision import (
     refined_answer_decision,
 )
-from onyx.agents.agent_search.deep_search_a.main.nodes.refined_sub_question_creation import (
+from onyx.agents.agent_search.deep_search_a.main__graph.nodes.refined_sub_question_creation import (
     refined_sub_question_creation,
 )
-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.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 (
+    answer_refined_query_graph_builder,
+)
 from onyx.agents.agent_search.orchestration.nodes.basic_use_tool_response import (
     basic_use_tool_response,
 )

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

@@ -3,11 +3,13 @@ from typing import cast
 
 from langchain_core.runnables import RunnableConfig
 
-from onyx.agents.agent_search.deep_search_a.main.models import AgentAdditionalMetrics
-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.deep_search_a.main__graph.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.models import AgentSearchConfig
 from onyx.agents.agent_search.shared_graph_utils.models import CombinedAgentMetrics
 from onyx.db.chat import log_agent_metrics

backend/onyx/agents/agent_search/deep_search_a/main__graph/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.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.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.models import AgentSearchConfig
 
 

backend/onyx/agents/agent_search/deep_search_a/main__graph/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.states import MainState
+from onyx.agents.agent_search.deep_search_a.main__graph.states import MainState
 
 
 def agent_path_routing(

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

@@ -3,9 +3,11 @@ from typing import cast
 
 from langchain_core.runnables import RunnableConfig
 
-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.states import MainState
+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 (
+    ExploratorySearchUpdate,
+)
+from onyx.agents.agent_search.deep_search_a.main__graph.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__graph/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.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.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.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__graph/nodes/direct_llm_handling.py

@@ -7,9 +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.operations import logger
-from onyx.agents.agent_search.deep_search_a.main.states import InitialAnswerUpdate
-from onyx.agents.agent_search.deep_search_a.main.states import MainState
+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 (
+    InitialAnswerUpdate,
+)
+from onyx.agents.agent_search.deep_search_a.main__graph.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__graph/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.operations import logger
-from onyx.agents.agent_search.deep_search_a.main.states import (
+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 (
     EntityTermExtractionUpdate,
 )
-from onyx.agents.agent_search.deep_search_a.main.states import MainState
+from onyx.agents.agent_search.deep_search_a.main__graph.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__graph/nodes/generate_initial_base_search_only_answer.py

@@ -4,9 +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.operations import logger
-from onyx.agents.agent_search.deep_search_a.main.states import InitialAnswerBASEUpdate
-from onyx.agents.agent_search.deep_search_a.main.states import MainState
+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 (
+    InitialAnswerBASEUpdate,
+)
+from onyx.agents.agent_search.deep_search_a.main__graph.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__graph/nodes/generate_refined_answer.py

@@ -7,14 +7,18 @@ 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.models import AgentRefinedMetrics
-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 (
+from onyx.agents.agent_search.deep_search_a.main__graph.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 (
     remove_document_citations,
 )
-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.deep_search_a.main__graph.states import MainState
+from onyx.agents.agent_search.deep_search_a.main__graph.states import (
+    RefinedAnswerUpdate,
+)
 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__graph/nodes/ingest_refined_answers.py

@@ -1,10 +1,12 @@
 from datetime import datetime
 
-from onyx.agents.agent_search.deep_search_a.answer_initial_sub_question.states import (
+from onyx.agents.agent_search.deep_search_a.initial__individual_sub_answer__subgraph.states import (
     AnswerQuestionOutput,
 )
-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.deep_search_a.main__graph.operations import logger
+from onyx.agents.agent_search.deep_search_a.main__graph.states import (
+    DecompAnswersUpdate,
+)
 from onyx.agents.agent_search.shared_graph_utils.operators import (
     dedup_inference_sections,
 )

backend/onyx/agents/agent_search/deep_search_a/main__graph/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.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 (
+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 (
     RequireRefinedAnswerUpdate,
 )
 from onyx.agents.agent_search.models import AgentSearchConfig

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

@@ -6,13 +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.models import FollowUpSubQuestion
-from onyx.agents.agent_search.deep_search_a.main.operations import dispatch_subquestion
-from onyx.agents.agent_search.deep_search_a.main.operations import logger
-from onyx.agents.agent_search.deep_search_a.main.states import (
+from onyx.agents.agent_search.deep_search_a.main__graph.models import (
+    FollowUpSubQuestion,
+)
+from onyx.agents.agent_search.deep_search_a.main__graph.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 (
     FollowUpSubQuestionsUpdate,
 )
-from onyx.agents.agent_search.deep_search_a.main.states import MainState
+from onyx.agents.agent_search.deep_search_a.main__graph.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__graph/states.py

@@ -6,12 +6,16 @@ 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.expanded_retrieval.models import (
+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 (
+    AgentRefinedMetrics,
+)
+from onyx.agents.agent_search.deep_search_a.main__graph.models import (
+    FollowUpSubQuestion,
+)
+from onyx.agents.agent_search.deep_search_a.util__expanded_retrieval__subgraph.models import (
     ExpandedRetrievalResult,
 )
-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.models import FollowUpSubQuestion
 from onyx.agents.agent_search.orchestration.states import ToolCallUpdate
 from onyx.agents.agent_search.orchestration.states import ToolChoiceInput
 from onyx.agents.agent_search.orchestration.states import ToolChoiceUpdate

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

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

backend/onyx/agents/agent_search/deep_search_a/refinement__consolidate_sub_answers__subgraph/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.answer_initial_sub_question.nodes.answer_check import (
+from onyx.agents.agent_search.deep_search_a.initial__individual_sub_answer__subgraph.nodes.answer_check import (
     answer_check,
 )
-from onyx.agents.agent_search.deep_search_a.answer_initial_sub_question.nodes.answer_generation import (
+from onyx.agents.agent_search.deep_search_a.initial__individual_sub_answer__subgraph.nodes.answer_generation import (
     answer_generation,
 )
-from onyx.agents.agent_search.deep_search_a.answer_initial_sub_question.nodes.format_answer import (
+from onyx.agents.agent_search.deep_search_a.initial__individual_sub_answer__subgraph.nodes.format_answer import (
     format_answer,
 )
-from onyx.agents.agent_search.deep_search_a.answer_initial_sub_question.nodes.ingest_retrieval import (
+from onyx.agents.agent_search.deep_search_a.initial__individual_sub_answer__subgraph.nodes.ingest_retrieval import (
     ingest_retrieval,
 )
-from onyx.agents.agent_search.deep_search_a.answer_initial_sub_question.states import (
+from onyx.agents.agent_search.deep_search_a.initial__individual_sub_answer__subgraph.states import (
     AnswerQuestionInput,
 )
-from onyx.agents.agent_search.deep_search_a.answer_initial_sub_question.states import (
+from onyx.agents.agent_search.deep_search_a.initial__individual_sub_answer__subgraph.states import (
     AnswerQuestionOutput,
 )
-from onyx.agents.agent_search.deep_search_a.answer_initial_sub_question.states import (
+from onyx.agents.agent_search.deep_search_a.initial__individual_sub_answer__subgraph.states import (
     AnswerQuestionState,
 )
-from onyx.agents.agent_search.deep_search_a.answer_refinement_sub_question.edges import (
+from onyx.agents.agent_search.deep_search_a.refinement__consolidate_sub_answers__subgraph.edges import (
     send_to_expanded_refined_retrieval,
 )
-from onyx.agents.agent_search.deep_search_a.expanded_retrieval.graph_builder import (
+from onyx.agents.agent_search.deep_search_a.util__expanded_retrieval__subgraph.graph_builder import (
     expanded_retrieval_graph_builder,
 )
 from onyx.utils.logger import setup_logger

backend/onyx/agents/agent_search/deep_search_a/util__expanded_retrieval__subgraph/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.expanded_retrieval.states import (
+from onyx.agents.agent_search.deep_search_a.util__expanded_retrieval__subgraph.states import (
     ExpandedRetrievalState,
 )
-from onyx.agents.agent_search.deep_search_a.expanded_retrieval.states import (
+from onyx.agents.agent_search.deep_search_a.util__expanded_retrieval__subgraph.states import (
     RetrievalInput,
 )
 from onyx.agents.agent_search.models import AgentSearchConfig

backend/onyx/agents/agent_search/deep_search_a/util__expanded_retrieval__subgraph/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.expanded_retrieval.edges import (
+from onyx.agents.agent_search.deep_search_a.util__expanded_retrieval__subgraph.edges import (
     parallel_retrieval_edge,
 )
-from onyx.agents.agent_search.deep_search_a.expanded_retrieval.nodes.doc_reranking import (
+from onyx.agents.agent_search.deep_search_a.util__expanded_retrieval__subgraph.nodes.doc_reranking import (
     doc_reranking,
 )
-from onyx.agents.agent_search.deep_search_a.expanded_retrieval.nodes.doc_retrieval import (
+from onyx.agents.agent_search.deep_search_a.util__expanded_retrieval__subgraph.nodes.doc_retrieval import (
     doc_retrieval,
 )
-from onyx.agents.agent_search.deep_search_a.expanded_retrieval.nodes.doc_verification import (
+from onyx.agents.agent_search.deep_search_a.util__expanded_retrieval__subgraph.nodes.doc_verification import (
     doc_verification,
 )
-from onyx.agents.agent_search.deep_search_a.expanded_retrieval.nodes.dummy import (
+from onyx.agents.agent_search.deep_search_a.util__expanded_retrieval__subgraph.nodes.dummy import (
     dummy,
 )
-from onyx.agents.agent_search.deep_search_a.expanded_retrieval.nodes.expand_queries import (
+from onyx.agents.agent_search.deep_search_a.util__expanded_retrieval__subgraph.nodes.expand_queries import (
     expand_queries,
 )
-from onyx.agents.agent_search.deep_search_a.expanded_retrieval.nodes.format_results import (
+from onyx.agents.agent_search.deep_search_a.util__expanded_retrieval__subgraph.nodes.format_results import (
     format_results,
 )
-from onyx.agents.agent_search.deep_search_a.expanded_retrieval.nodes.verification_kickoff import (
+from onyx.agents.agent_search.deep_search_a.util__expanded_retrieval__subgraph.nodes.verification_kickoff import (
     verification_kickoff,
 )
-from onyx.agents.agent_search.deep_search_a.expanded_retrieval.states import (
+from onyx.agents.agent_search.deep_search_a.util__expanded_retrieval__subgraph.states import (
     ExpandedRetrievalInput,
 )
-from onyx.agents.agent_search.deep_search_a.expanded_retrieval.states import (
+from onyx.agents.agent_search.deep_search_a.util__expanded_retrieval__subgraph.states import (
     ExpandedRetrievalOutput,
 )
-from onyx.agents.agent_search.deep_search_a.expanded_retrieval.states import (
+from onyx.agents.agent_search.deep_search_a.util__expanded_retrieval__subgraph.states import (
     ExpandedRetrievalState,
 )
 from onyx.agents.agent_search.shared_graph_utils.utils import get_test_config

backend/onyx/agents/agent_search/deep_search_a/util__expanded_retrieval__subgraph/nodes/doc_reranking.py

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

backend/onyx/agents/agent_search/deep_search_a/util__expanded_retrieval__subgraph/nodes/doc_retrieval.py

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

backend/onyx/agents/agent_search/deep_search_a/util__expanded_retrieval__subgraph/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.expanded_retrieval.states import (
+from onyx.agents.agent_search.deep_search_a.util__expanded_retrieval__subgraph.states import (
     DocVerificationInput,
 )
-from onyx.agents.agent_search.deep_search_a.expanded_retrieval.states import (
+from onyx.agents.agent_search.deep_search_a.util__expanded_retrieval__subgraph.states import (
     DocVerificationUpdate,
 )
 from onyx.agents.agent_search.models import AgentSearchConfig

backend/onyx/agents/agent_search/deep_search_a/util__expanded_retrieval__subgraph/nodes/dummy.py

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

backend/onyx/agents/agent_search/deep_search_a/util__expanded_retrieval__subgraph/nodes/expand_queries.py

@@ -5,14 +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.expanded_retrieval.operations import (
+from onyx.agents.agent_search.deep_search_a.util__expanded_retrieval__subgraph.operations import (
     dispatch_subquery,
 )
-from onyx.agents.agent_search.deep_search_a.expanded_retrieval.operations import logger
-from onyx.agents.agent_search.deep_search_a.expanded_retrieval.states import (
+from onyx.agents.agent_search.deep_search_a.util__expanded_retrieval__subgraph.operations import (
+    logger,
+)
+from onyx.agents.agent_search.deep_search_a.util__expanded_retrieval__subgraph.states import (
     ExpandedRetrievalInput,
 )
-from onyx.agents.agent_search.deep_search_a.expanded_retrieval.states import (
+from onyx.agents.agent_search.deep_search_a.util__expanded_retrieval__subgraph.states import (
     QueryExpansionUpdate,
 )
 from onyx.agents.agent_search.models import AgentSearchConfig

backend/onyx/agents/agent_search/deep_search_a/util__expanded_retrieval__subgraph/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.expanded_retrieval.models import (
+from onyx.agents.agent_search.deep_search_a.util__expanded_retrieval__subgraph.models import (
     ExpandedRetrievalResult,
 )
-from onyx.agents.agent_search.deep_search_a.expanded_retrieval.operations import (
+from onyx.agents.agent_search.deep_search_a.util__expanded_retrieval__subgraph.operations import (
     calculate_sub_question_retrieval_stats,
 )
-from onyx.agents.agent_search.deep_search_a.expanded_retrieval.states import (
+from onyx.agents.agent_search.deep_search_a.util__expanded_retrieval__subgraph.states import (
     ExpandedRetrievalState,
 )
-from onyx.agents.agent_search.deep_search_a.expanded_retrieval.states import (
+from onyx.agents.agent_search.deep_search_a.util__expanded_retrieval__subgraph.states import (
     ExpandedRetrievalUpdate,
 )
 from onyx.agents.agent_search.models import AgentSearchConfig

backend/onyx/agents/agent_search/deep_search_a/util__expanded_retrieval__subgraph/nodes/verification_kickoff.py

@@ -5,10 +5,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.expanded_retrieval.states import (
+from onyx.agents.agent_search.deep_search_a.util__expanded_retrieval__subgraph.states import (
     DocVerificationInput,
 )
-from onyx.agents.agent_search.deep_search_a.expanded_retrieval.states import (
+from onyx.agents.agent_search.deep_search_a.util__expanded_retrieval__subgraph.states import (
     ExpandedRetrievalState,
 )
 from onyx.agents.agent_search.models import AgentSearchConfig

backend/onyx/agents/agent_search/deep_search_a/util__expanded_retrieval__subgraph/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.expanded_retrieval.models import (
+from onyx.agents.agent_search.deep_search_a.util__expanded_retrieval__subgraph.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,12 @@ 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_builder import (
+from onyx.agents.agent_search.deep_search_a.main__graph.graph_builder import (
     main_graph_builder as main_graph_builder_a,
 )
-from onyx.agents.agent_search.deep_search_a.main.states import MainInput as MainInput_a
+from onyx.agents.agent_search.deep_search_a.main__graph.states import (
+    MainInput as MainInput_a,
+)
 from onyx.agents.agent_search.models import AgentSearchConfig
 from onyx.agents.agent_search.shared_graph_utils.utils import get_test_config
 from onyx.chat.models import AgentAnswerPiece

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

@@ -2,10 +2,14 @@ from typing import Literal
 
 from pydantic import BaseModel
 
-from onyx.agents.agent_search.deep_search_a.main.models import AgentAdditionalMetrics
-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.models import AgentTimings
+from onyx.agents.agent_search.deep_search_a.main__graph.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 (
+    AgentRefinedMetrics,
+)
+from onyx.agents.agent_search.deep_search_a.main__graph.models import AgentTimings
 from onyx.context.search.models import InferenceSection
 from onyx.tools.models import SearchQueryInfo