pydantic-graph-studio 0.1.0__py3-none-any.whl

pydantic_graph_studio/__init__.py

@@ -0,0 +1,55 @@
+"""Pydantic Graph Studio entrypoint."""
+
+from pydantic_graph_studio.cli import main
+from pydantic_graph_studio.introspection import build_graph_model, serialize_graph
+from pydantic_graph_studio.runtime import (
+    RunHooks,
+    instrument_graph_run,
+    iter_instrumented,
+    iter_run_events,
+    run_instrumented,
+    run_instrumented_sync,
+)
+from pydantic_graph_studio.schemas import (
+    EdgeTakenEvent,
+    ErrorEvent,
+    Event,
+    EventBase,
+    GraphEdge,
+    GraphModel,
+    GraphNode,
+    NodeEndEvent,
+    NodeStartEvent,
+    RunEndEvent,
+    event_schema,
+    export_schemas,
+    graph_schema,
+)
+from pydantic_graph_studio.server import RunRegistry, create_app
+
+__all__ = [
+    "EdgeTakenEvent",
+    "ErrorEvent",
+    "Event",
+    "EventBase",
+    "GraphEdge",
+    "GraphModel",
+    "GraphNode",
+    "NodeEndEvent",
+    "NodeStartEvent",
+    "RunEndEvent",
+    "RunHooks",
+    "RunRegistry",
+    "build_graph_model",
+    "create_app",
+    "event_schema",
+    "export_schemas",
+    "graph_schema",
+    "instrument_graph_run",
+    "iter_instrumented",
+    "iter_run_events",
+    "main",
+    "run_instrumented",
+    "run_instrumented_sync",
+    "serialize_graph",
+]

pydantic_graph_studio/cli.py

