almost working custom tools

Author: Orbital-Web

backend/onyx/agents/agent_search/dr/dr_prompt_builder.py

@@ -19,22 +19,23 @@
 def get_dr_prompt_orchestration_templates(
     purpose: DRPromptPurpose,
     research_type: ResearchType,
-    available_tools: list[OrchestratorTool],
+    available_tools: dict[str, OrchestratorTool],
     entity_types_string: str | None = None,
     relationship_types_string: str | None = None,
     reasoning_result: str | None = None,
     tool_calls_string: str | None = None,
 ) -> PromptTemplate:
-    available_tools = available_tools or []
-    tool_names = [tool.llm_path for tool in available_tools]
+    available_tools = available_tools or {}
+    tool_names = list(available_tools.keys())
     tool_description_str = "\n\n".join(
-        f"- {tool.llm_path}: {tool.description}" for tool in available_tools
+        f"- {tool_name}: {tool.description}"
+        for tool_name, tool in available_tools.items()
     )
     tool_cost_str = "\n".join(
-        f"{tool.llm_path}: {tool.cost}" for tool in available_tools
+        f"{tool_name}: {tool.cost}" for tool_name, tool in available_tools.items()
     )
 
-    available_tool_paths = [tool.path for tool in available_tools]
+    available_tool_paths = [tool.path for tool in available_tools.values()]
 
     tool_differentiations: list[str] = []
     for tool_1 in available_tool_paths:

backend/onyx/agents/agent_search/dr/models.py

@@ -4,6 +4,7 @@
 
 from onyx.agents.agent_search.dr.enums import DRPath
 from onyx.context.search.models import InferenceSection
+from onyx.tools.tool import Tool
 
 
 class OrchestratorStep(BaseModel):
@@ -59,6 +60,10 @@ class OrchestratorTool(BaseModel):
     description: str
     metadata: dict[str, str]
     cost: float
+    tool_object: Tool | None = None  # None for CLOSER
+
+    class Config:
+        arbitrary_types_allowed = True
 
 
 class IterationInstructions(BaseModel):
@@ -71,11 +76,10 @@ class IterationInstructions(BaseModel):
 class GenericToolAnswer(BaseModel):
     reasoning: str
     answer: str
-    background_info: str
 
 
 class IterationAnswer(BaseModel):
-    tool: DRPath
+    tool: str
     tool_id: int
     iteration_nr: int
     parallelization_nr: int

backend/onyx/agents/agent_search/dr/nodes/dr_a0_clarification.py

@@ -67,9 +67,9 @@ def _format_tool_name(tool_name: str) -> str:
 
 def _get_available_tools(
     graph_config: GraphConfig, kg_enabled: bool
-) -> list[OrchestratorTool]:
+) -> dict[str, OrchestratorTool]:
 
