From f68b74ff4af97da07c96c052a2c38be19c2ab1be Mon Sep 17 00:00:00 2001
From: joachim-danswer <joachim@danswer.ai>
Date: Wed, 30 Apr 2025 12:51:08 -0700
Subject: [PATCH] disable Agent Search  refinement by default (#4638)

- created env variable  AGENT_ALLOW_REFINEMENT  with default "". Must be set to true to enable Refinement.
 - added an environment variable for the upper limit of docs that can be sent to verification
---
 .../deep_search/main/nodes/decide_refinement_need.py  |  2 +-
 .../expanded_retrieval/nodes/kickoff_verification.py  |  3 ++-
 backend/onyx/agents/agent_search/run_graph.py         |  4 ++--
 backend/onyx/configs/agent_configs.py                 | 11 ++++++++---
 4 files changed, 13 insertions(+), 7 deletions(-)

diff --git a/backend/onyx/agents/agent_search/deep_search/main/nodes/decide_refinement_need.py b/backend/onyx/agents/agent_search/deep_search/main/nodes/decide_refinement_need.py
index 731db32fdb3..92fd7ae4e0a 100644
--- a/backend/onyx/agents/agent_search/deep_search/main/nodes/decide_refinement_need.py
+++ b/backend/onyx/agents/agent_search/deep_search/main/nodes/decide_refinement_need.py
@@ -26,7 +26,7 @@ def decide_refinement_need(
 
     graph_config = cast(GraphConfig, config["metadata"]["config"])
 
-    decision = True  # TODO: just for current testing purposes
+    decision = graph_config.behavior.allow_refinement
 
     if state.answer_error:
         return RequireRefinemenEvalUpdate(
diff --git a/backend/onyx/agents/agent_search/deep_search/shared/expanded_retrieval/nodes/kickoff_verification.py b/backend/onyx/agents/agent_search/deep_search/shared/expanded_retrieval/nodes/kickoff_verification.py
index 68c5597fb6d..cb092e0bbe5 100644
--- a/backend/onyx/agents/agent_search/deep_search/shared/expanded_retrieval/nodes/kickoff_verification.py
+++ b/backend/onyx/agents/agent_search/deep_search/shared/expanded_retrieval/nodes/kickoff_verification.py
@@ -10,6 +10,7 @@ from onyx.agents.agent_search.deep_search.shared.expanded_retrieval.states impor
 from onyx.agents.agent_search.deep_search.shared.expanded_retrieval.states import (
     ExpandedRetrievalState,
 )
+from onyx.configs.agent_configs import AGENT_MAX_VERIFICATION_HITS
 
 
 def kickoff_verification(
@@ -22,7 +23,7 @@ def kickoff_verification(
     are done here, so this could be replaced with an edge. But we may choose to make state
     updates later.)
     """
-    retrieved_documents = state.retrieved_documents
+    retrieved_documents = state.retrieved_documents[:AGENT_MAX_VERIFICATION_HITS]
     verification_question = state.question
 
     sub_question_id = state.sub_question_id
diff --git a/backend/onyx/agents/agent_search/run_graph.py b/backend/onyx/agents/agent_search/run_graph.py
index add734e1e18..2f993e93c0d 100644
--- a/backend/onyx/agents/agent_search/run_graph.py
+++ b/backend/onyx/agents/agent_search/run_graph.py
@@ -30,7 +30,7 @@ from onyx.chat.models import StreamStopInfo
 from onyx.chat.models import SubQueryPiece
 from onyx.chat.models import SubQuestionPiece
 from onyx.chat.models import ToolResponse
-from onyx.configs.agent_configs import ALLOW_REFINEMENT
+from onyx.configs.agent_configs import AGENT_ALLOW_REFINEMENT
 from onyx.configs.agent_configs import INITIAL_SEARCH_DECOMPOSITION_ENABLED
 from onyx.context.search.models import SearchRequest
 from onyx.db.engine import get_session_context_manager
@@ -104,7 +104,7 @@ def run_graph(
     config.behavior.perform_initial_search_decomposition = (
         INITIAL_SEARCH_DECOMPOSITION_ENABLED
     )
-    config.behavior.allow_refinement = ALLOW_REFINEMENT
+    config.behavior.allow_refinement = AGENT_ALLOW_REFINEMENT
 
     for event in manage_sync_streaming(
         compiled_graph=compiled_graph, config=config, graph_input=input
diff --git a/backend/onyx/configs/agent_configs.py b/backend/onyx/configs/agent_configs.py
index 791e124418c..bd6c4051e79 100644
--- a/backend/onyx/configs/agent_configs.py
+++ b/backend/onyx/configs/agent_configs.py
@@ -1,7 +1,5 @@
 import os
 
-INITIAL_SEARCH_DECOMPOSITION_ENABLED = True
-ALLOW_REFINEMENT = True
 
 AGENT_DEFAULT_RETRIEVAL_HITS = 15
 AGENT_DEFAULT_RERANKING_HITS = 10
@@ -19,10 +17,10 @@ AGENT_DEFAULT_MAX_ANSWER_CONTEXT_DOCS = 10
 AGENT_DEFAULT_MAX_STATIC_HISTORY_WORD_LENGTH = 2000
 
 INITIAL_SEARCH_DECOMPOSITION_ENABLED = True
-ALLOW_REFINEMENT = True
 
 AGENT_DEFAULT_RETRIEVAL_HITS = 15
 AGENT_DEFAULT_RERANKING_HITS = 10
+AGENT_DEFAULT_MAX_VERIFIVATION_HITS = 30
 AGENT_DEFAULT_SUB_QUESTION_MAX_CONTEXT_HITS = 8
 AGENT_DEFAULT_NUM_DOCS_FOR_INITIAL_DECOMPOSITION = 3
 AGENT_DEFAULT_NUM_DOCS_FOR_REFINED_DECOMPOSITION = 5
@@ -31,6 +29,9 @@ AGENT_DEFAULT_MIN_ORIG_QUESTION_DOCS = 3
 AGENT_DEFAULT_MAX_ANSWER_CONTEXT_DOCS = 10
 AGENT_DEFAULT_MAX_STATIC_HISTORY_WORD_LENGTH = 2000
 
+
+AGENT_ALLOW_REFINEMENT = os.environ.get("AGENT_ALLOW_REFINEMENT", "").lower() == "true"
+
 AGENT_ANSWER_GENERATION_BY_FAST_LLM = (
     os.environ.get("AGENT_ANSWER_GENERATION_BY_FAST_LLM", "").lower() == "true"
 )
@@ -40,6 +41,10 @@ AGENT_RETRIEVAL_STATS = (
 ) or True  # default True
 
 
+AGENT_MAX_VERIFICATION_HITS = int(
+    os.environ.get("AGENT_MAX_VERIFICATION_HITS") or AGENT_DEFAULT_MAX_VERIFIVATION_HITS
+)  # 30
+
 AGENT_MAX_QUERY_RETRIEVAL_RESULTS = int(
     os.environ.get("AGENT_MAX_QUERY_RETRIEVAL_RESULTS") or AGENT_DEFAULT_RETRIEVAL_HITS
 )  # 15