generic_internal tools

Author: joachim-danswer

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

@@ -25,5 +25,6 @@ class DRPath(str, Enum):
     KNOWLEDGE_GRAPH = "Knowledge Graph"
     INTERNET_SEARCH = "Internet Search"
     IMAGE_GENERATION = "Image Generation"
+    GENERIC_INTERNAL_TOOL = "Generic Internal Tool"
     CLOSER = "Closer"
     END = "End"

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

@@ -15,6 +15,9 @@
 from onyx.agents.agent_search.dr.sub_agents.custom_tool.dr_custom_tool_graph_builder import (
     dr_custom_tool_graph_builder,
 )
+from onyx.agents.agent_search.dr.sub_agents.generic_internal_tool.dr_generic_internal_tool_graph_builder import (
+    dr_generic_internal_tool_graph_builder,
+)
 from onyx.agents.agent_search.dr.sub_agents.image_generation.dr_image_generation_graph_builder import (
     dr_image_generation_graph_builder,
 )
@@ -59,6 +62,9 @@ def dr_graph_builder() -> StateGraph:
     custom_tool_graph = dr_custom_tool_graph_builder().compile()
     graph.add_node(DRPath.GENERIC_TOOL, custom_tool_graph)
 
+    generic_internal_tool_graph = dr_generic_internal_tool_graph_builder().compile()
+    graph.add_node(DRPath.GENERIC_INTERNAL_TOOL, generic_internal_tool_graph)
+
     graph.add_node(DRPath.CLOSER, closer)
 
     ### Add edges ###
@@ -74,6 +80,7 @@ def dr_graph_builder() -> StateGraph:
     graph.add_edge(start_key=DRPath.INTERNET_SEARCH, end_key=DRPath.ORCHESTRATOR)
     graph.add_edge(start_key=DRPath.IMAGE_GENERATION, end_key=DRPath.ORCHESTRATOR)
     graph.add_edge(start_key=DRPath.GENERIC_TOOL, end_key=DRPath.ORCHESTRATOR)
+    graph.add_edge(start_key=DRPath.GENERIC_INTERNAL_TOOL, end_key=DRPath.ORCHESTRATOR)
 
     graph.add_conditional_edges(DRPath.CLOSER, completeness_router)
 

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

@@ -90,6 +90,9 @@ class IterationAnswer(BaseModel):
     background_info: str | None = None
     claims: list[str] | None = None
     additional_data: dict[str, str] | None = None
+    response_type: str | None = None
+    data: dict | list | str | int | float | bool | None = None
+    file_ids: list[str] | None = None
 
 
 class AggregatedDRContext(BaseModel):

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

@@ -6,6 +6,7 @@
 from langchain_core.messages import merge_content
 from langchain_core.runnables import RunnableConfig
 from langgraph.types import StreamWriter
+from sqlalchemy.orm import Session
 
 from onyx.agents.agent_search.basic.utils import process_llm_stream
 from onyx.agents.agent_search.dr.constants import AVERAGE_TOOL_COSTS
@@ -38,6 +39,8 @@
 from onyx.configs.constants import DocumentSourceDescription
 from onyx.configs.constants import TMP_DRALPHA_PERSONA_NAME
 from onyx.db.connector import fetch_unique_document_sources
+from onyx.db.models import Tool
+from onyx.db.tools import get_tools
 from onyx.file_store.models import ChatFileType
 from onyx.file_store.models import InMemoryChatFile
 from onyx.kg.utils.extraction_utils import get_entity_types_str
@@ -57,7 +60,6 @@
 from onyx.server.query_and_chat.streaming_models import MessageStart
 from onyx.server.query_and_chat.streaming_models import OverallStop
 from onyx.server.query_and_chat.streaming_models import SectionEnd
-from onyx.tools.tool_implementations.custom.custom_tool import CustomTool
 from onyx.tools.tool_implementations.images.image_generation_tool import (
     ImageGenerationTool,
 )
