From dbdbb39936ec48abbe45e4c3cef1772079cf15d2 Mon Sep 17 00:00:00 2001 From: joachim-danswer Date: Tue, 26 Aug 2025 15:03:11 -0700 Subject: [PATCH 1/4] fix budget calculation --- .../onyx/agents/agent_search/dr/nodes/dr_a1_orchestrator.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/backend/onyx/agents/agent_search/dr/nodes/dr_a1_orchestrator.py b/backend/onyx/agents/agent_search/dr/nodes/dr_a1_orchestrator.py index 94d26392958..9e1bbce7411 100644 --- a/backend/onyx/agents/agent_search/dr/nodes/dr_a1_orchestrator.py +++ b/backend/onyx/agents/agent_search/dr/nodes/dr_a1_orchestrator.py @@ -163,7 +163,6 @@ def orchestrator( all_relationship_types = get_relationship_types_str(active=True) # default to closer - next_tool = DRPath.CLOSER.value query_list = ["Answer the question with the information you have."] decision_prompt = None @@ -310,7 +309,7 @@ def orchestrator( logger.error(f"Error in approach extraction: {e}") raise e - remaining_time_budget -= available_tools[next_tool].cost + remaining_time_budget -= available_tools[next_tool_name].cost else: if iteration_nr == 1 and not plan_of_record: # by default, we start a new iteration, but if there is a feedback request, @@ -434,7 +433,7 @@ def orchestrator( next_step = orchestrator_action.next_step next_tool_name = next_step.tool - next_tool = available_tools[next_tool_name].path + available_tools[next_tool_name].path query_list = [q for q in (next_step.questions or [])] reasoning_result = orchestrator_action.reasoning From 6ae28e38a883958e04950671b3674cc071bfaab2 Mon Sep 17 00:00:00 2001 From: joachim-danswer Date: Tue, 26 Aug 2025 16:47:14 -0700 Subject: [PATCH 2/4] Internal custom tool fix + Okta special casing --- .../dr_generic_internal_tool_2_act.py | 14 +++++-- .../dr_generic_internal_tool_3_reduce.py | 2 + backend/onyx/prompts/dr_prompts.py | 40 +++++++++++++++++-- 3 files changed, 49 insertions(+), 7 deletions(-) diff --git a/backend/onyx/agents/agent_search/dr/sub_agents/generic_internal_tool/dr_generic_internal_tool_2_act.py b/backend/onyx/agents/agent_search/dr/sub_agents/generic_internal_tool/dr_generic_internal_tool_2_act.py index 91ebea033b4..8931a38bf66 100644 --- a/backend/onyx/agents/agent_search/dr/sub_agents/generic_internal_tool/dr_generic_internal_tool_2_act.py +++ b/backend/onyx/agents/agent_search/dr/sub_agents/generic_internal_tool/dr_generic_internal_tool_2_act.py @@ -15,6 +15,7 @@ ) from onyx.prompts.dr_prompts import CUSTOM_TOOL_PREP_PROMPT from onyx.prompts.dr_prompts import CUSTOM_TOOL_USE_PROMPT +from onyx.prompts.dr_prompts import OKTA_TOOL_USE_SPECIAL_PROMPT from onyx.utils.logger import setup_logger logger = setup_logger() @@ -97,12 +98,17 @@ def generic_internal_tool_act( tool_result_str = json.dumps(final_data, ensure_ascii=False) tool_str = ( - f"Tool used: {generic_internal_tool_name}\n" + f"Tool used: {generic_internal_tool.display_name}\n" f"Description: {generic_internal_tool_info.description}\n" f"Result: {tool_result_str}" ) - tool_summary_prompt = CUSTOM_TOOL_USE_PROMPT.build( + if generic_internal_tool.display_name == "Okta Profile": + tool_prompt = OKTA_TOOL_USE_SPECIAL_PROMPT + else: + tool_prompt = CUSTOM_TOOL_USE_PROMPT + + tool_summary_prompt = tool_prompt.build( query=branch_query, base_question=base_question, tool_response=tool_str ) answer_string = str( @@ -118,7 +124,7 @@ def generic_internal_tool_act( return BranchUpdate( branch_iteration_responses=[ IterationAnswer( - tool=generic_internal_tool_name, + tool=generic_internal_tool.llm_name, tool_id=generic_internal_tool_info.tool_id, iteration_nr=iteration_nr, parallelization_nr=parallelization_nr, @@ -128,7 +134,7 @@ def generic_internal_tool_act( cited_documents={}, reasoning="", additional_data=None, - response_type="string", + response_type="text", data=answer_string, ) ], diff --git a/backend/onyx/agents/agent_search/dr/sub_agents/generic_internal_tool/dr_generic_internal_tool_3_reduce.py b/backend/onyx/agents/agent_search/dr/sub_agents/generic_internal_tool/dr_generic_internal_tool_3_reduce.py index 901c461d99c..4a51ac48722 100644 --- a/backend/onyx/agents/agent_search/dr/sub_agents/generic_internal_tool/dr_generic_internal_tool_3_reduce.py +++ b/backend/onyx/agents/agent_search/dr/sub_agents/generic_internal_tool/dr_generic_internal_tool_3_reduce.py @@ -46,6 +46,7 @@ def generic_internal_tool_reducer( write_custom_event( current_step_nr, CustomToolStart( + type="custom_tool_start", tool_name=new_update.tool, ), writer, @@ -54,6 +55,7 @@ def generic_internal_tool_reducer( write_custom_event( current_step_nr, CustomToolDelta( + type="custom_tool_delta", tool_name=new_update.tool, response_type=new_update.response_type, data=new_update.data, diff --git a/backend/onyx/prompts/dr_prompts.py b/backend/onyx/prompts/dr_prompts.py index c7ac6dcca83..761610d6f03 100644 --- a/backend/onyx/prompts/dr_prompts.py +++ b/backend/onyx/prompts/dr_prompts.py @@ -835,6 +835,40 @@ ) +OKTA_TOOL_USE_SPECIAL_PROMPT = PromptTemplate( + f"""\ +You are great at formatting the response from Okta and also provide a short reasoning and answer \ +in natural language to answer the specific task query, if possible. + +Here is the specific task query: +{SEPARATOR_LINE} +---query--- +{SEPARATOR_LINE} + +Here is the base question that ultimately needs to be answered: +{SEPARATOR_LINE} +---base_question--- +{SEPARATOR_LINE} + +Here is the tool response: +{SEPARATOR_LINE} +---tool_response--- +{SEPARATOR_LINE} + +Approach: + - start your answer by formatting the raw response from Okta in a readable format. + - then try to answer very concise and specifically to the specific task query, if possible. \ +If the Okta information appears not to be relevant, simply say that the Okta \ +information does not appear to relate to the specific task query. + +Guidelines: + - ONLY base any answer DIRECTLY on the Okta response. Do NOT DRAW on your own internal knowledge! + +ANSWER: +""" +) + + CUSTOM_TOOL_USE_PROMPT = PromptTemplate( f"""\ You are great at formatting the response from a tool into a short reasoning and answer \ @@ -863,10 +897,10 @@ - while the base question is important, really focus on answering the specific task query. \ That is your task. -Please respond with a short sentence explaining what the tool does and provide a concise answer to the \ +Please respond with a concise answer to the \ specific task query using the tool response. -If the tool definition and response did not provide information relevant to the specific context mentioned \ -in the query, start out with a short statement highlighting this (e.g., I was not able to find information \ +If the tool definition and response did not provide information relevant to the specific task query mentioned \ +, start out with a short statement highlighting this (e.g., I was not able to find information \ about yellow curry specifically, but I found information about curry...). ANSWER: From 764fe1f6dc4b2e8cf4fb867c784fd169b0c24a18 Mon Sep 17 00:00:00 2001 From: joachim-danswer Date: Tue, 26 Aug 2025 17:07:00 -0700 Subject: [PATCH 3/4] nits --- .../agent_search/dr/nodes/dr_a1_orchestrator.py | 12 ++++++++++-- backend/onyx/prompts/dr_prompts.py | 6 +++--- 2 files changed, 13 insertions(+), 5 deletions(-) diff --git a/backend/onyx/agents/agent_search/dr/nodes/dr_a1_orchestrator.py b/backend/onyx/agents/agent_search/dr/nodes/dr_a1_orchestrator.py index 9e1bbce7411..fa275940f55 100644 --- a/backend/onyx/agents/agent_search/dr/nodes/dr_a1_orchestrator.py +++ b/backend/onyx/agents/agent_search/dr/nodes/dr_a1_orchestrator.py @@ -309,7 +309,11 @@ def orchestrator( logger.error(f"Error in approach extraction: {e}") raise e - remaining_time_budget -= available_tools[next_tool_name].cost + if next_tool_name in available_tools.keys(): + remaining_time_budget -= available_tools[next_tool_name].cost + else: + logger.warning(f"Tool {next_tool_name} not found in available tools") + remaining_time_budget -= 1.0 else: if iteration_nr == 1 and not plan_of_record: # by default, we start a new iteration, but if there is a feedback request, @@ -443,7 +447,11 @@ def orchestrator( logger.error(f"Error in approach extraction: {e}") raise e - remaining_time_budget -= available_tools[next_tool_name].cost + if next_tool_name in available_tools.keys(): + remaining_time_budget -= available_tools[next_tool_name].cost + else: + logger.warning(f"Tool {next_tool_name} not found in available tools") + remaining_time_budget -= 1.0 else: reasoning_result = "Time to wrap up." diff --git a/backend/onyx/prompts/dr_prompts.py b/backend/onyx/prompts/dr_prompts.py index 761610d6f03..f80a84c4c6d 100644 --- a/backend/onyx/prompts/dr_prompts.py +++ b/backend/onyx/prompts/dr_prompts.py @@ -548,8 +548,6 @@ ---question--- {SEPARATOR_LINE} -The current iteration is ---iteration_nr---: - Here is the high-level plan: {SEPARATOR_LINE} ---current_plan_of_record_string--- @@ -838,7 +836,7 @@ OKTA_TOOL_USE_SPECIAL_PROMPT = PromptTemplate( f"""\ You are great at formatting the response from Okta and also provide a short reasoning and answer \ -in natural language to answer the specific task query, if possible. +in natural language to answer the specific task query (not the base question!), if possible. Here is the specific task query: {SEPARATOR_LINE} @@ -862,6 +860,8 @@ information does not appear to relate to the specific task query. Guidelines: + - only use the base question for context, but don't try to answer it. Try to answer \ +the 'specific task query', if possible. - ONLY base any answer DIRECTLY on the Okta response. Do NOT DRAW on your own internal knowledge! ANSWER: From 500d3b83fc0fc4d0c5fac81c53a14087123246a8 Mon Sep 17 00:00:00 2001 From: joachim-danswer Date: Tue, 26 Aug 2025 20:27:25 -0700 Subject: [PATCH 4/4] CW comments --- backend/onyx/agents/agent_search/dr/nodes/dr_a1_orchestrator.py | 2 -- .../generic_internal_tool/dr_generic_internal_tool_2_act.py | 2 +- .../generic_internal_tool/dr_generic_internal_tool_3_reduce.py | 2 -- backend/onyx/db/chat.py | 2 ++ 4 files changed, 3 insertions(+), 5 deletions(-) diff --git a/backend/onyx/agents/agent_search/dr/nodes/dr_a1_orchestrator.py b/backend/onyx/agents/agent_search/dr/nodes/dr_a1_orchestrator.py index fa275940f55..b9d49c558d4 100644 --- a/backend/onyx/agents/agent_search/dr/nodes/dr_a1_orchestrator.py +++ b/backend/onyx/agents/agent_search/dr/nodes/dr_a1_orchestrator.py @@ -437,8 +437,6 @@ def orchestrator( next_step = orchestrator_action.next_step next_tool_name = next_step.tool - available_tools[next_tool_name].path - query_list = [q for q in (next_step.questions or [])] reasoning_result = orchestrator_action.reasoning diff --git a/backend/onyx/agents/agent_search/dr/sub_agents/generic_internal_tool/dr_generic_internal_tool_2_act.py b/backend/onyx/agents/agent_search/dr/sub_agents/generic_internal_tool/dr_generic_internal_tool_2_act.py index 8931a38bf66..b9a86ff4536 100644 --- a/backend/onyx/agents/agent_search/dr/sub_agents/generic_internal_tool/dr_generic_internal_tool_2_act.py +++ b/backend/onyx/agents/agent_search/dr/sub_agents/generic_internal_tool/dr_generic_internal_tool_2_act.py @@ -134,7 +134,7 @@ def generic_internal_tool_act( cited_documents={}, reasoning="", additional_data=None, - response_type="text", + response_type="text", # TODO: convert all response types to enums data=answer_string, ) ], diff --git a/backend/onyx/agents/agent_search/dr/sub_agents/generic_internal_tool/dr_generic_internal_tool_3_reduce.py b/backend/onyx/agents/agent_search/dr/sub_agents/generic_internal_tool/dr_generic_internal_tool_3_reduce.py index 4a51ac48722..901c461d99c 100644 --- a/backend/onyx/agents/agent_search/dr/sub_agents/generic_internal_tool/dr_generic_internal_tool_3_reduce.py +++ b/backend/onyx/agents/agent_search/dr/sub_agents/generic_internal_tool/dr_generic_internal_tool_3_reduce.py @@ -46,7 +46,6 @@ def generic_internal_tool_reducer( write_custom_event( current_step_nr, CustomToolStart( - type="custom_tool_start", tool_name=new_update.tool, ), writer, @@ -55,7 +54,6 @@ def generic_internal_tool_reducer( write_custom_event( current_step_nr, CustomToolDelta( - type="custom_tool_delta", tool_name=new_update.tool, response_type=new_update.response_type, data=new_update.data, diff --git a/backend/onyx/db/chat.py b/backend/onyx/db/chat.py index aeb237640e7..73d14797974 100644 --- a/backend/onyx/db/chat.py +++ b/backend/onyx/db/chat.py @@ -1505,6 +1505,7 @@ def translate_db_message_to_packets( step_nr += 1 elif tool_name == "OktaProfileTool": + # TODO: convert all response types to enums packet_list.extend( create_custom_tool_packets( tool_name=tool_name, @@ -1516,6 +1517,7 @@ def translate_db_message_to_packets( step_nr += 1 else: + # TODO: convert all response types to enums packet_list.extend( create_custom_tool_packets( tool_name=tool_name,