feat: add extra metadata for tools (#608)

Author: ionut-mihalache-uipath

pyproject.toml

@@ -1,13 +1,13 @@
 [project]
 name = "uipath-langchain"
-version = "0.7.6"
+version = "0.7.7"
 description = "Python SDK that enables developers to build and deploy LangGraph agents to the UiPath Cloud Platform"
 readme = { file = "README.md", content-type = "text/markdown" }
 requires-python = ">=3.11"
 dependencies = [
-    "uipath>=2.9.10, <2.10.0",
+    "uipath>=2.9.12, <2.10.0",
     "uipath-core>=0.5.2, <0.6.0",
-    "uipath-platform>=0.0.1, < 0.1.0",
+    "uipath-platform>=0.0.4, < 0.1.0",
     "uipath-runtime>=0.9.1, <0.10.0",
     "langgraph>=1.0.0, <2.0.0",
     "langchain-core>=1.2.11, <2.0.0",

src/uipath_langchain/agent/tools/context_tool.py

@@ -123,6 +123,8 @@ async def context_tool_fn(query: Optional[str] = None) -> dict[str, Any]:
         metadata={
             "tool_type": "context",
             "display_name": resource.name,
+            "index_name": resource.index_name,
+            "context_retrieval_mode": resource.settings.retrieval_mode,
         },
     )
 
@@ -202,6 +204,8 @@ async def create_deep_rag():
         metadata={
             "tool_type": "context",
             "display_name": resource.name,
+            "index_name": resource.index_name,
+            "context_retrieval_mode": resource.settings.retrieval_mode,
         },
     )
 
@@ -335,6 +339,8 @@ async def context_batch_transform_wrapper(
         metadata={
             "tool_type": "context",
             "display_name": resource.name,
+            "index_name": resource.index_name,
+            "context_retrieval_mode": resource.settings.retrieval_mode,
         },
     )
     tool.set_tool_wrappers(awrapper=context_batch_transform_wrapper)

src/uipath_langchain/agent/tools/escalation_tool.py

@@ -154,6 +154,8 @@ class EscalationToolOutput(BaseModel):
         action: Literal["approve", "reject"]
         data: output_model
 
+    _bts_context: dict[str, Any] = {}
+
     async def escalation_tool_fn(**kwargs: Any) -> dict[str, Any]:
         recipient: TaskRecipient | None = (
             await resolve_recipient_value(channel.recipients[0])
@@ -191,6 +193,10 @@ async def create_escalation_task():
                     is_actionable_message_enabled=channel.properties.is_actionable_message_enabled,
                     actionable_message_metadata=channel.properties.actionable_message_meta_data,
                 )
+
+                if created_task.id is not None:
+                    _bts_context["task_key"] = str(created_task.id)
+
                 return WaitEscalation(
                     action=created_task,
                     app_folder_path=channel.properties.folder_name,
@@ -283,6 +289,7 @@ async def escalation_wrapper(
             "recipient": None,
             "args_schema": input_model,
             "output_schema": output_model,
+            "_bts_context": _bts_context,
         },
     )
     tool.set_tool_wrappers(awrapper=escalation_wrapper)

src/uipath_langchain/agent/tools/integration_tool.py

@@ -201,6 +201,8 @@ async def integration_tool_wrapper(
         metadata={
             "tool_type": "integration",
             "display_name": resource.name,
+            "connector_key": resource.properties.connection.id,
+            "connector_name": resource.properties.connection.name,
         },
     )
     tool.set_tool_wrappers(awrapper=integration_tool_wrapper)

src/uipath_langchain/agent/tools/process_tool.py

@@ -5,7 +5,7 @@
 from langchain.tools import BaseTool
 from langchain_core.messages import ToolCall
 from langchain_core.tools import StructuredTool
-from uipath.agent.models.agent import AgentProcessToolResourceConfig
+from uipath.agent.models.agent import AgentProcessToolResourceConfig, AgentToolType
 from uipath.eval.mocks import mockable
 from uipath.platform import UiPath
 from uipath.platform.common import WaitJob
