wip

Author: ionmincu

src/uipath_langchain/chat/models.py

@@ -165,13 +165,33 @@ def _generate(
                   downstream and understand why generation stopped.
             run_manager: A run manager with callbacks for the LLM.
         """
-        if kwargs.get("tools"):
-            kwargs["tools"] = [tool["function"] for tool in kwargs["tools"]]
-        if "tool_choice" in kwargs and kwargs["tool_choice"]["type"] == "function":
-            kwargs["tool_choice"] = {
-                "type": "tool",
-                "name": kwargs["tool_choice"]["function"]["name"],
-            }
+        if tools := kwargs.get("tools"):
+            kwargs["tools"] = [tool["function"] for tool in tools]
+            tool_names = [tool["function"]["name"] for tool in tools]
+            if tool_choice := kwargs.get("tool_choice"):
+                if isinstance(tool_choice, str):
+                    if tool_choice in ["auto", "required"]:
+                        kwargs["tool_choice"] = {"type": tool_choice}
+                    elif tool_choice in tool_names:
+                        kwargs["tool_choice"] = {"type": "tool", "name": tool_choice}
+                    else:
+                        kwargs["tool_choice"] = "required"
+                elif (
+                    isinstance(tool_choice, dict)
+                    and "type" in tool_choice
+                    and tool_choice["type"] == "function"
+                    and "function" in tool_choice
+                    and isinstance(tool_choice["function"], dict)
+                    and "name" in tool_choice["function"]
+                    and tool_choice["function"]["name"] in tool_names
+                ):
+                    kwargs["tool_choice"] = {
+                        "type": "tool",
+                        "name": tool_choice["function"]["name"],
+                    }
+                else:
+                    kwargs["tool_choice"] = {"type": "required"}
+
         payload = self._get_request_payload(messages, stop=stop, **kwargs)
 
         response = self._call(self.url, payload, self.auth_headers)
@@ -199,13 +219,33 @@ async def _agenerate(
                   downstream and understand why generation stopped.
             run_manager: A run manager with callbacks for the LLM.
         """
-        if kwargs.get("tools"):
-            kwargs["tools"] = [tool["function"] for tool in kwargs["tools"]]
-        if "tool_choice" in kwargs and kwargs["tool_choice"]["type"] == "function":
-            kwargs["tool_choice"] = {
-                "type": "tool",
-                "name": kwargs["tool_choice"]["function"]["name"],
-            }
+        if tools := kwargs.get("tools"):
+            kwargs["tools"] = [tool["function"] for tool in tools]
+            tool_names = [tool["function"]["name"] for tool in tools]
+            if tool_choice := kwargs.get("tool_choice"):
+                if isinstance(tool_choice, str):
+                    if tool_choice in ["auto", "required"]:
+                        kwargs["tool_choice"] = {"type": tool_choice}
+                    elif tool_choice in tool_names:
+                        kwargs["tool_choice"] = {"type": "tool", "name": tool_choice}
+                    else:
+                        kwargs["tool_choice"] = "required"
+                elif (
+                    isinstance(tool_choice, dict)
+                    and "type" in tool_choice
+                    and tool_choice["type"] == "function"
+                    and "function" in tool_choice
+                    and isinstance(tool_choice["function"], dict)
+                    and "name" in tool_choice["function"]
+                    and tool_choice["function"]["name"] in tool_names
+                ):
+                    kwargs["tool_choice"] = {
+                        "type": "tool",
+                        "name": tool_choice["function"]["name"],
+                    }
+                else:
+                    kwargs["tool_choice"] = {"type": "required"}
+
         payload = self._get_request_payload(messages, stop=stop, **kwargs)
 
         response = await self._acall(self.url, payload, self.auth_headers)

tests/cli_run/samples/1-simple-graph/langgraph.json

@@ -0,0 +1,9 @@
+{
+  "dependencies": [
+    "."
+  ],
+  "graphs": {
+    "agent": "./main.py:graph"
+  },
+  "env": ".env"
+}

tests/cli_run/samples/1-simple-graph/main.py

@@ -0,0 +1,59 @@
+import random
+from typing import Literal
+
+from langgraph.checkpoint.memory import MemorySaver
+from langgraph.graph import END, START, StateGraph
+from langgraph.types import interrupt
+from typing_extensions import TypedDict
+
+
+# State
+class State(TypedDict):
+    graph_state: str
+
+
+# Conditional edge
+def decide_mood(state) -> Literal["node_2", "node_3"]:
+    # Often, we will use state to decide on the next node to visit
+    user_input = state["graph_state"]
+
+    # Here, let's just do a 50 / 50 split between nodes 2, 3
+    if random.random() < 0.5:
+        # 50% of the time, we return Node 2
+        return "node_2"
+
+    # 50% of the time, we return Node 3
+    return "node_3"
+
+
+# Nodes
+def node_1(state):
+    print("---Node 1---")
+
+    return {"graph_state": state["graph_state"] + " I am"}
+
+
+def node_2(state):
+    print("---Node 2---")
+    return {"graph_state": state["graph_state"] + " happy!"}
+
+
+def node_3(state):
+    print("---Node 3---")
+    return {"graph_state": state["graph_state"] + " sad!"}
+
+
+builder = StateGraph(State)
+builder.add_node("node_1", node_1)
+builder.add_node("node_2", node_2)
+builder.add_node("node_3", node_3)
+
+builder.add_edge(START, "node_1")
+builder.add_edge("node_1", "node_2")
+builder.add_edge("node_2", END)
+builder.add_edge("node_3", END)
+
+
+memory = MemorySaver()
+
+graph = builder.compile(checkpointer=memory)

tests/cli_run/samples/1-simple-graph/pyproject.toml

@@ -0,0 +1,15 @@
+[project]
+name = "c-host-in-uipath"
+version = "0.1.0"
+description = "Add your description here"
+readme = "README.md"
+authors = [{ name = "Eduard Stanculet", email = "[email protected]" }]
+requires-python = ">=3.13"
+dependencies = [
+    "langchain-anthropic>=0.3.10",
+    "langchain-community>=0.3.21",
+    "langgraph>=0.3.29",
+    "tavily-python>=0.5.4",
+    "uipath>=2.0.8",
+    "uipath-langchain>=0.0.88",
+]

tests/cli_run/samples/1-simple-graph/uipath.json

@@ -0,0 +1,18 @@
+{
+  "entryPoints": [
+    {
+      "filePath": "agent",
+      "uniqueId": "dcc7a309-fbcc-4999-af4f-2a75a844b49a",
+      "type": "agent",
+      "input": {
+        "type": "string",
+        "title": "graph_state"
+      },
+      "output": {}
+    }
+  ],
+  "bindings": {
+    "version": "2.0",
+    "resources": []
+  }
+}

tests/cli_run/test_run_sample.py

@@ -0,0 +1,23 @@
+import os
+import sys
+
+from dotenv import load_dotenv
+
+from uipath_langchain._cli.cli_run import langgraph_run_middleware
+
+load_dotenv()
+
+
+def test_dummy():
+    test_folder_path = os.path.dirname(os.path.abspath(__file__))
+    sample_path = os.path.join(test_folder_path, "samples", "1-simple-graph")
+
+    sys.path.append(sample_path)
+    os.chdir(sample_path)
+    result = langgraph_run_middleware(
+        entrypoint=None,
+        input='{ "graph_state": "GET Assets API does not enforce proper permissions Assets.View" }',
+        resume=False,
+    )
+
+    assert result.error_message is None