-    available_tools: list[OrchestratorTool] = []
+    available_tools: dict[str, OrchestratorTool] = {}
     for tool in graph_config.tooling.tools:
         tool_info = OrchestratorTool(
             tool_id=tool.id,
@@ -79,6 +79,7 @@ def _get_available_tools(
             description=tool.description,
             metadata={},
             cost=1.0,
+            tool_object=tool,
         )
 
         if isinstance(tool, CustomTool):
@@ -107,29 +108,28 @@ def _get_available_tools(
         tool_info.description = TOOL_DESCRIPTION.get(tool_info.path, tool.description)
         tool_info.cost = AVERAGE_TOOL_COSTS[tool_info.path]
 
-        available_tools.append(tool_info)
+        # TODO: handle custom tools with same name as other tools (e.g., CLOSER)
+        available_tools[tool_info.llm_path] = tool_info
 
     # make sure KG isn't enabled without internal search
-    available_paths = [tool.path for tool in available_tools]
     if (
-        DRPath.KNOWLEDGE_GRAPH in available_paths
-        and DRPath.INTERNAL_SEARCH not in available_paths
+        DRPath.KNOWLEDGE_GRAPH.value in available_tools
+        and DRPath.INTERNAL_SEARCH.value not in available_tools
     ):
         raise ValueError(
             "The Knowledge Graph is not supported without internal search tool"
         )
 
     # add CLOSER tool, which is always available
-    available_tools.append(
-        OrchestratorTool(
-            tool_id=-1,
-            name="closer",
-            llm_path=DRPath.CLOSER.value,
-            path=DRPath.CLOSER,
-            description=TOOL_DESCRIPTION[DRPath.CLOSER],
-            metadata={},
-            cost=0.0,
-        )
+    available_tools[DRPath.CLOSER.value] = OrchestratorTool(
+        tool_id=-1,
+        name="closer",
+        llm_path=DRPath.CLOSER.value,
+        path=DRPath.CLOSER,
+        description=TOOL_DESCRIPTION[DRPath.CLOSER],
+        metadata={},
+        cost=0.0,
+        tool_object=None,
     )
 
     return available_tools

backend/onyx/agents/agent_search/dr/nodes/dr_a1_orchestrator.py

@@ -74,8 +74,7 @@ def orchestrator(
             ).context
             or "(No answer history yet available)"
         )
-    available_tools = state.available_tools or []
-    available_tool_map = {tool.llm_path: tool for tool in available_tools}
+    available_tools = state.available_tools or {}
 
     questions = [
         f"{iteration_response.tool}: {iteration_response.question}"
@@ -217,7 +216,7 @@ def orchestrator(
                 logger.error(f"Error in approach extraction: {e}")
                 raise e
 
-            remaining_time_budget -= available_tool_map[next_tool].cost
+            remaining_time_budget -= available_tools[next_tool].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,
@@ -304,7 +303,7 @@ def orchestrator(
                 logger.error(f"Error in approach extraction: {e}")
                 raise e
 
-            remaining_time_budget -= available_tool_map[next_tool].cost
+            remaining_time_budget -= available_tools[next_tool].cost
         else:
             reasoning_result = "Time to wrap up."
 

backend/onyx/agents/agent_search/dr/states.py

@@ -26,13 +26,11 @@ class OrchestrationUpdate(LoggerUpdate):
     tools_used: Annotated[list[str], add] = []
     query_list: list[str] = []
     iteration_nr: int = 0
-    plan_of_record: OrchestrationPlan | None = None  # None for FAST TimeBudget
+    plan_of_record: OrchestrationPlan | None = None  # None for Thoughtful
     remaining_time_budget: float = 2.0  # set by default to about 2 searches
     clarification: OrchestrationClarificationInfo | None = None
-    available_tools: list[OrchestratorTool] | None = None
-    num_closer_suggestions: int = (
-        0  # how many times the closer was suggested. (Closer can send back now.)
-    )
+    available_tools: dict[str, OrchestratorTool] | None = None
+    num_closer_suggestions: int = 0  # how many times the closer was suggested
     gaps: list[str] = (
         []
     )  # gaps that may be identified by the closer before being able to answer the question.

backend/onyx/agents/agent_search/dr/sub_agents/basic_search/dr_basic_search_2_act.py

@@ -5,7 +5,6 @@
 from langchain_core.runnables import RunnableConfig
 from langgraph.types import StreamWriter
 
-from onyx.agents.agent_search.dr.enums import DRPath
 from onyx.agents.agent_search.dr.enums import ResearchType
 from onyx.agents.agent_search.dr.models import BaseSearchProcessingResponse
 from onyx.agents.agent_search.dr.models import IterationAnswer
@@ -30,6 +29,7 @@
     SEARCH_RESPONSE_SUMMARY_ID,
 )
 from onyx.tools.tool_implementations.search.search_tool import SearchResponseSummary
+from onyx.tools.tool_implementations.search.search_tool import SearchTool
 from onyx.utils.logger import setup_logger
 
 logger = setup_logger()
@@ -56,13 +56,17 @@ def basic_search(
     base_question = graph_config.inputs.prompt_builder.raw_user_query
     research_type = graph_config.behavior.research_type
 
-    search_tool = graph_config.tooling.search_tool
+    if not state.available_tools:
+        raise ValueError("available_tools is not set")
 
-    if search_tool is None:
-        raise ValueError("search_tool must be provided for agentic search")
+    search_tool_info = state.available_tools[state.tools_used[-1]]
+    search_tool = cast(SearchTool, search_tool_info.tool_object)
 
-    # rewrite query and identify source types
+    # sanity check
+    if search_tool != graph_config.tooling.search_tool:
+        raise ValueError("search_tool does not match the configured search tool")
 
+    # rewrite query and identify source types
     active_source_types_str = ", ".join(
         [source.value for source in state.active_source_types or []]
     )
@@ -228,8 +232,8 @@ def basic_search(
     return AnswerUpdate(
         iteration_responses=[
             IterationAnswer(
-                tool=DRPath.INTERNAL_SEARCH,
-                tool_id=search_tool.id,
+                tool=search_tool_info.llm_path,
+                tool_id=search_tool_info.tool_id,
                 iteration_nr=iteration_nr,
                 parallelization_nr=parallelization_nr,
                 question=branch_query,
@@ -242,8 +246,8 @@ def basic_search(
         ],
         log_messages=[
             get_langgraph_node_log_string(
-                graph_component="main",
-                node_name="search",
+                graph_component="basic_search",
+                node_name="searching",
                 node_start_time=node_start_time,
             )
         ],

backend/onyx/agents/agent_search/dr/sub_agents/basic_search/dr_basic_search_3_reduce.py

@@ -36,7 +36,7 @@ def is_reducer(
         iteration_responses=new_updates,
         log_messages=[
             get_langgraph_node_log_string(
-                graph_component="internet_search",
+                graph_component="basic_search",
                 node_name="consolidation",
                 node_start_time=node_start_time,
             )

backend/onyx/agents/agent_search/dr/sub_agents/basic_search/dr_basic_search_conditional_edges.py

@@ -17,6 +17,8 @@ def branching_router(state: SubAgentInput) -> list[Send | Hashable]:
                 branch_question=query,
                 context="",
                 active_source_types=state.active_source_types,
+                tools_used=state.tools_used,
+                available_tools=state.available_tools,
             ),
         )
         for parallelization_nr, query in enumerate(

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

@@ -3,12 +3,8 @@
 from langchain_core.runnables import RunnableConfig
 from langgraph.types import StreamWriter
 
-from onyx.agents.agent_search.dr.sub_agents.custom_tool.dr_custom_tool_states import (
-    CustomToolSubAgentInput,
-)
-from onyx.agents.agent_search.dr.sub_agents.custom_tool.dr_custom_tool_states import (
-    CustomToolSubAgentPrepareUpdate,
-)
+from onyx.agents.agent_search.dr.states import LoggerUpdate
+from onyx.agents.agent_search.dr.sub_agents.states import SubAgentInput
 from onyx.agents.agent_search.shared_graph_utils.utils import (
     get_langgraph_node_log_string,
 )
@@ -18,32 +14,21 @@
 
 
 def custom_tool_branch(
-    state: CustomToolSubAgentInput,
-    config: RunnableConfig,
-    writer: StreamWriter = lambda _: None,
-) -> CustomToolSubAgentPrepareUpdate:
+    state: SubAgentInput, config: RunnableConfig, writer: StreamWriter = lambda _: None
+) -> LoggerUpdate:
     """
     LangGraph node to perform a generic tool call as part of the DR process.
     """
 
     node_start_time = datetime.now()
-    tool_name = state.query_path[-1]
-
-    if not state.available_tools:
-        raise ValueError("available_tools is not set")
+    iteration_nr = state.iteration_nr
 
-    tool_dict: dict[str, str] = {}
-    for available_tool_dict in state.available_tools:
-        if available_tool_dict["name"] == tool_name:
-            tool_dict = available_tool_dict
-            break
+    logger.debug(f"Search start for Generic Tool {iteration_nr} at {datetime.now()}")
 
-    return CustomToolSubAgentPrepareUpdate(
-        tool_name=tool_name,
-        tool_dict=tool_dict,
+    return LoggerUpdate(
         log_messages=[
             get_langgraph_node_log_string(
-                graph_component="custom_tool_sub_agent",
+                graph_component="custom_tool",
                 node_name="branching",
                 node_start_time=node_start_time,
             )