(WIP) (DO NOT MERGE) feat(Mockable): enable toolruntime access to all tools

src/uipath_langchain/agent/tools/escalation_tool.py

@@ -43,31 +43,34 @@ def create_escalation_tool(resource: AgentEscalationResourceConfig) -> Structure
         else None
     )
 
-    @mockable(
-        name=resource.name,
-        description=resource.description,
-        input_schema=input_model.model_json_schema(),
-        output_schema=output_model.model_json_schema(),
-    )
     async def escalation_tool_fn(
         runtime: ToolRuntime, **kwargs: Any
     ) -> Command[Any] | Any:
-        task_title = channel.task_title or "Escalation Task"
-
-        result = interrupt(
-            CreateEscalation(
-                title=task_title,
-                data=kwargs,
-                assignee=assignee,
-                app_name=channel.properties.app_name,
-                app_folder_path=channel.properties.folder_name,
-                app_version=channel.properties.app_version,
-                priority=channel.priority,
-                labels=channel.labels,
-                is_actionable_message_enabled=channel.properties.is_actionable_message_enabled,
-                actionable_message_metadata=channel.properties.actionable_message_meta_data,
-            )
+        @mockable(
+            name=resource.name,
+            description=resource.description,
+            input_schema=input_model.model_json_schema(),
+            output_schema=output_model.model_json_schema(),
+            example_calls=[],  # TODO: pass these in from runtime.
         )
+        async def escalation_tool_fn_impl(**inner_kwargs: Any) -> Any:
+            task_title = channel.task_title or "Escalation Task"
+            return interrupt(
+                CreateEscalation(
+                    title=task_title,
+                    data=inner_kwargs,
+                    assignee=assignee,
+                    app_name=channel.properties.app_name,
+                    app_folder_path=channel.properties.folder_name,
+                    app_version=channel.properties.app_version,
+                    priority=channel.priority,
+                    labels=channel.labels,
+                    is_actionable_message_enabled=channel.properties.is_actionable_message_enabled,
+                    actionable_message_metadata=channel.properties.actionable_message_meta_data,
+                )
+            )
+
+        result = await escalation_tool_fn_impl(**kwargs)
 
         escalation_action = getattr(result, "action", None)
         escalation_output = getattr(result, "data", {})

src/uipath_langchain/agent/tools/integration_tool.py

@@ -142,33 +142,34 @@ def create_integration_tool(
         else create_model({"type": "object", "properties": {}})
     )
 
-    sdk = UiPath()
-
-    @mockable(
-        name=resource.name,
-        description=resource.description,
-        input_schema=input_model.model_json_schema(),
-        output_schema=output_model.model_json_schema(),
-    )
     async def integration_tool_fn(runtime: ToolRuntime, **kwargs: Any):
-        try:
-            # we manually validating here and not passing input_model to StructuredTool
-            # because langchain itself will block their own injected arguments (like runtime) if the model is strict
-            val_args = input_model.model_validate(kwargs)
-            args = handle_static_args(
-                resource=resource,
-                runtime=runtime,
-                input_args=val_args.model_dump(),
-            )
-            result = await sdk.connections.invoke_activity_async(
-                activity_metadata=activity_metadata,
-                connection_id=connection_id,
-                activity_input=args,
-            )
-        except Exception:
-            raise
-
-        return result
+        @mockable(
+            name=resource.name,
+            description=resource.description,
+            input_schema=input_model.model_json_schema(),
+            output_schema=output_model.model_json_schema(),
+            example_calls=[],  # TODO: pass these in from runtime.
+        )
+        async def integration_tool_fn_impl(**inner_kwargs: Any):
+            try:
+                # we manually validating here and not passing input_model to StructuredTool
+                # because langchain itself will block their own injected arguments (like runtime) if the model is strict
+                sdk = UiPath()
+                val_args = input_model.model_validate(inner_kwargs)
+                args = handle_static_args(
+                    resource=resource,
+                    runtime=runtime,
+                    input_args=val_args.model_dump(),
+                )
+                return await sdk.connections.invoke_activity_async(
+                    activity_metadata=activity_metadata,
+                    connection_id=connection_id,
+                    activity_input=args,
+                )
+            except Exception:
+                raise
+
+        return await integration_tool_fn_impl(**kwargs)
 
     tool = StructuredToolWithOutputType(
         name=tool_name,

src/uipath_langchain/agent/tools/process_tool.py

@@ -4,6 +4,7 @@
 
 from jsonschema_pydantic_converter import transform as create_model
 from langchain_core.tools import StructuredTool
+from langgraph.prebuilt import ToolRuntime
 from langgraph.types import interrupt
 from uipath.agent.models.agent import AgentProcessToolResourceConfig
 from uipath.eval.mocks import mockable
@@ -22,21 +23,25 @@ def create_process_tool(resource: AgentProcessToolResourceConfig) -> StructuredT
     input_model: Any = create_model(resource.input_schema)
     output_model: Any = create_model(resource.output_schema)
 
-    @mockable(
-        name=resource.name,
-        description=resource.description,
-        input_schema=input_model.model_json_schema(),
-        output_schema=output_model.model_json_schema(),
-    )
-    async def process_tool_fn(**kwargs: Any):
-        return interrupt(
-            InvokeProcess(
-                name=process_name,
-                input_arguments=kwargs,
-                process_folder_path=folder_path,
-                process_folder_key=None,
-            )
+    async def process_tool_fn(runtime: ToolRuntime, **kwargs: Any):
+        @mockable(
+            name=resource.name,
+            description=resource.description,
+            input_schema=input_model.model_json_schema(),
+            output_schema=output_model.model_json_schema(),
+            example_calls=[],  # TODO: pass these in from runtime.
         )
+        async def process_tool_impl(**inner_kwargs: Any):
+            return interrupt(
+                InvokeProcess(
+                    name=process_name,
+                    input_arguments=inner_kwargs,
+                    process_folder_path=folder_path,
+                    process_folder_key=None,
+                )
+            )
+
+        return await process_tool_impl(**kwargs)
 
     tool = StructuredToolWithOutputType(
         name=tool_name,