diff --git a/.claude/settings.local.json b/.claude/settings.local.json index b2f408f..d646372 100644 --- a/.claude/settings.local.json +++ b/.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)" ] } } diff --git a/docs/api-reference.md b/docs/api-reference.md index f9a1fd5..3ffb660 100644 --- a/docs/api-reference.md +++ b/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. diff --git a/docs/changelog.md b/docs/changelog.md index b237daf..5772fc5 100644 --- a/docs/changelog.md +++ b/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 diff --git a/docs/index.md b/docs/index.md index 52e46c0..4525ccd 100644 --- a/docs/index.md +++ b/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) - ``` - --- diff --git a/docs/installation.md b/docs/installation.md index cb8fa1e..469f7d4 100644 --- a/docs/installation.md +++ b/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: diff --git a/docs/integrations/langgraph.md b/docs/integrations/langgraph.md index 579f27e..df27f9a 100644 --- a/docs/integrations/langgraph.md +++ b/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 diff --git a/docs/quickstart.md b/docs/quickstart.md index f7b2a3f..9c06e85 100644 --- a/docs/quickstart.md +++ b/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}") ``` diff --git a/examples/langgraph_demo.py b/examples/langgraph_demo.py index c960e28..4c0250a 100644 --- a/examples/langgraph_demo.py +++ b/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,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[langgraph,openai] typing_extensions") + print("Run: pip install shekel[openai] langgraph typing_extensions") return api_key = os.environ.get("OPENAI_API_KEY") @@ -29,6 +29,8 @@ def main() -> None: 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,11 +52,11 @@ 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}") @@ -62,22 +64,10 @@ def call_llm(state: State) -> State: 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"}, diff --git a/pyproject.toml b/pyproject.toml index cad4d95..cc7e2f6 100644 --- a/pyproject.toml +++ b/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" diff --git a/shekel/integrations/__init__.py b/shekel/integrations/__init__.py index c263b9d..7015b49 100644 --- a/shekel/integrations/__init__.py +++ b/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") diff --git a/shekel/integrations/langgraph.py b/shekel/integrations/langgraph.py deleted file mode 100644 index 9514352..0000000 --- a/shekel/integrations/langgraph.py +++ /dev/null @@ -1,50 +0,0 @@ -"""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 deleted file mode 100644 index 31c055a..0000000 --- a/tests/integrations/test_langgraph_integration.py +++ /dev/null @@ -1,77 +0,0 @@ -"""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