Artifex/graph/langgraph_runtime.py

from __future__ import annotations

from typing import Any, TypedDict
from uuid import UUID

from graph.models import GraphRun
from graph.native_runtime import NativeGraphRuntime
from graph.registry import NodeHandlerRegistry
from graph.runtime import GraphRuntime
from graph.spec import ExecutionGraphSpec


class LangGraphState(TypedDict, total=False):
    graph_run_id: int
    current_node: str
    status: str
    edge_result: str


class LangGraphRuntime(GraphRuntime):
    """LangGraph adapter behind Artifex's runtime-neutral graph boundary."""

    def __init__(self, registry: NodeHandlerRegistry | None = None) -> None:
        self.registry = registry

    async def start(self, project_id: UUID | None = None, **kwargs: Any) -> str:
        graph_run = kwargs.get("graph_run")
        if graph_run is not None:
            self.run_until_terminal_or_paused(graph_run)
            return str(graph_run.id)
        return f"project-{project_id}"

    async def pause(self, run_id: str) -> None:
        return None

    async def resume(self, run_id: str) -> None:
        if self.registry is None:
            return None
        graph_run = GraphRun.objects.get(id=run_id)
        self.run_until_terminal_or_paused(graph_run)

    async def cancel(self, run_id: str) -> None:
        return None

    async def signal(self, run_id: str, event: dict[str, Any]) -> None:
        return None

    def run_until_terminal_or_paused(self, graph_run: GraphRun) -> GraphRun:
        if self.registry is None:
            raise RuntimeError("LangGraphRuntime requires a NodeHandlerRegistry")
        try:
            from langgraph.graph import END, StateGraph
        except ImportError as exc:
            raise RuntimeError("LangGraphRuntime requires the langgraph package") from exc

        spec = ExecutionGraphSpec.from_dict(graph_run.execution_graph_version.graph_spec)
        native = NativeGraphRuntime(self.registry)

        workflow = StateGraph(LangGraphState)
        for node_id in spec.nodes:
            workflow.add_node(node_id, self._node_runner(native, graph_run, node_id))
        workflow.set_entry_point(graph_run.current_node or spec.entry)
        for node_id in spec.nodes:
            if node_id in spec.terminal_nodes:
                workflow.add_edge(node_id, END)
                continue
            edges = [edge for edge in spec.edges if edge.source == node_id]
            if not edges:
                workflow.add_edge(node_id, END)
                continue
            workflow.add_conditional_edges(
                node_id,
                lambda state: str(state.get("edge_result", "success")),
                {edge.condition or "success": edge.target for edge in edges},
            )
        compiled = workflow.compile()
        compiled.invoke({"graph_run_id": graph_run.id, "current_node": graph_run.current_node or spec.entry})
        graph_run.refresh_from_db()
        return graph_run

    def _node_runner(self, native: NativeGraphRuntime, graph_run: GraphRun, node_id: str):
        def run_node(state: LangGraphState) -> LangGraphState:
            if graph_run.current_node != node_id:
                graph_run.current_node = node_id
                graph_run.save(update_fields=["current_node", "updated_at"])
            native.run_until_terminal_or_paused(graph_run, interrupt_after=node_id)
            graph_run.refresh_from_db()
            return {
                "graph_run_id": graph_run.id,
                "current_node": graph_run.current_node,
                "status": graph_run.status,
                "edge_result": str(graph_run.metadata.get("last_edge_result", "success")),
            }

        return run_node