loopshield 1.0.0__py3-none-any.whl

loopshield/shield.py

@@ -0,0 +1,279 @@
+"""Middleware wrapper that inspects each agent step and blocks looping tools."""
+
+from __future__ import annotations
+
+from collections.abc import Callable, Sequence
+from functools import wraps
+from typing import Any, ParamSpec, TypeVar
+
+from loopshield.adapters import as_embed_fn
+from loopshield.graph_analyzer import GraphAnalyzer, tool_signature
+from loopshield.mitigator import Mitigator
+from loopshield.models import AgentStep, LoopVerdict
+from loopshield.semantic_analyzer import EmbedFn, SemanticAnalyzer
+
+P = ParamSpec("P")
+R = TypeVar("R")
+
+_INTERVENTIONS: dict[str, str] = {
+    "graph_cycle": (
+        "LoopShield detected a repeated tool cycle. "
+        "Do not call the same tools with the same arguments again. "
+        "Use the observations you already have, try a different approach, "
+        "or give the final answer."
+    ),
+    "semantic": (
+        "LoopShield detected a repeated line of thought. "
+        "Stop rephrasing the same idea. "
+        "Use the observations you already have, try a different approach, "
+        "or give the final answer."
+    ),
+    "both": (
+        "LoopShield detected a repeated tool cycle and a repeated line of thought. "
+        "Do not repeat the same action or reasoning. "
+        "Use the observations you already have, try a different approach, "
+        "or give the final answer."
+    ),
+}
+
+
+class LoopShield:
+    """Watch an agent loop, block a tool that starts repeating, and prune context.
+
+    ``step`` records one :class:`AgentStep` and returns the verdict for the
+    trajectory so far. ``guard`` wraps the function an agent calls on each
+    iteration: a looping call is not executed and the verdict is returned
+    instead. A guarded function that receives the whole ``list[AgentStep]``
+    is refused before it runs when any step in that list closes a loop.
+
+    Pass an initialized client straight through ``embedding``::
+
+        LoopShield(embedding=OpenAI())
+        LoopShield(embedding=Client(host="http://127.0.0.1:11434"))
+    """
+
+    def __init__(
+        self,
+        *,
+        max_cycles: int = 1,
+        similarity_threshold: float = 0.92,
+        max_semantic_repeats: int = 2,
+        embed_fn: EmbedFn | None = None,
+        embedding: Any | None = None,
+        embedding_model: str | None = None,
+        max_tokens_estimate: int | None = None,
+    ) -> None:
+        if max_tokens_estimate is not None and max_tokens_estimate < 0:
+            raise ValueError("max_tokens_estimate must be zero or greater")
+        if embed_fn is not None and embedding is not None:
+            raise ValueError("Pass either embed_fn or embedding, not both.")
+        if embedding is None and embedding_model is not None:
+            raise ValueError("embedding_model is only used with an OpenAI or Ollama client.")
+        resolved_embed_fn = embed_fn if embedding is None else as_embed_fn(embedding, model=embedding_model)
+        self.graph_analyzer = GraphAnalyzer(max_cycles=max_cycles)
+        self.semantic_analyzer = SemanticAnalyzer(
+            resolved_embed_fn,
+            similarity_threshold=similarity_threshold,
+            max_semantic_repeats=max_semantic_repeats,
+        )
+        self.mitigator = Mitigator()
+        self.max_tokens_estimate = max_tokens_estimate
+        self._history: list[AgentStep] = []
+
+    @property
+    def history(self) -> tuple[AgentStep, ...]:
+        """Steps recorded by :meth:`step` and :meth:`guard`, in call order."""
+
+        return tuple(self._history)
+
+    def step(self, agent_step: AgentStep) -> LoopVerdict:
+        """Record ``agent_step`` and judge the trajectory collected so far.
+
+        A loop locks ``agent_step.tool_name``. The verdict's
+        ``intervention_prompt`` is the system nudge that tells the model to
+        change strategy.
+        """
+
+        self._history.append(agent_step)
+        verdict = self._verdict_for_latest(self._history)
+        if not verdict.is_loop:
+            return verdict
+        nudge = self.mitigator.generate_nudge(verdict, agent_step.tool_name)
+        return verdict.model_copy(update={"intervention_prompt": nudge["content"]})
+
+    def guard(self, func: Callable[P, R]) -> Callable[P, R | LoopVerdict]:
+        """Decorate an agent step or an agent loop so a detected loop is not executed.
+
+        Per-step form::
+
+            @shield.guard
+            def execute(agent_step: AgentStep) -> str:
+                return run_tool(agent_step)
+
+        Loop form::
+
+            @shield.guard
+            def run(steps: list[AgentStep]) -> list[str]:
+                return [run_tool(step) for step in steps]
+        """
+
+        @wraps(func)
+        def wrapper(*args: P.args, **kwargs: P.kwargs) -> R | LoopVerdict:
+            agent_step = _find_step(args, kwargs)
+            if agent_step is not None:
+                verdict = self.step(agent_step)
+                if verdict.is_loop:
+                    return verdict
+                return self._maybe_prune(func(*args, **kwargs))
+
+            sequence = _find_step_sequence(args, kwargs)
+            if sequence is not None:
+                for item in sequence:
+                    verdict = self.step(item)
+                    if verdict.is_loop:
+                        return verdict
+                return self._maybe_prune(func(*args, **kwargs))
+
+            result = func(*args, **kwargs)
+            if isinstance(result, AgentStep):
+                verdict = self.step(result)
+                if verdict.is_loop:
+                    return verdict
+            return self._maybe_prune(result)
+
+        return wrapper
+
+    def inspect(self, steps: Sequence[AgentStep]) -> LoopVerdict:
+        """Return a verdict for ``steps``. The given trajectory is not modified."""
+
+        return _assemble(
+            self.graph_analyzer.detect(steps),
+            self.semantic_analyzer.detect(steps),
+        )
+
+    def _verdict_for_latest(self, steps: Sequence[AgentStep]) -> LoopVerdict:
+        """Flag a loop only when the newest step is still inside it.
+
+        A tool that is not part of the repeated cycle can run after a nudge,
+        which is what lets the agent leave the loop and finish the task.
+        """
+
+        if not steps:
+            return _assemble(False, False)
+        latest = tool_signature(steps[-1].tool_name, steps[-1].tool_args)
+        graph_loop = latest in self.graph_analyzer.repeated_signatures(steps)
+        semantic_loop = self.semantic_analyzer.detect(steps)
+        return _assemble(graph_loop, semantic_loop)
+
+    def protect(
+        self,
+        steps: Sequence[AgentStep],
+    ) -> tuple[LoopVerdict, list[AgentStep]]:
+        """Inspect ``steps`` and drop older copies of a repeated tool call.
+
+        When a loop is present, ``intervention_prompt`` is the tool-lock nudge
+        for the latest step.
+        """
+
+        verdict = self.inspect(steps)
+        if not verdict.is_loop or not steps:
+            return verdict, list(steps)
+        nudge = self.mitigator.generate_nudge(verdict, steps[-1].tool_name)
+        verdict = verdict.model_copy(update={"intervention_prompt": nudge["content"]})
+        return verdict, self.mitigator.prune(steps)
+
+    def __call__(
+        self,
+        steps: Sequence[AgentStep],
+    ) -> tuple[LoopVerdict, list[AgentStep]]:
+        """Run :meth:`protect` so the shield can be used as middleware."""
+
+        return self.protect(steps)
+
+    def evaluate(self, steps: Sequence[AgentStep]) -> LoopVerdict:
+        """Judge a complete trajectory by its newest step, without recording it.
+
+        A loop is reported only while the newest tool call is still inside the
+        repeated cycle, so an agent that switches tools after a nudge is let
+        through. Framework integrations call this on every message pipeline pass.
+        """
+
+        return self._verdict_for_latest(steps)
+
+    def langchain_middleware(self, config: Any | None = None) -> Any:
+        """Return a LangChain ``AgentMiddleware`` that screens model and tool calls.
+
+        ``config`` is a :class:`loopshield.frameworks.IntegrationConfig` or a
+        mapping of its fields. Pass the result in ``create_agent(middleware=[...])``.
+        """
+
+        from loopshield.frameworks import LoopShieldMiddleware
+
+        return LoopShieldMiddleware(self, config)
+
+    def crewai_agent(self) -> Callable[[Any], Any]:
+        """Return a decorator for CrewAI's ``Agent.execute_task`` or an ``Agent`` class."""
+
+        from loopshield.frameworks import crewai_agent
+
+        return crewai_agent(self)
+
+    def langgraph_node(
+        self,
+        node: Callable[..., Any] | None = None,
+        config: Any | None = None,
+    ) -> Any:
+        """Wrap a LangGraph node so a flagged ``messages`` state gets the system nudge.
+
+        Use ``@shield.langgraph_node()`` or ``shield.langgraph_node(node)``.
+        """
+
+        from loopshield.frameworks import langgraph_node
+
+        return langgraph_node(self, node, config)
+
+    def _maybe_prune(self, result: R) -> R:
+        if self.max_tokens_estimate is None or not _is_message_list(result):
+            return result
+        return self.mitigator.prune_context(result, self.max_tokens_estimate)  # type: ignore[return-value]
+
+
+def _assemble(graph_loop: bool, semantic_loop: bool) -> LoopVerdict:
+    if graph_loop and semantic_loop:
+        loop_type = "both"
+    elif graph_loop:
+        loop_type = "graph_cycle"
+    elif semantic_loop:
+        loop_type = "semantic"
+    else:
+        return LoopVerdict(is_loop=False, loop_type=None, intervention_prompt=None)
+    return LoopVerdict(
+        is_loop=True,
+        loop_type=loop_type,
+        intervention_prompt=_INTERVENTIONS[loop_type],
+    )
+
+
+def _find_step(args: tuple[Any, ...], kwargs: dict[str, Any]) -> AgentStep | None:
+    for value in (*args, *kwargs.values()):
+        if isinstance(value, AgentStep):
+            return value
+    return None
+
+
+def _find_step_sequence(
+    args: tuple[Any, ...],
+    kwargs: dict[str, Any],
+) -> list[AgentStep] | None:
+    for value in (*args, *kwargs.values()):
+        if (
+            isinstance(value, (list, tuple))
+            and value
+            and all(isinstance(item, AgentStep) for item in value)
+        ):
+            return list(value)
+    return None
+
+
+def _is_message_list(value: object) -> bool:
+    return isinstance(value, list) and bool(value) and all(isinstance(item, dict) for item in value)

