From ada4cd778673e5164a82146c6eabf55fb3d56d30 Mon Sep 17 00:00:00 2001 From: Brijmohan Siyag Date: Mon, 22 Sep 2025 02:43:45 +0530 Subject: [PATCH] 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. --- .../custom_tool/dr_custom_tool_2_act.py | 28 +++++++++++++------ backend/onyx/agents/agent_search/models.py | 2 ++ backend/onyx/chat/answer.py | 2 ++ backend/onyx/configs/tool_configs.py | 5 ++++ 4 files changed, 28 insertions(+), 9 deletions(-) diff --git a/backend/onyx/agents/agent_search/dr/sub_agents/custom_tool/dr_custom_tool_2_act.py b/backend/onyx/agents/agent_search/dr/sub_agents/custom_tool/dr_custom_tool_2_act.py index ebccbb0b7b..c08f2f932a 100644 --- a/backend/onyx/agents/agent_search/dr/sub_agents/custom_tool/dr_custom_tool_2_act.py +++ b/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( diff --git a/backend/onyx/agents/agent_search/models.py b/backend/onyx/agents/agent_search/models.py index 510d75bbaf..57e4f151fa 100644 --- a/backend/onyx/agents/agent_search/models.py +++ b/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 diff --git a/backend/onyx/chat/answer.py b/backend/onyx/chat/answer.py index ed10e3d246..0bd1bc78d2 100644 --- a/backend/onyx/chat/answer.py +++ b/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, diff --git a/backend/onyx/configs/tool_configs.py b/backend/onyx/configs/tool_configs.py index 78da1fc201..ad81d6cf7d 100644 --- a/backend/onyx/configs/tool_configs.py +++ b/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" +)