Move out dm-tree and networkx from main dependencies (#324)

Moved these do ldp[scg] subpackage

Author: whitead

.github/workflows/tests.yml

@@ -59,6 +59,12 @@ jobs:
         env:
           OPENAI_API_KEY: ${{ secrets.OPENAI_API_KEY }}
           ANTHROPIC_API_KEY: ${{ secrets.ANTHROPIC_API_KEY }}
+      - run:
+          | # this is to check that we can run README examples with lazily imported dependencies (the [scg] subpackage)
+          uv pip uninstall dm-tree usearch networkx && uv run --no-sync pytest --noconftest tests/test_readme.py
+        env:
+          OPENAI_API_KEY: ${{ secrets.OPENAI_API_KEY }}
+          ANTHROPIC_API_KEY: ${{ secrets.ANTHROPIC_API_KEY }}
   test-lmi:
     runs-on: ubuntu-latest
     steps:

README.md

@@ -185,7 +185,15 @@ which just converts the action constructed by plain python into a node in a comp
 
 For more advanced use-cases, LDP features a stochastic computation graph [^2]
 which enables differentiatiation with respect to agent parameters
-(including the weights of the LLM). The example computation graph below illustrates the functionality
+(including the weights of the LLM).
+
+You should install the `scg` subpackage to work with it:
+
+```bash
+pip install ldp[scg]
+```
+
+The example computation graph below illustrates the functionality
 
 ```py
 from ldp.graph import FxnOp, LLMCallOp, PromptOp, compute_graph

pyproject.toml

@@ -23,18 +23,15 @@ classifiers = [
 ]
 dependencies = [
     "aiofiles",
-    "dm-tree",
     "fhaviary>=0.20.0",  # For Environment.get_id
     "fhlmi",  # For new LLM client interface
     "httpx",
-    "networkx[default]~=3.4",  # Pin for pydot fix
     "numpy>=1.20",  # For numpy.typing
     "pydantic~=2.0",
     "tenacity",
     "tiktoken",
     "tqdm",
     "typing-extensions; python_version <= '3.11'",  # for typing.override
-    "usearch>=2.13",  # For py.typed
 ]
 description = "Agent framework for constructing language model agents and training on constructive tasks."
 dynamic = ["version"]
@@ -66,6 +63,7 @@ dev = [
     "pytest-xdist",
     "pytest>=8",  # Pin to keep recent
     "refurb>=2",  # Pin to keep recent
+    "usearch>=2.13",  # For py.typed
     "vcrpy>=6",  # Pin for https://github.com/kevin1024/vcrpy/issues/884
 ]
 monitor = [
@@ -76,6 +74,7 @@ nn = [
     "dask-cuda>=24.8.2",
     "dask-jobqueue",
     "dask[distributed]",
+    "ldp[scg]",
     "tokenizers>0.20",
     "torch>=2.5,<2.6",  # Temporarily pin <2.6 until someone fixes our CI with newer versions
     "transformers>=4.46,<4.51",  # Temporarily pin <4.51 until someone fixes our CI with newer versions, and also https://github.com/huggingface/transformers/issues/37339
@@ -85,6 +84,10 @@ rich = [
     "rich",
     "tqdm>=4.56",  # When tqdm.rich was created
 ]
+scg = [
+    "dm-tree",
+    "networkx[default]~=3.4",  # Pin for pydot fix
+]
 server = [
     "fastapi>=0.109",  # For Python 3.12 support
 ]

src/ldp/alg/algorithms.py

@@ -4,13 +4,15 @@
 import itertools
 import random
 from collections.abc import Awaitable, Callable, Hashable, Iterable, Sequence
-from typing import Any, Literal, TypeVar
+from typing import TYPE_CHECKING, Any, Literal, TypeVar
 
-import networkx as nx
+if TYPE_CHECKING:
+    import networkx as nx
 import numpy as np
 from aviary.core import Message, Tool, ToolRequestMessage, join
 
 from ldp.graph import OpResult
+from ldp.graph.op_utils import _lazy_import_networkx
 from ldp.graph.ops import GradOutType
 
 
@@ -36,6 +38,7 @@ def to_network(  # noqa: C901
     Returns:
         Populated a NetworkX multi-edge directed graph.
     """
+    nx = _lazy_import_networkx()
 
     def gvizify(x: Any) -> str:
         """Stringify and then escape colons for Graphviz labels."""

src/ldp/data_structures.py

@@ -9,16 +9,17 @@
 from uuid import UUID
 
 import aiofiles
-import networkx as nx
 from aviary.core import Message, ToolRequestMessage, ToolResponseMessage, join
 from pydantic import BaseModel, ConfigDict, Field, JsonValue, field_validator
 
 from ldp.graph import OpResult
+from ldp.graph.op_utils import _lazy_import_networkx
 from ldp.utils import discounted_returns
 
 if TYPE_CHECKING:
     from aviary.core import Environment
 
+
 logger = logging.getLogger(__name__)
 
 
@@ -180,6 +181,7 @@ def __init__(self, root_id: str | UUID):
         """
         self.root_id = str(root_id)
 
+        nx = _lazy_import_networkx()
         self.tree = nx.DiGraph()  # the actual tree
         self.rev_tree = nx.DiGraph()  # the same as self.tree, but with reversed edges
 
@@ -292,7 +294,7 @@ def assign_mc_value_estimates(self, discount_factor: float = 1.0) -> None:
             discount_factor: The discount factor to use when computing cumulative
                 future rewards.
         """
-        for step_id in nx.topological_sort(self.rev_tree):
+        for step_id in _lazy_import_networkx().topological_sort(self.rev_tree):
             step: Transition | None = self.tree.nodes[step_id]["transition"]
             if step is None:
                 continue
@@ -324,7 +326,9 @@ def compute_advantages(self) -> None:
         """
         state_values: dict[str, float] = {}
 
-        for step_id in cast("Iterable[str]", nx.topological_sort(self.tree)):
+        for step_id in cast(
+            "Iterable[str]", _lazy_import_networkx().topological_sort(self.tree)
+        ):
             # topological sort means we will update a parent node in-place before
             # descending to its children
 
@@ -412,7 +416,7 @@ def merge_identical_nodes(
         # old step ID -> new step ID
         node_remap: dict[str, str] = {self.root_id: self.root_id}
 
-        for step_id in nx.topological_sort(self.tree):
+        for step_id in _lazy_import_networkx().topological_sort(self.tree):
             step: Transition | None = self.tree.nodes[step_id]["transition"]
             if step is None:
                 continue

src/ldp/graph/common_ops.py

@@ -8,11 +8,10 @@
 import logging
 from collections.abc import Awaitable, Callable
 from functools import lru_cache
-from typing import Generic, TypeVar, cast, overload
+from typing import TYPE_CHECKING, Generic, TypeVar, cast, overload
 
 import numpy as np
 import tenacity
-import tree
 from aviary.core import Message, Tool, ToolRequestMessage, is_coroutine_callable
 from lmi import (
     EmbeddingModel,
@@ -24,11 +23,13 @@
 from lmi import LiteLLMModel as LLMModel
 from pydantic import BaseModel
 
-from .gradient_estimators import assign_constant_grads
 from .memory import Memory, MemoryModel, UIndexMemoryModel
-from .op_utils import CallID, get_call_id, get_training_mode
+from .op_utils import CallID, _lazy_import_tree, get_call_id, get_training_mode
 from .ops import GradInType, Op, OpCtx, ResultOrValue, TOutput_co
 
+if TYPE_CHECKING:
+    import tree
+
 logger = logging.getLogger(__name__)
 
 
@@ -76,6 +77,8 @@ def backward(
         grad_output: tree.Structure,
         call_id: CallID,
     ) -> GradInType:
+        from .gradient_estimators import assign_constant_grads
+
         return assign_constant_grads(input_args, input_kwargs, None)
 
 
@@ -101,7 +104,7 @@ def backward(
         call_id: CallID,
     ) -> GradInType:
         # Check that the grad_output structure is consistent with our config
-        tree.assert_same_structure(
+        _lazy_import_tree().assert_same_structure(
             grad_output, ctx.get(call_id, "output").value, check_types=False
         )
 
@@ -186,6 +189,8 @@ def backward(
         grad_output: tree.Structure,
         call_id: CallID,
     ) -> GradInType:
+        from .gradient_estimators import assign_constant_grads
+
         return assign_constant_grads(input_args, input_kwargs, 0.0)
 
 
@@ -397,7 +402,9 @@ def backward(
         # but not necessarily each message or tool.
 
         # tree.map_structure allows us to assign a gradient of 0 to all fields of config
-        grad_config = tree.map_structure(lambda _: 0.0, input_kwargs["config"])
+        grad_config = _lazy_import_tree().map_structure(
+            lambda _: 0.0, input_kwargs["config"]
+        )
         grad_kwargs = {"config": grad_config}
         for arg in ("msgs", "tools", "tool_choice"):
             if arg in input_kwargs:
@@ -472,6 +479,8 @@ def backward(
         call_id: CallID,
     ) -> GradInType:
         """Backward pass for memory retrieval - goes back to item."""
+        from .gradient_estimators import assign_constant_grads
+
         return assign_constant_grads(input_args, input_kwargs, 0.0)
 
 

src/ldp/graph/memory.py

@@ -17,9 +17,10 @@
     field_validator,
     model_validator,
 )
-from usearch.index import Index
 
 if TYPE_CHECKING:
+    from usearch.index import Index  # noqa: F401
+
     from .common_ops import MemoryOp
     from .op_utils import CallID
     from .ops import Op, OpResult, TOutput_co
@@ -164,13 +165,20 @@ async def _search_index(
         """Search the internal Index, returning a 'matches' amount of Memories."""
 
 
-class UIndexMemoryModel(MemoryModel[Index]):
+class UIndexMemoryModel(MemoryModel["Index"]):
     """Memory model using a U-Search index."""
 
     def __init__(self, **kwargs):
         super().__init__(**kwargs)
         if not self.embedding_model.ndim:
             raise TypeError("Specify dimensions to the embedding model.")
+        try:
+            from usearch.index import Index
+        except ImportError as e:
+            raise ImportError(
+                "U-Search library not found. Unable to use UIndexMemoryModel."
+                " To install U-Search dependencies, please run `pip install usearch`."
+            ) from e
         self._index = Index(ndim=self.embedding_model.ndim)
 
     async def _add_to_index(self, embedding: np.ndarray) -> int:

src/ldp/graph/op_utils.py

@@ -1,12 +1,35 @@
 import contextvars
 from collections.abc import AsyncIterator
 from contextlib import asynccontextmanager
+from types import ModuleType
 from uuid import UUID, uuid4
 
 from aviary.core import is_coroutine_callable
 from pydantic import BaseModel, field_serializer, field_validator
 
 
+def _lazy_import_networkx() -> ModuleType:
+    try:
+        import networkx as nx
+    except ImportError as e:
+        raise ImportError(
+            "networkx is required for compute graph operations. "
+            "Please install it with: pip install ldp[scg]"
+        ) from e
+    return nx
+
+
+def _lazy_import_tree() -> ModuleType:
+    try:
+        import tree
+    except ImportError as e:
+        raise ImportError(
+            "tree is required for compute graph operations. "
+            "Please install it with: pip install ldp[scg]"
+        ) from e
+    return tree
+
+
 class CallID(BaseModel):
     run_id: UUID
     fwd_id: UUID

src/ldp/graph/ops.py

@@ -9,23 +9,34 @@
 from abc import ABC, abstractmethod
 from collections import defaultdict
 from collections.abc import Callable, Collection, Iterable, Iterator, Mapping, Sequence
-from typing import Any, ClassVar, Generic, TypeAlias, TypeVar
+from typing import TYPE_CHECKING, Any, ClassVar, Generic, TypeAlias, TypeVar
 from uuid import UUID
 
-import networkx as nx
-import tree
+if TYPE_CHECKING:
+    import networkx as nx
+    import tree
 from pydantic import BaseModel, Field
 
-from .op_utils import CallID, compute_graph, get_call_id, get_training_mode, op_call
+from .op_utils import (
+    CallID,
+    _lazy_import_networkx,
+    _lazy_import_tree,
+    compute_graph,
+    get_call_id,
+    get_training_mode,
+    op_call,
+)
 
 logger = logging.getLogger(__name__)
 
 
-GradOutType: TypeAlias = tree.Structure | None  # None means the gradient has terminated
+GradOutType: TypeAlias = (
+    "tree.Structure | None"  # None means the gradient has terminated
+)
 GradInType: TypeAlias = tuple[Sequence[GradOutType], Mapping[str, GradOutType]]
 BackwardsType: TypeAlias = Callable[
     # Call signature of Op.backward
-    ["OpCtx", list, dict, tree.Structure, CallID],
+    ["OpCtx", list, dict, "tree.Structure", CallID],
     GradInType,
 ]
 TOutput_co = TypeVar("TOutput_co", covariant=True)
@@ -93,8 +104,11 @@ def compute_grads(
             (a) define the backward computation
             (b) store internal gradients for optimizer updates.
         """
-        # call ID -> [d op(x) / d x] for each op that consumes x
-        grad_outputs: dict[CallID, list[tree.Structure]] = defaultdict(list)
+        tree = _lazy_import_tree()
+
+        # call ID -> [d op(x) / d x] for each op that consumes x[
+        # due to interaction between ruff and mypy, we need type ignore
+        grad_outputs: dict[CallID, list[tree.Structure]] = defaultdict(list)  # type: ignore[name-defined]
 
         # grad_outputs stores a list of output grads (corresponding to each consuming op call).
         # Since the root node is not consumed by any other node, we create a singleton list here.
@@ -212,7 +226,7 @@ def add_edges(graph: nx.DiGraph, node: OpResult) -> None:
                     graph.add_edge(*edge)
                     add_edges(graph, x)
 
-        graph = nx.DiGraph()
+        graph = _lazy_import_networkx().DiGraph()
         graph.add_node(self)
         add_edges(graph, self)
 
@@ -251,7 +265,7 @@ def traverse(
         """
         if topological_order:
             G = self.get_compute_graph()
-            for node in nx.topological_sort(G):
+            for node in _lazy_import_networkx().topological_sort(G):
                 if filter_fn(node):
                     yield node