diff --git a/redbox-core/redbox/graph/root.py b/redbox-core/redbox/graph/root.py index 5aee48e6..8eb66522 100644 --- a/redbox-core/redbox/graph/root.py +++ b/redbox-core/redbox/graph/root.py @@ -47,7 +47,11 @@ def self_route_question_is_unanswerable(llm_response: str): return "unanswerable" in llm_response # Processes - builder.add_node("p_condense_question", build_chat_pattern(prompt_set=PromptSet.CondenseQuestion)) + builder.add_node( + "p_condense_question", + build_chat_pattern(prompt_set=PromptSet.CondenseQuestion), + retry=RetryPolicy(max_attempts=3), + ) builder.add_node( "p_retrieve_docs", build_retrieve_pattern( @@ -55,6 +59,7 @@ def self_route_question_is_unanswerable(llm_response: str): structure_func=structure_documents_by_file_name, final_source_chain=False, ), + retry=RetryPolicy(max_attempts=3), ) builder.add_node( "p_answer_question_or_decide_unanswerable", @@ -76,8 +81,13 @@ def self_route_question_is_unanswerable(llm_response: str): true_condition_state_update={"route_name": ChatRoute.chat_with_docs_map_reduce}, false_condition_state_update={"route_name": ChatRoute.search}, ), + retry=RetryPolicy(max_attempts=3), + ) + builder.add_node( + "p_clear_documents", + clear_documents_process, + retry=RetryPolicy(max_attempts=3), ) - builder.add_node("p_clear_documents", clear_documents_process) # Edges builder.add_edge(START, "p_condense_question") @@ -104,10 +114,15 @@ def get_chat_graph( builder = StateGraph(RedboxState) # Processes - builder.add_node("p_set_chat_route", build_set_route_pattern(route=ChatRoute.chat)) + builder.add_node( + "p_set_chat_route", + build_set_route_pattern(route=ChatRoute.chat), + retry=RetryPolicy(max_attempts=3), + ) builder.add_node( "p_chat", build_chat_pattern(prompt_set=PromptSet.Chat, final_response_chain=True), + retry=RetryPolicy(max_attempts=3), ) # Edges @@ -129,8 +144,16 @@ def get_search_graph( builder = StateGraph(RedboxState) # Processes - builder.add_node("p_set_search_route", build_set_route_pattern(route=ChatRoute.search)) - builder.add_node("p_condense_question", build_chat_pattern(prompt_set=PromptSet.CondenseQuestion)) + builder.add_node( + "p_set_search_route", + build_set_route_pattern(route=ChatRoute.search), + retry=RetryPolicy(max_attempts=3), + ) + builder.add_node( + "p_condense_question", + build_chat_pattern(prompt_set=PromptSet.CondenseQuestion), + retry=RetryPolicy(max_attempts=3), + ) builder.add_node( "p_retrieve_docs", build_retrieve_pattern( @@ -138,6 +161,7 @@ def get_search_graph( structure_func=structure_documents_by_group_and_indices, final_source_chain=final_sources, ), + retry=RetryPolicy(max_attempts=3), ) builder.add_node( "p_stuff_docs", @@ -165,7 +189,11 @@ def get_agentic_search_graph(tools: List[StructuredTool], debug: bool = False) - # agent_tools: list[StructuredTool] = [tools[tool_name] for tool_name in agent_tool_names] # Processes - builder.add_node("p_set_agentic_search_route", build_set_route_pattern(route=ChatRoute.gadget)) + builder.add_node( + "p_set_agentic_search_route", + build_set_route_pattern(route=ChatRoute.gadget), + retry=RetryPolicy(max_attempts=3), + ) builder.add_node( "p_search_agent", build_stuff_pattern( @@ -180,13 +208,18 @@ def get_agentic_search_graph(tools: List[StructuredTool], debug: bool = False) - builder.add_node( "p_retrieval_tools", ToolNode(tools=tools), + retry=RetryPolicy(max_attempts=3), ) builder.add_node( "p_give_up_agent", build_stuff_pattern(prompt_set=PromptSet.GiveUpAgentic, final_response_chain=True), retry=RetryPolicy(max_attempts=3), ) - builder.add_node("p_report_sources", report_sources_process) + builder.add_node( + "p_report_sources", + report_sources_process, + retry=RetryPolicy(max_attempts=3), + ) # Log builder.add_node( @@ -197,13 +230,22 @@ def get_agentic_search_graph(tools: List[StructuredTool], debug: bool = False) - for tool_state_entry in s.last_message.tool_calls ] ), + retry=RetryPolicy(max_attempts=3), ) # Decisions - builder.add_node("d_x_steps_left_or_less", empty_process) + builder.add_node( + "d_x_steps_left_or_less", + empty_process, + retry=RetryPolicy(max_attempts=3), + ) # Sends - builder.add_node("s_tool", empty_process) + builder.add_node( + "s_tool", + empty_process, + retry=RetryPolicy(max_attempts=3), + ) # Edges builder.add_edge(START, "p_set_agentic_search_route") @@ -235,11 +277,20 @@ def get_chat_with_documents_graph( builder = StateGraph(RedboxState) # Processes - builder.add_node("p_pass_question_to_text", build_passthrough_pattern()) - builder.add_node("p_set_chat_docs_route", build_set_route_pattern(route=ChatRoute.chat_with_docs)) + builder.add_node( + "p_pass_question_to_text", + build_passthrough_pattern(), + retry=RetryPolicy(max_attempts=3), + ) + builder.add_node( + "p_set_chat_docs_route", + build_set_route_pattern(route=ChatRoute.chat_with_docs), + retry=RetryPolicy(max_attempts=3), + ) builder.add_node( "p_set_chat_docs_map_reduce_route", build_set_route_pattern(route=ChatRoute.chat_with_docs_map_reduce), + retry=RetryPolicy(max_attempts=3), ) builder.add_node( "p_summarise_each_document", @@ -259,17 +310,23 @@ def get_chat_with_documents_graph( ), retry=RetryPolicy(max_attempts=3), ) - builder.add_node("p_clear_documents", clear_documents_process) + builder.add_node( + "p_clear_documents", + clear_documents_process, + retry=RetryPolicy(max_attempts=3), + ) builder.add_node( "p_too_large_error", build_error_pattern( text="These documents are too large to work with.", route_name=ErrorRoute.files_too_large, ), + retry=RetryPolicy(max_attempts=3), ) builder.add_node( "p_answer_or_decide_route", get_self_route_graph(parameterised_retriever, PromptSet.SelfRoute), + retry=RetryPolicy(max_attempts=3), ) builder.add_node( "p_retrieve_all_chunks", @@ -278,24 +335,58 @@ def get_chat_with_documents_graph( structure_func=structure_documents_by_file_name, final_source_chain=True, ), + retry=RetryPolicy(max_attempts=3), ) builder.add_node( "p_activity_log_tool_decision", build_activity_log_node(lambda state: RedboxActivityEvent(message=f"Using _{state.route_name}_")), + retry=RetryPolicy(max_attempts=3), ) # Decisions - builder.add_node("d_request_handler_from_total_tokens", empty_process) - builder.add_node("d_single_doc_summaries_bigger_than_context", empty_process) - builder.add_node("d_doc_summaries_bigger_than_context", empty_process) - builder.add_node("d_groups_have_multiple_docs", empty_process) - builder.add_node("d_self_route_is_enabled", empty_process) + builder.add_node( + "d_request_handler_from_total_tokens", + empty_process, + retry=RetryPolicy(max_attempts=3), + ) + builder.add_node( + "d_single_doc_summaries_bigger_than_context", + empty_process, + retry=RetryPolicy(max_attempts=3), + ) + builder.add_node( + "d_doc_summaries_bigger_than_context", + empty_process, + retry=RetryPolicy(max_attempts=3), + ) + builder.add_node( + "d_groups_have_multiple_docs", + empty_process, + retry=RetryPolicy(max_attempts=3), + ) + builder.add_node( + "d_self_route_is_enabled", + empty_process, + retry=RetryPolicy(max_attempts=3), + ) # Sends - builder.add_node("s_chunk", empty_process) - builder.add_node("s_group_1", empty_process) - builder.add_node("s_group_2", empty_process) + builder.add_node( + "s_chunk", + empty_process, + retry=RetryPolicy(max_attempts=3), + ) + builder.add_node( + "s_group_1", + empty_process, + retry=RetryPolicy(max_attempts=3), + ) + builder.add_node( + "s_group_2", + empty_process, + retry=RetryPolicy(max_attempts=3), + ) # Edges builder.add_edge(START, "p_pass_question_to_text") @@ -391,9 +482,18 @@ def get_retrieve_metadata_graph(metadata_retriever: VectorStoreRetriever, debug: retriever=metadata_retriever, structure_func=structure_documents_by_file_name, ), + retry=RetryPolicy(max_attempts=3), + ) + builder.add_node( + "p_set_metadata", + build_set_metadata_pattern(), + retry=RetryPolicy(max_attempts=3), + ) + builder.add_node( + "p_clear_metadata_documents", + clear_documents_process, + retry=RetryPolicy(max_attempts=3), ) - builder.add_node("p_set_metadata", build_set_metadata_pattern()) - builder.add_node("p_clear_metadata_documents", clear_documents_process) # Edges builder.add_edge(START, "p_retrieve_metadata") @@ -428,12 +528,36 @@ def get_root_graph( new_route = build_new_graph(all_chunks_retriever, parameterised_retriever, metadata_retriever, tools, debug) # Processes - builder.add_node("p_search", rag_subgraph) - builder.add_node("p_search_agentic", agent_subgraph) - builder.add_node("p_chat", chat_subgraph) - builder.add_node("p_chat_with_documents", cwd_subgraph) - builder.add_node("p_retrieve_metadata", metadata_subgraph) - builder.add_node("p_new_route", new_route) + builder.add_node( + "p_search", + rag_subgraph, + retry=RetryPolicy(max_attempts=3), + ) + builder.add_node( + "p_search_agentic", + agent_subgraph, + retry=RetryPolicy(max_attempts=3), + ) + builder.add_node( + "p_chat", + chat_subgraph, + retry=RetryPolicy(max_attempts=3), + ) + builder.add_node( + "p_chat_with_documents", + cwd_subgraph, + retry=RetryPolicy(max_attempts=3), + ) + builder.add_node( + "p_retrieve_metadata", + metadata_subgraph, + retry=RetryPolicy(max_attempts=3), + ) + builder.add_node( + "p_new_route", + new_route, + retry=RetryPolicy(max_attempts=3), + ) # Log builder.add_node( @@ -447,11 +571,20 @@ def get_root_graph( else "You selected no files", ] ), + retry=RetryPolicy(max_attempts=3), ) # Decisions - builder.add_node("d_keyword_exists", empty_process) - builder.add_node("d_docs_selected", empty_process) + builder.add_node( + "d_keyword_exists", + empty_process, + retry=RetryPolicy(max_attempts=3), + ) + builder.add_node( + "d_docs_selected", + empty_process, + retry=RetryPolicy(max_attempts=3), + ) # Edges builder.add_edge(START, "p_activity_log_user_request") @@ -501,9 +634,17 @@ def build_new_graph( # Subgraphs/may need to convert into tools metadata_subgraph = get_retrieve_metadata_graph(metadata_retriever=metadata_retriever, debug=debug) - builder.add_node("p_strip_route", strip_route) + builder.add_node( + "p_strip_route", + strip_route, + retry=RetryPolicy(max_attempts=3), + ) # Nodes - builder.add_node("p_retrieve_metadata", metadata_subgraph) + builder.add_node( + "p_retrieve_metadata", + metadata_subgraph, + retry=RetryPolicy(max_attempts=3), + ) # add back when move to this graph # Show activity logs from user request # builder.add_node( @@ -525,8 +666,13 @@ def build_new_graph( structure_func=structure_documents_by_file_name, final_source_chain=False, ), + retry=RetryPolicy(max_attempts=3), + ) + builder.add_node( + "d_docs_selected", + empty_process, + retry=RetryPolicy(max_attempts=3), ) - builder.add_node("d_docs_selected", empty_process) builder.add_node( "p_search_agent", build_stuff_pattern( @@ -536,17 +682,24 @@ def build_new_graph( format_instructions=format_instructions, final_response_chain=True, # Output parser handles streaming ), + retry=RetryPolicy(max_attempts=3), ) builder.add_node( "p_retrieval_tools", ToolNode(tools=tools), + retry=RetryPolicy(max_attempts=3), ) builder.add_node( "p_give_up_agent", build_stuff_pattern(prompt_set=PromptSet.GiveUpAgentic, final_response_chain=True), + retry=RetryPolicy(max_attempts=3), + ) + builder.add_node( + "p_report_sources", + report_sources_process, + retry=RetryPolicy(max_attempts=3), ) - builder.add_node("p_report_sources", report_sources_process) builder.add_node( "p_activity_log_retrieval_tool_calls", @@ -556,9 +709,18 @@ def build_new_graph( for tool_state_entry in s.last_message.tool_calls ] ), + retry=RetryPolicy(max_attempts=3), + ) + builder.add_node( + "s_tool", + empty_process, + retry=RetryPolicy(max_attempts=3), + ) + builder.add_node( + "d_x_steps_left_or_less", + empty_process, + retry=RetryPolicy(max_attempts=3), ) - builder.add_node("s_tool", empty_process) - builder.add_node("d_x_steps_left_or_less", empty_process) # Edges # builder.add_edge(START, "p_activity_log_user_request") # add back when move to this graph