diff --git a/backend/onyx/agents/agent_search/deep_search_a/initial/generate_initial_answer/graph_builder.py b/backend/onyx/agents/agent_search/deep_search_a/initial/generate_initial_answer/graph_builder.py index 44ea6243bb29..dd01659eab3c 100644 --- a/backend/onyx/agents/agent_search/deep_search_a/initial/generate_initial_answer/graph_builder.py +++ b/backend/onyx/agents/agent_search/deep_search_a/initial/generate_initial_answer/graph_builder.py @@ -2,9 +2,6 @@ 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.graph_builder import ( - consolidate_sub_answers_graph_builder, -) from onyx.agents.agent_search.deep_search_a.initial.generate_initial_answer.nodes.generate_initial_answer import ( generate_initial_answer, ) @@ -17,6 +14,9 @@ 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.generate_sub_answers.graph_builder import ( + generate_sub_answers_graph_builder, +) from onyx.agents.agent_search.deep_search_a.initial.retrieve_orig_question_docs.graph_builder import ( retrieve_orig_question_docs_graph_builder, ) @@ -31,15 +31,15 @@ def generate_initial_answer_graph_builder(test_mode: bool = False) -> StateGraph input=SearchSQInput, ) - consolidate_sub_answers = consolidate_sub_answers_graph_builder().compile() + generate_sub_answers = generate_sub_answers_graph_builder().compile() graph.add_node( - node="consolidate_sub_answers_subgraph", - action=consolidate_sub_answers, + node="generate_sub_answers_subgraph", + action=generate_sub_answers, ) retrieve_orig_question_docs = retrieve_orig_question_docs_graph_builder().compile() graph.add_node( - node="retrieve_orig_question_docs_subgraph", + node="retrieve_orig_question_docs_subgraph_wrapper", action=retrieve_orig_question_docs, ) @@ -62,18 +62,18 @@ def generate_initial_answer_graph_builder(test_mode: bool = False) -> StateGraph graph.add_edge( start_key=START, - end_key="retrieve_orig_question_docs_subgraph", + end_key="retrieve_orig_question_docs_subgraph_wrapper", ) graph.add_edge( start_key=START, - end_key="consolidate_sub_answers_subgraph", + end_key="generate_sub_answers_subgraph", ) graph.add_edge( start_key=[ - "retrieve_orig_question_docs_subgraph", - "consolidate_sub_answers_subgraph", + "retrieve_orig_question_docs_subgraph_wrapper", + "generate_sub_answers_subgraph", ], end_key="generate_initial_answer", ) diff --git a/backend/onyx/agents/agent_search/deep_search_a/initial/consolidate_sub_answers/edges.py b/backend/onyx/agents/agent_search/deep_search_a/initial/generate_sub_answers/edges.py similarity index 97% rename from backend/onyx/agents/agent_search/deep_search_a/initial/consolidate_sub_answers/edges.py rename to backend/onyx/agents/agent_search/deep_search_a/initial/generate_sub_answers/edges.py index f7426c3fc028..fd2f8b684f22 100644 --- a/backend/onyx/agents/agent_search/deep_search_a/initial/consolidate_sub_answers/edges.py +++ b/backend/onyx/agents/agent_search/deep_search_a/initial/generate_sub_answers/edges.py @@ -32,7 +32,7 @@ def parallelize_initial_sub_question_answering( return [ Send( - "answer_sub_question_subgraph", + "answer_sub_question_subgraphs", AnswerQuestionInput( question=question, question_id=make_question_id(0, question_nr + 1), diff --git a/backend/onyx/agents/agent_search/deep_search_a/initial/consolidate_sub_answers/graph_builder.py b/backend/onyx/agents/agent_search/deep_search_a/initial/generate_sub_answers/graph_builder.py similarity index 69% rename from backend/onyx/agents/agent_search/deep_search_a/initial/consolidate_sub_answers/graph_builder.py rename to backend/onyx/agents/agent_search/deep_search_a/initial/generate_sub_answers/graph_builder.py index f5683208235c..b1ebcfd204e2 100644 --- a/backend/onyx/agents/agent_search/deep_search_a/initial/consolidate_sub_answers/graph_builder.py +++ b/backend/onyx/agents/agent_search/deep_search_a/initial/generate_sub_answers/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.edges import ( - parallelize_initial_sub_question_answering, -) -from onyx.agents.agent_search.deep_search_a.initial.consolidate_sub_answers.nodes.decompose_orig_question import ( - decompose_orig_question, -) -from onyx.agents.agent_search.deep_search_a.initial.consolidate_sub_answers.nodes.format_initial_sub_answers import ( - format_initial_sub_answers, -) -from onyx.agents.agent_search.deep_search_a.initial.consolidate_sub_answers.states import ( - SQInput, -) -from onyx.agents.agent_search.deep_search_a.initial.consolidate_sub_answers.states import ( - SQState, -) from onyx.agents.agent_search.deep_search_a.initial.generate_individual_sub_answer.graph_builder import ( answer_query_graph_builder, ) +from onyx.agents.agent_search.deep_search_a.initial.generate_sub_answers.edges import ( + parallelize_initial_sub_question_answering, +) +from onyx.agents.agent_search.deep_search_a.initial.generate_sub_answers.nodes.decompose_orig_question import ( + decompose_orig_question, +) +from onyx.agents.agent_search.deep_search_a.initial.generate_sub_answers.nodes.format_initial_sub_answers import ( + format_initial_sub_answers, +) +from onyx.agents.agent_search.deep_search_a.initial.generate_sub_answers.states import ( + SQInput, +) +from onyx.agents.agent_search.deep_search_a.initial.generate_sub_answers.states import ( + SQState, +) from onyx.utils.logger import setup_logger logger = setup_logger() @@ -27,7 +27,7 @@ logger = setup_logger() test_mode = False -def consolidate_sub_answers_graph_builder() -> StateGraph: +def generate_sub_answers_graph_builder() -> StateGraph: graph = StateGraph( state_schema=SQState, input=SQInput, @@ -37,10 +37,10 @@ def consolidate_sub_answers_graph_builder() -> StateGraph: node="decompose_orig_question", action=decompose_orig_question, ) - answer_sub_question_subgraph = answer_query_graph_builder().compile() + answer_sub_question_subgraphs = answer_query_graph_builder().compile() graph.add_node( - node="answer_sub_question_subgraph", - action=answer_sub_question_subgraph, + node="answer_sub_question_subgraphs", + action=answer_sub_question_subgraphs, ) graph.add_node( @@ -75,10 +75,10 @@ def consolidate_sub_answers_graph_builder() -> StateGraph: graph.add_conditional_edges( source="decompose_orig_question", path=parallelize_initial_sub_question_answering, - path_map=["answer_sub_question_subgraph"], + path_map=["answer_sub_question_subgraphs"], ) graph.add_edge( - start_key="answer_sub_question_subgraph", + start_key=["answer_sub_question_subgraphs"], end_key="format_initial_sub_question_answers", ) diff --git a/backend/onyx/agents/agent_search/deep_search_a/initial/consolidate_sub_answers/nodes/decompose_orig_question.py b/backend/onyx/agents/agent_search/deep_search_a/initial/generate_sub_answers/nodes/decompose_orig_question.py similarity index 100% rename from backend/onyx/agents/agent_search/deep_search_a/initial/consolidate_sub_answers/nodes/decompose_orig_question.py rename to backend/onyx/agents/agent_search/deep_search_a/initial/generate_sub_answers/nodes/decompose_orig_question.py diff --git a/backend/onyx/agents/agent_search/deep_search_a/initial/consolidate_sub_answers/nodes/format_initial_sub_answers.py b/backend/onyx/agents/agent_search/deep_search_a/initial/generate_sub_answers/nodes/format_initial_sub_answers.py similarity index 100% rename from backend/onyx/agents/agent_search/deep_search_a/initial/consolidate_sub_answers/nodes/format_initial_sub_answers.py rename to backend/onyx/agents/agent_search/deep_search_a/initial/generate_sub_answers/nodes/format_initial_sub_answers.py diff --git a/backend/onyx/agents/agent_search/deep_search_a/initial/consolidate_sub_answers/states.py b/backend/onyx/agents/agent_search/deep_search_a/initial/generate_sub_answers/states.py similarity index 100% rename from backend/onyx/agents/agent_search/deep_search_a/initial/consolidate_sub_answers/states.py rename to backend/onyx/agents/agent_search/deep_search_a/initial/generate_sub_answers/states.py diff --git a/backend/onyx/agents/agent_search/deep_search_a/initial/retrieve_orig_question_docs/graph_builder.py b/backend/onyx/agents/agent_search/deep_search_a/initial/retrieve_orig_question_docs/graph_builder.py index 656f7a0f63c3..698492722b48 100644 --- a/backend/onyx/agents/agent_search/deep_search_a/initial/retrieve_orig_question_docs/graph_builder.py +++ b/backend/onyx/agents/agent_search/deep_search_a/initial/retrieve_orig_question_docs/graph_builder.py @@ -38,7 +38,7 @@ def retrieve_orig_question_docs_graph_builder() -> StateGraph: expanded_retrieval = expanded_retrieval_graph_builder().compile() graph.add_node( - node="expanded_retrieval_base_search", + node="retrieve_orig_question_docs_subgraph", action=expanded_retrieval, ) graph.add_node( @@ -52,10 +52,10 @@ def retrieve_orig_question_docs_graph_builder() -> StateGraph: graph.add_edge( start_key="format_orig_question_search_input", - end_key="expanded_retrieval_base_search", + end_key="retrieve_orig_question_docs_subgraph", ) graph.add_edge( - start_key="expanded_retrieval_base_search", + start_key="retrieve_orig_question_docs_subgraph", end_key="format_orig_question_search_output", ) diff --git a/backend/onyx/agents/agent_search/run_graph.py b/backend/onyx/agents/agent_search/run_graph.py index c1f38dcab85b..ea644ae1979e 100644 --- a/backend/onyx/agents/agent_search/run_graph.py +++ b/backend/onyx/agents/agent_search/run_graph.py @@ -204,7 +204,7 @@ def run_basic_graph( if __name__ == "__main__": from onyx.llm.factory import get_default_llms - for _ in range(1): + for _ in range(2): now_start = datetime.now() logger.debug(f"Start at {now_start}") graph = main_graph_builder_a()