uipath-langchain 0.0.133__py3-none-any.whl → 0.1.28__py3-none-any.whl

uipath_langchain/runtime/schema.py

@@ -0,0 +1,386 @@
+from dataclasses import dataclass
+from typing import Any, Callable, TypeVar
+
+from langchain_core.language_models.base import BaseLanguageModel
+from langchain_core.language_models.chat_models import BaseChatModel
+from langchain_core.runnables.base import Runnable
+from langchain_core.runnables.graph import Graph, Node
+from langgraph.graph.state import CompiledStateGraph
+from langgraph.prebuilt import ToolNode
+from langgraph.pregel._read import PregelNode
+from langgraph.pregel._write import ChannelWrite, ChannelWriteEntry
+from uipath.runtime.schema import (
+    UiPathRuntimeEdge,
+    UiPathRuntimeGraph,
+    UiPathRuntimeNode,
+)
+
+try:
+    from langgraph._internal._runnable import RunnableCallable
+except ImportError:
+    RunnableCallable = None  # type: ignore
+
+T = TypeVar("T")
+
+
+@dataclass
+class SchemaDetails:
+    schema: dict[str, Any]
+    has_input_circular_dependency: bool
+    has_output_circular_dependency: bool
+
+
+def _unwrap_runnable_callable(
+    runnable: Runnable[Any, Any], target_type: type[T]
+) -> T | None:
+    """Unwrap a RunnableCallable to find an instance of the target type.
+
+    Args:
+        runnable: The runnable to unwrap
+        target_type: The type to search for (e.g., BaseChatModel)
+
+    Returns:
+        Instance of target_type if found in the closure, None otherwise
+    """
+    if isinstance(runnable, target_type):
+        return runnable
+
+    if RunnableCallable is not None and isinstance(runnable, RunnableCallable):
+        func: Callable[..., Any] | None = getattr(runnable, "func", None)
+        if func is not None and hasattr(func, "__closure__") and func.__closure__:
+            for cell in func.__closure__:
+                if hasattr(cell, "cell_contents"):
+                    content = cell.cell_contents
+                    if isinstance(content, target_type):
+                        return content
+
+    return None
+
+
+def _get_node_type(node: Node) -> str:
+    """Determine the type of a LangGraph node using strongly-typed isinstance checks.
+
+    Args:
+        node: A Node object from the graph
+
+    Returns:
+        String representing the node type
+    """
+    if node.id in ("__start__", "__end__"):
+        return node.id
+
+    if node.data is None:
+        return "node"
+
+    if not isinstance(node.data, Runnable):
+        return "node"
+
+    tool_node = _unwrap_runnable_callable(node.data, ToolNode)
+    if tool_node is not None:
+        return "tool"
+
+    chat_model = _unwrap_runnable_callable(node.data, BaseChatModel)  # type: ignore[type-abstract]
+    if chat_model is not None:
+        return "model"
+
+    language_model = _unwrap_runnable_callable(node.data, BaseLanguageModel)  # type: ignore[type-abstract]
+    if language_model is not None:
+        return "model"
+
+    return "node"
+
+
+def _get_node_metadata(node: Node) -> dict[str, Any]:
+    """Extract metadata from a node in a type-safe manner.
+
+    Args:
+        node: A Node object from the graph
+
+    Returns:
+        Dictionary containing node metadata
+    """
+    if node.data is None:
+        return {}
+
+    # Early return if data is not a Runnable
+    if not isinstance(node.data, Runnable):
+        return {}
+
+    metadata: dict[str, Any] = {}
+
+    tool_node = _unwrap_runnable_callable(node.data, ToolNode)
+    if tool_node is not None:
+        if hasattr(tool_node, "_tools_by_name"):
+            tools_by_name = tool_node._tools_by_name
+            metadata["tool_names"] = list(tools_by_name.keys())
+            metadata["tool_count"] = len(tools_by_name)
+        return metadata
+
+    chat_model = _unwrap_runnable_callable(node.data, BaseChatModel)  # type: ignore[type-abstract]
+    if chat_model is not None:
+        if hasattr(chat_model, "model") and isinstance(chat_model.model, str):
+            metadata["model_name"] = chat_model.model
+        elif hasattr(chat_model, "model_name") and chat_model.model_name:
+            metadata["model_name"] = chat_model.model_name
+
+        if hasattr(chat_model, "temperature") and chat_model.temperature is not None:
+            metadata["temperature"] = chat_model.temperature
+
+        if hasattr(chat_model, "max_tokens") and chat_model.max_tokens is not None:
+            metadata["max_tokens"] = chat_model.max_tokens
+        elif (
+            hasattr(chat_model, "max_completion_tokens")
+            and chat_model.max_completion_tokens is not None
+        ):
+            metadata["max_tokens"] = chat_model.max_completion_tokens
+
+    return metadata
+
+
+def _convert_graph_to_uipath(graph: Graph) -> UiPathRuntimeGraph:
+    """Helper to convert a LangGraph Graph object to UiPathRuntimeGraph.
+
+    Args:
+        graph: A LangGraph Graph object (from get_graph() call)
+
+    Returns:
+        UiPathRuntimeGraph with nodes and edges
+    """
+    nodes: list[UiPathRuntimeNode] = []
+    for _, node in graph.nodes.items():
+        nodes.append(
+            UiPathRuntimeNode(
+                id=node.id,
+                name=node.name or node.id,
+                type=_get_node_type(node),
+                metadata=_get_node_metadata(node),
+                subgraph=None,
+            )
+        )
+
+    edges: list[UiPathRuntimeEdge] = []
+    for edge in graph.edges:
+        edges.append(
+            UiPathRuntimeEdge(
+                source=edge.source,
+                target=edge.target,
+                label=getattr(edge, "data", None) or getattr(edge, "label", None),
+            )
+        )
+
+    return UiPathRuntimeGraph(nodes=nodes, edges=edges)
+
+
+def get_graph_schema(
+    compiled_graph: CompiledStateGraph[Any, Any, Any, Any], xray: int = 1
+) -> UiPathRuntimeGraph:
+    """Convert a compiled LangGraph to UiPathRuntimeGraph structure.
+
+    Args:
+        compiled_graph: A compiled LangGraph (Pregel instance)
+        xray: Depth of subgraph expansion (0 = no subgraphs, 1 = one level, etc.)
+
+    Returns:
+        UiPathRuntimeGraph with hierarchical subgraph structure
+    """
+    graph: Graph = compiled_graph.get_graph(xray=0)  # Keep parent at xray=0
+
+    subgraphs_dict: dict[str, UiPathRuntimeGraph] = {}
+    if xray:
+        for name, subgraph_pregel in compiled_graph.get_subgraphs():
+            next_xray: int = xray - 1 if isinstance(xray, int) and xray > 0 else 0
+            subgraph_graph: Graph = subgraph_pregel.get_graph(xray=next_xray)
+            subgraphs_dict[name] = _convert_graph_to_uipath(subgraph_graph)
+
+    nodes: list[UiPathRuntimeNode] = []
+    for node_id, node in graph.nodes.items():
+        subgraph: UiPathRuntimeGraph | None = subgraphs_dict.get(node_id)
+        nodes.append(
+            UiPathRuntimeNode(
+                id=node.id,
+                name=node.name or node.id,
+                type=_get_node_type(node),
+                metadata=_get_node_metadata(node),
+                subgraph=subgraph,
+            )
+        )
+
+    # Use a set to track unique edges (source, target)
+    seen_edges: set[tuple[str, str]] = set()
+    edges: list[UiPathRuntimeEdge] = []
+
+    # First, add edges from graph.edges (static edges)
+    for edge in graph.edges:
+        edge_tuple = (edge.source, edge.target)
+        if edge_tuple not in seen_edges:
+            seen_edges.add(edge_tuple)
+            edges.append(
+                UiPathRuntimeEdge(
+                    source=edge.source,
+                    target=edge.target,
+                    label=getattr(edge, "data", None) or getattr(edge, "label", None),
+                )
+            )
+
+    # Build a map of channel -> target node
+    channel_to_target: dict[str, str] = {}
+    node_spec: PregelNode
+    for node_name, node_spec in compiled_graph.nodes.items():
+        for trigger in node_spec.triggers:
+            if isinstance(trigger, str) and trigger.startswith("branch:to:"):
+                channel_to_target[trigger] = node_name
+
+    # Extract edges by looking at what each node writes to (dynamic edges from Command)
+    for source_name, node_spec in compiled_graph.nodes.items():
+        writer: Runnable[Any, Any]
+        for writer in node_spec.writers:
+            if isinstance(writer, ChannelWrite):
+                write_entry: Any
+                for write_entry in writer.writes:
+                    if isinstance(write_entry, ChannelWriteEntry) and isinstance(
+                        write_entry.channel, str
+                    ):
+                        target = channel_to_target.get(write_entry.channel)
+                        if target:
+                            edge_tuple = (source_name, target)
+                            if edge_tuple not in seen_edges:
+                                seen_edges.add(edge_tuple)
+                                edges.append(
+                                    UiPathRuntimeEdge(
+                                        source=source_name, target=target, label=None
+                                    )
+                                )
+
+    return UiPathRuntimeGraph(nodes=nodes, edges=edges)
+
+
+def get_entrypoints_schema(
+    graph: CompiledStateGraph[Any, Any, Any],
+) -> SchemaDetails:
+    """Extract input/output schema from a LangGraph graph"""
+    input_circular_dependency = False
+    output_circular_dependency = False
+    schema = {
+        "input": {"type": "object", "properties": {}, "required": []},
+        "output": {"type": "object", "properties": {}, "required": []},
+    }
+
+    if hasattr(graph, "input_schema"):
+        if hasattr(graph.input_schema, "model_json_schema"):
+            input_schema = graph.input_schema.model_json_schema()
+            unpacked_ref_def_properties, input_circular_dependency = _resolve_refs(
+                input_schema
+            )
+
+            # Process the schema to handle nullable types
+            processed_properties = _process_nullable_types(
+                unpacked_ref_def_properties.get("properties", {})
+            )
+
+            schema["input"]["properties"] = processed_properties
+            schema["input"]["required"] = unpacked_ref_def_properties.get(
+                "required", []
+            )
+
+    if hasattr(graph, "output_schema"):
+        if hasattr(graph.output_schema, "model_json_schema"):
+            output_schema = graph.output_schema.model_json_schema()
+            unpacked_ref_def_properties, output_circular_dependency = _resolve_refs(
+                output_schema
+            )
+
+            # Process the schema to handle nullable types
+            processed_properties = _process_nullable_types(
+                unpacked_ref_def_properties.get("properties", {})
+            )
+
+            schema["output"]["properties"] = processed_properties
+            schema["output"]["required"] = unpacked_ref_def_properties.get(
+                "required", []
+            )
+
+    return SchemaDetails(schema, input_circular_dependency, output_circular_dependency)
+
+
+def _resolve_refs(schema, root=None, visited=None):
+    """Recursively resolves $ref references in a JSON schema, handling circular references.
+
+    Returns:
+        tuple: (resolved_schema, has_circular_dependency)
+    """
+    if root is None:
+        root = schema
+
+    if visited is None:
+        visited = set()
+
+    has_circular = False
+
+    if isinstance(schema, dict):
+        if "$ref" in schema:
+            ref_path = schema["$ref"]
+
+            if ref_path in visited:
+                # Circular dependency detected
+                return {
+                    "type": "object",
+                    "description": f"Circular reference to {ref_path}",
+                }, True
+
+            visited.add(ref_path)
+
+            # Resolve the reference
+            ref_parts = ref_path.lstrip("#/").split("/")
+            ref_schema = root
+            for part in ref_parts:
+                ref_schema = ref_schema.get(part, {})
+
+            result, circular = _resolve_refs(ref_schema, root, visited)
+            has_circular = has_circular or circular
+
+            # Remove from visited after resolution (allows the same ref in different branches)
+            visited.discard(ref_path)
+
+            return result, has_circular
+
+        resolved_dict = {}
+        for k, v in schema.items():
+            resolved_value, circular = _resolve_refs(v, root, visited)
+            resolved_dict[k] = resolved_value
+            has_circular = has_circular or circular
+        return resolved_dict, has_circular
+
+    elif isinstance(schema, list):
+        resolved_list = []
+        for item in schema:
+            resolved_item, circular = _resolve_refs(item, root, visited)
+            resolved_list.append(resolved_item)
+            has_circular = has_circular or circular
+        return resolved_list, has_circular
+
+    return schema, False
+
+
+def _process_nullable_types(
+    schema: dict[str, Any] | list[Any] | Any,
+) -> dict[str, Any] | list[Any]:
+    """Process the schema to handle nullable types by removing anyOf with null and keeping the base type."""
+    if isinstance(schema, dict):
+        if "anyOf" in schema and len(schema["anyOf"]) == 2:
+            types = [t.get("type") for t in schema["anyOf"]]
+            if "null" in types:
+                non_null_type = next(
+                    t for t in schema["anyOf"] if t.get("type") != "null"
+                )
+                return non_null_type
+
+        return {k: _process_nullable_types(v) for k, v in schema.items()}
+    elif isinstance(schema, list):
+        return [_process_nullable_types(item) for item in schema]
+    return schema
+
+
+__all__ = [
+    "get_graph_schema",
+    "get_entrypoints_schema",
+]