loopshield/transcript.py

@@ -0,0 +1,205 @@
+"""Convert chat transcripts from any framework into :class:`AgentStep` trajectories."""
+
+from __future__ import annotations
+
+import json
+from collections.abc import Iterable, Mapping, Sequence
+from typing import Any
+
+from loopshield.mitigator import estimate_tokens, generate_nudge
+from loopshield.models import AgentStep, LoopVerdict
+
+REPLY_TOOL = "reply"
+
+_ROLE_ALIASES = {
+    "human": "user",
+    "ai": "assistant",
+    "chat": "assistant",
+    "model": "assistant",
+}
+
+
+def content_text(content: Any) -> str:
+    """Read plain text or a list of content parts as one string."""
+
+    if content is None:
+        return ""
+    if isinstance(content, str):
+        return content
+    if isinstance(content, list):
+        parts = [
+            content_text(part.get("text", part.get("content", "")))
+            if isinstance(part, Mapping)
+            else content_text(part)
+            for part in content
+        ]
+        return "\n".join(part for part in parts if part)
+    return str(content)
+
+
+def argument_dict(arguments: Any) -> dict[str, Any]:
+    """Return tool arguments as a dict, parsing a JSON string when needed."""
+
+    if isinstance(arguments, str):
+        try:
+            parsed = json.loads(arguments) if arguments.strip() else {}
+        except json.JSONDecodeError:
+            return {"raw": arguments}
+        arguments = parsed
+    if isinstance(arguments, dict):
+        return arguments
+    if arguments is None:
+        return {}
+    return {"value": arguments}
+
+
+def describe_call(name: str, arguments: Mapping[str, Any]) -> str:
+    """Describe a tool call as text for steps that carry no written thought."""
+
+    return f"{name} {json.dumps(dict(arguments), sort_keys=True, default=str)}"
+
+
+def coerce_messages(raw: Iterable[Any]) -> list[dict[str, Any]]:
+    """Return OpenAI-style message dicts for dicts, role tuples, or message objects.
+
+    LangChain-style objects (``type``, ``content``, ``tool_calls``) are read by
+    attribute, so no framework import is needed.
+    """
+
+    return [coerce_message(item) for item in raw]
+
+
+def coerce_message(item: Any) -> dict[str, Any]:
+    """Return one message as an OpenAI-style dict. The input is not modified."""
+
+    if isinstance(item, Mapping):
+        data = dict(item)
+        data["role"] = _role(data.get("role") or data.get("type"))
+        _carry_marker(data, item.get("additional_kwargs"))
+        return data
+    if isinstance(item, str):
+        return {"role": "user", "content": item}
+    if isinstance(item, (tuple, list)) and len(item) == 2:
+        return {"role": _role(item[0]), "content": item[1]}
+
+    data = {
+        "role": _role(getattr(item, "role", None) or getattr(item, "type", None)),
+        "content": getattr(item, "content", ""),
+    }
+    for key in ("name", "tool_call_id", "status"):
+        value = getattr(item, key, None)
+        if isinstance(value, str):
+            data[key] = value
+    calls = getattr(item, "tool_calls", None)
+    if isinstance(calls, list) and calls:
+        data["tool_calls"] = [_openai_call(call) for call in calls]
+    _carry_marker(data, getattr(item, "additional_kwargs", None))
+    return data
+
+
+def messages_to_steps(messages: Sequence[Mapping[str, Any]]) -> list[AgentStep]:
+    """Turn assistant turns into steps: one per tool call, or one text reply.
+
+    A text-only reply is a ``reply`` step whose arguments carry its position,
+    so replies never form a tool cycle on their own.
+    """
+
+    steps: list[AgentStep] = []
+    for index, message in enumerate(messages):
+        if str(message.get("role", "")).casefold() != "assistant":
+            continue
+        thought = content_text(message.get("content")).strip()
+        calls = tool_calls(message)
+        tokens_used = estimate_tokens([message])
+        if calls:
+            for call in calls:
+                name, arguments = parse_call(call)
+                steps.append(
+                    AgentStep(
+                        thought=thought or describe_call(name, arguments),
+                        tool_name=name,
+                        tool_args=arguments,
+                        tokens_used=tokens_used,
+                    )
+                )
+        elif thought:
+            steps.append(
+                AgentStep(
+                    thought=thought,
+                    tool_name=REPLY_TOOL,
+                    tool_args={"index": index},
+                    tokens_used=tokens_used,
+                )
+            )
+    return steps
+
+
+def tool_calls(message: Mapping[str, Any]) -> list[Mapping[str, Any]]:
+    """Return OpenAI ``tool_calls`` and Anthropic ``tool_use`` blocks of a message."""
+
+    calls: list[Mapping[str, Any]] = []
+    raw_calls = message.get("tool_calls")
+    if isinstance(raw_calls, list):
+        calls.extend(call for call in raw_calls if isinstance(call, Mapping))
+    content = message.get("content")
+    if isinstance(content, list):
+        calls.extend(
+            part for part in content if isinstance(part, Mapping) and part.get("type") == "tool_use"
+        )
+    return calls
+
+
+def parse_call(call: Mapping[str, Any]) -> tuple[str, dict[str, Any]]:
+    """Return the tool name and argument dict of one call in any supported shape."""
+
+    function = call.get("function")
+    payload = function if isinstance(function, Mapping) else call
+    name = str(payload.get("name") or "").strip() or "tool"
+    arguments = payload.get("arguments", payload.get("input", payload.get("args", {})))
+    return name, argument_dict(arguments)
+
+
+def nudge_for(verdict: LoopVerdict, steps: Sequence[AgentStep]) -> str:
+    """Return the intervention text for a flagged trajectory."""
+
+    tool = steps[-1].tool_name if steps else "tool"
+    if tool != REPLY_TOOL:
+        return str(generate_nudge(verdict, tool)["content"])
+    if verdict.intervention_prompt:
+        return str(verdict.intervention_prompt)
+    return "LoopShield detected a loop. Choose a different strategy immediately."
+
+
+def is_intervention(message: Mapping[str, Any]) -> bool:
+    """Return True for a message that already carries a LoopShield nudge."""
+
+    if message.get("loopshield_intervention") is True:
+        return True
+    text = content_text(message.get("content"))
+    return text.startswith("LoopShield detected") or (
+        "is temporarily locked" in text and "different strategy" in text
+    )
+
+
+def _role(value: Any) -> str:
+    role = str(value or "").strip().casefold()
+    return _ROLE_ALIASES.get(role, role)
+
+
+def _carry_marker(data: dict[str, Any], extras: Any) -> None:
+    if isinstance(extras, Mapping) and extras.get("loopshield") is True:
+        data["loopshield_intervention"] = True
+
+
+def _openai_call(call: Any) -> Any:
+    if not isinstance(call, Mapping) or isinstance(call.get("function"), Mapping):
+        return call
+    arguments = call.get("args", call.get("arguments", call.get("input", {})))
+    return {
+        "id": call.get("id"),
+        "type": "function",
+        "function": {
+            "name": call.get("name"),
+            "arguments": arguments if isinstance(arguments, str) else json.dumps(arguments, default=str),
+        },
+    }