Merge pull request #36 from Floe-Labs/feat/pre-request-call

feat:request-sized pre-call estimates and mid-stream budget enforceme…

Author: rajbhensdadiya

CHANGELOG.md

@@ -10,6 +10,30 @@ both packages adhere to [Semantic Versioning](https://semver.org/).
 
 ## Unreleased
 
+## py 0.4.0 — 2026-07-15
+
+### Added (py)
+
+- **Request-sized pre-call estimates**: `BudgetGuard.estimate_call(model,
+  prompt_tokens, max_completion_tokens)` prices the actual incoming request
+  from the cost map, so `reserve(est)` / `check(est)` block an oversized call
+  — including the very first one, which the last-cost prediction is blind to.
+  The LiteLLM adapter reserves request-sized automatically (prompt via
+  `litellm.token_counter`, cap from `max_tokens`); the LangChain handler sizes
+  its pre-call `check()` from the serialized model config and a ~4 chars/token
+  prompt heuristic. Anything unpriceable/unsized falls back to the previous
+  last-cost behaviour.
+- **Mid-stream enforcement**: `StreamGuard` / `guard_stream()` re-price a
+  streaming response chunk-by-chunk and raise `BudgetExceeded`
+  mid-generation when the running call would cross the ceiling — the partial
+  spend is settled (and lands in `spend_log`) instead of the whole overshoot
+  being discovered post-mortem. `finish(prompt_tokens=…, completion_tokens=…)`
+  reconciles the chunk heuristic to provider-reported usage; unpriceable
+  models fail closed before the stream starts; parallel streams count each
+  other's in-flight accrual, so unreserved streams share the ceiling instead
+  of each spending it in full. Demo: `examples/streaming_guard.py` (no API
+  key).
+
 ## py 0.3.0 / js 0.3.0 — 2026-07-14
 
 ### Added (py + js)

README.md

@@ -175,6 +175,45 @@ cost_usd, label?, reserved?}` — identical to the TS package's `exportLog()`, s
 every agent produces the same shape regardless of stack and the streams can be
 concatenated and analysed together.
 
+## Request-sized estimates and mid-stream enforcement
+
+Two gaps in last-cost prediction, closed in 0.4.0 (Python):
+
+**The oversized first call.** `check()`/`reserve()` predict from the *last*
+call — blind on call #1, wrong for a call much bigger than the previous one.
+`estimate_call()` prices the **actual incoming request** so even a first call
+that alone would cross the cap blocks pre-flight:
+
+```python
+est = guard.estimate_call("gpt-4o", prompt_tokens=12_000, max_completion_tokens=4_096)
+handle = guard.reserve(est)   # raises BudgetExceeded NOW if this call can't fit
+```
+
+The LiteLLM adapter does this automatically (prompt tokens via
+`litellm.token_counter`, output cap from `max_tokens`), and the LangChain
+handler sizes its pre-call `check()` the same way. Unpriceable or unsized
+requests fall back to the old last-cost prediction — the wiring only ever
+tightens enforcement.
+
+**The stream that runs long.** `record()` meters a *completed* response — too
+late for a generation that starts cheap and keeps going. `guard_stream()` (or
+the underlying `StreamGuard`) re-prices the call on every chunk and cuts the
+stream off **mid-generation**, settling the tokens actually consumed instead of
+recording a big overshoot after the fact:
+
+```python
+from floe_guard import guard_stream
+
+for chunk in guard_stream(guard, "gpt-4o", stream, prompt_tokens=1_000):
+    print(chunk, end="")   # raises BudgetExceeded mid-stream at the ceiling
+```
+
+Chunk sizes are estimated at ~4 chars/token (pass `count_tokens=` for a real
+tokenizer); the final accrual reconciles to provider-reported usage via
+`StreamGuard.finish(...)`. See
+[`examples/streaming_guard.py`](examples/streaming_guard.py) for a runnable
+demo (no API key).
+
 ## Framework adapters (optional extras)
 
 ### CrewAI
@@ -320,10 +359,13 @@ vendored cost map *inside your process*:
 - It only sees the vendors you instrument.
 - A determined agent or a bug could route around an in-process check.
 - Under heavy or cold-start concurrency it bounds steady-state spend, not the
-  first parallel wave. Reservations size from the last call's cost (`0` until the
-  first `record()`), so the opening fan-out has nothing to estimate from. Pass a
-  known per-call max to `reserve()` to bound it, or use hosted Floe for a hard cap
+  first parallel wave. Reservations default to the last call's cost (`0` until
+  the first `record()`) — size them to the real request with `estimate_call()`
+  (the LiteLLM adapter does this for you), or use hosted Floe for a hard cap
   under arbitrary concurrency.
+- Mid-stream enforcement (`guard_stream`) prices chunks by a ~4 chars/token
+  heuristic unless you supply a tokenizer, so the cut-off point is approximate;
+  the final accrual reconciles to provider-reported usage.
 
 It's genuinely useful on its own, and it's honest about its limits. No inflated
 metrics, no "zero defaults" claims — it's a free local stop, not a vault.

examples/streaming_guard.py

@@ -0,0 +1,67 @@
+"""Mid-stream kill-switch demo — runs with NO API key and NO account.
+
+Two enforcement gaps this demo closes (both new in 0.4.0):
+
+1. **The oversized first call.** ``check()``/``reserve()`` predict the next
+   call from the LAST one — on call #1 there is no baseline, so a huge request
+   sails through and the overshoot is only discovered after the money is spent.
+   ``estimate_call()`` prices the ACTUAL request (real prompt size + output
+   cap) so the very first call blocks pre-flight if it alone would cross.
+
+2. **The stream that runs long.** ``record()`` meters a COMPLETED response —
+   too late to stop a generation that starts cheap and keeps going.
+   ``guard_stream()`` re-prices the call on every chunk and cuts the stream
+   off mid-generation, settling only the tokens actually consumed.
+
+Run it::
+
+    python examples/streaming_guard.py
+
+The "LLM" is a stub generator — no network, no key. Costs are computed offline
+from the bundled cost map, exactly as for a real ``gpt-4o`` call.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Iterator
+
+from floe_guard import BudgetExceeded, BudgetGuard, guard_stream
+
+MODEL = "gpt-4o"
+
+
+def stub_llm_stream() -> Iterator[str]:
+    """A fake streaming LLM that never stops talking."""
+    while True:
+        yield "and another thing... "  # ~5 tokens per chunk by the heuristic
+
+
+def main() -> None:
+    guard = BudgetGuard(limit_usd=0.01)  # ≙ 1_000 gpt-4o output tokens
+
+    # ── 1. the oversized FIRST call is blocked pre-flight ──────────────────────
+    print(f"Budget: ${guard.limit_usd:.2f}\n")
+    print("1) First call asks for 100k output tokens (≈ $1.00):")
+    estimate = guard.estimate_call(MODEL, prompt_tokens=1_000, max_completion_tokens=100_000)
+    try:
+        guard.reserve(estimate)  # request-sized, so call #1 has no free pass
+    except BudgetExceeded:
+        print("   blocked BEFORE the request was sent — $0.00 spent.\n")
+
+    # ── 2. a stream that starts cheap but runs long is cut off mid-flight ──────
+    print("2) Streaming a response that never wants to stop:")
+    chunks = 0
+    try:
+        for _ in guard_stream(guard, MODEL, stub_llm_stream()):
+            chunks += 1
+    except BudgetExceeded:
+        print(f"   stream cut off mid-generation after {chunks} chunks.")
+        print(f"   spent ${guard.spent_usd:.4f} of the ${guard.limit_usd:.2f} ceiling — "
+              "the partial spend is settled, not lost:")
+        for event in guard.spend_log:
+            print(f"     ledger: {event.kind} {event.model_or_tool} "
+                  f"{event.completion_tokens} tokens → ${event.cost_usd:.4f}")
+
+
+if __name__ == "__main__":
+    main()

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "floe-guard"
-version = "0.3.0"
+version = "0.4.0"
 description = "Local budget guardrail for AI agents — hard-stops a runaway loop before its next LLM call crosses a spend ceiling. No account, no network."
 readme = "README.md"
 requires-python = ">=3.10"

src/floe_guard/__init__.py

@@ -23,13 +23,16 @@
 from .guard import BudgetAdvisory, BudgetGuard, SpendEvent
 from .hosted import hosted_enforcement_available, hosted_remaining_usd
 from .pricing import ManualPrice, PricedModel, price_tokens, resolve_price
+from .stream import StreamGuard, guard_stream
 
-__version__ = "0.3.0"  # keep in lockstep with pyproject.toml
+__version__ = "0.4.0"  # keep in lockstep with pyproject.toml
 
 __all__ = [
     "BudgetGuard",
     "BudgetAdvisory",
     "SpendEvent",
+    "StreamGuard",
+    "guard_stream",
     "BudgetExceeded",
     "FloeGuardError",
     "HostedEnforcementError",

src/floe_guard/guard.py

@@ -193,6 +193,10 @@ def __init__(
             )
         # Per-call ledger; deque(maxlen=None) is unbounded, otherwise a ring buffer.
         self._spend_log: deque[SpendEvent] = deque(maxlen=max_log_events)
+        # Active streams' (accrued_usd, reserved_usd), keyed by registry token —
+        # see _stream_register(). Lets parallel streams count each other's
+        # in-flight accrual against the ceiling before anything settles.
+        self._stream_costs: dict[object, tuple[float, float]] = {}
         self._lock = threading.Lock()
 
     # ── enforcement ───────────────────────────────────────────────────────────
@@ -214,6 +218,36 @@ def check(self, estimated_next_cost: float | None = None) -> None:
         if self._would_cross(estimated_next_cost):
             self._block()
 
+    def estimate_call(
+        self,
+        model: str,
+        prompt_tokens: int,
+        max_completion_tokens: int = 0,
+        *,
+        price: ManualPrice | None = None,
+    ) -> float | None:
+        """Price the ACTUAL incoming request, for a request-sized reserve()/check().
+
+        :meth:`check` and :meth:`reserve` default to predicting the next call
+        from the LAST call's cost — which is blind on the first call and wrong
+        for a call much larger than the previous one. Feed this the request you
+        are about to send (its real prompt size and output cap) and pass the
+        result straight through::
+
+            est = guard.estimate_call("gpt-4o", prompt_tokens, max_completion_tokens=1024)
+            handle = guard.reserve(est)   # blocks NOW if this call alone would cross
+
+        The estimate is worst-case on output (the model may stop well short of
+        ``max_completion_tokens``); the hold is corrected to actual cost at
+        :meth:`settle`. Returns ``None`` when the model is unpriceable — and
+        ``reserve(None)`` / ``check(None)`` fall back to the last-cost
+        prediction, so the wiring degrades gracefully instead of failing.
+        """
+        priced = self._resolve(model, price)
+        if priced is None:
+            return None
+        return price_tokens(priced, prompt_tokens, max_completion_tokens)
+
     def reserve(self, estimated_cost: float | None = None) -> float:
         """Atomically check the ceiling AND hold the estimated cost in-flight.
 
@@ -454,6 +488,41 @@ def _validate_estimate(self, estimated: float | None) -> None:
         if estimated is not None and not math.isfinite(estimated):
             raise ValueError(f"estimated cost must be a finite number, got {estimated!r}")
 
+    def _stream_register(self, reserved: float) -> object:
+        """Register an active stream (see :class:`~floe_guard.stream.StreamGuard`)
+        and return its registry key. Active streams' accrued-but-unsettled costs
+        count against the ceiling for each OTHER stream, so parallel unreserved
+        streams share the budget instead of each spending the full ceiling."""
+        key = object()
+        with self._lock:
+            self._stream_costs[key] = (0.0, max(0.0, reserved))
+        return key
+
+    def _stream_unregister(self, key: object) -> None:
+        """Drop a stream's registry entry once its cost is settled (settle()
+        moves the accrual into ``spent_usd``, so keeping it would double-count)."""
+        with self._lock:
+            self._stream_costs.pop(key, None)
+
+    def _stream_would_cross(self, key: object, cumulative_call_cost: float) -> bool:
+        """Atomically record stream ``key``'s cumulative cost so far and answer:
+        would it cross the ceiling? Counted against the limit: settled spend,
+        other calls' reservations (this stream's own hold is excluded — its real
+        accrued cost replaces the estimate), and each OTHER active stream's
+        accrual beyond its own reservation (the reservation part is already
+        inside ``_reserved``). Used by :class:`~floe_guard.stream.StreamGuard`.
+        """
+        with self._lock:
+            own_reserved = self._stream_costs.get(key, (0.0, 0.0))[1]
+            self._stream_costs[key] = (cumulative_call_cost, own_reserved)
+            other_overage = sum(
+                max(0.0, accrued - held)
+                for k, (accrued, held) in self._stream_costs.items()
+                if k is not key
+            )
+            others = self.spent_usd + max(0.0, self._reserved - own_reserved) + other_overage
+            return others + cumulative_call_cost > self.limit_usd + _EPS
+
     def _would_cross(self, estimated_next_cost: float | None) -> bool:
         with self._lock:
             estimate = (

src/floe_guard/integrations/langchain.py

@@ -23,6 +23,7 @@
 from typing import Any
 
 from ..guard import BudgetGuard
+from ..stream import approx_tokens
 
 
 def _require_langchain() -> Any:
@@ -36,6 +37,43 @@ def _require_langchain() -> Any:
     return langchain_core
 
 
+def _estimate_start(guard: BudgetGuard, serialized: Any, texts: list[str]) -> float | None:
+    """Request-sized cost estimate for the pre-call check.
+
+    ``check()`` defaults to predicting from the LAST call's cost — blind on the
+    first call and wrong for a much larger one. Size the prediction to the
+    request instead: model id and ``max_tokens`` from the serialized model
+    config, prompt tokens via the ~4 chars/token heuristic (langchain-core has
+    no tokenizer; a rough request-sized figure still beats a stale or zero
+    baseline). Returns ``None`` when the model is unknown/unpriceable —
+    ``check(None)`` keeps the old behaviour.
+    """
+    model_kwargs = serialized.get("kwargs") if isinstance(serialized, dict) else None
+    if not isinstance(model_kwargs, dict):
+        return None
+    model = model_kwargs.get("model_name") or model_kwargs.get("model")
+    if not model:
+        return None
+    prompt_tokens = sum(approx_tokens(t) for t in texts if isinstance(t, str))
+    max_out = model_kwargs.get("max_tokens") or model_kwargs.get("max_completion_tokens") or 0
+    try:
+        max_out = max(0, int(max_out))
+    except (TypeError, ValueError):
+        max_out = 0
+    return guard.estimate_call(str(model), prompt_tokens, max_out)
+
+
+def _chat_texts(messages: Any) -> list[str]:
+    """Flatten LangChain's batched chat messages to their text contents."""
+    texts: list[str] = []
+    for batch in messages or []:
+        for message in batch or []:
+            content = getattr(message, "content", None)
+            if isinstance(content, str):
+                texts.append(content)
+    return texts
+
+
 def _model_from_result(response: Any) -> str:
     """Pull the model name from a LangChain ``LLMResult``."""
     llm_output = getattr(response, "llm_output", None)
@@ -112,10 +150,13 @@ def __init__(self) -> None:
             self.guard = guard
 
         def on_llm_start(self, serialized: Any, prompts: Any, **kwargs: Any) -> None:
-            self.guard.check()
+            # Request-sized when derivable (see _estimate_start); check(None)
+            # falls back to the last-cost prediction.
+            texts = [p for p in (prompts or []) if isinstance(p, str)]
+            self.guard.check(_estimate_start(self.guard, serialized, texts))
 
         def on_chat_model_start(self, serialized: Any, messages: Any, **kwargs: Any) -> None:
-            self.guard.check()
+            self.guard.check(_estimate_start(self.guard, serialized, _chat_texts(messages)))
 
         def on_llm_end(self, response: Any, **kwargs: Any) -> None:
             _record_result(self.guard, response)