fix: small DR/Thoughtful mode fixes

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,11 @@ def orchestrator(
                 logger.error(f"Error in approach extraction: {e}")
                 raise e
 
-            remaining_time_budget -= available_tools[next_tool].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,
@@ -434,8 +437,6 @@ def orchestrator(
                 next_step = orchestrator_action.next_step
                 next_tool_name = next_step.tool
 
-                next_tool = available_tools[next_tool_name].path
-
                 query_list = [q for q in (next_step.questions or [])]
                 reasoning_result = orchestrator_action.reasoning
 
@@ -444,7 +445,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."
 

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",  # TODO: convert all response types to enums
                 data=answer_string,
             )
         ],

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,

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---
@@ -835,6 +833,42 @@
 )
 
 
+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 (not the base question!), 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 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:
+"""
+)
+
+
 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: