feat: add LangGraph integration helper (v0.2.6)

docs/api-reference.md

@@ -251,6 +251,79 @@ print(w.tree())
 
 ---
 
+## `budgeted_graph()`
+
+Convenience context manager for running LangGraph graphs inside a shekel budget. Requires `pip install shekel[langgraph]`.
+
+### Signature
+
+```python
+from shekel.integrations.langgraph import budgeted_graph
+
+def budgeted_graph(
+    max_usd: float,
+    **budget_kwargs: Any,
+) -> Generator[Budget, None, None]
+```
+
+### Parameters
+
+| Parameter | Type | Description |
+|-----------|------|-------------|
+| `max_usd` | `float` | Maximum spend in USD before `BudgetExceededError` is raised. |
+| `**budget_kwargs` | `Any` | Forwarded to `budget()` — e.g. `name`, `warn_at`, `fallback`, `max_llm_calls`. |
+
+### Yields
+
+The active `Budget` instance.
+
+### Examples
+
+#### Basic Usage
+
+```python
+from shekel.integrations.langgraph import budgeted_graph
+
+app = graph.compile()
+
+with budgeted_graph(max_usd=0.50) as b:
+    result = app.invoke({"question": "What is 2+2?", "answer": ""})
+    print(f"Cost: ${b.spent:.4f}")
+```
+
+#### With Budget Options
+
+```python
+with budgeted_graph(
+    max_usd=1.00,
+    name="research-graph",
+    warn_at=0.8,
+    fallback={"at_pct": 0.8, "model": "gpt-4o-mini"},
+) as b:
+    result = app.invoke(state)
+    print(f"Spent: ${b.spent:.4f} / ${b.limit:.2f}")
+    if b.model_switched:
+        print(f"Switched to fallback at ${b.switched_at_usd:.4f}")
+```
+
+#### Equivalent to `budget()`
+
+`budgeted_graph()` is a thin wrapper — these are identical:
+
+```python
+from shekel.integrations.langgraph import budgeted_graph
+from shekel import budget
+
+# These are equivalent
+with budgeted_graph(max_usd=0.50) as b:
+    app.invoke(state)
+
+with budget(max_usd=0.50) as b:
+    app.invoke(state)
+```
+
+---
+
 ## `@with_budget`
 
 Decorator that wraps functions with a budget context.

docs/installation.md

@@ -43,6 +43,20 @@ For access to OpenAI, Anthropic, Gemini, Cohere, Ollama, Azure, Bedrock, and 90+
 pip install shekel[litellm]
 ```
 
+### LangGraph
+
+For the `budgeted_graph()` convenience helper with LangGraph:
+
+```bash
+pip install shekel[langgraph]
+```
+
+Or combined with your LLM provider:
+
+```bash
+pip install shekel[langgraph,openai]
+```
+
 ### Extended Model Support (400+ Models)
 
 For support of 400+ models via [tokencost](https://github.com/AgentOps-AI/tokencost):
@@ -107,6 +121,7 @@ Shekel has zero required dependencies beyond the Python standard library. The Op
 | `openai>=1.0.0` | Optional | Track OpenAI API costs |
 | `anthropic>=0.7.0` | Optional | Track Anthropic API costs |
 | `litellm>=1.0.0` | Optional | Track costs via LiteLLM (100+ providers) |
+| `langgraph>=0.1.0` | Optional | `budgeted_graph()` convenience helper for LangGraph |
 | `tokencost>=0.1.0` | Optional | Support 400+ models |
 | `click>=8.0.0` | Optional | CLI tools |
 
@@ -136,6 +151,14 @@ If you see this error, install LiteLLM:
 pip install shekel[litellm]
 ```
 
+### ImportError: No module named 'langgraph'
+
+If you see this error when using `budgeted_graph()`, install LangGraph:
+
+```bash
+pip install shekel[langgraph]
+```
+
 ### Model pricing not found
 
 For models not in shekel's built-in pricing table:

docs/integrations/langgraph.md

@@ -5,7 +5,13 @@ Shekel works seamlessly with [LangGraph](https://github.com/langchain-ai/langgra
 ## Installation
 
 ```bash
-pip install shekel[openai] "langgraph>=0.2"
+pip install shekel[langgraph]
+```
+
+Or combined with your LLM provider:
+
+```bash
+pip install shekel[langgraph,openai]
 ```
 
 ## Convenience Helper

examples/langgraph_demo.py

@@ -1,4 +1,4 @@
-# Requires: pip install shekel[openai] "langgraph>=0.2"
+# Requires: pip install shekel[langgraph,openai]
 """
 LangGraph demo: budget enforcement with the budgeted_graph() helper.
 
@@ -21,7 +21,7 @@ def main() -> None:
         from typing_extensions import TypedDict
     except ImportError as e:
         print(f"Missing dependency: {e}")
-        print("Run: pip install shekel[openai] 'langgraph>=0.2' typing_extensions")
+        print("Run: pip install shekel[langgraph,openai] typing_extensions")
         return
 
     api_key = os.environ.get("OPENAI_API_KEY")

pyproject.toml

@@ -62,8 +62,9 @@ dependencies = []
 openai = ["openai>=1.0.0"]
 anthropic = ["anthropic>=0.7.0"]
 langfuse = ["langfuse>=2.0.0"]
+langgraph = ["langgraph>=0.1.0"]
 litellm = ["litellm>=1.0.0"]
-all = ["openai>=1.0.0", "anthropic>=0.7.0", "langfuse>=2.0.0", "litellm>=1.0.0"]
+all = ["openai>=1.0.0", "anthropic>=0.7.0", "langfuse>=2.0.0", "langgraph>=0.1.0", "litellm>=1.0.0"]
 all-models = ["openai>=1.0.0", "anthropic>=0.7.0", "langfuse>=2.0.0", "tokencost>=0.1.0"]
 cli = ["click>=8.0.0"]
 dev = [
@@ -123,6 +124,7 @@ python_version = "3.9"
 strict = true
 warn_return_any = true
 warn_unused_configs = true
+exclude = ["/_pytest/"]
 
 [[tool.mypy.overrides]]
 module = "tokencost"
@@ -132,6 +134,15 @@ ignore_missing_imports = true
 module = "langfuse"
 ignore_missing_imports = true
 
+[[tool.mypy.overrides]]
+module = "langgraph"
+ignore_missing_imports = true
+
 [[tool.mypy.overrides]]
 module = "litellm"
 ignore_missing_imports = true
+
+[[tool.mypy.overrides]]
+module = ["_pytest", "_pytest.*"]
+follow_imports = "skip"
+ignore_errors = true

shekel/integrations/__init__.py

@@ -18,3 +18,8 @@
 except ImportError:
     # langfuse is an optional dependency
     pass
+
+# LangGraph integration helper
+from shekel.integrations import langgraph  # noqa: F401
+
+__all__.append("langgraph")

shekel/integrations/langgraph.py

@@ -0,0 +1,50 @@
+"""LangGraph integration helper for Shekel LLM cost tracking.
+
+Shekel works with LangGraph out of the box — existing OpenAI/Anthropic patches
+automatically intercept all LLM calls inside graph nodes. This module provides
+a convenience context manager so you don't need to import budget directly.
+
+Example:
+    >>> from shekel.integrations.langgraph import budgeted_graph
+    >>>
+    >>> with budgeted_graph(max_usd=0.10) as b:
+    ...     result = app.invoke({"question": "What is 2+2?"})
+    ...     print(f"Spent: ${b.spent:.4f}")
+"""
+
+from __future__ import annotations
+
+from collections.abc import Generator
+from contextlib import contextmanager
+from typing import Any
+
+from shekel import budget as shekel_budget
+from shekel._budget import Budget
+
+
+@contextmanager
+def budgeted_graph(max_usd: float, **budget_kwargs: Any) -> Generator[Budget, None, None]:
+    """Run a LangGraph graph inside a shekel budget.
+
+    A convenience wrapper around :func:`shekel.budget` for LangGraph workflows.
+    All LLM calls made within LangGraph nodes are automatically tracked via
+    shekel's existing OpenAI/Anthropic patches — no extra configuration needed.
+
+    Args:
+        max_usd: Maximum spend in USD before :class:`shekel.BudgetExceededError` is raised.
+        **budget_kwargs: Additional keyword arguments forwarded to :func:`shekel.budget`
+            (e.g. ``name``, ``warn_at``, ``fallback``, ``max_llm_calls``).
+
+    Yields:
+        The active :class:`shekel._budget.Budget` instance.
+
+    Raises:
+        shekel.BudgetExceededError: If spend exceeds ``max_usd``.
+
+    Example:
+        >>> with budgeted_graph(max_usd=0.50, name="research-graph") as b:
+        ...     result = app.invoke(state)
+        ...     print(f"Spent: ${b.spent:.4f} / ${b.limit:.2f}")
+    """
+    with shekel_budget(max_usd=max_usd, **budget_kwargs) as b:
+        yield b

tests/integrations/test_langgraph_integration.py

@@ -0,0 +1,77 @@
+"""TDD tests for LangGraph integration helper."""
+
+from __future__ import annotations
+
+from unittest.mock import MagicMock, patch
+
+
+class TestBudgetedGraphExists:
+    def test_module_importable(self):
+        from shekel.integrations import langgraph  # noqa: F401
+
+    def test_budgeted_graph_is_callable(self):
+        from shekel.integrations.langgraph import budgeted_graph
+
+        assert callable(budgeted_graph)
+
+
+class TestBudgetedGraphContextManager:
+    def test_yields_budget_object(self):
+        from shekel.integrations.langgraph import budgeted_graph
+
+        with budgeted_graph(max_usd=1.0) as b:
+            assert b is not None
+            assert hasattr(b, "spent")
+            assert hasattr(b, "limit")
+
+    def test_budget_limit_is_set(self):
+        from shekel.integrations.langgraph import budgeted_graph
+
+        with budgeted_graph(max_usd=0.50) as b:
+            assert b.limit == 0.50
+
+    def test_budget_kwargs_forwarded(self):
+        from shekel.integrations.langgraph import budgeted_graph
+
+        with budgeted_graph(max_usd=1.0, name="my-graph") as b:
+            assert b.name == "my-graph"
+
+    def test_spent_starts_at_zero(self):
+        from shekel.integrations.langgraph import budgeted_graph
+
+        with budgeted_graph(max_usd=1.0) as b:
+            assert b.spent == 0.0
+
+    def test_budget_exceeded_error_propagates(self):
+        from shekel import BudgetExceededError
+        from shekel.integrations.langgraph import budgeted_graph
+
+        with patch("shekel._patch._originals", {"openai_sync": MagicMock()}):
+            try:
+                with budgeted_graph(max_usd=0.000001) as b:
+                    # Force a spend by directly recording
+                    b._record_spend(1.0, "gpt-4o", {"input": 1000, "output": 500})
+            except BudgetExceededError:
+                pass  # Expected — budget exceeded propagates correctly
+
+
+class TestBudgetedGraphRecordsCost:
+    def test_openai_call_inside_records_spend(self):
+        """OpenAI calls inside budgeted_graph are tracked via existing patches."""
+        from shekel.integrations.langgraph import budgeted_graph
+
+        mock_response = MagicMock()
+        mock_response.usage.prompt_tokens = 100
+        mock_response.usage.completion_tokens = 50
+        mock_response.model = "gpt-4o-mini"
+
+        original_fn = MagicMock(return_value=mock_response)
+
+        with patch("shekel._patch._originals", {"openai_sync": original_fn}):
+
+            with budgeted_graph(max_usd=1.0) as b:
+                # Simulate what the openai wrapper does
+                from shekel._patch import _record
+
+                _record(100, 50, "gpt-4o-mini")
+                assert b.spent > 0