Merge pull request #3 from Floe-Labs/feat/langchain-adapter

feat: add LangChain budget-guard integration

Author: achris7

README.md

@@ -108,10 +108,29 @@ response = guarded_completion(guard, model="gpt-4o", messages=[...])
 Prefer the LiteLLM-native callback? Register `budget_guard_callback(guard)` on
 `litellm.callbacks`.
 
+### LangChain
+
+```bash
+pip install floe-guard[langchain] langchain-openai   # langchain-openai only for the ChatOpenAI example below
+```
+
+```python
+from langchain_openai import ChatOpenAI
+from floe_guard import BudgetGuard
+from floe_guard.integrations.langchain import budget_guard_callback_handler
+
+guard = BudgetGuard(limit_usd=1.00)
+llm = ChatOpenAI(model="gpt-4o", callbacks=[budget_guard_callback_handler(guard)])
+llm.invoke("hello")            # checks budget before the call, records spend after
+```
+
+The handler checks the budget on LLM start (raising `BudgetExceeded` aborts the
+call before it runs) and records token usage on LLM end.
+
 ### Coming next
 
-LangChain (callback) and the Vercel AI SDK (TypeScript middleware) are next. Open
-an issue if you want one sooner.
+The Vercel AI SDK (TypeScript middleware) is next. Open an issue if you want one
+sooner.
 
 ## Honest about what this is
 

pyproject.toml

@@ -10,7 +10,7 @@ readme = "README.md"
 requires-python = ">=3.10"
 license = { file = "LICENSE" }
 authors = [{ name = "Floe Labs" }]
-keywords = ["llm", "agents", "budget", "guardrail", "crewai", "litellm", "cost", "openai", "anthropic"]
+keywords = ["llm", "agents", "budget", "guardrail", "crewai", "litellm", "langchain", "cost", "openai", "anthropic"]
 classifiers = [
     "Development Status :: 4 - Beta",
     "Intended Audience :: Developers",
@@ -26,6 +26,7 @@ dependencies = []
 [project.optional-dependencies]
 litellm = ["litellm>=1.0"]
 crewai = ["crewai>=0.30", "litellm>=1.0"]
+langchain = ["langchain-core>=0.3"]
 dev = ["pytest>=7.0", "ruff>=0.4"]
 
 [project.urls]

src/floe_guard/integrations/__init__.py

@@ -4,4 +4,5 @@
 
     pip install floe-guard[litellm]
     pip install floe-guard[crewai]
+    pip install floe-guard[langchain]
 """

src/floe_guard/integrations/langchain.py

@@ -0,0 +1,126 @@
+"""LangChain adapter (optional extra: ``pip install floe-guard[langchain]``).
+
+:func:`budget_guard_callback_handler` builds a LangChain callback handler you pass
+to any LLM or chat model. It checks the budget *before* the call
+(``on_llm_start`` / ``on_chat_model_start``) — raising
+:class:`~floe_guard.BudgetExceeded` aborts the call so it never runs — and records
+spend *after* the response (``on_llm_end``).
+
+    from langchain_openai import ChatOpenAI
+    from floe_guard import BudgetGuard
+    from floe_guard.integrations.langchain import budget_guard_callback_handler
+
+    guard = BudgetGuard(limit_usd=1.00)
+    llm = ChatOpenAI(model="gpt-4o", callbacks=[budget_guard_callback_handler(guard)])
+    llm.invoke("hello")
+
+Every priced response routes through the same :class:`~floe_guard.BudgetGuard` as
+the other adapters, so token usage is priced via the bundled cost map.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+from ..guard import BudgetGuard
+
+
+def _require_langchain() -> Any:
+    try:
+        import langchain_core  # noqa: F401
+    except ImportError as e:  # pragma: no cover - exercised only without the extra
+        raise ImportError(
+            "The LangChain adapter requires langchain-core. "
+            "Install with: pip install floe-guard[langchain]"
+        ) from e
+    return langchain_core
+
+
+def _model_from_result(response: Any) -> str:
+    """Pull the model name from a LangChain ``LLMResult``."""
+    llm_output = getattr(response, "llm_output", None)
+    if isinstance(llm_output, dict):
+        model = llm_output.get("model_name") or llm_output.get("model")
+        if model:
+            return str(model)
+    return ""
+
+
+def _usage_from_result(response: Any) -> tuple[int, int]:
+    """Pull (prompt_tokens, completion_tokens) from a LangChain ``LLMResult``.
+
+    Handles both shapes LangChain emits: the provider ``token_usage`` block in
+    ``llm_output`` (e.g. OpenAI: ``prompt_tokens``/``completion_tokens``) and the
+    standardized ``usage_metadata`` on a message (``input_tokens``/``output_tokens``).
+    """
+    llm_output = getattr(response, "llm_output", None)
+    if isinstance(llm_output, dict):
+        usage = llm_output.get("token_usage") or llm_output.get("usage")
+        if isinstance(usage, dict):
+            prompt = int(usage.get("prompt_tokens", 0) or 0)
+            completion = int(usage.get("completion_tokens", 0) or 0)
+            if prompt > 0 or completion > 0:
+                return prompt, completion
+
+    # Fall back to per-message usage_metadata (input_tokens/output_tokens), the
+    # provider-agnostic shape newer chat models attach to each generation.
+    prompt = completion = 0
+    for batch in getattr(response, "generations", None) or []:
+        for gen in batch:
+            meta = getattr(getattr(gen, "message", None), "usage_metadata", None)
+            if isinstance(meta, dict):
+                prompt += int(meta.get("input_tokens", 0) or 0)
+                completion += int(meta.get("output_tokens", 0) or 0)
+    return prompt, completion
+
+
+def _record_result(guard: BudgetGuard, response: Any) -> None:
+    model = _model_from_result(response)
+    prompt_tokens, completion_tokens = _usage_from_result(response)
+    if prompt_tokens <= 0 and completion_tokens <= 0:
+        # No tokens were spent — nothing to meter.
+        return
+    # There IS spend to account for. Route it through record() even when the model
+    # id is missing, so the guard's configured policy applies (fail-closed → warn +
+    # raise; fail-open → warn + skip). Silently skipping here would let a real,
+    # completed call go unmetered and skew the next check().
+    guard.record(model, prompt_tokens, completion_tokens)
+
+
+def budget_guard_callback_handler(guard: BudgetGuard) -> Any:
+    """Build a LangChain callback handler that enforces ``guard`` on every call.
+
+    Pass it to any LLM or chat model::
+
+        llm = ChatOpenAI(model="gpt-4o", callbacks=[budget_guard_callback_handler(guard)])
+
+    ``on_llm_start``/``on_chat_model_start`` run ``guard.check()`` (raising
+    :class:`~floe_guard.BudgetExceeded` to abort), and ``on_llm_end`` records the
+    response's token cost.
+    """
+    _require_langchain()
+    from langchain_core.callbacks import BaseCallbackHandler
+
+    class BudgetGuardCallbackHandler(BaseCallbackHandler):  # type: ignore[misc]
+        # LangChain swallows exceptions raised inside callbacks by default, which
+        # would let a blocked call run anyway. raise_error=True propagates
+        # BudgetExceeded so the call is actually aborted — the whole point.
+        raise_error = True
+
+        def __init__(self) -> None:
+            super().__init__()
+            self.guard = guard
+
+        def on_llm_start(self, serialized: Any, prompts: Any, **kwargs: Any) -> None:
+            self.guard.check()
+
+        def on_chat_model_start(self, serialized: Any, messages: Any, **kwargs: Any) -> None:
+            self.guard.check()
+
+        def on_llm_end(self, response: Any, **kwargs: Any) -> None:
+            _record_result(self.guard, response)
+
+    return BudgetGuardCallbackHandler()
+
+
+__all__ = ["budget_guard_callback_handler"]

tests/test_langchain_adapter.py

@@ -0,0 +1,134 @@
+"""LangChain adapter tests.
+
+The parsing helpers (``_model_from_result``/``_usage_from_result``/
+``_record_result``) are duck-typed over an ``LLMResult`` and need no langchain
+install, so they run in CI without the optional extra. The two handler tests
+call the factory, which hard-imports ``langchain_core`` — they are skipped
+unless that extra is installed. When available they exercise the real callback
+and prove the hard-stop: a call under budget is allowed and accrued, and a call
+that would cross the ceiling raises ``BudgetExceeded`` in ``on_llm_start`` —
+before the call runs.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+
+import pytest
+
+from floe_guard import BudgetExceeded, BudgetGuard, UnpriceableModelError, UnpriceableModelWarning
+from floe_guard.integrations.langchain import (
+    _model_from_result,
+    _record_result,
+    _usage_from_result,
+    budget_guard_callback_handler,
+)
+
+
+@dataclass
+class _Msg:
+    usage_metadata: dict
+
+
+@dataclass
+class _Gen:
+    message: _Msg
+
+
+@dataclass
+class _Result:
+    """Stand-in for a LangChain ``LLMResult``."""
+
+    llm_output: dict | None = None
+    generations: list = field(default_factory=list)
+
+
+def _openai_result(prompt: int, completion: int, model: str = "gpt-4o") -> _Result:
+    return _Result(
+        llm_output={
+            "model_name": model,
+            "token_usage": {"prompt_tokens": prompt, "completion_tokens": completion},
+        }
+    )
+
+
+def test_model_from_result_reads_llm_output() -> None:
+    assert _model_from_result(_openai_result(1, 1)) == "gpt-4o"
+
+
+def test_model_from_result_missing_is_empty() -> None:
+    assert _model_from_result(_Result(llm_output=None)) == ""
+
+
+def test_usage_from_token_usage_block() -> None:
+    assert _usage_from_result(_openai_result(5, 7)) == (5, 7)
+
+
+def test_usage_from_usage_metadata_fallback() -> None:
+    # No token_usage in llm_output — fall back to per-message usage_metadata.
+    result = _Result(
+        llm_output={"model_name": "gpt-4o"},
+        generations=[[_Gen(_Msg({"input_tokens": 5, "output_tokens": 7}))]],
+    )
+    assert _usage_from_result(result) == (5, 7)
+
+
+def test_record_result_accrues() -> None:
+    guard = BudgetGuard(limit_usd=1.0)
+    _record_result(guard, _openai_result(1_000, 1_000))
+    assert guard.spent_usd == pytest.approx(0.0125)
+
+
+def test_usage_present_but_model_missing_fails_closed() -> None:
+    # Tokens were spent but the model id is missing. This MUST go through record()
+    # (fail-closed → raise), not be silently skipped unmetered.
+    guard = BudgetGuard(limit_usd=1.0)  # fail_closed defaults to True
+    result = _Result(
+        llm_output={"token_usage": {"prompt_tokens": 1_000, "completion_tokens": 1_000}}
+    )
+    with pytest.warns(UnpriceableModelWarning):
+        with pytest.raises(UnpriceableModelError):
+            _record_result(guard, result)
+    assert guard.spent_usd == 0.0
+
+
+def test_usage_present_but_model_missing_fail_open_warns_and_skips() -> None:
+    guard = BudgetGuard(limit_usd=1.0, fail_closed=False)
+    result = _Result(
+        llm_output={"token_usage": {"prompt_tokens": 1_000, "completion_tokens": 1_000}}
+    )
+    with pytest.warns(UnpriceableModelWarning):
+        _record_result(guard, result)
+    assert guard.spent_usd == 0.0
+
+
+def test_no_usage_response_is_a_noop() -> None:
+    guard = BudgetGuard(limit_usd=1.0)
+    _record_result(guard, _Result(llm_output={"model_name": "gpt-4o"}))
+    _record_result(guard, _openai_result(0, 0))
+    assert guard.spent_usd == 0.0
+
+
+def test_handler_allows_under_budget_and_records() -> None:
+    pytest.importorskip("langchain_core")
+    guard = BudgetGuard(limit_usd=1.0)
+    handler = budget_guard_callback_handler(guard)
+
+    handler.on_llm_start({}, ["hello"])  # under budget — no raise
+    handler.on_llm_end(_openai_result(1_000, 1_000))
+    assert guard.spent_usd == pytest.approx(0.0125)
+
+
+def test_handler_blocks_before_crossing() -> None:
+    pytest.importorskip("langchain_core")
+    # First call costs 0.0125 and primes _last_cost; the next call's projection
+    # (0.025) crosses the 0.02 ceiling, so on_llm_start raises BEFORE it runs.
+    guard = BudgetGuard(limit_usd=0.02)
+    handler = budget_guard_callback_handler(guard)
+
+    handler.on_llm_start({}, ["hello"])
+    handler.on_llm_end(_openai_result(1_000, 1_000))
+    assert guard.spent_usd == pytest.approx(0.0125)
+
+    with pytest.raises(BudgetExceeded):
+        handler.on_llm_start({}, ["hello"])