fix: trace resource overrides

Author: radu-mocanu

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "uipath"
-version = "2.4.3"
+version = "2.4.4"
 description = "Python SDK and CLI for UiPath Platform, enabling programmatic interaction with automation services, process management, and deployment tools."
 readme = { file = "README.md", content-type = "text/markdown" }
 requires-python = ">=3.11"

src/uipath/_utils/__init__.py

@@ -1,4 +1,4 @@
-from ._bindings import get_inferred_bindings_names, resource_override
+from ._bindings import resource_override
 from ._endpoint import Endpoint
 from ._logs import setup_logging
 from ._request_override import header_folder
@@ -12,7 +12,6 @@
     "setup_logging",
     "RequestSpec",
     "header_folder",
-    "get_inferred_bindings_names",
     "resource_override",
     "header_user_agent",
     "user_agent_value",

src/uipath/_utils/_bindings.py

@@ -139,11 +139,29 @@ def resource_override(
     resource_identifier: str = "name",
     folder_identifier: str = "folder_path",
 ) -> Callable[..., Any]:
+    """Decorator for applying resource overrides for an overridable resource.
+
+    It checks the current ContextVar to identify the requested overrides and, if any key matches, it invokes the decorated function
+    with the extracted resource and folder identifiers.
+
+    Args:
+        resource_type: Type of resource to check for overrides (e.g., "model", "dataset")
+        resource_identifier: Key name for the resource ID in override data (default: "name")
+        folder_identifier: Key name for the folder path in override data (default: "folder_path")
+
+    Returns:
+        Decorated function that receives overridden resource identifiers when applicable
+
+    Note:
+        Must be applied BEFORE the @traced decorator to ensure proper execution order.
+    """
+
     def decorator(func: Callable[..., Any]):
-        @functools.wraps(func)
-        def wrapper(*args, **kwargs):
+        sig = inspect.signature(func)
+
+        def process_args(args, kwargs) -> dict[str, Any]:
+            """Process arguments and apply resource overrides if applicable."""
             # convert both args and kwargs to single dict
-            sig = inspect.signature(func)
             bound = sig.bind_partial(*args, **kwargs)
             bound.apply_defaults()
             all_args = dict(bound.arguments)
@@ -177,22 +195,23 @@ def wrapper(*args, **kwargs):
                             matched_overwrite.folder_identifier
                         )
 
-            return func(**all_args)
+            return all_args
 
-        wrapper._should_infer_bindings = True  # type: ignore[attr-defined] # probably a better way to do this
-        wrapper._infer_bindings_mappings = {  # type: ignore[attr-defined] # probably a better way to do this
-            "name": resource_identifier,
-            "folder_path": folder_identifier,
-        }
-        return wrapper
+        if inspect.iscoroutinefunction(func):
 
-    return decorator
+            @functools.wraps(func)
+            async def async_wrapper(*args, **kwargs):
+                all_args = process_args(args, kwargs)
+                return await func(**all_args)
 
+            return async_wrapper
+        else:
 
-def get_inferred_bindings_names(cls: T):
-    inferred_bindings = {}
-    for name, method in inspect.getmembers(cls, inspect.isfunction):
-        if hasattr(method, "_should_infer_bindings") and method._should_infer_bindings:
-            inferred_bindings[name] = method._infer_bindings_mappings  # type: ignore # probably a better way to do this
+            @functools.wraps(func)
+            def wrapper(*args, **kwargs):
+                all_args = process_args(args, kwargs)
+                return func(**all_args)
 
-    return inferred_bindings
+            return wrapper
+
+    return decorator

src/uipath/platform/action_center/_tasks_service.py

@@ -166,12 +166,12 @@ def __init__(
     ) -> None:
         super().__init__(config=config, execution_context=execution_context)
 
-    @traced(name="tasks_create", run_type="uipath")
     @resource_override(
         resource_type="app",
         resource_identifier="app_name",
         folder_identifier="app_folder_path",
     )
