Merge pull request #298 from UiPath/akshaya/update_output_schema

feat(ReturnDicts): tools to return dicts

Author: akshaylive

pyproject.toml

@@ -1,12 +1,12 @@
 [project]
 name = "uipath-langchain"
-version = "0.1.8"
+version = "0.1.9"
 description = "UiPath Langchain"
 readme = { file = "README.md", content-type = "text/markdown" }
 requires-python = ">=3.11"
 dependencies = [
     "uipath-runtime>=0.1.2, <0.2.0",
-    "uipath>=2.2.8, <2.3.0",
+    "uipath>=2.2.11, <2.3.0",
     "langgraph>=1.0.0, <2.0.0",
     "langchain-core>=1.0.0, <2.0.0",
     "langgraph-checkpoint-sqlite>=3.0.0, <4.0.0",

src/uipath_langchain/agent/tools/context_tool.py

@@ -2,6 +2,8 @@
 
 from __future__ import annotations
 
+from typing import Any
+
 from langchain_core.documents import Document
 from langchain_core.tools import StructuredTool
 from pydantic import BaseModel, Field
@@ -40,11 +42,8 @@ class ContextOutputSchemaModel(BaseModel):
         input_schema=input_model.model_json_schema(),
         output_schema=output_model.model_json_schema(),
     )
-    async def context_tool_fn(query: str) -> ContextOutputSchemaModel:
-        documents = await retriever.ainvoke(query)
-        return ContextOutputSchemaModel(
-            documents=documents,
-        )
+    async def context_tool_fn(query: str) -> dict[str, Any]:
+        return {"documents": await retriever.ainvoke(query)}
 
     return StructuredTool(
         name=tool_name,

src/uipath_langchain/agent/tools/integration_tool.py

@@ -6,7 +6,6 @@
 
 from jsonschema_pydantic_converter import transform as create_model
 from langchain_core.tools import StructuredTool
-from pydantic import TypeAdapter
 from uipath.agent.models.agent import AgentIntegrationToolResourceConfig
 from uipath.eval.mocks import mockable
 from uipath.platform import UiPath
@@ -137,17 +136,12 @@ def sanitize_for_serialization(args: dict[str, Any]) -> dict[str, Any]:
         input_schema=input_model.model_json_schema(),
         output_schema=output_model.model_json_schema(),
     )
-    async def integration_tool_fn(**kwargs: Any) -> output_model:
-        try:
-            result = await sdk.connections.invoke_activity_async(
-                activity_metadata=activity_metadata,
-                connection_id=connection_id,
-                activity_input=sanitize_for_serialization(kwargs),
-            )
-        except Exception:
-            raise
-
-        return TypeAdapter(output_model).validate_python(result)
+    async def integration_tool_fn(**kwargs: Any):
+        return await sdk.connections.invoke_activity_async(
+            activity_metadata=activity_metadata,
+            connection_id=connection_id,
+            activity_input=sanitize_for_serialization(kwargs),
+        )
 
     tool = StructuredTool(
         name=tool_name,
@@ -156,6 +150,4 @@ async def integration_tool_fn(**kwargs: Any) -> output_model:
         coroutine=integration_tool_fn,
     )
 
-    tool.__dict__["OutputType"] = output_model
-
     return tool

src/uipath_langchain/agent/tools/process_tool.py

@@ -6,8 +6,6 @@
 
 from jsonschema_pydantic_converter import transform as create_model
 from langchain_core.tools import StructuredTool
-from langgraph.types import interrupt
-from pydantic import TypeAdapter
 from uipath.agent.models.agent import AgentProcessToolResourceConfig
 from uipath.eval.mocks import mockable
 from uipath.platform.common import InvokeProcess
@@ -30,25 +28,19 @@ def create_process_tool(resource: AgentProcessToolResourceConfig) -> StructuredT
         input_schema=input_model.model_json_schema(),
         output_schema=output_model.model_json_schema(),
     )
-    async def process_tool_fn(**kwargs: Any) -> output_model:
-        result = interrupt(
-            InvokeProcess(
-                name=process_name,
-                input_arguments=kwargs,
-                process_folder_path=folder_path,
-                process_folder_key=None,
-            )
+    async def process_tool_fn(**kwargs: Any):
+        return InvokeProcess(
+            name=process_name,
+            input_arguments=kwargs,
+            process_folder_path=folder_path,
+            process_folder_key=None,
         )
 
-        return TypeAdapter(output_model).validate_python(result)
-
     tool = StructuredTool(
         name=tool_name,
         description=resource.description,
         args_schema=input_model,
         coroutine=process_tool_fn,
     )
 
-    tool.__dict__["OutputType"] = output_model
-
     return tool