@@ -0,0 +1,201 @@
+"""Command-line entrypoint for launching the studio server."""
+
+from __future__ import annotations
+
+import argparse
+import importlib
+import importlib.util
+import sys
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any
+
+from pydantic_graph import Graph
+from pydantic_graph.nodes import BaseNode
+
+from pydantic_graph_studio.introspection import build_graph_model
+from pydantic_graph_studio.server import create_app
+
+BetaGraph: type[Any] | None = None
+try:  # pragma: no cover - optional beta support
+    from pydantic_graph.beta.graph import Graph as _BetaGraph
+except ModuleNotFoundError:  # pragma: no cover
+    pass
+else:
+    BetaGraph = _BetaGraph
+
+
+@dataclass(slots=True)
+class GraphRef:
+    target: str
+    attribute: str
+
+
+class CLIError(RuntimeError):
+    """Raised for user-facing CLI errors."""
+
+
+def main(argv: list[str] | None = None) -> None:
+    """CLI entrypoint for the `pgraph` launcher."""
+
+    args = _parse_args(argv)
+    try:
+        graph = _load_graph(args.graph_ref)
+        start_node = _resolve_start_node(graph, args.start)
+        _run_server(graph, start_node, host=args.host, port=args.port)
+    except CLIError as exc:
+        print(f"error: {exc}", file=sys.stderr)
+        raise SystemExit(2) from exc
+
+
+def _parse_args(argv: list[str] | None) -> argparse.Namespace:
+    parser = argparse.ArgumentParser(
+        prog="pgraph",
+        description="Launch the local Pydantic Graph Studio for a graph reference.",
+    )
+    parser.add_argument(
+        "graph_ref",
+        help="Graph reference in the form module:var or path.py:var",
+    )
+    parser.add_argument(
+        "--host",
+        default="127.0.0.1",
+        help="Host to bind the local server (default: 127.0.0.1)",
+    )
+    parser.add_argument(
+        "--port",
+        type=int,
+        default=8000,
+        help="Port to bind the local server (default: 8000)",
+    )
+    parser.add_argument(
+        "--start",
+        help="Explicit node id to use as the entry point",
+    )
+    return parser.parse_args(argv)
+
+
+def _load_graph(graph_ref: str) -> Any:
+    parsed = _parse_graph_ref(graph_ref)
+    module = _load_module(parsed.target)
+    graph = _resolve_attribute(module, parsed.attribute)
+    if not isinstance(graph, Graph) and not _is_beta_graph(graph):
+        raise CLIError(
+            "Graph reference did not resolve to a Graph instance. "
+            "Ensure the reference points to a pydantic_graph.Graph object."
+        )
+    return graph
+
+
+def _parse_graph_ref(graph_ref: str) -> GraphRef:
+    if ":" not in graph_ref:
+        raise CLIError("Graph reference must be in the form module:var or path.py:var")
+    target, attribute = graph_ref.rsplit(":", 1)
+    if not target or not attribute:
+        raise CLIError("Graph reference must include both target and attribute")
+    return GraphRef(target=target, attribute=attribute)
+
+
+def _load_module(target: str) -> Any:
+    path = Path(target)
+    if _looks_like_path(target):
+        if not path.exists():
+            raise CLIError(f"File not found: {path}")
+        if path.suffix != ".py":
+            raise CLIError("File reference must point to a .py file")
+        return _load_module_from_file(path)
+    try:
+        return importlib.import_module(target)
+    except ModuleNotFoundError as exc:
+        raise CLIError(f"Module not found: {target}") from exc
+
+
+def _looks_like_path(target: str) -> bool:
+    return "/" in target or "\\" in target or target.endswith(".py")
+
+
+def _load_module_from_file(path: Path) -> Any:
+    module_name = f"pgraph_user_{path.stem}_{abs(hash(path))}"
+    spec = importlib.util.spec_from_file_location(module_name, path)
+    if spec is None or spec.loader is None:
+        raise CLIError(f"Unable to load module from file: {path}")
+    module = importlib.util.module_from_spec(spec)
+    sys.modules[module_name] = module
+    spec.loader.exec_module(module)
+    return module
+
+
+def _resolve_attribute(module: Any, attribute: str) -> Any:
+    current: Any = module
+    for segment in attribute.split("."):
+        if not hasattr(current, segment):
+            raise CLIError(f"Attribute '{segment}' not found while resolving '{attribute}'")
+        current = getattr(current, segment)
+    return current
+
+
+def _resolve_start_node(
+    graph: Any,
+    start_node_id: str | None,
+) -> BaseNode[Any, Any, Any] | None:
+    if _is_beta_graph(graph):
+        if start_node_id:
+            raise CLIError("Beta graphs use a fixed start node; --start is not supported.")
+        return None
+
+    node_defs = graph.node_defs
+    if not node_defs:
+        raise CLIError("Graph contains no nodes")
+
+    if start_node_id:
+        node_def = node_defs.get(start_node_id)
+        if node_def is None:
+            available = ", ".join(sorted(node_defs.keys()))
+            raise CLIError(f"Unknown start node '{start_node_id}'. Available nodes: {available}")
+    else:
+        entry_nodes = build_graph_model(graph).entry_nodes
+        if not entry_nodes:
+            raise CLIError("Unable to infer an entry node. Use --start to specify one.")
+        if len(entry_nodes) > 1:
+            entries = ", ".join(entry_nodes)
+            raise CLIError(f"Multiple entry nodes found: {entries}. Use --start to choose one.")
+        node_def = node_defs[entry_nodes[0]]
+
+    node_cls_or_instance = node_def.node
+    if isinstance(node_cls_or_instance, BaseNode):
+        return node_cls_or_instance
+    if isinstance(node_cls_or_instance, type) and issubclass(node_cls_or_instance, BaseNode):
+        try:
+            instance = node_cls_or_instance()
+        except TypeError as exc:
+            raise CLIError(
+                "Failed to instantiate the start node. "
+                "Ensure it can be constructed with no arguments or provide a different entry node."
+            ) from exc
+        return instance
+
+    raise CLIError("Start node did not resolve to a BaseNode instance")
+
+
+def _run_server(
+    graph: Any,
+    start_node: BaseNode[Any, Any, Any] | None,
+    *,
+    host: str,
+    port: int,
+) -> None:
+    if port <= 0 or port > 65535:
+        raise CLIError("Port must be between 1 and 65535")
+
+    app = create_app(graph, start_node)
+    try:
+        import uvicorn
+    except ModuleNotFoundError as exc:
+        raise CLIError("uvicorn is required to run the server") from exc
+
+    print(f"Studio running at http://{host}:{port}")
+    uvicorn.run(app, host=host, port=port, log_level="info")
+
+
+def _is_beta_graph(graph: Any) -> bool:
+    return BetaGraph is not None and isinstance(graph, BetaGraph)

