feat: LangGraph integration docs + mypy CI fix

.claude/settings.local.json

@@ -103,7 +103,15 @@
       "Bash(python -m pytest tests/test_fallback.py::test_fallback_model_rewritten_in_kwargs tests/test_fallback.py::test_fallback_spent_tracks_separately -xvs 2>&1 | tail -20)",
       "Bash(python -m pytest tests/test_fallback.py::test_fallback_model_rewritten_in_kwargs -xvs 2>&1 | tail -20)",
       "Bash(python -m pytest tests/ -q --tb=no 2>&1 | tail -10)",
-      "Bash(python -m pytest tests/test_fallback.py::test_fallback_small_call_within_budget -xvs 2>&1 | grep -A 5 \"assert\")"
+      "Bash(python -m pytest tests/test_fallback.py::test_fallback_small_call_within_budget -xvs 2>&1 | grep -A 5 \"assert\")",
+      "Bash(GIT_EDITOR=true git merge --continue)",
+      "Bash(python -m ruff check . 2>&1)",
+      "Bash(python -m mypy shekel/ 2>&1)",
+      "Bash(python -m black --check . 2>&1)",
+      "Bash(python -m isort --check-only . 2>&1)",
+      "Bash(python -m pytest --cov=shekel --cov-report=term-missing 2>&1)",
+      "Bash(python -c \"\nimport yaml, os\nwith open\\('/mnt/c/Users/elish/code/shekel/mkdocs.yml'\\) as f:\n    cfg = yaml.safe_load\\(f\\)\n\ndef extract_paths\\(nav\\):\n    for item in nav:\n        if isinstance\\(item, dict\\):\n            for k, v in item.items\\(\\):\n                if isinstance\\(v, str\\):\n                    yield v\n                elif isinstance\\(v, list\\):\n                    yield from extract_paths\\(v\\)\n\nmissing = []\nfor path in extract_paths\\(cfg['nav']\\):\n    full = f'/mnt/c/Users/elish/code/shekel/docs/{path}'\n    if not os.path.exists\\(full\\):\n        missing.append\\(path\\)\n\nif missing:\n    print\\('MISSING:', missing\\)\nelse:\n    print\\('All nav files exist.'\\)\n\")",
+      "Bash(python -m mypy shekel/ --verbose 2>&1 | grep \"_pytest\" | head -10)"
     ]
   }
 }

docs/api-reference.md

@@ -251,79 +251,6 @@ 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/changelog.md

@@ -48,12 +48,6 @@ All notable changes to this project are documented here. For detailed informatio
 - Tracks costs across all 100+ providers LiteLLM supports (Gemini, Cohere, Ollama, Azure, Bedrock, Mistral, and more)
 - Model names with provider prefix (e.g. `gemini/gemini-1.5-flash`) pass through to the pricing engine
 
-**LangGraph integration helper**
-
-- `from shekel.integrations.langgraph import budgeted_graph`
-- `budgeted_graph(max_usd, **kwargs)` — convenience context manager wrapping `budget()` for LangGraph workflows
-- Install with `pip install shekel[langgraph]`
-
 ## [0.2.5] - 2026-03-11
 
 ### 🔧 Extensible Provider Architecture

docs/index.md

@@ -255,18 +255,6 @@ print(f"Remaining: ${b.remaining:.4f}")
     pip install shekel[litellm]
     ```
 
--   :material-graph:{ .lg .middle } **[LangGraph Helper](integrations/langgraph.md)**
-
-    ---
-
-    New `budgeted_graph()` context manager for cleaner LangGraph integration.
-
-    ```python
-    from shekel.integrations.langgraph import budgeted_graph
-    with budgeted_graph(max_usd=0.50) as b:
-        result = app.invoke(state)
-    ```
-
 </div>
 
 ---

docs/installation.md

@@ -43,20 +43,6 @@ 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):
@@ -121,7 +107,6 @@ 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 |
 
@@ -151,14 +136,6 @@ 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,35 +5,14 @@ Shekel works seamlessly with [LangGraph](https://github.com/langchain-ai/langgra
 ## Installation
 
 ```bash
-pip install shekel[langgraph]
+pip install shekel[openai]   # or shekel[anthropic], shekel[litellm]
 ```
 
-Or combined with your LLM provider:
-
-```bash
-pip install shekel[langgraph,openai]
-```
-
-## Convenience Helper
-
-Shekel provides a `budgeted_graph()` context manager so you don't need to import `budget` directly:
-
-```python
-from shekel.integrations.langgraph import budgeted_graph
-
-app = graph.compile()
-
-with budgeted_graph(max_usd=0.50, name="research-graph") as b:
-    result = app.invoke({"question": "What is 2+2?", "answer": ""})
-    print(f"Answer: {result['answer']}")
-    print(f"Cost: ${b.spent:.4f}")
-```
-
-It accepts the same keyword arguments as `budget()` (`name`, `warn_at`, `fallback`, `max_llm_calls`, etc.) and yields the active budget object.
+No special LangGraph extra needed — shekel works with LangGraph out of the box by intercepting all LLM calls inside graph nodes automatically.
 
 ## Basic Integration
 
-You can also use `budget()` directly — they are equivalent:
+Use `budget()` directly:
 
 ```python
 from langgraph.graph import StateGraph, END

