refined search + question answering as sub-graphs

This commit is contained in:
joachim-danswer
2025-01-26 12:03:28 -08:00
committed by Evan Lohn
parent de5ce8a613
commit 7487b15522

View File

@@ -2,6 +2,7 @@ from langgraph.graph import END
from langgraph.graph import START from langgraph.graph import START
from langgraph.graph import StateGraph 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.base_raw_search.graph_builder import (
base_raw_search_graph_builder, base_raw_search_graph_builder,
) )
@@ -20,6 +21,7 @@ from onyx.agents.agent_search.deep_search_a.initial_search_sq_subgraph.states im
from onyx.agents.agent_search.deep_search_a.initial_search_sq_subgraph.states import ( from onyx.agents.agent_search.deep_search_a.initial_search_sq_subgraph.states import (
SearchSQState, SearchSQState,
) )
from onyx.agents.agent_search.deep_search_a.initial_sub_question_answering.graph_builder import ( from onyx.agents.agent_search.deep_search_a.initial_sub_question_answering.graph_builder import (
initial_sq_subgraph_builder, initial_sq_subgraph_builder,
) )
@@ -36,6 +38,7 @@ def initial_search_sq_subgraph_builder(test_mode: bool = False) -> StateGraph:
input=SearchSQInput, input=SearchSQInput,
) )
# graph.add_node( # graph.add_node(
# node="initial_sub_question_creation", # node="initial_sub_question_creation",
# action=initial_sub_question_creation, # action=initial_sub_question_creation,
@@ -59,20 +62,12 @@ def initial_search_sq_subgraph_builder(test_mode: bool = False) -> StateGraph:
action=base_raw_search_subgraph, action=base_raw_search_subgraph,
) )
# graph.add_node(
# node="ingest_initial_retrieval",
# action=ingest_initial_base_retrieval,
# )
graph.add_node( graph.add_node(
node="retrieval_consolidation", node="retrieval_consolidation",
action=retrieval_consolidation, action=retrieval_consolidation,
) )
# graph.add_node(
# node="ingest_initial_sub_question_answers",
# action=ingest_initial_sub_question_answers,
# )
graph.add_node( graph.add_node(
node="generate_initial_answer", node="generate_initial_answer",
action=generate_initial_answer, action=generate_initial_answer,
@@ -92,11 +87,6 @@ def initial_search_sq_subgraph_builder(test_mode: bool = False) -> StateGraph:
end_key="base_raw_search_subgraph", end_key="base_raw_search_subgraph",
) )
# graph.add_edge(
# start_key="base_raw_search_subgraph",
# end_key="ingest_initial_retrieval",
# )
# graph.add_edge( # graph.add_edge(
# start_key="agent_search_start", # start_key="agent_search_start",
# end_key="entity_term_extraction_llm", # end_key="entity_term_extraction_llm",
@@ -104,12 +94,13 @@ def initial_search_sq_subgraph_builder(test_mode: bool = False) -> StateGraph:
graph.add_edge( graph.add_edge(
start_key=START, start_key=START,
end_key="sub_question_answering_subgraph", end_key="sub_question_answering_subgraph",
) )
graph.add_edge( graph.add_edge(
start_key=["base_raw_search_subgraph", "sub_question_answering_subgraph"], start_key=["base_raw_search_subgraph", "sub_question_answering_subgraph"],
end_key="retrieval_consolidation", end_key="retrieval_consolidation",
) )
graph.add_edge( graph.add_edge(
@@ -127,15 +118,6 @@ def initial_search_sq_subgraph_builder(test_mode: bool = False) -> StateGraph:
# end_key="initial_sub_question_creation", # end_key="initial_sub_question_creation",
# ) # )
# graph.add_conditional_edges(
# source="initial_sub_question_creation",
# path=parallelize_initial_sub_question_answering,
# path_map=["answer_query_subgraph"],
# )
# graph.add_edge(
# start_key="answer_query_subgraph",
# end_key="ingest_initial_sub_question_answers",
# )
graph.add_edge( graph.add_edge(
start_key="retrieval_consolidation", start_key="retrieval_consolidation",