feat: add configurable LLM summarization for DR custom tools

Add DR_CUSTOM_TOOL_USE_LLM_SUMMARY config to control whether custom tool
responses are summarized by LLM in DR agent. Defaults to true for backward
compatibility. Set to false to return raw tool results directly.

Author: brijsiyag-meesho

backend/onyx/agents/agent_search/dr/sub_agents/custom_tool/dr_custom_tool_2_act.py

@@ -121,15 +121,25 @@ def custom_tool_act(
         f"Result: {tool_result_str}"
     )
 
-    tool_summary_prompt = CUSTOM_TOOL_USE_PROMPT.build(
-        query=branch_query, base_question=base_question, tool_response=tool_str
-    )
-    answer_string = str(
-        graph_config.tooling.primary_llm.invoke(
-            tool_summary_prompt, timeout_override=TF_DR_TIMEOUT_SHORT
-        ).content
-    ).strip()
-
+    # Use LLM summary if configured, otherwise use raw tool result
+    if graph_config.tooling.dr_custom_tool_use_llm_summary:
+        tool_summary_prompt = CUSTOM_TOOL_USE_PROMPT.build(
+            query=branch_query, base_question=base_question, tool_response=tool_str
+        )
+        answer_string = str(
+            graph_config.tooling.primary_llm.invoke(
+                tool_summary_prompt, timeout_override=TF_DR_TIMEOUT_SHORT
+            ).content
+        ).strip()
+    else:
+        answer_string = tool_result_str
+        # Format JSON response for better readability when not using LLM summary
+        if response_summary.response_type == "json":
+            try:
+                parsed_json = json.loads(tool_result_str)
+                answer_string = json.dumps(parsed_json, indent=2, ensure_ascii=False)
+            except json.JSONDecodeError:
+                answer_string = tool_result_str
     # get file_ids:
     file_ids = None
     if response_summary.response_type in {"image", "csv"} and hasattr(

backend/onyx/agents/agent_search/models.py

@@ -40,6 +40,8 @@ class GraphTooling(BaseModel):
     # force tool args IF the tool is used
     force_use_tool: ForceUseTool
     using_tool_calling_llm: bool = False
+    # Whether to use LLM to summarize custom tool responses in DR agent
+    dr_custom_tool_use_llm_summary: bool = True
 
     class Config:
         arbitrary_types_allowed = True

backend/onyx/chat/answer.py

@@ -20,6 +20,7 @@
 from onyx.configs.agent_configs import AGENT_ALLOW_REFINEMENT
 from onyx.configs.agent_configs import INITIAL_SEARCH_DECOMPOSITION_ENABLED
 from onyx.configs.agent_configs import TF_DR_DEFAULT_FAST
+from onyx.configs.tool_configs import DR_CUSTOM_TOOL_USE_LLM_SUMMARY
 from onyx.context.search.models import RerankingDetails
 from onyx.db.kg_config import get_kg_config_settings
 from onyx.db.models import Persona
@@ -105,6 +106,7 @@ def __init__(
             tools=tools or [],
             force_use_tool=force_use_tool,
             using_tool_calling_llm=using_tool_calling_llm,
+            dr_custom_tool_use_llm_summary=DR_CUSTOM_TOOL_USE_LLM_SUMMARY,
         )
         self.graph_persistence = GraphPersistence(
             db_session=db_session,

backend/onyx/configs/tool_configs.py

@@ -24,3 +24,8 @@
         logger.error(
             "Failed to parse CUSTOM_TOOL_PASS_THROUGH_HEADERS, must be a valid JSON object"
         )
+
+# Whether to use LLM to summarize custom tool responses in DR agent
+DR_CUSTOM_TOOL_USE_LLM_SUMMARY = (
+    os.environ.get("DR_CUSTOM_TOOL_USE_LLM_SUMMARY", "true").lower() == "true"
+)