pydantic_graph_studio/introspection.py

@@ -0,0 +1,209 @@
+from __future__ import annotations
+
+from collections.abc import Iterable, Mapping
+from typing import Any
+
+from pydantic_graph import Graph
+from pydantic_graph.nodes import NodeDef
+
+from pydantic_graph_studio.schemas import GraphEdge, GraphModel, GraphNode
+
+BetaGraph: type[Any] | None = None
+try:  # pragma: no cover - optional beta support
+    from pydantic_graph.beta.decision import Decision as BetaDecision
+    from pydantic_graph.beta.graph import Graph as _BetaGraph
+    from pydantic_graph.beta.join import Join as BetaJoin
+    from pydantic_graph.beta.node import EndNode as BetaEndNode
+    from pydantic_graph.beta.node import Fork as BetaFork
+    from pydantic_graph.beta.node import StartNode as BetaStartNode
+    from pydantic_graph.beta.paths import DestinationMarker as BetaDestinationMarker
+    from pydantic_graph.beta.paths import Path as BetaPath
+    from pydantic_graph.beta.step import NodeStep as BetaNodeStep
+    from pydantic_graph.beta.step import Step as BetaStep
+except ModuleNotFoundError:  # pragma: no cover
+    pass
+else:
+    BetaGraph = _BetaGraph
+
+
+def build_graph_model(graph: Any) -> GraphModel:
+    """Build a GraphModel payload from a pydantic_graph.Graph instance."""
+
+    if _is_beta_graph(graph):
+        return _build_beta_graph_model(graph)
+
+    node_defs = _sorted_node_defs(graph.node_defs)
+    nodes = _build_nodes(node_defs)
+    edges = _build_edges(node_defs)
+    entry_nodes = _infer_entry_nodes(node_defs, edges)
+    terminal_nodes = _infer_terminal_nodes(node_defs)
+    return GraphModel(
+        nodes=nodes,
+        edges=edges,
+        entry_nodes=entry_nodes,
+        terminal_nodes=terminal_nodes,
+    )
+
+
+def serialize_graph(graph: Graph[Any, Any, Any]) -> dict[str, Any]:
+    """Serialize a graph into a JSON-safe dict payload."""
+
+    return build_graph_model(graph).model_dump(mode="json")
+
+
+def _sorted_node_defs(node_defs: Mapping[str, NodeDef[Any, Any, Any]]) -> list[NodeDef[Any, Any, Any]]:
+    return [node_defs[node_id] for node_id in sorted(node_defs.keys())]
+
+
+def _build_nodes(node_defs: Iterable[NodeDef[Any, Any, Any]]) -> list[GraphNode]:
+    return [
+        GraphNode(
+            node_id=node_def.node_id,
+            label=node_def.node.__name__,
+        )
+        for node_def in node_defs
+    ]
+
+
+def _build_edges(node_defs: Iterable[NodeDef[Any, Any, Any]]) -> list[GraphEdge]:
+    edges: list[GraphEdge] = []
+    for node_def in node_defs:
+        for target_id in sorted(node_def.next_node_edges.keys()):
+            edges.append(
+                GraphEdge(
+                    source_node_id=node_def.node_id,
+                    target_node_id=target_id,
+                    dynamic=False,
+                )
+            )
+        if node_def.returns_base_node:
+            edges.append(
+                GraphEdge(
+                    source_node_id=node_def.node_id,
+                    target_node_id=None,
+                    dynamic=True,
+                )
+            )
+    return edges
+
+
+def _infer_entry_nodes(
+    node_defs: Iterable[NodeDef[Any, Any, Any]],
+    edges: Iterable[GraphEdge],
+) -> list[str]:
+    node_ids = {node_def.node_id for node_def in node_defs}
+    inbound: set[str] = set()
+    for edge in edges:
+        if edge.target_node_id is not None:
+            inbound.add(edge.target_node_id)
+    entry_nodes = sorted(node_ids - inbound)
+    return entry_nodes
+
+
+def _infer_terminal_nodes(node_defs: Iterable[NodeDef[Any, Any, Any]]) -> list[str]:
+    terminal_nodes = sorted(node_def.node_id for node_def in node_defs if node_def.end_edge is not None)
+    return terminal_nodes
+
+
+def _is_beta_graph(graph: Any) -> bool:
+    return BetaGraph is not None and isinstance(graph, BetaGraph)
+
+
+def _build_beta_graph_model(graph: Any) -> GraphModel:
+    nodes = _build_beta_nodes(graph)
+    edges = _build_beta_edges(graph.edges_by_source, graph.nodes)
+    entry_nodes = _infer_beta_entry_nodes(graph)
+    terminal_nodes = _infer_beta_terminal_nodes(graph)
+    return GraphModel(
+        nodes=nodes,
+        edges=edges,
+        entry_nodes=entry_nodes,
+        terminal_nodes=terminal_nodes,
+    )
+
+
+def _build_beta_nodes(graph: Any) -> list[GraphNode]:
+    items = sorted(graph.nodes.items(), key=lambda item: str(item[0]))
+    return [
+        GraphNode(
+            node_id=str(node_id),
+            label=_beta_node_label(node),
+        )
+        for node_id, node in items
+    ]
+
+
+def _beta_node_label(node: Any) -> str | None:
+    if BetaGraph is None:
+        return None
+    if isinstance(node, BetaStartNode):
+        return "Start"
+    if isinstance(node, BetaEndNode):
+        return "Done"
+    if isinstance(node, BetaFork):
+        return "Map Fork" if node.is_map else "Fork"
+    if isinstance(node, BetaJoin):
+        return "Join"
+    if isinstance(node, BetaDecision):
+        return "Decision"
+    if isinstance(node, BetaNodeStep):
+        return node.node_type.__name__
+    if isinstance(node, BetaStep):
+        return node.label or str(node.id)
+    raw = getattr(node, "label", None) or getattr(node, "id", None)
+    return str(raw) if raw is not None else None
+
+
+def _build_beta_edges(
+    edges_by_source: Mapping[Any, list[Any]], nodes: Mapping[Any, Any] | None = None
+) -> list[GraphEdge]:
+    edge_pairs: set[tuple[str, str]] = set()
+
+    for source_id, paths in edges_by_source.items():
+        for path in paths:
+            for target_id in _beta_path_destinations(path):
+                edge_pairs.add((str(source_id), str(target_id)))
+
+    if nodes:
+        for node_id, node in nodes.items():
+            if BetaGraph is None or not isinstance(node, BetaDecision):
+                continue
+            for branch in node.branches:
+                for target_id in _beta_path_destinations(branch.path):
+                    edge_pairs.add((str(node_id), str(target_id)))
+
+    edges = [
+        GraphEdge(
+            source_node_id=source_id,
+            target_node_id=target_id,
+            dynamic=False,
+        )
+        for source_id, target_id in sorted(edge_pairs)
+    ]
+    return edges
+
+
+def _beta_path_destinations(path: Any) -> list[Any]:
+    if BetaGraph is None:
+        return []
+    if not isinstance(path, BetaPath):
+        return []
+    destinations: list[Any] = []
+    for item in path.items:
+        if isinstance(item, BetaDestinationMarker):
+            destinations.append(item.destination_id)
+    return destinations
+
+
+def _infer_beta_entry_nodes(graph: Any) -> list[str]:
+    if BetaGraph is None:
+        return []
+    start_id = str(BetaStartNode.id)
+    return [start_id] if start_id in {str(node_id) for node_id in graph.nodes.keys()} else []
+
+
+def _infer_beta_terminal_nodes(graph: Any) -> list[str]:
+    if BetaGraph is None:
+        return []
+    end_id = str(BetaEndNode.id)
+    return [end_id] if end_id in {str(node_id) for node_id in graph.nodes.keys()} else []