Add optional extras to slim the core dependency set

.github/workflows/validate.yaml

@@ -23,7 +23,7 @@ jobs:
 
       - name: Install Python dependencies
         run: |
-          uv sync
+          uv sync --all-extras
 
       - name: Run automated validation checks
         run: |

docs/getting_started.rst

@@ -28,6 +28,25 @@ Install using `uv <https://docs.astral.sh/uv/>`__:
 
     uv add python-flexeval
 
+Optional extras
+~~~~~~~~~~~~~~~
+
+Some functionality depends on optional packages that most users do not need.
+
+To load `LangGraph <https://langchain-ai.github.io/langgraph/>`__ checkpoint
+SQLite databases (the ``langgraph_sqlite`` data source), install the
+``langgraph`` extra:
+
+.. code-block:: bash
+
+    pip install 'python-flexeval[langgraph]'
+
+To use `LiteLLM <https://docs.litellm.ai/>`__ completions via :func:`~flexeval.configuration.completion_functions.litellm_completion`, install the ``litellm`` extra:
+
+.. code-block:: bash
+
+    pip install 'python-flexeval[litellm]'
+
 .. _getting-started-usage:
 
 Usage

pyproject.toml

@@ -26,11 +26,6 @@ requires-python = ">=3.10"
 dependencies = [
     "flatten-json>=0.1.14",
     "jsonschema>=4.23.0",
-    "langchain>=0.3.20",
-    "langchain-openai>=0.3.8",
-    "langgraph>=1.0.0",
-    "langgraph-checkpoint-sqlite>=3.0.0",
-    "litellm>=1.74.3",
     "msgpack>=1.1.0",
     "networkx>=3.4.2",
     "openai>=1.66.0",
@@ -40,11 +35,20 @@ dependencies = [
     "python-dotenv>=1.0.1",
     "pyyaml>=6.0.2",
     "requests>=2.32.3",
-    "sympy>=1.13.3",
     "textstat>=0.7.5",
     "typer>=0.16.0",
 ]
 
+[project.optional-dependencies]
+# Optional, user-facing extras. Install with e.g. `pip install 'python-flexeval[litellm]'`.
+# litellm enables `litellm_completion`, giving access to many non-OpenAI providers.
+litellm = ["litellm>=1.74.3"]
+# langgraph enables loading langgraph checkpoint SQLite databases (the `langgraph_sqlite` data source). 
+# Users who don't load langgraph checkpoints do not need this.
+langgraph = ["langgraph-checkpoint>=4.0.0"]
+# viz enables the optional visualization helpers in flexeval.helpers.
+viz = ["matplotlib>=3.10.3"]
+
 [project.urls]
 Homepage = "https://digitalharborfoundation.github.io/FlexEval/"
 GitHub = "https://github.com/DigitalHarborFoundation/FlexEval"
@@ -54,9 +58,20 @@ Issues = "https://github.com/DigitalHarborFoundation/FlexEval/issues"
 dev = [
     "hatch>=1.14.1",
     "jupyter>=1.1.1",
+    # langgraph packages are needed to build/read checkpoint DBs in the tests.
+    # langgraph-checkpoint (the langgraph extra) comes in transitively via langgraph.
+    # langchain-openai is used only to generate integration-test fixtures.
+    "langchain-openai>=0.3.8",
+    "langgraph>=1.0.0",
+    "langgraph-checkpoint-sqlite>=3.0.0",
+    # litellm is needed for the litellm optional dependency.
+    "litellm>=1.74.3",
+    # matplotlib is needed for the viz optional dependency.
     "matplotlib>=3.10.3",
     "pre-commit>=4.2.0",
     "ruff>=0.12.2",
+    # sympy is only used by the integration test fixtures (langgraph_data.py).
+    "sympy>=1.13.3",
 ]
 docs = [
     "autodoc-pydantic>=2.2.0",

src/flexeval/__about__.py

@@ -1 +1 @@
-__version__ = "0.4.3"
+__version__ = "0.5.0"

src/flexeval/configuration/completion_functions.py

@@ -28,7 +28,6 @@
 import os
 from typing import Any, Dict, List
 
-import litellm
 import requests
 import tiktoken
 from openai import OpenAI
@@ -85,6 +84,13 @@ def litellm_completion(
     Returns:
         Dict[str, Any]: The response.
     """
+    try:
+        import litellm
+    except ImportError as e:
+        raise ImportError(
+            "litellm is required for litellm_completion but is not installed. "
+            "Install it with: pip install 'python-flexeval[litellm]'"
+        ) from e
     response = litellm.completion(
         messages=conversation_history,
         model=model,

src/flexeval/data_loader.py

@@ -6,8 +6,6 @@
 import random as rd
 import sqlite3
 
-from langgraph.checkpoint.serde.jsonplus import JsonPlusSerializer
-
 from flexeval.classes.dataset import Dataset
 from flexeval.classes.message import Message
 from flexeval.classes.thread import Thread
@@ -180,6 +178,14 @@ def load_langgraph_sqlite(
     Reads the final checkpoint for each thread and extracts the cumulative
     message list from channel_values.messages. Compatible with langgraph >= 1.0.
     """
+    try:
+        from langgraph.checkpoint.serde.jsonplus import JsonPlusSerializer
+    except ImportError as e:
+        raise ImportError(
+            "langgraph-checkpoint is required to load langgraph SQLite checkpoint "
+            "databases but is not installed. "
+            "Install it with: pip install 'python-flexeval[langgraph]'"
+        ) from e
     serializer = JsonPlusSerializer()
 
     with sqlite3.connect(filename) as conn:

src/flexeval/helpers.py

@@ -21,32 +21,53 @@ def generate_hash():
     return full_hash[:8]
 
 
-def visualize_graph(graph: nx.DiGraph, output_path: str | None = None):
+def _graph_node_label(om) -> str:
+    """Build the display label for a single node in a metric dependency graph."""
+    metric = om.metric
+    return (
+        f"{om.object.__class__.__name__} {om.object.id}\n"
+        f"{metric.get('id')}("
+        f"l={metric.get('metric_level')},"
+        f"r={metric.get('kwargs', {}).get('response')})"
+    )
+
+
+def visualize_graph(graph: nx.DiGraph, ax=None, output_path: str | None = None):
     """Visualize graphs produced by :class:`~flexeval.compute_metrics.MetricGraphBuilder`.
 
     Args:
-        graph (nx.DiGraph): The graph
-        output_path (str | None, optional): If not None, will save the graph as an image using :meth:`matplotlib.pyplot.Figure.savefig`.
+        graph (nx.DiGraph): The graph.
+        ax (matplotlib.axes.Axes | None, optional): An existing Axes to draw into.
+            If None, a new figure and axes are created.
+        output_path (str | None, optional): If not None, save the figure using
+            :meth:`matplotlib.pyplot.Figure.savefig`.
+
+    Returns:
+        tuple: The ``(fig, ax)`` the graph was drawn into.
 
     Raises:
         ImportError: If matplotlib is not installed.
     """
     try:
         import matplotlib.pyplot as plt
-    except ImportError:
-        raise ImportError("matplotlib must be installed to use this helper function.")
-    fig, ax = plt.subplots(1, 1, figsize=(12, 5))
+    except ImportError as e:
+        raise ImportError(
+            "matplotlib is required for visualize_graph but is not installed. "
+            "Install it with: pip install 'python-flexeval[viz]'"
+        ) from e
+    if ax is None:
+        fig, ax = plt.subplots(1, 1, figsize=(12, 5))
+    else:
+        fig = ax.figure
     pos = nx.spring_layout(graph)
     nx.draw(graph, ax=ax, pos=pos)
     nx.draw_networkx_labels(
         graph,
         font_size=8,
         ax=ax,
         pos=pos,
-        labels={
-            om: f"{om.object.__class__.__name__} {om.object.id}\n{om.metric.get('id')}(l={om.metric.get('metric_level')},r={om.metric.get('kwargs', {}).get('response')})"
-            for om in graph
-        },
+        labels={om: _graph_node_label(om) for om in graph},
     )
     if output_path is not None:
         fig.savefig(output_path)
+    return fig, ax