needlepath-langchain 0.1.0__py3-none-any.whl

needlepath_langchain/__init__.py

@@ -0,0 +1,23 @@
+"""Needlepath context selection for LangChain agents and LangGraph.
+
+    from langchain.agents import create_agent
+    from needlepath_langchain import NeedlepathMiddleware
+
+    agent = create_agent(
+        model, tools,
+        middleware=[NeedlepathMiddleware(operating_point="np-2026-07-r2")],
+    )
+
+One line. Tool results are selected against the current step, the history is
+selected before each model call, and every error path passes the original
+context through unmodified with the reason recorded in message metadata.
+
+LangGraph users who cannot move to ``create_agent`` get
+:mod:`needlepath_langchain.langgraph` — the ``pre_model_hook`` recipe, in this
+same package rather than a second distribution.
+"""
+
+from .middleware import USER_AGENT_SUFFIX, NeedlepathMiddleware
+from .version import __version__
+
+__all__ = ["NeedlepathMiddleware", "USER_AGENT_SUFFIX", "__version__"]

needlepath_langchain/_messages.py

@@ -0,0 +1,377 @@
+"""Message ↔ record mapping, and the safety rules that govern rewriting.
+
+Everything structural about this integration lives here, because the sharp edges
+are structural. Two of them are load-bearing enough to state as invariants that
+the rest of the package is built to satisfy:
+
+**Invariant 1 — a message is never removed, only rewritten.**
+LangGraph's ``_validate_chat_history`` raises when an ``AIMessage`` carrying
+``tool_calls`` loses its matching ``ToolMessage``, and every major provider
+rejects the same shape at the API boundary::
+
+    ValueError: Found AIMessages with tool_calls that do not have a
+    corresponding ToolMessage.
+
+A selector that drops messages has to reason about tool-call pairing on every
+path and get it right every time. A selector that only ever *rewrites content
+in place* cannot break the pairing at all — the message objects, their ids,
+their ``tool_call_id``\\ s and their order are exactly what they were. We take
+the second option, and give up the token savings of deleting a message envelope
+to get it.
+
+**Invariant 2 — a rewrite never grows the context and never empties a message.**
+An excerpt longer than what it replaces, or an empty excerpt applied to a
+non-empty message, is a regression however the service arrived at it. Both are
+refused here rather than trusted upstream.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+from typing import Any
+
+from needlepath import ContextRecord, RecordKind
+
+try:  # pragma: no cover - exercised by the installed-langchain test matrix
+    from langchain_core.messages import AIMessage, BaseMessage, SystemMessage, ToolMessage
+except ImportError as exc:  # pragma: no cover
+    raise ImportError(
+        "needlepath-langchain requires langchain (>=1.2.9,<2). "
+        "Install it with: pip install 'needlepath-langchain'"
+    ) from exc
+
+
+#: Marker written into ``response_metadata`` on every message this package
+#: rewrites, mirroring what LangChain's own ``ContextEditingMiddleware`` does.
+#: It makes a rewrite visible in a trace instead of looking like model output.
+METADATA_KEY = "needlepath"
+
+#: What replaces a message the service did not select. Deliberately short,
+#: deliberately explicit that content was withheld rather than absent — a model
+#: that sees "[…]" cannot tell "there was nothing" from "there was something".
+DEFAULT_PLACEHOLDER = "[needlepath: this tool result was not selected for the current step]"
+
+
+def message_text(message: BaseMessage) -> str:
+    """Text of a message, whatever shape its content takes.
+
+    ``content`` is ``str | list[str | dict]``. ``BaseMessage.text`` handles both
+    and returns a ``TextAccessor``, not a ``str`` — hence the cast, which is
+    load-bearing the moment the value is serialized into a JSON request body.
+    """
+    return str(message.text)
+
+
+def approx_tokens(text: str) -> int:
+    """A cheap, dependency-free token estimate for the engage/stand-down trigger.
+
+    Deliberately the same ``len // 4`` shape the published ``estimate_tokens``
+    falls back to, so a trigger set from a documented figure means roughly what
+    the reader expects. It is **never** compared against a billed quantity: the
+    meter is server-side and authoritative, and this is only ever used to decide
+    whether calling is worth it at all.
+    """
+    return len(text) // 4 + 1
+
+
+def messages_tokens(messages: Sequence[BaseMessage]) -> int:
+    return sum(approx_tokens(message_text(m)) for m in messages)
+
+
+def has_tool_calls(message: BaseMessage) -> bool:
+    return bool(getattr(message, "tool_calls", None))
+
+
+def kind_for(message: BaseMessage) -> str:
+    """Map a LangChain message onto the published record-role enum."""
+    if isinstance(message, ToolMessage):
+        return (
+            RecordKind.ERROR
+            if getattr(message, "status", None) == "error"
+            else (RecordKind.TOOL_RESULT)
+        )
+    if isinstance(message, AIMessage):
+        return RecordKind.TOOL_CALL if has_tool_calls(message) else RecordKind.LLM_RESPONSE
+    if isinstance(message, SystemMessage):
+        # Unreachable in practice — system messages are excluded before this is
+        # called — but mapped rather than left to fall through to a default that
+        # would silently mislabel instructions as retrieved data.
+        return RecordKind.USER_INPUT
+    return RecordKind.USER_INPUT
+
+
+def is_prunable(
+    message: BaseMessage,
+    *,
+    include_ai_messages: bool = False,
+    include_human_messages: bool = False,
+) -> bool:
+    """Whether this message's content may be rewritten.
+
+    The default is **tool results only**, and that is a product decision as much
+    as a safety one. Tool results are where the tokens are in an agent thread,
+    they are the seam §8.1 calls our sweet spot, and they are the one message
+    kind whose content is machine-produced — so replacing it with an excerpt of
+    itself cannot silently rewrite something a human wrote or the model
+    committed to.
+
+    Opening this up to AI and Human messages is supported and off by default.
+
+    A ``SystemMessage`` is never prunable, on any setting: it carries the
+    instructions the whole run depends on, it is small, and on LangChain's
+    ``wrap_model_call`` seam it is not even in ``request.messages``.
+    """
+    if isinstance(message, SystemMessage):
+        return False
+    if isinstance(message, ToolMessage):
+        return True
+    if isinstance(message, AIMessage):
+        # An AIMessage carrying tool_calls is structure, not prose: its content
+        # is usually empty and the part that matters is the call list. Nothing
+        # to save, and every reason not to touch it.
+        return include_ai_messages and not has_tool_calls(message)
+    return include_human_messages
+
+
+def record_id_for(index: int) -> str:
+    """A short, stable per-call record id.
+
+    Positional rather than derived from ``message.id``: it is guaranteed present
+    and unique inside one request, it is inside the published identifier
+    charset, and it is two or three bytes rather than a UUID — which matters
+    because every byte counts against the 6 MB guard on a long thread.
+    """
+    return f"m{index}"
+
+
+def build_records(
+    messages: Sequence[BaseMessage],
+    *,
+    protected_indices: set | None = None,
+    include_ai_messages: bool = False,
+    include_human_messages: bool = False,
+) -> tuple[list[ContextRecord], dict[str, int]]:
+    """Records for every prunable, unprotected, non-empty message.
+
+    Returns the records and a ``record_id -> message index`` map. Anything not
+    in that map is passed through untouched, by construction.
+    """
+    protected = protected_indices or set()
+    records: list[ContextRecord] = []
+    index_by_id: dict[str, int] = {}
+    total = len(messages)
+    for index, message in enumerate(messages):
+        if index in protected:
+            continue
+        if not is_prunable(
+            message,
+            include_ai_messages=include_ai_messages,
+            include_human_messages=include_human_messages,
+        ):
+            continue
+        text = message_text(message)
+        if not text.strip():
+            continue
+        record_id = record_id_for(index)
+        index_by_id[record_id] = index
+        records.append(
+            ContextRecord(
+                text=text,
+                kind=kind_for(message),
+                id=record_id,
+                title=_title_for(message),
+                # A recency prior the engine can use, without this package
+                # deciding recency policy on its own. Newest → 1.0, oldest → ~0.
+                importance=(index + 1) / total if total else 0.0,
+            )
+        )
+    return records, index_by_id
+
+
+def _title_for(message: BaseMessage) -> str | None:
+    """A short human-readable label. For a tool result, the tool's name."""
+    name = getattr(message, "name", None)
+    if isinstance(name, str) and name:
+        return name[:120]
+    return None
+
+
+def latest_user_prompt(messages: Sequence[BaseMessage]) -> str:
+    """The most recent human turn — what the selected context has to serve.
+
+    Falls back to the last non-empty message of any kind, because a thread whose
+    latest turn is a tool result still has a task; using an empty prompt would
+    hand the service no query at all and make every selection meaningless.
+    """
+    for message in reversed(messages):
+        if message.type == "human":
+            text = message_text(message)
+            if text.strip():
+                return text
+    for message in reversed(messages):
+        text = message_text(message)
+        if text.strip():
+            return text
+    return ""
+
+
+def prior_user_prompts(messages: Sequence[BaseMessage], limit: int = 4) -> list[str]:
+    """Earlier human turns, newest first — the contract's drift signal."""
+    found: list[str] = []
+    seen_latest = False
+    for message in reversed(messages):
+        if message.type != "human":
+            continue
+        text = message_text(message)
+        if not text.strip():
+            continue
+        if not seen_latest:
+            seen_latest = True
+            continue
+        found.append(text)
+        if len(found) >= limit:
+            break
+    return found
+
+
+def protected_tail(messages: Sequence[BaseMessage], preserve_recent: int) -> set:
+    """Indices of the trailing messages that are never rewritten.
+
+    Belt and braces alongside ``task.required_record_ids``: the contract lets us
+    *ask* the service to keep them, and this guarantees it regardless of what
+    comes back. A guard that depends on the far side honouring a request is not
+    a guard.
+    """
+    if preserve_recent <= 0:
+        return set()
+    return set(range(max(0, len(messages) - preserve_recent), len(messages)))
+
+
+def rewrite_message(
+    message: BaseMessage,
+    new_text: str,
+    *,
+    reason: str,
+    request_id: str,
+) -> BaseMessage:
+    """A copy of ``message`` with new content and a traceable marker.
+
+    ``model_copy`` rather than reconstruction: reconstructing a ``ToolMessage``
+    silently drops ``artifact``, ``id``, ``response_metadata`` and any provider
+    metadata, and the whole safety argument here rests on the envelope being
+    untouched.
+    """
+    metadata = dict(getattr(message, "response_metadata", {}) or {})
+    existing = metadata.get(METADATA_KEY)
+    merged = dict(existing) if isinstance(existing, dict) else {}
+    merged.update(
+        {
+            "rewritten": True,
+            # `rewrite_reason`, not `reason`: the selection result's own
+            # metadata also carries a `reason`, the two are merged into this
+            # same dict, and a collision would silently overwrite one of them.
+            "rewrite_reason": reason,
+            "request_id": request_id,
+            "original_tokens": approx_tokens(message_text(message)),
+            "new_tokens": approx_tokens(new_text),
+        }
+    )
+    metadata[METADATA_KEY] = merged
+    return message.model_copy(update={"content": new_text, "response_metadata": metadata})
+
+
+def apply_selection(
+    messages: Sequence[BaseMessage],
+    *,
+    excerpts: dict[str, str],
+    index_by_id: dict[str, int],
+    placeholder: str,
+    request_id: str,
+) -> tuple[list[BaseMessage], int, int]:
+    """Rewrite the selected and unselected messages. Returns (messages, before, after).
+
+    Enforces invariant 2 message by message: an excerpt that is empty, or that
+    would make a message *longer*, is refused and the original kept. Both are
+    regressions however the service arrived at them, and refusing one message's
+    rewrite is cheaper than failing the whole selection open.
+    """
+    out: list[BaseMessage] = list(messages)
+    before = 0
+    after = 0
+    for record_id, index in index_by_id.items():
+        original = out[index]
+        original_text = message_text(original)
+        before += approx_tokens(original_text)
+        excerpt = excerpts.get(record_id)
+        if excerpt is None and not placeholder.strip():
+            # Invariant 2, at the mutation boundary. A blank custom placeholder
+            # would rewrite a non-empty message to nothing — the pairing
+            # survives, the content does not. Refused rather than trusted.
+            after += approx_tokens(original_text)
+            continue
+        if excerpt is None:
+            # Not selected: collapse to the placeholder. This is the only place
+            # content is discarded, and it is the difference between selection
+            # and a no-op.
+            new_text = placeholder
+            reason = "not_selected"
+        elif not excerpt.strip():
+            # Selected with nothing in it. Applying this would delete content
+            # the service said to keep — refuse it and keep the original.
+            after += approx_tokens(original_text)
+            continue
+        elif len(excerpt) >= len(original_text):
+            # An "excerpt" that is not shorter buys nothing and may cost.
+            after += approx_tokens(original_text)
+            continue
+        else:
+            new_text = excerpt
+            reason = "excerpt"
+        if len(new_text) >= len(original_text) and reason == "not_selected":
+            # A placeholder longer than the message it replaces. Rare, but a
+            # short tool result plus a long custom placeholder gets here.
+            after += approx_tokens(original_text)
+            continue
+        out[index] = rewrite_message(original, new_text, reason=reason, request_id=request_id)
+        after += approx_tokens(new_text)
+    return out, before, after
+
+
+def tool_call_keywords(args: Any, limit: int = 6, max_len: int = 120) -> list[str]:
+    """Short string values from a tool call's arguments, as selection keywords.
+
+    A tool call's arguments are the closest thing to a query the service will
+    ever see for a tool result, so they are worth sending. Bounded on both count
+    and length because they are model-generated and therefore unbounded.
+    """
+    if not isinstance(args, dict):
+        return []
+    out: list[str] = []
+    for value in args.values():
+        if isinstance(value, str) and value.strip():
+            out.append(value.strip()[:max_len])
+        elif isinstance(value, (int, float, bool)):
+            out.append(str(value))
+        if len(out) >= limit:
+            break
+    return out
+
+
+__all__ = [
+    "METADATA_KEY",
+    "DEFAULT_PLACEHOLDER",
+    "message_text",
+    "approx_tokens",
+    "messages_tokens",
+    "kind_for",
+    "is_prunable",
+    "has_tool_calls",
+    "record_id_for",
+    "build_records",
+    "latest_user_prompt",
+    "prior_user_prompts",
+    "protected_tail",
+    "rewrite_message",
+    "apply_selection",
+    "tool_call_keywords",
+]