@@ -82,6 +84,7 @@ def _format_tool_name(tool_name: str) -> str:
 
 
 def _get_available_tools(
+    db_session: Session,
     graph_config: GraphConfig,
     kg_enabled: bool,
     active_source_types: list[DocumentSource],
@@ -97,49 +100,59 @@ def _get_available_tools(
     else:
         include_kg = False
 
+    tool_dict: dict[int, Tool] = {tool.id: tool for tool in get_tools(db_session)}
+
     for tool in graph_config.tooling.tools:
-        tool_info = OrchestratorTool(
-            tool_id=tool.id,
-            name=tool.name,
-            llm_path=_format_tool_name(tool.name),
-            path=DRPath.GENERIC_TOOL,
-            description=tool.description,
-            metadata={},
-            cost=1.0,
-            tool_object=tool,
-        )
 
-        if isinstance(tool, CustomTool):
-            # tool_info.metadata["summary_signature"] = CUSTOM_TOOL_RESPONSE_ID
-            pass
-        elif isinstance(tool, InternetSearchTool):
-            # tool_info.metadata["summary_signature"] = (
-            #     INTERNET_SEARCH_RESPONSE_SUMMARY_ID
-            # )
-            tool_info.llm_path = DRPath.INTERNET_SEARCH.value
-            tool_info.path = DRPath.INTERNET_SEARCH
+        tool_db_info = tool_dict.get(tool.id)
+        if tool_db_info:
+            incode_tool_id = tool_db_info.in_code_tool_id
+        else:
+            raise ValueError(f"Tool {tool.name} is not found in the database")
+
+        if isinstance(tool, InternetSearchTool):
+            llm_path = DRPath.INTERNET_SEARCH.value
+            path = DRPath.INTERNET_SEARCH
         elif isinstance(tool, SearchTool) and len(active_source_types) > 0:
             # tool_info.metadata["summary_signature"] = SEARCH_RESPONSE_SUMMARY_ID
-            tool_info.llm_path = DRPath.INTERNAL_SEARCH.value
-            tool_info.path = DRPath.INTERNAL_SEARCH
+            llm_path = DRPath.INTERNAL_SEARCH.value
+            path = DRPath.INTERNAL_SEARCH
         elif (
             isinstance(tool, KnowledgeGraphTool)
             and include_kg
             and len(active_source_types) > 0
         ):
-            tool_info.llm_path = DRPath.KNOWLEDGE_GRAPH.value
-            tool_info.path = DRPath.KNOWLEDGE_GRAPH
+            llm_path = DRPath.KNOWLEDGE_GRAPH.value
+            path = DRPath.KNOWLEDGE_GRAPH
         elif isinstance(tool, ImageGenerationTool):
-            tool_info.llm_path = DRPath.IMAGE_GENERATION.value
-            tool_info.path = DRPath.IMAGE_GENERATION
+            llm_path = DRPath.IMAGE_GENERATION.value
+            path = DRPath.IMAGE_GENERATION
+        elif incode_tool_id:
+            # if incode tool id is found, it is a generic internal tool
+            llm_path = DRPath.GENERIC_INTERNAL_TOOL.value
+            path = DRPath.GENERIC_INTERNAL_TOOL
         else:
-            logger.warning(
-                f"Tool {tool.name} ({type(tool)}) is not supported/available"
-            )
-            continue
+            # otherwise it is a custom tool
+            llm_path = DRPath.GENERIC_TOOL.value
+            path = DRPath.GENERIC_TOOL
+
+        if path not in {DRPath.GENERIC_INTERNAL_TOOL, DRPath.GENERIC_TOOL}:
+            description = TOOL_DESCRIPTION.get(path, tool.description)
+            cost = AVERAGE_TOOL_COSTS[path]
+        else:
+            description = tool.description
+            cost = 1.0
 
-        tool_info.description = TOOL_DESCRIPTION.get(tool_info.path, tool.description)
-        tool_info.cost = AVERAGE_TOOL_COSTS[tool_info.path]
+        tool_info = OrchestratorTool(
+            tool_id=tool.id,
+            name=tool.name,
+            llm_path=llm_path,
+            path=path,
+            description=description,
+            metadata={},
+            cost=cost,
+            tool_object=tool,
+        )
 
         # TODO: handle custom tools with same name as other tools (e.g., CLOSER)
         available_tools[tool_info.llm_path] = tool_info
@@ -236,8 +249,9 @@ def _get_existing_clarification_request(
     "type": "function",
     "function": {
         "name": "run_any_knowledge_retrieval_and_any_action_tool",
-        "description": "Use this tool to get any external information \
-that is relevant to the question, or for any action to be taken, including image generation.",
+        "description": "Use this tool to get ANY external information \
+that is relevant to the question, or for any action to be taken, including image generation. In fact, \
+ANY tool mentioned can be accessed through this generic tool.",
         "parameters": {
             "type": "object",
             "properties": {
@@ -292,7 +306,7 @@ def clarifier(
     active_source_types = fetch_unique_document_sources(db_session)
 
     available_tools = _get_available_tools(
-        graph_config, kg_enabled, active_source_types
+        db_session, graph_config, kg_enabled, active_source_types
     )
 
     available_tool_descriptions_str = "\n -" + "\n -".join(

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

@@ -32,6 +32,7 @@
 from onyx.agents.agent_search.utils import create_question_prompt
 from onyx.kg.utils.extraction_utils import get_entity_types_str
 from onyx.kg.utils.extraction_utils import get_relationship_types_str
+from onyx.prompts.dr_prompts import DEFAULLT_DECISION_PROMPT
 from onyx.prompts.dr_prompts import SUFFICIENT_INFORMATION_STRING
 from onyx.server.query_and_chat.streaming_models import ReasoningDelta
 from onyx.server.query_and_chat.streaming_models import ReasoningStart
@@ -62,9 +63,15 @@ def orchestrator(
     clarification = state.clarification
     assistant_system_prompt = state.assistant_system_prompt
 
-    decision_system_prompt = _DECISION_SYSTEM_PROMPT_PREFIX + assistant_system_prompt
+    if assistant_system_prompt:
+        decision_system_prompt: str = (
+            DEFAULLT_DECISION_PROMPT
+            + _DECISION_SYSTEM_PROMPT_PREFIX
+            + assistant_system_prompt
+        )
+    else:
+        decision_system_prompt = DEFAULLT_DECISION_PROMPT
 
-    state.assistant_task_prompt
     iteration_nr = state.iteration_nr + 1
     current_step_nr = state.current_step_nr
 

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

@@ -102,6 +102,9 @@ def custom_tool_act(
         raise ValueError("Custom tool did not return a valid response summary")
 
     # summarise tool result
+    if not response_summary.response_type:
+        raise ValueError("Response type is not returned.")
+
     if response_summary.response_type == "json":
         tool_result_str = json.dumps(response_summary.tool_result, ensure_ascii=False)
     elif response_summary.response_type in {"image", "csv"}:
@@ -124,6 +127,13 @@ def custom_tool_act(
         ).content
     ).strip()
 
+    # get file_ids:
+    file_ids = None
+    if response_summary.response_type in {"image", "csv"} and hasattr(
+        response_summary.tool_result, "file_ids"
+    ):
+        file_ids = response_summary.tool_result.file_ids
+
     logger.debug(
         f"Tool call end for {custom_tool_name} {iteration_nr}.{parallelization_nr} at {datetime.now()}"
     )
@@ -141,6 +151,9 @@ def custom_tool_act(
                 cited_documents={},
                 reasoning="",
                 additional_data=None,
+                response_type=response_summary.response_type,
+                data=response_summary.tool_result,
+                file_ids=file_ids,
             )
         ],
         log_messages=[

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

@@ -8,6 +8,10 @@
 from onyx.agents.agent_search.shared_graph_utils.utils import (
     get_langgraph_node_log_string,
 )
+from onyx.agents.agent_search.shared_graph_utils.utils import write_custom_event
+from onyx.server.query_and_chat.streaming_models import CustomToolDelta
+from onyx.server.query_and_chat.streaming_models import CustomToolStart
+from onyx.server.query_and_chat.streaming_models import SectionEnd
 from onyx.utils.logger import setup_logger
 
 
@@ -25,13 +29,47 @@ def custom_tool_reducer(
 
     node_start_time = datetime.now()
 
+    current_step_nr = state.current_step_nr
+
     branch_updates = state.branch_iteration_responses
     current_iteration = state.iteration_nr
 
     new_updates = [
         update for update in branch_updates if update.iteration_nr == current_iteration
     ]
 
+    for new_update in new_updates:
+
+        if not new_update.response_type:
+            raise ValueError("Response type is not returned.")
+
+        write_custom_event(
+            current_step_nr,
+            CustomToolStart(
+                tool_name=new_update.tool,
+            ),
+            writer,
+        )
+
+        write_custom_event(
+            current_step_nr,
+            CustomToolDelta(
+                tool_name=new_update.tool,
+                response_type=new_update.response_type,
+                data=new_update.data,
+                file_ids=new_update.file_ids,
+            ),
+            writer,
+        )
+
+        write_custom_event(
+            current_step_nr,
+            SectionEnd(),
+            writer,
+        )
+
+        current_step_nr += 1
+
     return SubAgentUpdate(
         iteration_responses=new_updates,
         log_messages=[

backend/onyx/agents/agent_search/dr/sub_agents/generic_internal_tool/dr_generic_internal_tool_1_branch.py

@@ -0,0 +1,36 @@
+from datetime import datetime
+
+from langchain_core.runnables import RunnableConfig
+from langgraph.types import StreamWriter
+
+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,
+)
+from onyx.utils.logger import setup_logger
+
+logger = setup_logger()
+
+
+def generic_internal_tool_branch(
+    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()
+    iteration_nr = state.iteration_nr
+
+    logger.debug(f"Search start for Generic Tool {iteration_nr} at {datetime.now()}")
+
+    return LoggerUpdate(
+        log_messages=[
+            get_langgraph_node_log_string(
+                graph_component="generic_internal_tool",
+                node_name="branching",
+                node_start_time=node_start_time,
+            )
+        ],
+    )