+    @traced(name="tasks_create", run_type="uipath")
     async def create_async(
         self,
         title: str,
@@ -234,12 +234,12 @@ async def create_async(
             )
         return Task.model_validate(json_response)
 
-    @traced(name="tasks_create", run_type="uipath")
     @resource_override(
         resource_type="app",
         resource_identifier="app_name",
         folder_identifier="app_folder_path",
     )
+    @traced(name="tasks_create", run_type="uipath")
     def create(
         self,
         title: str,
@@ -302,12 +302,12 @@ def create(
             )
         return Task.model_validate(json_response)
 
-    @traced(name="tasks_retrieve", run_type="uipath")
     @resource_override(
         resource_type="app",
         resource_identifier="app_name",
         folder_identifier="app_folder_path",
     )
+    @traced(name="tasks_retrieve", run_type="uipath")
     def retrieve(
         self,
         action_key: str,
@@ -336,12 +336,12 @@ def retrieve(
 
         return Task.model_validate(response.json())
 
-    @traced(name="tasks_retrieve", run_type="uipath")
     @resource_override(
         resource_type="app",
         resource_identifier="app_name",
         folder_identifier="app_folder_path",
     )
+    @traced(name="tasks_retrieve", run_type="uipath")
     async def retrieve_async(
         self,
         action_key: str,

src/uipath/platform/connections/_connections_service.py

@@ -37,12 +37,12 @@ def __init__(
         super().__init__(config=config, execution_context=execution_context)
         self._folders_service = folders_service
 
+    @resource_override("connection", resource_identifier="key")
     @traced(
         name="connections_retrieve",
         run_type="uipath",
         hide_output=True,
     )
-    @resource_override("connection", resource_identifier="key")
     def retrieve(self, key: str) -> Connection:
         """Retrieve connection details by its key.
 
@@ -248,12 +248,12 @@ async def list_async(
 
         return self._parse_and_validate_list_response(response)
 
+    @resource_override("connection", resource_identifier="key")
     @traced(
         name="connections_retrieve",
         run_type="uipath",
         hide_output=True,
     )
-    @resource_override("connection", resource_identifier="key")
     async def retrieve_async(self, key: str) -> Connection:
         """Asynchronously retrieve connection details by its key.
 

src/uipath/platform/context_grounding/_context_grounding_service.py

@@ -74,8 +74,8 @@ def __init__(
         super().__init__(config=config, execution_context=execution_context)
 
     # 2.3.0 prefix trace name with contextgrounding
-    @traced(name="add_to_index", run_type="uipath")
     @resource_override(resource_type="index")
+    @traced(name="add_to_index", run_type="uipath")
     def add_to_index(
         self,
         name: str,
@@ -130,8 +130,8 @@ def add_to_index(
             self.ingest_data(index, folder_key=folder_key, folder_path=folder_path)
 
     # 2.3.0 prefix trace name with contextgrounding
-    @traced(name="add_to_index", run_type="uipath")
     @resource_override(resource_type="index")
+    @traced(name="add_to_index", run_type="uipath")
     async def add_to_index_async(
         self,
         name: str,
@@ -189,8 +189,8 @@ async def add_to_index_async(
                 index, folder_key=folder_key, folder_path=folder_path
             )
 
-    @traced(name="contextgrounding_retrieve", run_type="uipath")
     @resource_override(resource_type="index")
+    @traced(name="contextgrounding_retrieve", run_type="uipath")
     def retrieve(
         self,
         name: str,
@@ -231,8 +231,8 @@ def retrieve(
         except StopIteration as e:
             raise Exception("ContextGroundingIndex not found") from e
 
-    @traced(name="contextgrounding_retrieve", run_type="uipath")
     @resource_override(resource_type="index")
+    @traced(name="contextgrounding_retrieve", run_type="uipath")
     async def retrieve_async(
         self,
         name: str,
@@ -343,8 +343,8 @@ async def retrieve_by_id_async(
 
         return response.json()
 
-    @traced(name="contextgrounding_create_index", run_type="uipath")
     @resource_override(resource_type="index")
+    @traced(name="contextgrounding_create_index", run_type="uipath")
     def create_index(
         self,
         name: str,
@@ -398,8 +398,8 @@ def create_index(
 
         return ContextGroundingIndex.model_validate(response.json())
 
-    @traced(name="contextgrounding_create_index", run_type="uipath")
     @resource_override(resource_type="index")
+    @traced(name="contextgrounding_create_index", run_type="uipath")
     async def create_index_async(
         self,
         name: str,
@@ -453,8 +453,8 @@ async def create_index_async(
 
         return ContextGroundingIndex.model_validate(response.json())
 
-    @traced(name="contextgrounding_retrieve_deep_rag", run_type="uipath")
     @resource_override(resource_type="index", resource_identifier="index_name")
+    @traced(name="contextgrounding_retrieve_deep_rag", run_type="uipath")
     def retrieve_deep_rag(
         self,
         id: str,
@@ -482,8 +482,8 @@ def retrieve_deep_rag(
         )
         return DeepRagResponse.model_validate(response.json())
 
-    @traced(name="contextgrounding_retrieve_deep_rag_async", run_type="uipath")
     @resource_override(resource_type="index", resource_identifier="index_name")
+    @traced(name="contextgrounding_retrieve_deep_rag_async", run_type="uipath")
     async def retrieve_deep_rag_async(
         self,
         id: str,
@@ -513,8 +513,8 @@ async def retrieve_deep_rag_async(
 
         return DeepRagResponse.model_validate(response.json())
 
-    @traced(name="contextgrounding_start_batch_transform", run_type="uipath")
     @resource_override(resource_type="index", resource_identifier="index_name")
+    @traced(name="contextgrounding_start_batch_transform", run_type="uipath")
     def start_batch_transform(
         self,
         name: str,
@@ -571,8 +571,8 @@ def start_batch_transform(
         )
         return BatchTransformCreationResponse.model_validate(response.json())
 
-    @traced(name="contextgrounding_start_batch_transform_async", run_type="uipath")
     @resource_override(resource_type="index", resource_identifier="index_name")
+    @traced(name="contextgrounding_start_batch_transform_async", run_type="uipath")
     async def start_batch_transform_async(
         self,
         name: str,
@@ -629,8 +629,8 @@ async def start_batch_transform_async(
         )
         return BatchTransformCreationResponse.model_validate(response.json())
 
-    @traced(name="contextgrounding_retrieve_batch_transform", run_type="uipath")
     @resource_override(resource_type="index", resource_identifier="index_name")
+    @traced(name="contextgrounding_retrieve_batch_transform", run_type="uipath")
     def retrieve_batch_transform(
         self,
         id: str,
@@ -655,8 +655,8 @@ def retrieve_batch_transform(
         )
         return BatchTransformResponse.model_validate(response.json())
 
-    @traced(name="contextgrounding_retrieve_batch_transform_async", run_type="uipath")
     @resource_override(resource_type="index", resource_identifier="index_name")
+    @traced(name="contextgrounding_retrieve_batch_transform_async", run_type="uipath")
     async def retrieve_batch_transform_async(
         self,
         id: str,
@@ -681,8 +681,8 @@ async def retrieve_batch_transform_async(
         )
         return BatchTransformResponse.model_validate(response.json())
 
-    @traced(name="contextgrounding_download_batch_transform_result", run_type="uipath")
     @resource_override(resource_type="index", resource_identifier="index_name")
+    @traced(name="contextgrounding_download_batch_transform_result", run_type="uipath")
     def download_batch_transform_result(
         self,
         id: str,
@@ -727,10 +727,10 @@ def download_batch_transform_result(
                 file_content = client.get(uri_response.uri).content
             file.write(file_content)
 
+    @resource_override(resource_type="index", resource_identifier="index_name")
     @traced(
         name="contextgrounding_download_batch_transform_result_async", run_type="uipath"
     )
-    @resource_override(resource_type="index", resource_identifier="index_name")
     async def download_batch_transform_result_async(
         self,
         id: str,
@@ -777,8 +777,8 @@ async def download_batch_transform_result_async(
         with open(destination_path, "wb") as file:
             file.write(file_content)
 
-    @traced(name="contextgrounding_start_deep_rag", run_type="uipath")
     @resource_override(resource_type="index", resource_identifier="index_name")
+    @traced(name="contextgrounding_start_deep_rag", run_type="uipath")
     def start_deep_rag(
         self,
         name: str,
@@ -829,8 +829,8 @@ def start_deep_rag(
 
         return DeepRagCreationResponse.model_validate(response.json())
 
-    @traced(name="contextgrounding_start_deep_rag_async", run_type="uipath")
     @resource_override(resource_type="index", resource_identifier="index_name")
+    @traced(name="contextgrounding_start_deep_rag_async", run_type="uipath")
     async def start_deep_rag_async(
         self,
         name: str,
@@ -882,8 +882,8 @@ async def start_deep_rag_async(
 
         return DeepRagCreationResponse.model_validate(response.json())
 
-    @traced(name="contextgrounding_search", run_type="uipath")
     @resource_override(resource_type="index")
+    @traced(name="contextgrounding_search", run_type="uipath")
     def search(
         self,
         name: str,
@@ -931,8 +931,8 @@ def search(
             response.json()
         )
 
-    @traced(name="contextgrounding_search", run_type="uipath")
     @resource_override(resource_type="index")
+    @traced(name="contextgrounding_search", run_type="uipath")
     async def search_async(
         self,
         name: str,