uipath_langchain/runtime/storage.py

@@ -0,0 +1,115 @@
+"""SQLite implementation of UiPathResumableStorageProtocol."""
+
+import json
+from typing import cast
+
+from langgraph.checkpoint.sqlite.aio import AsyncSqliteSaver
+from pydantic import BaseModel
+from uipath.runtime import (
+    UiPathApiTrigger,
+    UiPathResumeTrigger,
+    UiPathResumeTriggerName,
+    UiPathResumeTriggerType,
+)
+
+
+class SqliteResumableStorage:
+    """SQLite storage for resume triggers."""
+
+    def __init__(
+        self, memory: AsyncSqliteSaver, table_name: str = "__uipath_resume_triggers"
+    ):
+        self.memory = memory
+        self.table_name = table_name
+        self._initialized = False
+
+    async def _ensure_table(self) -> None:
+        """Create table if needed."""
+        if self._initialized:
+            return
+
+        await self.memory.setup()
+        async with self.memory.lock, self.memory.conn.cursor() as cur:
+            await cur.execute(f"""
+                CREATE TABLE IF NOT EXISTS {self.table_name} (
+                    id INTEGER PRIMARY KEY AUTOINCREMENT,
+                    type TEXT NOT NULL,
+                    name TEXT NOT NULL,
+                    key TEXT,
+                    folder_key TEXT,
+                    folder_path TEXT,
+                    payload TEXT,
+                    timestamp DATETIME DEFAULT (strftime('%Y-%m-%d %H:%M:%S', 'now', 'utc'))
+                )
+            """)
+            await self.memory.conn.commit()
+            self._initialized = True
+
+    async def save_trigger(self, trigger: UiPathResumeTrigger) -> None:
+        """Save resume trigger to database."""
+        await self._ensure_table()
+
+        trigger_key = (
+            trigger.api_resume.inbox_id if trigger.api_resume else trigger.item_key
+        )
+        payload = trigger.payload
+        if payload:
+            payload = (
+                (
+                    payload.model_dump()
+                    if isinstance(payload, BaseModel)
+                    else json.dumps(payload)
+                )
+                if isinstance(payload, dict)
+                else str(payload)
+            )
+
+        async with self.memory.lock, self.memory.conn.cursor() as cur:
+            await cur.execute(
+                f"INSERT INTO {self.table_name} (type, key, name, payload, folder_path, folder_key) VALUES (?, ?, ?, ?, ?, ?)",
+                (
+                    trigger.trigger_type.value,
+                    trigger_key,
+                    trigger.trigger_name.value,
+                    payload,
+                    trigger.folder_path,
+                    trigger.folder_key,
+                ),
+            )
+            await self.memory.conn.commit()
+
+    async def get_latest_trigger(self) -> UiPathResumeTrigger | None:
+        """Get most recent trigger from database."""
+        await self._ensure_table()
+
+        async with self.memory.lock, self.memory.conn.cursor() as cur:
+            await cur.execute(f"""
+                SELECT type, key, name, folder_path, folder_key, payload
+                FROM {self.table_name}
+                ORDER BY timestamp DESC
+                LIMIT 1
+            """)
+            result = await cur.fetchone()
+
+            if not result:
+                return None
+
+            trigger_type, key, name, folder_path, folder_key, payload = cast(
+                tuple[str, str, str, str, str, str], tuple(result)
+            )
+
+            resume_trigger = UiPathResumeTrigger(
+                trigger_type=UiPathResumeTriggerType(trigger_type),
+                trigger_name=UiPathResumeTriggerName(name),
+                item_key=key,
+                folder_path=folder_path,
+                folder_key=folder_key,
+                payload=payload,
+            )
+
+            if resume_trigger.trigger_type == UiPathResumeTriggerType.API:
+                resume_trigger.api_resume = UiPathApiTrigger(
+                    inbox_id=resume_trigger.item_key, request=resume_trigger.payload
+                )
+
+            return resume_trigger