diff --git a/src/uipath_langchain/agent/tools/context_tool.py b/src/uipath_langchain/agent/tools/context_tool.py index d4364481..95a4e4a7 100644 --- a/src/uipath_langchain/agent/tools/context_tool.py +++ b/src/uipath_langchain/agent/tools/context_tool.py @@ -77,6 +77,7 @@ class ContextOutputSchemaModel(BaseModel): input_schema=input_model, output_schema=output_model.model_json_schema(), example_calls=[], # Examples cannot be provided for context. + recording=False, ) async def context_tool_fn() -> dict[str, Any]: return {"documents": await retriever.ainvoke(static_query_value)} @@ -96,6 +97,7 @@ class ContextInputSchemaModel(BaseModel): input_schema=input_model.model_json_schema(), output_schema=output_model.model_json_schema(), example_calls=[], # Examples cannot be provided for context. + recording=False, ) async def context_tool_fn(query: str) -> dict[str, Any]: return {"documents": await retriever.ainvoke(query)} @@ -107,8 +109,15 @@ async def context_tool_fn(query: str) -> dict[str, Any]: coroutine=context_tool_fn, output_type=output_model, metadata={ - "tool_type": "context", + "tool_type": "context_grounding", "display_name": resource.name, + "retrieval_mode": "SemanticSearch", + "number_of_results": resource.settings.result_count, + **( + {"static_query": static_query_value} + if is_static_query(resource) + else {} + ), }, ) @@ -145,6 +154,7 @@ def handle_deep_rag( input_schema=input_model, output_schema=output_model.model_json_schema(), example_calls=[], # Examples cannot be provided for context. + recording=False, ) async def context_tool_fn() -> dict[str, Any]: # TODO: add glob pattern support @@ -173,6 +183,7 @@ class DeepRagInputSchemaModel(BaseModel): input_schema=input_model.model_json_schema(), output_schema=output_model.model_json_schema(), example_calls=[], # Examples cannot be provided for context. + recording=False, ) async def context_tool_fn(query: str) -> dict[str, Any]: # TODO: add glob pattern support @@ -192,8 +203,11 @@ async def context_tool_fn(query: str) -> dict[str, Any]: coroutine=context_tool_fn, output_type=output_model, metadata={ - "tool_type": "context", + "tool_type": "context_grounding", "display_name": resource.name, + "retrieval_mode": "DeepRag", + "citation_mode": citation_mode.value, + **({"static_query": static_prompt} if is_static_query(resource) else {}), }, ) @@ -254,6 +268,7 @@ class StaticBatchTransformSchemaModel(BaseModel): input_schema=input_model.model_json_schema(), output_schema=output_model.model_json_schema(), example_calls=[], # Examples cannot be provided for context. + recording=False, ) async def context_tool_fn( destination_path: str = "output.csv", @@ -291,6 +306,7 @@ class DynamicBatchTransformSchemaModel(BaseModel): input_schema=input_model.model_json_schema(), output_schema=output_model.model_json_schema(), example_calls=[], # Examples cannot be provided for context. + recording=False, ) async def context_tool_fn( query: str, destination_path: str = "output.csv" @@ -315,8 +331,15 @@ async def context_tool_fn( coroutine=context_tool_fn, output_type=output_model, metadata={ - "tool_type": "context", + "tool_type": "context_grounding", "display_name": resource.name, + "retrieval_mode": "BatchTransform", + "output_columns": [ + {"name": col.name, "description": col.description} + for col in batch_transform_output_columns + ], + "web_search_grounding": enable_web_search_grounding, + **({"static_query": static_prompt} if is_static_query(resource) else {}), }, ) diff --git a/src/uipath_langchain/agent/tools/internal_tools/batch_transform_tool.py b/src/uipath_langchain/agent/tools/internal_tools/batch_transform_tool.py index 76b9e68b..ab17f731 100644 --- a/src/uipath_langchain/agent/tools/internal_tools/batch_transform_tool.py +++ b/src/uipath_langchain/agent/tools/internal_tools/batch_transform_tool.py @@ -36,6 +36,38 @@ from uipath_langchain.agent.tools.tool_node import ToolWrapperReturnType from uipath_langchain.agent.tools.utils import sanitize_tool_name +# Define the output schema with job-attachment +BATCH_TRANSFORM_OUTPUT_SCHEMA = { + "type": "object", + "properties": { + "result": { + "$ref": "#/definitions/job-attachment", + "description": "The transformed result file as an attachment", + } + }, + "required": ["result"], + "definitions": { + "job-attachment": { + "type": "object", + "properties": { + "ID": {"type": "string", "description": "Orchestrator attachment key"}, + "FullName": {"type": "string", "description": "File name"}, + "MimeType": { + "type": "string", + "description": "The MIME type of the content", + }, + "Metadata": { + "type": "object", + "description": "Dictionary of metadata", + "additionalProperties": {"type": "string"}, + }, + }, + "required": ["ID", "FullName", "MimeType"], + "x-uipath-resource-kind": "JobAttachment", + } + }, +} + def create_batch_transform_tool( resource: AgentInternalToolResourceConfig, llm: BaseChatModel @@ -84,7 +116,8 @@ def create_batch_transform_tool( # Create input model from modified schema input_model = create_model(input_schema) - output_model = create_model(resource.output_schema) + # Create output model from schema with job-attachment definition + output_model = create_model(BATCH_TRANSFORM_OUTPUT_SCHEMA) async def batch_transform_tool_fn(**kwargs: Any) -> dict[str, Any]: query = kwargs.get("query") if not is_query_static else static_query @@ -144,7 +177,10 @@ async def create_ephemeral_index(): ) ) - return await invoke_batch_transform() + result_attachment = await invoke_batch_transform() + + # The resume trigger returns the attachment info directly + return {"result": result_attachment} # Import here to avoid circular dependency from uipath_langchain.agent.wrappers import get_job_attachment_wrapper @@ -167,10 +203,17 @@ async def batch_transform_tool_wrapper( output_type=output_model, argument_properties=resource.argument_properties, metadata={ - "tool_type": resource.type.lower(), + "tool_type": "context_grounding", "display_name": tool_name, "args_schema": input_model, "output_schema": output_model, + "retrieval_mode": "BatchTransform", + "output_columns": [ + {"name": col.name, "description": col.description} + for col in batch_transform_output_columns + ], + "web_search_grounding": static_web_search, + **({"static_query": static_query} if is_query_static else {}), }, ) tool.set_tool_wrappers(awrapper=batch_transform_tool_wrapper) diff --git a/src/uipath_langchain/agent/tools/internal_tools/deeprag_tool.py b/src/uipath_langchain/agent/tools/internal_tools/deeprag_tool.py index 4259b1d8..e6db5919 100644 --- a/src/uipath_langchain/agent/tools/internal_tools/deeprag_tool.py +++ b/src/uipath_langchain/agent/tools/internal_tools/deeprag_tool.py @@ -96,6 +96,7 @@ async def deeprag_tool_fn(**kwargs: Any) -> dict[str, Any]: input_schema=input_model.model_json_schema() if input_model else None, output_schema=output_model.model_json_schema(), example_calls=[], # Examples cannot be provided for internal tools + recording=False, ) async def invoke_deeprag(): @task @@ -150,10 +151,13 @@ async def deeprag_tool_wrapper( output_type=output_model, argument_properties=resource.argument_properties, metadata={ - "tool_type": resource.type.lower(), + "tool_type": "context_grounding", "display_name": tool_name, "args_schema": input_model, "output_schema": output_model, + "retrieval_mode": "DeepRag", + "citation_mode": citation_mode.value, + **({"static_query": static_query} if is_query_static else {}), }, ) tool.set_tool_wrappers(awrapper=deeprag_tool_wrapper) diff --git a/tests/agent/wrappers/test_job_attachment_wrapper.py b/tests/agent/wrappers/test_job_attachment_wrapper.py index 7d1ec671..5d4e0b68 100644 --- a/tests/agent/wrappers/test_job_attachment_wrapper.py +++ b/tests/agent/wrappers/test_job_attachment_wrapper.py @@ -43,7 +43,7 @@ def assert_command_success( result: Command[Any], tool_call_id: str = "call_123", job_attachments: dict[str, Any] | None = None, - expected_content: str | None = "{'result': 'success'}", + expected_content: str | None = '{"result": "success"}', ): """Assert that result is a successful Command with expected structure.""" if job_attachments is None: