From 56283c69504e6f879a413dc197f938cef3ae8d81 Mon Sep 17 00:00:00 2001 From: arieradle Date: Thu, 12 Mar 2026 12:23:46 +0200 Subject: [PATCH 1/6] feat: add LangGraph integration helper (v0.2.6) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adds shekel.integrations.langgraph.budgeted_graph() — a convenience context manager wrapping shekel.budget() for LangGraph workflows. Existing OpenAI/Anthropic patches cover cost tracking automatically. Co-Authored-By: Claude Sonnet 4.6 --- pyproject.toml | 7 +- shekel/integrations/__init__.py | 5 ++ shekel/integrations/langgraph.py | 50 ++++++++++++ .../test_langgraph_integration.py | 76 +++++++++++++++++++ 4 files changed, 137 insertions(+), 1 deletion(-) create mode 100644 shekel/integrations/langgraph.py create mode 100644 tests/integrations/test_langgraph_integration.py diff --git a/pyproject.toml b/pyproject.toml index a53c723..058ef28 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -62,7 +62,8 @@ dependencies = [] openai = ["openai>=1.0.0"] anthropic = ["anthropic>=0.7.0"] langfuse = ["langfuse>=2.0.0"] -all = ["openai>=1.0.0", "anthropic>=0.7.0", "langfuse>=2.0.0"] +langgraph = ["langgraph>=0.1.0"] +all = ["openai>=1.0.0", "anthropic>=0.7.0", "langfuse>=2.0.0", "langgraph>=0.1.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 = [ @@ -127,3 +128,7 @@ ignore_missing_imports = true [[tool.mypy.overrides]] module = "langfuse" ignore_missing_imports = true + +[[tool.mypy.overrides]] +module = "langgraph" +ignore_missing_imports = true diff --git a/shekel/integrations/__init__.py b/shekel/integrations/__init__.py index 7015b49..c263b9d 100644 --- a/shekel/integrations/__init__.py +++ b/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") diff --git a/shekel/integrations/langgraph.py b/shekel/integrations/langgraph.py new file mode 100644 index 0000000..9514352 --- /dev/null +++ b/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 diff --git a/tests/integrations/test_langgraph_integration.py b/tests/integrations/test_langgraph_integration.py new file mode 100644 index 0000000..8b1d127 --- /dev/null +++ b/tests/integrations/test_langgraph_integration.py @@ -0,0 +1,76 @@ +"""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 From 3b925f63e8a65bcfededeb00b2bd339cf36c7bd7 Mon Sep 17 00:00:00 2001 From: arieradle Date: Thu, 12 Mar 2026 13:43:30 +0200 Subject: [PATCH 2/6] docs: update docs for LangGraph integration and budgeted_graph() - Add budgeted_graph() to api-reference.md with full signature and examples - Add shekel[langgraph] extra to installation.md (section, dep table, troubleshooting) - Fix install command in langgraph.md and langgraph_demo.py to use shekel[langgraph] Co-Authored-By: Claude Sonnet 4.6 --- docs/api-reference.md | 73 ++++++++++++++++++++++++++++++++++ docs/installation.md | 23 +++++++++++ docs/integrations/langgraph.md | 8 +++- examples/langgraph_demo.py | 4 +- 4 files changed, 105 insertions(+), 3 deletions(-) diff --git a/docs/api-reference.md b/docs/api-reference.md index 3ffb660..f9a1fd5 100644 --- a/docs/api-reference.md +++ b/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. diff --git a/docs/installation.md b/docs/installation.md index 469f7d4..cb8fa1e 100644 --- a/docs/installation.md +++ b/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: diff --git a/docs/integrations/langgraph.md b/docs/integrations/langgraph.md index a8a6b51..579f27e 100644 --- a/docs/integrations/langgraph.md +++ b/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 diff --git a/examples/langgraph_demo.py b/examples/langgraph_demo.py index 72f0e99..c960e28 100644 --- a/examples/langgraph_demo.py +++ b/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") From f7dabb4640bbdb0ecbbe1e0a3af0b71a3949f628 Mon Sep 17 00:00:00 2001 From: arieradle Date: Thu, 12 Mar 2026 13:52:34 +0200 Subject: [PATCH 3/6] style: reformat test_langgraph_integration.py with black Co-Authored-By: Claude Sonnet 4.6 --- tests/integrations/test_langgraph_integration.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/integrations/test_langgraph_integration.py b/tests/integrations/test_langgraph_integration.py index 8b1d127..31c055a 100644 --- a/tests/integrations/test_langgraph_integration.py +++ b/tests/integrations/test_langgraph_integration.py @@ -72,5 +72,6 @@ def test_openai_call_inside_records_spend(self): 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 From 91271b9bb66ef91010d45d3fbdd9e8cc788138b8 Mon Sep 17 00:00:00 2001 From: arieradle Date: Thu, 12 Mar 2026 13:57:55 +0200 Subject: [PATCH 4/6] fix: suppress mypy syntax errors in _pytest for python_version=3.9 mypy 1.19.1 follows imports into _pytest/terminal.py which uses match syntax (Python 3.10+), causing a hard [syntax] error when python_version = "3.9". Co-Authored-By: Claude Sonnet 4.6 --- pyproject.toml | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/pyproject.toml b/pyproject.toml index 3f76037..6df4952 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -140,3 +140,7 @@ ignore_missing_imports = true [[tool.mypy.overrides]] module = "litellm" ignore_missing_imports = true + +[[tool.mypy.overrides]] +module = "_pytest.*" +ignore_errors = true From ec7602ba9b425a79b18068e198648f751d3965ee Mon Sep 17 00:00:00 2001 From: arieradle Date: Thu, 12 Mar 2026 14:14:30 +0200 Subject: [PATCH 5/6] fix: exclude _pytest from mypy path scanning mypy 1.19.1 traverses into _pytest/terminal.py (site-packages) which uses match syntax incompatible with python_version = "3.9". Adding exclude = ["/_pytest/"] prevents the path from being scanned. Co-Authored-By: Claude Sonnet 4.6 --- pyproject.toml | 1 + 1 file changed, 1 insertion(+) diff --git a/pyproject.toml b/pyproject.toml index 6df4952..bdb5f80 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -124,6 +124,7 @@ python_version = "3.9" strict = true warn_return_any = true warn_unused_configs = true +exclude = ["/_pytest/"] [[tool.mypy.overrides]] module = "tokencost" From 41247fb16828b69e40b7b8c3ca82f06111a0b6f6 Mon Sep 17 00:00:00 2001 From: arieradle Date: Thu, 12 Mar 2026 14:26:20 +0200 Subject: [PATCH 6/6] fix: add follow_imports=skip for _pytest.* to fix mypy on Python 3.11 On Python 3.11, mypy 1.19.1 follows imports into _pytest/terminal.py which uses match syntax (Python 3.10+), causing a fatal [syntax] error. Python 3.12 can parse match natively so the issue only surfaces in CI. Using follow_imports="skip" prevents mypy from traversing into _pytest.* regardless of how the import is reached. Co-Authored-By: Claude Sonnet 4.6 --- pyproject.toml | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index bdb5f80..cad4d95 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -143,5 +143,6 @@ module = "litellm" ignore_missing_imports = true [[tool.mypy.overrides]] -module = "_pytest.*" +module = ["_pytest", "_pytest.*"] +follow_imports = "skip" ignore_errors = true