152 lines
6.8 KiB
Python
152 lines
6.8 KiB
Python
from __future__ import annotations
|
|
|
|
from typing import Any, TypedDict
|
|
from uuid import UUID
|
|
|
|
from control_plane.events.bus import EventBus
|
|
from graph.models import GraphEdgeTraversal, GraphRun, GraphRunStatus
|
|
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 GraphInterrupted(RuntimeError):
|
|
pass
|
|
|
|
|
|
class LangGraphRuntime(GraphRuntime):
|
|
"""LangGraph adapter behind Artifex's runtime-neutral graph boundary."""
|
|
|
|
def __init__(self, registry: NodeHandlerRegistry | None = None, bus: EventBus | None = None) -> None:
|
|
self.registry = registry
|
|
self.bus = bus or EventBus()
|
|
|
|
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, *, interrupt_after: str | None = None) -> 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, interrupt_after))
|
|
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
|
|
route_map = {edge.condition or "success": edge.target for edge in edges}
|
|
route_map["__end__"] = END
|
|
workflow.add_conditional_edges(node_id, self._route(graph_run, spec, node_id), route_map)
|
|
compiled = workflow.compile()
|
|
try:
|
|
compiled.invoke({"graph_run_id": graph_run.id, "current_node": graph_run.current_node or spec.entry})
|
|
except GraphInterrupted:
|
|
pass
|
|
graph_run.refresh_from_db()
|
|
return graph_run
|
|
|
|
def _node_runner(self, native: NativeGraphRuntime, graph_run: GraphRun, node_id: str, interrupt_after: str | None):
|
|
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()
|
|
if interrupt_after == node_id:
|
|
raise GraphInterrupted(node_id)
|
|
return {
|
|
"graph_run_id": graph_run.id,
|
|
"current_node": graph_run.current_node,
|
|
"status": graph_run.status,
|
|
"edge_result": str(graph_run.metadata.get("interrupted_edge_result", graph_run.metadata.get("last_edge_result", "success"))),
|
|
}
|
|
|
|
return run_node
|
|
|
|
def _route(self, graph_run: GraphRun, spec: ExecutionGraphSpec, node_id: str):
|
|
def route(state: LangGraphState) -> str:
|
|
graph_run.refresh_from_db()
|
|
if graph_run.status != GraphRunStatus.RUNNING:
|
|
return "__end__"
|
|
edge_result = str(state.get("edge_result") or graph_run.metadata.get("interrupted_edge_result") or "success")
|
|
target = None
|
|
for edge in spec.edges:
|
|
if edge.source == node_id and edge.condition == edge_result:
|
|
target = edge.target
|
|
break
|
|
if target is None:
|
|
for edge in spec.edges:
|
|
if edge.source == node_id and not edge.condition:
|
|
target = edge.target
|
|
edge_result = edge.condition or "success"
|
|
break
|
|
if target is None:
|
|
metadata = dict(graph_run.metadata)
|
|
metadata["final_failure_reason"] = f"No edge from {node_id} for {edge_result}"
|
|
graph_run.status = GraphRunStatus.FAILED
|
|
graph_run.failure_reason = metadata["final_failure_reason"]
|
|
graph_run.metadata = metadata
|
|
graph_run.save(update_fields=["status", "failure_reason", "metadata", "updated_at"])
|
|
try:
|
|
from agents.progeny import ProgenyService
|
|
|
|
node_run = graph_run.node_runs.filter(node_id=node_id).order_by("-visit_index").first()
|
|
ProgenyService(self.bus).create_graph_runtime_signal(graph_run, "GRAPH_RUNTIME_ERROR", graph_run.failure_reason, {"node_id": node_id, "edge_result": edge_result}, graph_node_run=node_run)
|
|
except Exception:
|
|
pass
|
|
return "__end__"
|
|
GraphEdgeTraversal.objects.create(graph_run=graph_run, source_node=node_id, target_node=target, condition=edge_result, result=edge_result)
|
|
metadata = dict(graph_run.metadata)
|
|
metadata["last_node_id"] = node_id
|
|
metadata["last_edge_result"] = edge_result
|
|
metadata["last_output"] = metadata.get("interrupted_output", {})
|
|
metadata.pop("interrupted_after_node", None)
|
|
metadata.pop("interrupted_edge_result", None)
|
|
metadata.pop("interrupted_output", None)
|
|
graph_run.current_node = target
|
|
graph_run.metadata = metadata
|
|
graph_run.save(update_fields=["current_node", "metadata", "updated_at"])
|
|
self.bus.publish("GRAPH_EDGE_TRAVERSED", project=graph_run.project, task=graph_run.task, payload={"graph_run_id": graph_run.id, "source": node_id, "target": target, "condition": edge_result})
|
|
return edge_result
|
|
|
|
return route
|