diff --git a/pyproject.toml b/pyproject.toml index 6eacebc0..975be820 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "uipath" -version = "2.0.62" +version = "2.0.63" 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.10" diff --git a/src/uipath/_services/assets_service.py b/src/uipath/_services/assets_service.py index 3112ed64..452e3adc 100644 --- a/src/uipath/_services/assets_service.py +++ b/src/uipath/_services/assets_service.py @@ -25,7 +25,13 @@ def __init__(self, config: Config, execution_context: ExecutionContext) -> None: self._base_url = "assets" @traced( - name="assets_retrieve", run_type="uipath", hide_input=True, hide_output=True + name="assets_retrieve", + run_type="uipath", + dependency={ + "targetName": lambda inputs: inputs["name"], + "targetType": "Asset", + "operationName": "GET Asset", + }, ) @infer_bindings(resource_type="asset") def retrieve( @@ -80,7 +86,15 @@ def retrieve( return Asset.model_validate(response.json()["value"][0]) @traced( - name="assets_retrieve", run_type="uipath", hide_input=True, hide_output=True + name="assets_retrieve", + run_type="uipath", + hide_input=False, + hide_output=False, + dependency={ + "targetName": lambda inputs: inputs["name"], + "targetType": "Asset", + "operationName": "GET Asset", + }, ) @infer_bindings(resource_type="asset") async def retrieve_async( @@ -126,7 +140,15 @@ async def retrieve_async( return Asset.model_validate(response.json()["value"][0]) @traced( - name="assets_credential", run_type="uipath", hide_input=True, hide_output=True + name="assets_credential", + run_type="uipath", + hide_input=False, + hide_output=True, + dependency={ + "targetName": lambda inputs: inputs["name"], + "targetType": "Asset", + "operationName": "GET Credential", + }, ) @infer_bindings(resource_type="asset") def retrieve_credential( @@ -179,7 +201,15 @@ def retrieve_credential( return user_asset.credential_password @traced( - name="assets_credential", run_type="uipath", hide_input=True, hide_output=True + name="assets_credential", + run_type="uipath", + hide_input=False, + hide_output=True, + dependency={ + "targetName": lambda inputs: inputs["name"], + "targetType": "Asset", + "operationName": "GET Credential", + }, ) @infer_bindings(resource_type="asset") async def retrieve_credential_async( @@ -231,7 +261,17 @@ async def retrieve_credential_async( return user_asset.credential_password - @traced(name="assets_update", run_type="uipath", hide_input=True, hide_output=True) + @traced( + name="assets_update", + run_type="uipath", + hide_input=False, + hide_output=True, + dependency={ + "targetName": lambda inputs: inputs["robot_asset"].name, + "targetType": "Asset", + "operationName": "UPDATE Asset", + }, + ) def update( self, robot_asset: UserAsset, @@ -273,7 +313,17 @@ def update( return response.json() - @traced(name="assets_update", run_type="uipath", hide_input=True, hide_output=True) + @traced( + name="assets_update", + run_type="uipath", + hide_input=False, + hide_output=True, + dependency={ + "targetName": lambda inputs: inputs["robot_asset"].name, + "targetType": "Asset", + "operationName": "UPDATE Asset", + }, + ) async def update_async( self, robot_asset: UserAsset, diff --git a/src/uipath/_services/buckets_service.py b/src/uipath/_services/buckets_service.py index a41bb59e..720bafda 100644 --- a/src/uipath/_services/buckets_service.py +++ b/src/uipath/_services/buckets_service.py @@ -28,7 +28,18 @@ def __init__(self, config: Config, execution_context: ExecutionContext) -> None: self.custom_client = httpx.Client() self.custom_client_async = httpx.AsyncClient() - @traced(name="buckets_download", run_type="uipath") + @traced( + name="buckets_download", + run_type="uipath", + hide_input=False, + hide_output=False, + dependency={ + "targetName": lambda inputs: inputs.get("name"), + "targetType": "Bucket", + "targetId": lambda inputs: inputs.get("key"), + "operationName": "DOWNLOAD File", + }, + ) @infer_bindings(resource_type="bucket") def download( self, @@ -84,7 +95,18 @@ def download( file_content = self.custom_client.get(read_uri, headers=headers).content file.write(file_content) - @traced(name="buckets_download", run_type="uipath") + @traced( + name="buckets_download_async", + run_type="uipath", + hide_input=False, + hide_output=False, + dependency={ + "targetName": lambda inputs: inputs.get("name"), + "targetType": "Bucket", + "targetId": lambda inputs: inputs.get("key"), + "operationName": "DOWNLOAD File", + }, + ) @infer_bindings(resource_type="bucket") async def download_async( self, @@ -146,7 +168,19 @@ async def download_async( ).content file.write(file_content) - @traced(name="buckets_upload", run_type="uipath") + @traced( + name="buckets_upload", + run_type="uipath", + input_processor=_upload_from_memory_input_processor, + hide_input=False, + hide_output=False, + dependency={ + "targetName": lambda inputs: inputs.get("name"), + "targetType": "Bucket", + "targetId": lambda inputs: inputs.get("key"), + "operationName": "UPLOAD File", + }, + ) @infer_bindings(resource_type="bucket") def upload( self, @@ -229,7 +263,19 @@ def upload( write_uri, headers=headers, files={"file": file} ) - @traced(name="buckets_upload", run_type="uipath") + @traced( + name="buckets_upload_async", + run_type="uipath", + input_processor=_upload_from_memory_input_processor, + hide_input=False, + hide_output=False, + dependency={ + "targetName": lambda inputs: inputs.get("name"), + "targetType": "Bucket", + "targetId": lambda inputs: inputs.get("key"), + "operationName": "UPLOAD File", + }, + ) @infer_bindings(resource_type="bucket") async def upload_async( self, @@ -317,7 +363,18 @@ async def upload_async( write_uri, headers=headers, files={"file": file} ) - @traced(name="buckets_retrieve", run_type="uipath") + @traced( + name="buckets_retrieve", + run_type="uipath", + hide_input=False, + hide_output=False, + dependency={ + "targetName": lambda inputs: inputs.get("name"), + "targetType": "Bucket", + "targetId": lambda inputs: inputs.get("key"), + "operationName": "GET Bucket", + }, + ) @infer_bindings(resource_type="bucket") def retrieve( self, @@ -365,7 +422,18 @@ def retrieve( raise Exception(f"Bucket with name '{name}' not found") from e return Bucket.model_validate(response) - @traced(name="buckets_retrieve", run_type="uipath") + @traced( + name="buckets_retrieve_async", + run_type="uipath", + hide_input=False, + hide_output=False, + dependency={ + "targetName": lambda inputs: inputs.get("name"), + "targetType": "Bucket", + "targetId": lambda inputs: inputs.get("key"), + "operationName": "GET Bucket", + }, + ) @infer_bindings(resource_type="bucket") async def retrieve_async( self, @@ -487,3 +555,33 @@ def _retrieve_by_key_spec( **header_folder(folder_key, folder_path), }, ) + + def retrieve_read_uri( + self, + *, + bucket_id: int, + blob_file_path: str, + folder_key: Optional[str] = None, + folder_path: Optional[str] = None, + ) -> dict: + """Get the read URI and headers for a file in a bucket. + + Args: + bucket_id (int): The ID of the bucket. + blob_file_path (str): The path to the file in the bucket. + folder_key (Optional[str]): The key of the folder where the bucket resides. + folder_path (Optional[str]): The path of the folder where the bucket resides. + + Returns: + dict: Contains 'Uri', 'Headers', and 'RequiresAuth'. + """ + spec = self._retrieve_readUri_spec( + bucket_id, blob_file_path, folder_key=folder_key, folder_path=folder_path + ) + result = self.request( + spec.method, + url=spec.endpoint, + params=spec.params, + headers=spec.headers, + ).json() + return result diff --git a/src/uipath/_services/jobs_service.py b/src/uipath/_services/jobs_service.py index 64ce5e24..3ab86989 100644 --- a/src/uipath/_services/jobs_service.py +++ b/src/uipath/_services/jobs_service.py @@ -38,7 +38,16 @@ def resume(self, *, inbox_id: str, payload: Any) -> None: ... @overload def resume(self, *, job_id: str, payload: Any) -> None: ... - @traced(name="jobs_resume", run_type="uipath") + @traced( + name="jobs_resume", + run_type="uipath", + dependency={ + "targetName": lambda inputs: inputs.get("job_id"), + "targetType": "Job", + "targetId": lambda inputs: inputs.get("job_id"), + "operationName": "RESUME Job", + }, + ) def resume( self, *, @@ -84,6 +93,16 @@ def resume( content=spec.content, ) + @traced( + name="jobs_resume", + run_type="uipath", + dependency={ + "targetName": lambda inputs: inputs.get("job_id"), + "targetType": "Job", + "targetId": lambda inputs: inputs.get("job_id"), + "operationName": "RESUME Job", + }, + ) async def resume_async( self, *, @@ -149,6 +168,16 @@ async def main(): # noqa: D103 def custom_headers(self) -> Dict[str, str]: return self.folder_headers + @traced( + name="jobs_retrieve", + run_type="uipath", + dependency={ + "targetName": lambda inputs: inputs.get("job_key"), + "targetType": "Job", + "targetId": lambda inputs: inputs.get("job_key"), + "operationName": "GET Job", + }, + ) def retrieve( self, job_key: str, @@ -185,6 +214,16 @@ def retrieve( return Job.model_validate(response.json()) + @traced( + name="jobs_retrieve", + run_type="uipath", + dependency={ + "targetName": lambda inputs: inputs.get("job_key"), + "targetType": "Job", + "targetId": lambda inputs: inputs.get("job_key"), + "operationName": "GET Job", + }, + ) async def retrieve_async( self, job_key: str, @@ -334,7 +373,16 @@ def _retrieve_spec( }, ) - @traced(name="jobs_list_attachments", run_type="uipath") + @traced( + name="jobs_list_attachments", + run_type="uipath", + dependency={ + "targetName": lambda inputs: inputs.get("job_key"), + "targetType": "Job", + "targetId": lambda inputs: inputs.get("job_key"), + "operationName": "LIST JobAttachments", + }, + ) def list_attachments( self, *, @@ -372,7 +420,16 @@ def list_attachments( return [item.get("attachmentId") for item in response] - @traced(name="jobs_list_attachments", run_type="uipath") + @traced( + name="jobs_list_attachments", + run_type="uipath", + dependency={ + "targetName": lambda inputs: inputs.get("job_key"), + "targetType": "Job", + "targetId": lambda inputs: inputs.get("job_key"), + "operationName": "LIST JobAttachments", + }, + ) async def list_attachments_async( self, *, @@ -427,7 +484,18 @@ async def main(): return [item.get("attachmentId") for item in response] - @traced(name="jobs_link_attachment", run_type="uipath") + @traced( + name="jobs_link_attachment", + run_type="uipath", + hide_input=True, + hide_output=True, + dependency={ + "targetName": lambda inputs: inputs.get("job_key"), + "targetType": "Job", + "targetId": lambda inputs: inputs.get("job_key"), + "operationName": "LINK JobAttachment", + }, + ) def link_attachment( self, *, @@ -465,7 +533,18 @@ def link_attachment( json=spec.json, ) - @traced(name="jobs_link_attachment", run_type="uipath") + @traced( + name="jobs_link_attachment", + run_type="uipath", + hide_input=True, + hide_output=True, + dependency={ + "targetName": lambda inputs: inputs.get("job_key"), + "targetType": "Job", + "targetId": lambda inputs: inputs.get("job_key"), + "operationName": "LINK JobAttachment", + }, + ) async def link_attachment_async( self, *, @@ -541,7 +620,18 @@ def _link_job_attachment_spec( }, ) - @traced(name="jobs_create_attachment", run_type="uipath") + @traced( + name="jobs_create_attachment", + run_type="uipath", + hide_input=True, + hide_output=True, + dependency={ + "targetName": lambda inputs: inputs.get("job_key"), + "targetType": "Job", + "targetId": lambda inputs: inputs.get("job_key"), + "operationName": "CREATE Attachment", + }, + ) def create_attachment( self, *, @@ -680,7 +770,18 @@ def create_attachment( # Return only the UUID return attachment_id - @traced(name="jobs_create_attachment", run_type="uipath") + @traced( + name="jobs_create_attachment", + run_type="uipath", + hide_input=True, + hide_output=True, + dependency={ + "targetName": lambda inputs: inputs.get("job_key"), + "targetType": "Job", + "targetId": lambda inputs: inputs.get("job_key"), + "operationName": "CREATE Attachment", + }, + ) async def create_attachment_async( self, *, @@ -820,3 +921,57 @@ async def main(): # Return only the UUID return attachment_id + + @traced( + name="jobs_list", + run_type="uipath", + dependency={ + "targetType": "Job", + "operationName": "LIST Jobs", + }, + ) + def list_jobs( + self, + *, + folder_key: Optional[str] = None, + folder_path: Optional[str] = None, + top: Optional[int] = None, + skip: Optional[int] = None, + filters: Optional[dict] = None, + ) -> List[Job]: + """List all jobs in a folder. + + Args: + folder_key (Optional[str]): The key of the folder to list jobs from. + folder_path (Optional[str]): The path of the folder to list jobs from. + top (Optional[int]): Max number of jobs to return. + skip (Optional[int]): Number of jobs to skip (for paging). + filters (Optional[dict]): Additional OData filters as a dict. + + Returns: + List[Job]: List of Job objects in the folder. + """ + params = {} + if top is not None: + params["$top"] = top + if skip is not None: + params["$skip"] = skip + if filters: + # Merge filters into OData $filter string + filter_str = " and ".join(f"{k} eq '{v}'" for k, v in filters.items()) + params["$filter"] = filter_str + headers = {**header_folder(folder_key, folder_path)} + spec = RequestSpec( + method="GET", + endpoint=Endpoint("/orchestrator_/odata/Jobs"), + params=params, + headers=headers, + ) + response = self.request( + spec.method, + url=spec.endpoint, + params=spec.params, + headers=spec.headers, + ) + jobs_data = response.json().get("value", []) + return [Job.model_validate(job) for job in jobs_data] diff --git a/src/uipath/_services/llm_gateway_service.py b/src/uipath/_services/llm_gateway_service.py index 1e63c211..8fa414e8 100644 --- a/src/uipath/_services/llm_gateway_service.py +++ b/src/uipath/_services/llm_gateway_service.py @@ -54,7 +54,18 @@ class UiPathOpenAIService(BaseService): def __init__(self, config: Config, execution_context: ExecutionContext) -> None: super().__init__(config=config, execution_context=execution_context) - @traced(name="llm_embeddings_usage", run_type="uipath") + @traced( + name="llm_embeddings_usage", + run_type="uipath", + hide_input=False, + hide_output=False, + dependency={ + "targetName": lambda inputs: f"LLMGateway:OpenAI:EmbeddingModel:{inputs['embedding_model']}", + "targetType": "LLMGateway", + "targetId": lambda inputs: inputs["embedding_model"], + "operationName": "GET Embedding Usage", + }, + ) async def embeddings_usage( self, input: str, embedding_model: str = EmbeddingModels.text_embedding_ada_002 ): @@ -81,7 +92,18 @@ async def embeddings_usage( return UsageInfo.model_validate(response.json()) - @traced(name="llm_embeddings", run_type="uipath") + @traced( + name="llm_embeddings", + run_type="uipath", + hide_input=False, + hide_output=False, + dependency={ + "targetName": lambda inputs: f"LLMGateway:OpenAI:EmbeddingModel:{inputs['embedding_model']}", + "targetType": "LLMGateway", + "targetId": lambda inputs: inputs["embedding_model"], + "operationName": "GET Embedding", + }, + ) async def embeddings( self, input: str, embedding_model: str = EmbeddingModels.text_embedding_ada_002 ): @@ -107,7 +129,18 @@ async def embeddings( return TextEmbedding.model_validate(response.json()) - @traced(name="llm_chat_completions", run_type="uipath") + @traced( + name="llm_chat_completions", + run_type="uipath", + hide_input=False, + hide_output=False, + dependency={ + "targetName": lambda inputs: f"LLMGateway:OpenAI:ChatModel:{inputs['model']}", + "targetType": "LLMGateway", + "targetId": lambda inputs: inputs["model"], + "operationName": "GET Chat Completion", + }, + ) async def chat_completions( self, messages: List[Dict[str, str]], @@ -157,7 +190,18 @@ async def chat_completions( return ChatCompletion.model_validate(response.json()) - @traced(name="llm_chat_completions_usage", run_type="uipath") + @traced( + name="llm_chat_completions_usage", + run_type="uipath", + hide_input=False, + hide_output=False, + dependency={ + "targetName": lambda inputs: f"LLMGateway:OpenAI:ChatModel:{inputs['model']}", + "targetType": "LLMGateway", + "targetId": lambda inputs: inputs["model"], + "operationName": "GET Chat Completion Usage", + }, + ) async def chat_completions_usage( self, messages: List[Dict[str, str]], @@ -216,7 +260,18 @@ class UiPathLlmChatService(BaseService): def __init__(self, config: Config, execution_context: ExecutionContext) -> None: super().__init__(config=config, execution_context=execution_context) - @traced(name="llm_chat_completions", run_type="uipath") + @traced( + name="llm_chat_completions", + run_type="uipath", + hide_input=True, + hide_output=True, + dependency={ + "targetName": lambda inputs: f"LLMGateway:Normalized:ChatModel:{inputs['model']}", + "targetType": "LLMGateway", + "targetId": lambda inputs: inputs["model"], + "operationName": "GET Normalized Chat Completion", + }, + ) async def chat_completions( self, messages: List[Dict[str, str]], diff --git a/src/uipath/_services/processes_service.py b/src/uipath/_services/processes_service.py index 6ee9ca44..44b0132c 100644 --- a/src/uipath/_services/processes_service.py +++ b/src/uipath/_services/processes_service.py @@ -23,7 +23,15 @@ class ProcessesService(FolderContext, BaseService): def __init__(self, config: Config, execution_context: ExecutionContext) -> None: super().__init__(config=config, execution_context=execution_context) - @traced(name="processes_invoke", run_type="uipath") + @traced( + name="processes_invoke", + run_type="uipath", + dependency={ + "targetName": lambda inputs: inputs.get("name"), + "targetType": "Process", + "operationName": "INVOKE Process", + }, + ) @infer_bindings(resource_type="process") def invoke( self, @@ -81,7 +89,15 @@ def invoke( return Job.model_validate(response.json()["value"][0]) - @traced(name="processes_invoke", run_type="uipath") + @traced( + name="processes_invoke", + run_type="uipath", + dependency={ + "targetName": lambda inputs: inputs.get("name"), + "targetType": "Process", + "operationName": "INVOKE Process", + }, + ) @infer_bindings(resource_type="process") async def invoke_async( self, diff --git a/src/uipath/_services/queues_service.py b/src/uipath/_services/queues_service.py index 2a6e8fb9..5788215d 100644 --- a/src/uipath/_services/queues_service.py +++ b/src/uipath/_services/queues_service.py @@ -11,6 +11,20 @@ from ._base_service import BaseService +def _get_queue_item_name_for_tracing( + inputs: Dict[str, Any], primary: str, secondary: str +) -> str: + item = inputs.get("item") + if not item: + return "UnknownItem" + name_val = ( + item.get(primary, item.get(secondary, "UnknownQueue")) + if isinstance(item, dict) + else getattr(item, primary, getattr(item, secondary, "UnknownQueue")) + ) + return name_val + + class QueuesService(FolderContext, BaseService): """Service for managing UiPath queues and queue items. @@ -21,7 +35,17 @@ class QueuesService(FolderContext, BaseService): def __init__(self, config: Config, execution_context: ExecutionContext) -> None: super().__init__(config=config, execution_context=execution_context) - @traced(name="queues_list_items", run_type="uipath") + @traced( + name="queues_list_items", + run_type="uipath", + hide_input=False, + hide_output=False, + dependency={ + "targetName": "All", + "targetType": "Queue", + "operationName": "LIST QueueItems", + }, + ) def list_items(self) -> Response: """Retrieves a list of queue items from the Orchestrator. @@ -33,7 +57,17 @@ def list_items(self) -> Response: return response.json() - @traced(name="queues_list_items", run_type="uipath") + @traced( + name="queues_list_items", + run_type="uipath", + hide_input=False, + hide_output=False, + dependency={ + "targetName": "All", + "targetType": "Queue", + "operationName": "LIST QueueItems", + }, + ) async def list_items_async(self) -> Response: """Asynchronously retrieves a list of queue items from the Orchestrator. @@ -44,7 +78,25 @@ async def list_items_async(self) -> Response: response = await self.request_async(spec.method, url=spec.endpoint) return response.json() - @traced(name="queues_create_item", run_type="uipath") + @traced( + name="queues_create_item", + run_type="uipath", + hide_input=False, + hide_output=False, + dependency={ + "targetName": lambda inputs: _get_queue_item_name_for_tracing( + inputs, "name", "Name" + ), + "targetType": "Queue", + "targetId": lambda inputs: ( + _get_queue_item_name_for_tracing(inputs, "key", "Key") + if _get_queue_item_name_for_tracing(inputs, "key", "Key") + != "UnknownQueue" + else _get_queue_item_name_for_tracing(inputs, "name", "Name") + ), + "operationName": "ADD QueueItem", + }, + ) def create_item(self, item: Union[Dict[str, Any], QueueItem]) -> Response: """Creates a new queue item in the Orchestrator. @@ -60,7 +112,25 @@ def create_item(self, item: Union[Dict[str, Any], QueueItem]) -> Response: response = self.request(spec.method, url=spec.endpoint, json=spec.json) return response.json() - @traced(name="queues_create_item", run_type="uipath") + @traced( + name="queues_create_item", + run_type="uipath", + hide_input=False, + hide_output=False, + dependency={ + "targetName": lambda inputs: _get_queue_item_name_for_tracing( + inputs, "name", "Name" + ), + "targetType": "Queue", + "targetId": lambda inputs: ( + _get_queue_item_name_for_tracing(inputs, "key", "Key") + if _get_queue_item_name_for_tracing(inputs, "key", "Key") + != "UnknownQueue" + else _get_queue_item_name_for_tracing(inputs, "name", "Name") + ), + "operationName": "ADD QueueItem", + }, + ) async def create_item_async( self, item: Union[Dict[str, Any], QueueItem] ) -> Response: @@ -80,7 +150,18 @@ async def create_item_async( ) return response.json() - @traced(name="queues_create_items", run_type="uipath") + @traced( + name="queues_create_items", + run_type="uipath", + hide_input=False, + hide_output=False, + dependency={ + "targetName": lambda inputs: inputs["queue_name"], # Changed + "targetType": "Queue", + "targetId": lambda inputs: inputs["queue_name"], + "operationName": "ADD QueueItems", + }, + ) def create_items( self, items: List[Union[Dict[str, Any], QueueItem]], @@ -101,7 +182,18 @@ def create_items( response = self.request(spec.method, url=spec.endpoint, json=spec.json) return response.json() - @traced(name="queues_create_items", run_type="uipath") + @traced( + name="queues_create_items", + run_type="uipath", + hide_input=False, + hide_output=False, + dependency={ + "targetName": lambda inputs: inputs["queue_name"], # Changed + "targetType": "Queue", + "targetId": lambda inputs: inputs["queue_name"], + "operationName": "ADD QueueItems", + }, + ) async def create_items_async( self, items: List[Union[Dict[str, Any], QueueItem]], @@ -124,7 +216,22 @@ async def create_items_async( ) return response.json() - @traced(name="queues_create_transaction_item", run_type="uipath") + @traced( + name="queues_create_transaction_item", + run_type="uipath", + hide_input=False, + hide_output=False, + dependency={ + "targetName": lambda inputs: _get_queue_item_name_for_tracing( + inputs, "name", "Name" + ), + "targetType": "Queue", + "targetId": lambda inputs: _get_queue_item_name_for_tracing( # Added + inputs, "name", "Name" + ), + "operationName": "ADD QueueItem", + }, + ) def create_transaction_item( self, item: Union[Dict[str, Any], TransactionItem], no_robot: bool = False ) -> Response: @@ -141,7 +248,22 @@ def create_transaction_item( response = self.request(spec.method, url=spec.endpoint, json=spec.json) return response.json() - @traced(name="queues_create_transaction_item", run_type="uipath") + @traced( + name="queues_create_transaction_item", + run_type="uipath", + hide_input=False, + hide_output=False, + dependency={ + "targetName": lambda inputs: _get_queue_item_name_for_tracing( + inputs, "name", "Name" + ), + "targetType": "Queue", # Added + "targetId": lambda inputs: _get_queue_item_name_for_tracing( # Added + inputs, "name", "Name" + ), + "operationName": "ADD QueueItem", + }, + ) async def create_transaction_item_async( self, item: Union[Dict[str, Any], TransactionItem], no_robot: bool = False ) -> Response: @@ -160,7 +282,16 @@ async def create_transaction_item_async( ) return response.json() - @traced(name="queues_update_progress_of_transaction_item", run_type="uipath") + @traced( + name="queues_update_progress_of_transaction_item", + run_type="uipath", + dependency={ + "targetName": "QueueItem", + "targetType": "QueueItem", # Added + "targetId": lambda inputs: inputs.get("transaction_key"), # Added + "operationName": "UPDATE QueueItem", + }, + ) def update_progress_of_transaction_item( self, transaction_key: str, progress: str ) -> Response: @@ -179,7 +310,16 @@ def update_progress_of_transaction_item( response = self.request(spec.method, url=spec.endpoint, json=spec.json) return response.json() - @traced(name="queues_update_progress_of_transaction_item", run_type="uipath") + @traced( + name="queues_update_progress_of_transaction_item", + run_type="uipath", + dependency={ + "targetName": "QueueItem", + "targetType": "QueueItem", # Added + "targetId": lambda inputs: inputs.get("transaction_key"), # Added + "operationName": "UPDATE QueueItem", + }, + ) async def update_progress_of_transaction_item_async( self, transaction_key: str, progress: str ) -> Response: @@ -200,7 +340,16 @@ async def update_progress_of_transaction_item_async( ) return response.json() - @traced(name="queues_complete_transaction_item", run_type="uipath") + @traced( + name="queues_complete_transaction_item", + run_type="uipath", + dependency={ + "targetName": lambda inputs: inputs.get("result").get("name"), + "targetType": "Queue", + "targetId": lambda inputs: inputs.get("transaction_key"), + "operationName": "UPDATE QueueItem", + }, + ) def complete_transaction_item( self, transaction_key: str, result: Union[Dict[str, Any], TransactionItemResult] ) -> Response: @@ -217,9 +366,18 @@ def complete_transaction_item( """ spec = self._complete_transaction_item_spec(transaction_key, result) response = self.request(spec.method, url=spec.endpoint, json=spec.json) - return response.json() - - @traced(name="queues_complete_transaction_item", run_type="uipath") + return response + + @traced( + name="queues_complete_transaction_item", + run_type="uipath", + dependency={ + "targetName": "QueueItem", + "targetType": "QueueItem", + "targetId": lambda inputs: inputs.get("transaction_key"), + "operationName": "UPDATE QueueItem", + }, + ) async def complete_transaction_item_async( self, transaction_key: str, result: Union[Dict[str, Any], TransactionItemResult] ) -> Response: @@ -240,6 +398,32 @@ async def complete_transaction_item_async( ) return response.json() + @traced( + name="queues_get_queue_item", + run_type="uipath", + dependency={ + "targetName": lambda inputs: inputs.get("queue_name", "UnknownQueueItem"), + "targetType": "Queue", + "targetId": lambda inputs: inputs.get("key"), + "operationName": "GET QueueItem", + }, + ) + def get_queue_item(self, key: int, queue_name: str) -> Response: + """Retrieves a queue item by its key from the Orchestrator. + + Args: + key: The unique integer key of the queue item. + + Returns: + Response: HTTP response containing the queue item details. + """ + spec = RequestSpec( + method="GET", + endpoint=Endpoint(f"/orchestrator_/odata/QueueItems({key})"), + ) + response = self.request(spec.method, url=spec.endpoint) + return response.json() + @property def custom_headers(self) -> Dict[str, str]: return self.folder_headers diff --git a/src/uipath/models/queues.py b/src/uipath/models/queues.py index 15e4e925..49421d83 100644 --- a/src/uipath/models/queues.py +++ b/src/uipath/models/queues.py @@ -127,6 +127,11 @@ class TransactionItemResult(BaseModel): json_encoders={datetime: lambda v: v.isoformat() if v else None}, ) + name: str = Field( + description="The name of the queue in which to search for the next item or in which to insert the item before marking it as InProgress and sending it to the robot.", + alias="Name", + ) + is_successful: Optional[bool] = Field( default=None, description="States if the processing was successful or not.", diff --git a/src/uipath/tracing/_traced.py b/src/uipath/tracing/_traced.py index 7fb74c80..27ef8e75 100644 --- a/src/uipath/tracing/_traced.py +++ b/src/uipath/tracing/_traced.py @@ -2,8 +2,9 @@ import inspect import json import logging +import os from functools import wraps -from typing import Any, Callable, List, Optional, Tuple +from typing import Any, Callable, List, Optional, Tuple, TypedDict, Union from opentelemetry import trace @@ -14,6 +15,17 @@ tracer = trace.get_tracer(__name__) +class DependencyInfo(TypedDict, total=False): + """Type definition for dependency tracking information.""" + + sourceName: Union[str, Callable[..., str]] + sourceType: Union[str, Callable[..., str]] # Added + targetName: Union[str, Callable[..., str]] + targetType: Union[str, Callable[..., str]] # Added + targetId: Union[str, Callable[..., str]] # Added + operationName: Union[str, Callable[..., str]] + + class TracingManager: """Static utility class to manage tracing implementations and decorated functions.""" @@ -129,6 +141,7 @@ def _opentelemetry_traced( span_type: Optional[str] = None, input_processor: Optional[Callable[..., Any]] = None, output_processor: Optional[Callable[..., Any]] = None, + dependency: Optional[DependencyInfo] = None, ): """Default tracer implementation using OpenTelemetry.""" @@ -147,14 +160,40 @@ def sync_wrapper(*args, **kwargs): span.set_attribute("run_type", run_type) # Format arguments for tracing - inputs = _SpanUtils.format_args_for_trace_json( + inputs_json_str = _SpanUtils.format_args_for_trace_json( # Store the raw JSON string of args inspect.signature(func), *args, **kwargs ) - # Apply input processor if provided + + # Determine the value for the 'inputs' attribute if input_processor is not None: - processed_inputs = input_processor(json.loads(inputs)) - inputs = json.dumps(processed_inputs, default=str) - span.set_attribute("inputs", inputs) + processed_inputs_val = input_processor(json.loads(inputs_json_str)) + inputs_attribute_value = json.dumps( + processed_inputs_val, default=str + ) + else: + inputs_attribute_value = ( + inputs_json_str # Use raw JSON string if no processor + ) + span.set_attribute("inputs", inputs_attribute_value) + + # Add dependency information as a JSON attribute if provided + if dependency is not None: + processed_dependency = {} + parsed_args_dict = None + # Check if any dependency value is callable to decide if we need to parse args + if any(callable(v) for v in dependency.values()): + parsed_args_dict = json.loads(inputs_json_str) + + for key, value in dependency.items(): + if callable(value): + # We've already ensured parsed_args_dict is populated if any value is callable + processed_dependency[key] = value(parsed_args_dict) + else: + processed_dependency[key] = value + + dependency_json = json.dumps(processed_dependency, default=str) + span.set_attribute("dependency", dependency_json) + try: result = func(*args, **kwargs) # Process output if processor is provided @@ -182,14 +221,38 @@ async def async_wrapper(*args, **kwargs): span.set_attribute("run_type", run_type) # Format arguments for tracing - inputs = _SpanUtils.format_args_for_trace_json( + inputs_json_str = _SpanUtils.format_args_for_trace_json( # Store the raw JSON string of args inspect.signature(func), *args, **kwargs ) - # Apply input processor if provided + + # Determine the value for the 'inputs' attribute if input_processor is not None: - processed_inputs = input_processor(json.loads(inputs)) - inputs = json.dumps(processed_inputs, default=str) - span.set_attribute("inputs", inputs) + processed_inputs_val = input_processor(json.loads(inputs_json_str)) + inputs_attribute_value = json.dumps( + processed_inputs_val, default=str + ) + else: + inputs_attribute_value = ( + inputs_json_str # Use raw JSON string if no processor + ) + span.set_attribute("inputs", inputs_attribute_value) + + # Add dependency information as a JSON attribute if provided + if dependency is not None: + processed_dependency = {} + parsed_args_dict = None + if any(callable(v) for v in dependency.values()): + parsed_args_dict = json.loads(inputs_json_str) + + for key, value in dependency.items(): + if callable(value): + processed_dependency[key] = value(parsed_args_dict) + else: + processed_dependency[key] = value + + dependency_json = json.dumps(processed_dependency, default=str) + span.set_attribute("dependency", dependency_json) + try: result = await func(*args, **kwargs) # Process output if processor is provided @@ -217,14 +280,38 @@ def generator_wrapper(*args, **kwargs): span.set_attribute("run_type", run_type) # Format arguments for tracing - inputs = _SpanUtils.format_args_for_trace_json( + inputs_json_str = _SpanUtils.format_args_for_trace_json( # Store the raw JSON string of args inspect.signature(func), *args, **kwargs ) - # Apply input processor if provided + + # Determine the value for the 'inputs' attribute if input_processor is not None: - processed_inputs = input_processor(json.loads(inputs)) - inputs = json.dumps(processed_inputs, default=str) - span.set_attribute("inputs", inputs) + processed_inputs_val = input_processor(json.loads(inputs_json_str)) + inputs_attribute_value = json.dumps( + processed_inputs_val, default=str + ) + else: + inputs_attribute_value = ( + inputs_json_str # Use raw JSON string if no processor + ) + span.set_attribute("inputs", inputs_attribute_value) + + # Add dependency information as a JSON attribute if provided + if dependency is not None: + processed_dependency = {} + parsed_args_dict = None + if any(callable(v) for v in dependency.values()): + parsed_args_dict = json.loads(inputs_json_str) + + for key, value in dependency.items(): + if callable(value): + processed_dependency[key] = value(parsed_args_dict) + else: + processed_dependency[key] = value + + dependency_json = json.dumps(processed_dependency, default=str) + span.set_attribute("dependency", dependency_json) + outputs = [] try: for item in func(*args, **kwargs): @@ -258,14 +345,38 @@ async def async_generator_wrapper(*args, **kwargs): span.set_attribute("run_type", run_type) # Format arguments for tracing - inputs = _SpanUtils.format_args_for_trace_json( + inputs_json_str = _SpanUtils.format_args_for_trace_json( # Store the raw JSON string of args inspect.signature(func), *args, **kwargs ) - # Apply input processor if provided + + # Determine the value for the 'inputs' attribute if input_processor is not None: - processed_inputs = input_processor(json.loads(inputs)) - inputs = json.dumps(processed_inputs, default=str) - span.set_attribute("inputs", inputs) + processed_inputs_val = input_processor(json.loads(inputs_json_str)) + inputs_attribute_value = json.dumps( + processed_inputs_val, default=str + ) + else: + inputs_attribute_value = ( + inputs_json_str # Use raw JSON string if no processor + ) + span.set_attribute("inputs", inputs_attribute_value) + + # Add dependency information as a JSON attribute if provided + if dependency is not None: + processed_dependency = {} + parsed_args_dict = None + if any(callable(v) for v in dependency.values()): + parsed_args_dict = json.loads(inputs_json_str) + + for key, value in dependency.items(): + if callable(value): + processed_dependency[key] = value(parsed_args_dict) + else: + processed_dependency[key] = value + + dependency_json = json.dumps(processed_dependency, default=str) + span.set_attribute("dependency", dependency_json) + outputs = [] try: async for item in func(*args, **kwargs): @@ -337,10 +448,12 @@ def traced( output_processor: Optional[Callable[..., Any]] = None, hide_input: bool = False, hide_output: bool = False, + dependency: Optional[DependencyInfo] = None, ): """Decorator that will trace function invocations. Args: + name: Optional name for the span (defaults to function name) run_type: Optional string to categorize the run type span_type: Optional string to categorize the span type input_processor: Optional function to process function inputs before recording @@ -349,6 +462,15 @@ def traced( Should accept the function output and return a processed value hide_input: If True, don't log any input data hide_output: If True, don't log any output data + dependency: Optional dictionary with dependency tracking information: + sourceName: The source system/component (str or callable returning str) + targetName: The target system/component (str or callable returning str) + operationName: The operation being performed (str or callable returning str) + + For sourceName, targetName, and operationName: + - If a string is provided, it's used directly + - If a callable is provided, it's called with the same args/kwargs as the + decorated function and should return a string """ # Apply default processors selectively based on hide flags if hide_input: @@ -356,6 +478,18 @@ def traced( if hide_output: output_processor = _default_output_processor + if dependency is not None: + if dependency.get("sourceName") is None: + # Updated to accept a single dictionary argument, though it's not used by default + dependency["sourceName"] = lambda _dep_args_dict: os.environ.get( + "UIPATH_PROCESS_KEY", "Unknown source" + ) + + if dependency.get("sourceType") is None: + dependency["sourceType"] = ( + "Agent" # Corrected from sourceName to sourceType + ) + # Store the parameters for later reapplication params = { "name": name, @@ -363,6 +497,7 @@ def traced( "span_type": span_type, "input_processor": input_processor, "output_processor": output_processor, + "dependency": dependency, } # Check for custom implementation first diff --git a/src/uipath/tracing/_utils.py b/src/uipath/tracing/_utils.py index 654a4a3b..cf7e69b7 100644 --- a/src/uipath/tracing/_utils.py +++ b/src/uipath/tracing/_utils.py @@ -11,10 +11,18 @@ from opentelemetry.sdk.trace import ReadableSpan from opentelemetry.trace import StatusCode +from pydantic import BaseModel logger = logging.getLogger(__name__) +def _custom_encoder(obj): + if isinstance(obj, BaseModel): + return obj.model_dump() # Convert Pydantic models to dictionaries + + return str(obj) # Fallback for other types + + @dataclass class UiPathSpan: """Represents a span in the UiPath tracing system.""" @@ -150,30 +158,10 @@ def otel_span_to_uipath_span(otel_span: ReadableSpan) -> UiPathSpan: status = 2 # Error attributes_dict["error"] = otel_span.status.description - original_inputs = attributes_dict.get("inputs", None) - original_outputs = attributes_dict.get("outputs", None) - - if original_inputs: - try: - if isinstance(original_inputs, str): - json_inputs = json.loads(original_inputs) - attributes_dict["inputs"] = json_inputs - else: - attributes_dict["inputs"] = original_inputs - except Exception as e: - print(f"Error parsing inputs: {e}") - attributes_dict["inputs"] = str(original_inputs) - - if original_outputs: - try: - if isinstance(original_outputs, str): - json_outputs = json.loads(original_outputs) - attributes_dict["outputs"] = json_outputs - else: - attributes_dict["outputs"] = original_outputs - except Exception as e: - print(f"Error parsing outputs: {e}") - attributes_dict["outputs"] = str(original_outputs) + # Parse JSON attributes using the helper method + for attr_name in ["inputs", "outputs", "dependency"]: + if attr_name in attributes_dict: + _SpanUtils._process_json_attribute(attributes_dict, attr_name) # Add events as additional attributes if they exist if otel_span.events: @@ -208,13 +196,22 @@ def otel_span_to_uipath_span(otel_span: ReadableSpan) -> UiPathSpan: ).isoformat() end_time_str = None + duration = None # Initialize duration as None + if otel_span.end_time is not None: end_time_str = datetime.fromtimestamp( (otel_span.end_time or 0) / 1e9 ).isoformat() + # Calculate duration in seconds if both start and end times are available + if otel_span.start_time is not None: + duration = (otel_span.end_time - otel_span.start_time) / 1e9 else: end_time_str = datetime.now().isoformat() + # Add duration to attributes if available + if duration is not None: + attributes_dict["duration"] = duration + return UiPathSpan( id=span_id, trace_id=trace_id, @@ -233,7 +230,7 @@ def format_args_for_trace_json( ) -> str: """Return a JSON string of inputs from the function signature.""" result = _SpanUtils.format_args_for_trace(signature, *args, **kwargs) - return json.dumps(result, default=str) + return json.dumps(result, default=_custom_encoder) @staticmethod def format_args_for_trace( @@ -271,3 +268,18 @@ def format_args_for_trace( f"Error formatting arguments for trace: {e}. Using args and kwargs directly." ) return {"args": args, "kwargs": kwargs} + + @staticmethod + def _process_json_attribute(attributes_dict, attr_name): + original_value = attributes_dict.get(attr_name) + + if not original_value: + return original_value + + try: + if isinstance(original_value, str): + json_data = json.loads(original_value) + attributes_dict[attr_name] = json_data + + except Exception as e: + logger.warning(f"Error parsing {attr_name}: {e}") diff --git a/tests/tracing/test_traced.py b/tests/tracing/test_traced.py index 252cf4af..62a00185 100644 --- a/tests/tracing/test_traced.py +++ b/tests/tracing/test_traced.py @@ -72,9 +72,83 @@ def sample_function(x, y): assert span.attributes["span_type"] == "function_call_sync" assert "inputs" in span.attributes assert "output" in span.attributes + # No dependency attribute by default + assert "dependency" not in span.attributes assert span.attributes["output"] == "5" +@traced( + dependency={ + "targetName": "TestTarget", + "targetType": "TestType", + "targetId": "TestId", + "operationName": "TestOperation", + } +) +def sample_function_with_dependency(x, y): + return x + y + + +def test_traced_sync_function_with_dependency(setup_tracer): + exporter, provider = setup_tracer + + result = sample_function_with_dependency(2, 3) + assert result == 5 + + provider.shutdown() + spans = exporter.get_exported_spans() + + assert len(spans) == 1 + span = spans[0] + assert span.name == "sample_function_with_dependency" + assert span.attributes["span_type"] == "function_call_sync" + assert "inputs" in span.attributes + assert "output" in span.attributes + assert span.attributes["output"] == "5" + assert "dependency" in span.attributes + dependency_attr = json.loads(span.attributes["dependency"]) + assert dependency_attr["targetName"] == "TestTarget" + assert dependency_attr["targetType"] == "TestType" + assert dependency_attr["targetId"] == "TestId" + assert dependency_attr["operationName"] == "TestOperation" + assert dependency_attr["sourceName"] == "Unknown source" # Default + assert dependency_attr["sourceType"] == "Agent" # Default + + +@traced( + dependency={ + "targetName": lambda inputs: inputs["target"], + "targetType": lambda inputs: inputs["type"], + "targetId": lambda inputs: inputs["id"], + "operationName": "DynamicOperation", + } +) +def sample_function_with_callable_dependency(target, type, id): + return f"{target}-{type}-{id}" + + +def test_traced_sync_function_with_callable_dependency(setup_tracer): + exporter, provider = setup_tracer + + result = sample_function_with_callable_dependency("MyTarget", "MyType", "MyId") + assert result == "MyTarget-MyType-MyId" + + provider.shutdown() + spans = exporter.get_exported_spans() + + assert len(spans) == 1 + span = spans[0] + assert span.name == "sample_function_with_callable_dependency" + assert "dependency" in span.attributes + dependency_attr = json.loads(span.attributes["dependency"]) + assert dependency_attr["targetName"] == "MyTarget" + assert dependency_attr["targetType"] == "MyType" + assert dependency_attr["targetId"] == "MyId" + assert dependency_attr["operationName"] == "DynamicOperation" + assert dependency_attr["sourceName"] == "Unknown source" + assert dependency_attr["sourceType"] == "Agent" + + @pytest.mark.asyncio async def test_traced_async_function(setup_tracer): exporter, provider = setup_tracer @@ -97,9 +171,47 @@ async def sample_async_function(x, y): assert span.attributes["span_type"] == "function_call_async" assert "inputs" in span.attributes assert "output" in span.attributes + # No dependency attribute by default + assert "dependency" not in span.attributes assert span.attributes["output"] == "6" +@traced( + dependency={ + "targetName": "AsyncTarget", + "targetType": "AsyncType", + "targetId": "AsyncId", + "operationName": "AsyncOperation", + } +) +async def sample_async_function_with_dependency(x, y): + return x * y + + +@pytest.mark.asyncio +async def test_traced_async_function_with_dependency(setup_tracer): + exporter, provider = setup_tracer + + result = await sample_async_function_with_dependency(2, 3) + assert result == 6 + + provider.shutdown() + await sleep(0.1) # allow time for export + spans = exporter.get_exported_spans() + + assert len(spans) == 1 + span = spans[0] + assert span.name == "sample_async_function_with_dependency" + assert "dependency" in span.attributes + dependency_attr = json.loads(span.attributes["dependency"]) + assert dependency_attr["targetName"] == "AsyncTarget" + assert dependency_attr["targetType"] == "AsyncType" + assert dependency_attr["targetId"] == "AsyncId" + assert dependency_attr["operationName"] == "AsyncOperation" + assert dependency_attr["sourceName"] == "Unknown source" + assert dependency_attr["sourceType"] == "Agent" + + def test_traced_generator_function(setup_tracer): exporter, provider = setup_tracer @@ -120,9 +232,46 @@ def sample_generator_function(n): assert span.attributes["span_type"] == "function_call_generator_sync" assert "inputs" in span.attributes assert "output" in span.attributes + # No dependency attribute by default + assert "dependency" not in span.attributes assert span.attributes["output"] == "[0, 1, 2]" +@traced( + dependency={ + "targetName": "GenTarget", + "targetType": "GenType", + "targetId": "GenId", + "operationName": "GenOperation", + } +) +def sample_generator_function_with_dependency(n): + for i in range(n): + yield i + + +def test_traced_generator_function_with_dependency(setup_tracer): + exporter, provider = setup_tracer + + results = list(sample_generator_function_with_dependency(3)) + assert results == [0, 1, 2] + + provider.shutdown() + spans = exporter.get_exported_spans() + + assert len(spans) == 1 + span = spans[0] + assert span.name == "sample_generator_function_with_dependency" + assert "dependency" in span.attributes + dependency_attr = json.loads(span.attributes["dependency"]) + assert dependency_attr["targetName"] == "GenTarget" + assert dependency_attr["targetType"] == "GenType" + assert dependency_attr["targetId"] == "GenId" + assert dependency_attr["operationName"] == "GenOperation" + assert dependency_attr["sourceName"] == "Unknown source" + assert dependency_attr["sourceType"] == "Agent" + + @pytest.mark.asyncio async def test_traced_async_generator_function(setup_tracer): exporter, provider = setup_tracer @@ -144,9 +293,49 @@ async def sample_async_generator_function(n): assert span.attributes["span_type"] == "function_call_generator_async" assert "inputs" in span.attributes assert "output" in span.attributes + # No dependency attribute by default + assert "dependency" not in span.attributes assert span.attributes["output"] == "[0, 1, 2]" +@traced( + dependency={ + "targetName": "AsyncGenTarget", + "targetType": "AsyncGenType", + "targetId": "AsyncGenId", + "operationName": "AsyncGenOperation", + } +) +async def sample_async_generator_function_with_dependency(n): + for i in range(n): + yield i + + +@pytest.mark.asyncio +async def test_traced_async_generator_function_with_dependency(setup_tracer): + exporter, provider = setup_tracer + + results = [ + item async for item in sample_async_generator_function_with_dependency(3) + ] + assert results == [0, 1, 2] + + provider.shutdown() + spans = exporter.get_exported_spans() + + assert len(spans) == 1 + span = spans[0] + assert span.name == "sample_async_generator_function_with_dependency" + assert "dependency" in span.attributes + dependency_attr = json.loads(span.attributes["dependency"]) + assert dependency_attr["targetName"] == "AsyncGenTarget" + assert dependency_attr["targetType"] == "AsyncGenType" + assert dependency_attr["targetId"] == "AsyncGenId" + assert dependency_attr["operationName"] == "AsyncGenOperation" + assert dependency_attr["sourceName"] == "Unknown source" + assert dependency_attr["sourceType"] == "Agent" + + def test_traced_with_basic_processors(setup_tracer): """Test traced decorator with basic input and output processors.""" exporter, provider = setup_tracer @@ -572,3 +761,170 @@ def fully_private_function(sensitive_input): output_json = span.attributes["output"] output = json.loads(output_json) assert output == {"redacted": "Output data not logged for privacy/security"} + + +def test_traced_with_static_dependency_info(setup_tracer): + """Test traced decorator with static dependency information.""" + exporter, provider = setup_tracer + + dependency_info = { + "sourceName": "TestSource", + "targetName": "TestTarget", + "operationName": "TestOperation", + } + + @traced(dependency=dependency_info) + def function_with_dependency(x, y): + return x + y + + result = function_with_dependency(10, 20) + assert result == 30 + + provider.shutdown() + spans = exporter.get_exported_spans() + + assert len(spans) == 1 + span = spans[0] + assert span.name == "function_with_dependency" + assert "dependency" in span.attributes + + dependency_attr = json.loads(span.attributes["dependency"]) + assert dependency_attr["sourceName"] == "TestSource" + assert dependency_attr["targetName"] == "TestTarget" + assert dependency_attr["operationName"] == "TestOperation" + + +def test_traced_with_callable_dependency_info(setup_tracer): + """Test traced decorator with callable dependency information.""" + exporter, provider = setup_tracer + + def get_source_name(inputs: Dict[str, Any]): + return f"Source_{inputs['x']}" + + def get_target_name(inputs: Dict[str, Any]): + return f"Target_{inputs['y']}" + + dependency_info = { + "sourceName": get_source_name, + "targetName": get_target_name, + "operationName": "CallableOperation", + } + + @traced(dependency=dependency_info) + def function_with_callable_dependency(x, y): + return x * y + + result = function_with_callable_dependency( + 5, y=4 + ) # Use keyword arg for get_target_name + assert result == 20 + + provider.shutdown() + spans = exporter.get_exported_spans() + + assert len(spans) == 1 + span = spans[0] + assert span.name == "function_with_callable_dependency" + assert "dependency" in span.attributes + + dependency_attr = json.loads(span.attributes["dependency"]) + assert dependency_attr["sourceName"] == "Source_5" + assert dependency_attr["targetName"] == "Target_4" + assert dependency_attr["operationName"] == "CallableOperation" + + +def test_traced_with_dependency_and_env_var(setup_tracer, monkeypatch): + """Test traced decorator with dependency using UIPATH_PROCESS_KEY env var.""" + exporter, provider = setup_tracer + monkeypatch.setenv("UIPATH_PROCESS_KEY", "EnvProcessKey") + + dependency_info = { + "targetName": "EnvTarget", + "operationName": "EnvOperation", + } # sourceName will be picked from env var + + @traced(dependency=dependency_info) + def function_with_env_dependency(a): + return f"Processed {a}" + + result = function_with_env_dependency("data") + assert result == "Processed data" + + provider.shutdown() + spans = exporter.get_exported_spans() + + assert len(spans) == 1 + span = spans[0] + assert span.name == "function_with_env_dependency" + assert "dependency" in span.attributes + + dependency_attr = json.loads(span.attributes["dependency"]) + assert dependency_attr["sourceName"] == "EnvProcessKey" + assert dependency_attr["targetName"] == "EnvTarget" + assert dependency_attr["operationName"] == "EnvOperation" + + +@pytest.mark.asyncio +async def test_traced_async_with_dependency_info(setup_tracer): + """Test traced decorator with dependency information for async functions.""" + exporter, provider = setup_tracer + + dependency_info = { + "sourceName": "AsyncSource", + "targetName": "AsyncTarget", + "operationName": lambda inputs: f"AsyncOp_{inputs['msg']}", + } + + @traced(dependency=dependency_info) + async def async_function_with_dependency(msg): + await sleep(0.01) + return f"Async says: {msg}" + + result = await async_function_with_dependency("hello") + assert result == "Async says: hello" + + provider.shutdown() + await sleep(0.1) # ensure spans are processed + spans = exporter.get_exported_spans() + + assert len(spans) == 1 + span = spans[0] + assert span.name == "async_function_with_dependency" + assert "dependency" in span.attributes + + dependency_attr = json.loads(span.attributes["dependency"]) + assert dependency_attr["sourceName"] == "AsyncSource" + assert dependency_attr["targetName"] == "AsyncTarget" + assert dependency_attr["operationName"] == "AsyncOp_hello" + + +def test_traced_generator_with_dependency_info(setup_tracer): + """Test traced decorator with dependency information for generator functions.""" + exporter, provider = setup_tracer + + dependency_info = { + "sourceName": "GeneratorSource", + "targetName": "GeneratorTarget", + "operationName": lambda inputs: f"GenOp_Count_{inputs['n_items']}", + } + + @traced(dependency=dependency_info) + def generator_with_dependency(n_items): + for i in range(n_items): + yield f"Item {i}" + + results = list(generator_with_dependency(2)) + assert results == ["Item 0", "Item 1"] + + provider.shutdown() + spans = exporter.get_exported_spans() + + assert len(spans) == 1 + span = spans[0] + assert span.name == "generator_with_dependency" + assert "dependency" in span.attributes + + dependency_attr = json.loads(span.attributes["dependency"]) + assert dependency_attr["sourceName"] == "GeneratorSource" + assert dependency_attr["targetName"] == "GeneratorTarget" + assert dependency_attr["operationName"] == "GenOp_Count_2"