@@ -38,6 +38,7 @@ def create_process_tool(resource: AgentProcessToolResourceConfig) -> StructuredT
     output_model: Any = create_model(resource.output_schema)
 
     _span_context: dict[str, Any] = {}
+    _bts_context: dict[str, Any] = {}
 
     async def process_tool_fn(**kwargs: Any):
         attachments = get_job_attachments(input_model, kwargs)
@@ -52,6 +53,7 @@ async def process_tool_fn(**kwargs: Any):
         )
         async def invoke_process(**_tool_kwargs: Any):
             parent_span_id = _span_context.pop("parent_span_id", None)
+            parent_operation_id = _bts_context.pop("parent_operation_id", None)
 
             @durable_interrupt
             async def start_job():
@@ -62,7 +64,17 @@ async def start_job():
                     folder_path=folder_path,
                     attachments=attachments,
                     parent_span_id=parent_span_id,
+                    parent_operation_id=parent_operation_id,
                 )
+
+                if job.key:
+                    bts_key = (
+                        "wait_for_agent_job_key"
+                        if resource.type == AgentToolType.AGENT
+                        else "wait_for_job_key"
+                    )
+                    _bts_context[bts_key] = str(job.key)
+
                 return WaitJob(job=job, process_folder_key=job.folder_key)
 
             return await start_job()
@@ -92,6 +104,7 @@ async def process_tool_wrapper(
             "args_schema": input_model,
             "output_schema": output_model,
             "_span_context": _span_context,
+            "_bts_context": _bts_context,
         },
         argument_properties=resource.argument_properties,
     )

tests/agent/tools/test_process_tool.py

@@ -121,6 +121,7 @@ async def test_invoke_calls_processes_invoke_async(
     ):
         """Test that invoking the tool calls client.processes.invoke_async."""
         mock_job = MagicMock(spec=Job)
+        mock_job.key = "job-key-123"
         mock_job.folder_key = "folder-key-123"
 
         mock_client = MagicMock()
@@ -138,6 +139,7 @@ async def test_invoke_calls_processes_invoke_async(
             folder_path="/Shared/MyFolder",
             attachments=[],
             parent_span_id=None,
+            parent_operation_id=None,
         )
 
     @pytest.mark.asyncio
@@ -148,6 +150,7 @@ async def test_invoke_interrupts_with_wait_job(
     ):
         """Test that after invoking, the tool interrupts with WaitJob."""
         mock_job = MagicMock(spec=Job)
+        mock_job.key = "job-key-456"
         mock_job.folder_key = "folder-key-456"
 
         mock_client = MagicMock()
@@ -173,6 +176,7 @@ async def test_invoke_passes_input_arguments(
     ):
         """Test that input arguments are correctly passed to invoke_async."""
         mock_job = MagicMock(spec=Job)
+        mock_job.key = "job-key"
         mock_job.folder_key = "folder-key"
 
         mock_client = MagicMock()
@@ -197,6 +201,7 @@ async def test_invoke_returns_interrupt_value(
     ):
         """Test that the tool returns the value from interrupt()."""
         mock_job = MagicMock(spec=Job)
+        mock_job.key = "job-key"
         mock_job.folder_key = "folder-key"
 
         mock_client = MagicMock()
@@ -222,6 +227,7 @@ async def test_span_context_parent_span_id_passed_to_invoke(
     ):
         """Test that parent_span_id from _span_context is forwarded to invoke_async."""
         mock_job = MagicMock(spec=Job)
+        mock_job.key = "job-key"
         mock_job.folder_key = "folder-key"
 
         mock_client = MagicMock()
@@ -249,6 +255,7 @@ async def test_span_context_consumed_after_invoke(
     ):
         """Test that parent_span_id is popped (consumed) from _span_context after use."""
         mock_job = MagicMock(spec=Job)
+        mock_job.key = "job-key"
         mock_job.folder_key = "folder-key"
 
         mock_client = MagicMock()
@@ -274,6 +281,7 @@ async def test_span_context_defaults_to_none_when_empty(
     ):
         """Test that parent_span_id defaults to None when _span_context is empty."""
         mock_job = MagicMock(spec=Job)
+        mock_job.key = "job-key"
         mock_job.folder_key = "folder-key"
 
         mock_client = MagicMock()