docs/quickstart.md

@@ -257,15 +257,15 @@ print(f"Cost: ${b.spent:.4f}")
 
 ### LangGraph
 
-Use `budget()` directly, or the convenience `budgeted_graph()` helper:
+Shekel works with LangGraph out of the box — just wrap with `budget()`:
 
 ```python
-from shekel.integrations.langgraph import budgeted_graph
+from shekel import budget
 
 # Your graph definition here
 app = graph.compile()
 
-with budgeted_graph(max_usd=0.50, name="my-graph") as b:
+with budget(max_usd=0.50, name="my-graph") as b:
     result = app.invoke({"question": "What is 2+2?"})
 print(f"Graph execution cost: ${b.spent:.4f}")
 ```

examples/langgraph_demo.py

@@ -1,18 +1,18 @@
-# Requires: pip install shekel[langgraph,openai]
+# Requires: pip install shekel[openai]
 """
-LangGraph demo: budget enforcement with the budgeted_graph() helper.
+LangGraph demo: budget enforcement with shekel.
+
+Shekel works with LangGraph out of the box — just wrap with budget().
+All LLM calls inside graph nodes are automatically tracked.
 
 Shows three patterns:
-1. budgeted_graph() convenience helper (recommended)
-2. budget() directly — equivalent but more verbose
-3. Fallback model when budget threshold is reached
+1. Basic budget enforcement
+2. Fallback model when budget threshold is reached
+3. Nested budgets for multi-node graphs
 """
 
 import os
 
-from shekel import BudgetExceededError, budget
-from shekel.integrations.langgraph import budgeted_graph
-
 
 def main() -> None:
     try:
@@ -21,14 +21,16 @@ def main() -> None:
         from typing_extensions import TypedDict
     except ImportError as e:
         print(f"Missing dependency: {e}")
-        print("Run: pip install shekel[langgraph,openai] typing_extensions")
+        print("Run: pip install shekel[openai] langgraph typing_extensions")
         return
 
     api_key = os.environ.get("OPENAI_API_KEY")
     if not api_key:
         print("Set OPENAI_API_KEY to run this demo.")
         return
 
+    from shekel import BudgetExceededError, budget
+
     client = openai.OpenAI(api_key=api_key)
 
     class State(TypedDict):
@@ -50,34 +52,22 @@ def call_llm(state: State) -> State:
     app = graph.compile()
 
     # ------------------------------------------------------------------
-    # 1. budgeted_graph() — recommended convenience helper
+    # 1. Basic budget enforcement
     # ------------------------------------------------------------------
-    print("=== budgeted_graph() helper ===")
+    print("=== Basic budget enforcement ===")
     try:
-        with budgeted_graph(max_usd=0.10, name="demo", warn_at=0.8) as b:
+        with budget(max_usd=0.10, name="demo", warn_at=0.8) as b:
             result = app.invoke({"question": "What is 2+2?", "answer": ""})
             print(f"Answer: {result['answer']}")
             print(f"Spent: ${b.spent:.4f} / ${b.limit:.2f}")
     except BudgetExceededError as e:
         print(f"Budget exceeded: {e}")
 
     # ------------------------------------------------------------------
-    # 2. budget() directly — same result, more explicit
-    # ------------------------------------------------------------------
-    print("\n=== budget() directly ===")
-    try:
-        with budget(max_usd=0.10) as b:
-            result = app.invoke({"question": "Name a planet.", "answer": ""})
-            print(f"Answer: {result['answer']}")
-            print(f"Spent: ${b.spent:.4f}")
-    except BudgetExceededError as e:
-        print(f"Budget exceeded: {e}")
-
-    # ------------------------------------------------------------------
-    # 3. Fallback model when threshold is reached
+    # 2. Fallback model when threshold is reached
     # ------------------------------------------------------------------
     print("\n=== Fallback model ===")
-    with budgeted_graph(
+    with budget(
         max_usd=0.001,
         name="fallback-demo",
         fallback={"at_pct": 0.5, "model": "gpt-4o-mini"},

pyproject.toml

@@ -62,9 +62,8 @@ 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", "langgraph>=0.1.0", "litellm>=1.0.0"]
+all = ["openai>=1.0.0", "anthropic>=0.7.0", "langfuse>=2.0.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 = [
@@ -134,9 +133,6 @@ ignore_missing_imports = true
 module = "langfuse"
 ignore_missing_imports = true
 
-[[tool.mypy.overrides]]
-module = "langgraph"
-ignore_missing_imports = true
 
 [[tool.mypy.overrides]]
 module = "litellm"

shekel/integrations/__init__.py

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