renamings and consolidation of formatting nodes in orig question retrieval

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

@@ -20,7 +20,7 @@ from onyx.agents.agent_search.deep_search_a.initial.generate_initial_answer.stat
 from onyx.agents.agent_search.deep_search_a.initial.generate_initial_answer.states import (
     SearchSQState,
 )
-from onyx.agents.agent_search.deep_search_a.initial.retrieve_orig_question_documents.graph_builder import (
+from onyx.agents.agent_search.deep_search_a.initial.retrieve_orig_question_docs.graph_builder import (
     base_raw_search_graph_builder,
 )
 from onyx.utils.logger import setup_logger

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

@@ -1,6 +1,6 @@
 from datetime import datetime
 
-from onyx.agents.agent_search.deep_search_a.initial.retrieve_orig_question_documents.states import (
+from onyx.agents.agent_search.deep_search_a.initial.retrieve_orig_question_docs.states import (
     BaseRawSearchOutput,
 )
 from onyx.agents.agent_search.deep_search_a.main.operations import logger
@@ -10,7 +10,7 @@ from onyx.agents.agent_search.deep_search_a.main.states import (
 from onyx.agents.agent_search.shared_graph_utils.models import AgentChunkStats
 
 
-def ingest_initial_base_retrieval(
+def ingest_retrieved_documents(
     state: BaseRawSearchOutput,
 ) -> ExpandedRetrievalUpdate:
     now_start = datetime.now()

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

@@ -2,22 +2,19 @@ from langgraph.graph import END
 from langgraph.graph import START
 from langgraph.graph import StateGraph
 
-from onyx.agents.agent_search.deep_search_a.initial.retrieve_orig_question_documents.nodes.format_raw_search_results import (
-    format_raw_search_results,
+from onyx.agents.agent_search.deep_search_a.initial.retrieve_orig_question_docs.nodes.format_orig_question_search_input import (
+    format_orig_question_search_input,
 )
-from onyx.agents.agent_search.deep_search_a.initial.retrieve_orig_question_documents.nodes.generate_raw_search_data import (
-    generate_raw_search_data,
+from onyx.agents.agent_search.deep_search_a.initial.retrieve_orig_question_docs.nodes.format_orig_question_search_output import (
+    format_orig_question_search_output,
 )
-from onyx.agents.agent_search.deep_search_a.initial.retrieve_orig_question_documents.nodes.ingest_initial_base_retrieval import (
-    ingest_initial_base_retrieval,
-)
-from onyx.agents.agent_search.deep_search_a.initial.retrieve_orig_question_documents.states import (
+from onyx.agents.agent_search.deep_search_a.initial.retrieve_orig_question_docs.states import (
     BaseRawSearchInput,
 )
-from onyx.agents.agent_search.deep_search_a.initial.retrieve_orig_question_documents.states import (
+from onyx.agents.agent_search.deep_search_a.initial.retrieve_orig_question_docs.states import (
     BaseRawSearchOutput,
 )
-from onyx.agents.agent_search.deep_search_a.initial.retrieve_orig_question_documents.states import (
+from onyx.agents.agent_search.deep_search_a.initial.retrieve_orig_question_docs.states import (
     BaseRawSearchState,
 )
 from onyx.agents.agent_search.deep_search_a.shared.expanded_retrieval.graph_builder import (
@@ -35,8 +32,8 @@ def base_raw_search_graph_builder() -> StateGraph:
     ### Add nodes ###
 
     graph.add_node(
-        node="generate_raw_search_data",
-        action=generate_raw_search_data,
+        node="format_orig_question_search_input",
+        action=format_orig_question_search_input,
     )
 
     expanded_retrieval = expanded_retrieval_graph_builder().compile()
@@ -45,40 +42,25 @@ def base_raw_search_graph_builder() -> StateGraph:
         action=expanded_retrieval,
     )
     graph.add_node(
-        node="format_raw_search_results",
-        action=format_raw_search_results,
-    )
-
-    graph.add_node(
-        node="ingest_initial_base_retrieval",
-        action=ingest_initial_base_retrieval,
+        node="format_orig_question_search_output",
+        action=format_orig_question_search_output,
     )
 
     ### Add edges ###
 
-    graph.add_edge(start_key=START, end_key="generate_raw_search_data")
+    graph.add_edge(start_key=START, end_key="format_orig_question_search_input")
 
     graph.add_edge(
-        start_key="generate_raw_search_data",
+        start_key="format_orig_question_search_input",
         end_key="expanded_retrieval_base_search",
     )
     graph.add_edge(
         start_key="expanded_retrieval_base_search",
-        end_key="format_raw_search_results",
-    )
-
-    # graph.add_edge(
-    #     start_key="expanded_retrieval_base_search",
-    #     end_key=END,
-    # )
-
-    graph.add_edge(
-        start_key="format_raw_search_results",
-        end_key="ingest_initial_base_retrieval",
+        end_key="format_orig_question_search_output",
     )
 
     graph.add_edge(
-        start_key="ingest_initial_base_retrieval",
+        start_key="format_orig_question_search_output",
         end_key=END,
     )
 

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

@@ -12,7 +12,7 @@ from onyx.utils.logger import setup_logger
 logger = setup_logger()
 
 
-def generate_raw_search_data(
+def format_orig_question_search_input(
     state: CoreState, config: RunnableConfig
 ) -> ExpandedRetrievalInput:
     logger.debug("generate_raw_search_data")

backend/onyx/agents/agent_search/deep_search_a/initial/retrieve_orig_question_docs/nodes/format_orig_question_search_output.py

@@ -0,0 +1,33 @@
+from onyx.agents.agent_search.deep_search_a.main.states import ExpandedRetrievalUpdate
+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
+from onyx.utils.logger import setup_logger
+
+logger = setup_logger()
+
+
+def format_orig_question_search_output(
+    state: ExpandedRetrievalOutput,
+) -> ExpandedRetrievalUpdate:
+    # return BaseRawSearchOutput(
+    #     base_expanded_retrieval_result=state.expanded_retrieval_result,
+    #     # base_retrieval_results=[state.expanded_retrieval_result],
+    #     # base_search_documents=[],
+    # )
+
+    sub_question_retrieval_stats = (
+        state.expanded_retrieval_result.sub_question_retrieval_stats
+    )
+    if sub_question_retrieval_stats is None:
+        sub_question_retrieval_stats = AgentChunkStats()
+    else:
+        sub_question_retrieval_stats = sub_question_retrieval_stats
+
+    return ExpandedRetrievalUpdate(
+        original_question_retrieval_results=state.expanded_retrieval_result.expanded_queries_results,
+        all_original_question_documents=state.expanded_retrieval_result.context_documents,
+        original_question_retrieval_stats=sub_question_retrieval_stats,
+        log_messages=[],
+    )

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

@@ -1,18 +0,0 @@
-from onyx.agents.agent_search.deep_search_a.initial.retrieve_orig_question_documents.states import (
-    BaseRawSearchOutput,
-)
-from onyx.agents.agent_search.deep_search_a.shared.expanded_retrieval.states import (
-    ExpandedRetrievalOutput,
-)
-from onyx.utils.logger import setup_logger
-
-logger = setup_logger()
-
-
-def format_raw_search_results(state: ExpandedRetrievalOutput) -> BaseRawSearchOutput:
-    logger.debug("format_raw_search_results")
-    return BaseRawSearchOutput(
-        base_expanded_retrieval_result=state.expanded_retrieval_result,
-        # base_retrieval_results=[state.expanded_retrieval_result],
-        # base_search_documents=[],
-    )

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

@@ -1,41 +0,0 @@
-from datetime import datetime
-
-from onyx.agents.agent_search.deep_search_a.initial.retrieve_orig_question_documents.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.shared_graph_utils.models import AgentChunkStats
-
-
-def ingest_initial_base_retrieval(
-    state: BaseRawSearchOutput,
-) -> ExpandedRetrievalUpdate:
-    now_start = datetime.now()
-
-    logger.info(f"--------{now_start}--------INGEST INITIAL RETRIEVAL---")
-
-    sub_question_retrieval_stats = (
-        state.base_expanded_retrieval_result.sub_question_retrieval_stats
-    )
-    if sub_question_retrieval_stats is None:
-        sub_question_retrieval_stats = AgentChunkStats()
-    else:
-        sub_question_retrieval_stats = sub_question_retrieval_stats
-
-    now_end = datetime.now()
-
-    logger.debug(
-        f"--------{now_end}--{now_end - now_start}--------INGEST INITIAL RETRIEVAL END---"
-    )
-
-    return ExpandedRetrievalUpdate(
-        original_question_retrieval_results=state.base_expanded_retrieval_result.expanded_queries_results,
-        all_original_question_documents=state.base_expanded_retrieval_result.context_documents,
-        original_question_retrieval_stats=sub_question_retrieval_stats,
-        log_messages=[
-            f"{now_start} -- Main - Ingestion base retrieval,  Time taken: {now_end - now_start}"
-        ],
-    )