dreadnode 1.15.3__tar.gz → 1.17.0__tar.gz

{dreadnode-1.15.3 → dreadnode-1.17.0}/.gitignore

@@ -156,6 +156,7 @@ venv.bak/
 
 # mkdocs documentation
 /site
+debug.html
 
 # mypy
 .mypy_cache/

{dreadnode-1.15.3 → dreadnode-1.17.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: dreadnode
-Version: 1.15.3
+Version: 1.17.0
 Summary: Dreadnode SDK
 Project-URL: Homepage, https://github.com/dreadnode/sdk
 Project-URL: Repository, https://github.com/dreadnode/sdk
@@ -208,21 +208,29 @@ License:                                  Apache License
         See the License for the specific language governing permissions and
         limitations under the License.
 License-File: LICENSE
+Classifier: License :: OSI Approved :: Apache Software License
+Classifier: Operating System :: OS Independent
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
 Requires-Python: <3.14,>=3.10
+Requires-Dist: aiofiles<25.0.0,>=24.1.0
 Requires-Dist: coolname<3.0.0,>=2.2.0
 Requires-Dist: cyclopts>=4.2.0
-Requires-Dist: fsspec[s3]<=2025.3.0,>=2023.1.0
+Requires-Dist: fsspec[s3]<=2025.12.0,>=2023.1.0
 Requires-Dist: httpx<1.0.0,>=0.28.0
 Requires-Dist: logfire<=3.20.0,>=3.5.3
 Requires-Dist: loguru>=0.7.3
-Requires-Dist: numpy<=2.2.6
+Requires-Dist: numpy<=2.3.5
 Requires-Dist: optuna<5.0.0,>=4.5.0
 Requires-Dist: pandas<3.0.0,>=2.2.3
 Requires-Dist: pydantic<3.0.0,>=2.9.2
 Requires-Dist: python-jsonpath>=2.0.1
 Requires-Dist: python-ulid<4.0.0,>=3.0.0
 Requires-Dist: pyyaml>=6.0.2
-Requires-Dist: rigging<4.0.0,>=3.2.1
+Requires-Dist: rigging>=3.3.4
 Requires-Dist: universal-pathlib<0.4.0,>=0.3.3
 Provides-Extra: all
 Requires-Dist: confusables<2.0.0,>=1.2.0; extra == 'all'
@@ -278,7 +286,7 @@ Dreadnode Strikes SDK
     <img alt="PyPI - Python Version" src="https://img.shields.io/pypi/pyversions/dreadnode">
     <img alt="PyPI - Version" src="https://img.shields.io/pypi/v/dreadnode">
     <img alt="GitHub License" src="https://img.shields.io/github/license/dreadnode/sdk">
-    <img alt="Tests" src="https://img.shields.io/github/actions/workflow/status/dreadnode/sdk/tests.yaml">
+    <img alt="Tests" src="https://img.shields.io/github/actions/workflow/status/dreadnode/sdk/test.yaml">
     <img alt="Pre-Commit" src="https://img.shields.io/github/actions/workflow/status/dreadnode/sdk/pre-commit.yaml">
     <img alt="Renovate" src="https://img.shields.io/github/actions/workflow/status/dreadnode/sdk/renovate.yaml">
 </h4>

{dreadnode-1.15.3 → dreadnode-1.17.0}/README.md

@@ -16,7 +16,7 @@ Dreadnode Strikes SDK
     <img alt="PyPI - Python Version" src="https://img.shields.io/pypi/pyversions/dreadnode">
     <img alt="PyPI - Version" src="https://img.shields.io/pypi/v/dreadnode">
     <img alt="GitHub License" src="https://img.shields.io/github/license/dreadnode/sdk">
-    <img alt="Tests" src="https://img.shields.io/github/actions/workflow/status/dreadnode/sdk/tests.yaml">
+    <img alt="Tests" src="https://img.shields.io/github/actions/workflow/status/dreadnode/sdk/test.yaml">
     <img alt="Pre-Commit" src="https://img.shields.io/github/actions/workflow/status/dreadnode/sdk/pre-commit.yaml">
     <img alt="Renovate" src="https://img.shields.io/github/actions/workflow/status/dreadnode/sdk/renovate.yaml">
 </h4>

{dreadnode-1.15.3 → dreadnode-1.17.0}/dreadnode/agent/agent.py

@@ -1,6 +1,8 @@
 import inspect
+import json
+import re
 import typing as t
-from contextlib import aclosing, asynccontextmanager
+from contextlib import AsyncExitStack, aclosing, asynccontextmanager
 from copy import deepcopy
 from textwrap import dedent
 
@@ -20,7 +22,6 @@ from dreadnode.agent.events import (
     AgentEventInStep,
     AgentStalled,
     AgentStart,
-    AgentStopReason,
     GenerationEnd,
     Reacted,
     StepStart,
@@ -37,7 +38,7 @@ from dreadnode.agent.reactions import (
     Retry,
     RetryWithFeedback,
 )
-from dreadnode.agent.result import AgentResult
+from dreadnode.agent.result import AgentResult, AgentStopReason
 from dreadnode.agent.stop import StopCondition, never
 from dreadnode.agent.thread import Thread
 from dreadnode.agent.tools import AnyTool, Tool, Toolset, discover_tools_on_obj
@@ -59,7 +60,6 @@ from dreadnode.util import (
 litellm.suppress_debug_info = True
 
 CommitBehavior = t.Literal["always", "on-success"]
-HookMap = dict[type[AgentEvent], list[Hook]]
 
 
 class AgentWarning(UserWarning):
@@ -250,6 +250,9 @@ class Agent(Model):
             new.scorers = scorers if scorers is not None else new.scorers
             new.assert_scores = assert_scores if assert_scores is not None else new.assert_scores
 
+        # Retrigger model_post_init functions to ensure consistency
+        new.model_post_init(None)
+
         return new
 
     def _get_transforms(self) -> list[rg.Transform]:
@@ -269,29 +272,11 @@ class Agent(Model):
                     transforms.append(rg.transform.tools_to_json_with_tag_transform)
                 case "json":
                     transforms.append(rg.transform.tools_to_json_transform)
+                case "pythonic":
+                    transforms.append(rg.transform.tools_to_pythonic_transform)
 
         return transforms
 
-    def _get_hooks(self) -> dict[type[AgentEvent], list[Hook]]:
-        hooks: dict[type[AgentEvent], list[Hook]] = {}
-        for hook in self.hooks:
-            sig = inspect.signature(hook)
-            if not (params := list(sig.parameters.values())):
-                continue
-            event_type = params[0].annotation
-
-            if hasattr(event_type, "__origin__") and event_type.__origin__ is t.Union:
-                union_args = event_type.__args__
-                for arg in union_args:
-                    if inspect.isclass(arg) and issubclass(arg, AgentEvent):
-                        hooks.setdefault(arg, []).append(hook)
-            elif inspect.isclass(event_type) and issubclass(event_type, AgentEvent):
-                hooks.setdefault(event_type, []).append(hook)
-            else:
-                hooks.setdefault(AgentEvent, []).append(hook)
-
-        return hooks
-
     async def _generate(
         self,
         messages: list[rg.Message],
@@ -352,7 +337,6 @@ class Agent(Model):
         self,
         thread: "Thread",
         messages: list[rg.Message],
-        hooks: HookMap,
         *,
         commit: CommitBehavior,
     ) -> t.AsyncGenerator[AgentEvent, None]:
@@ -369,26 +353,21 @@ class Agent(Model):
 
         # Event dispatcher
 
-        async def _dispatch(event: AgentEvent) -> t.AsyncIterator[AgentEvent]:
+        async def _dispatch(event: AgentEvent) -> t.AsyncIterator[AgentEvent]:  # noqa: PLR0912
             nonlocal messages, events
 
             yield event
 
             events.append(event)
 
-            # If we have no hooks, just return the event
-            applicable_hooks = list(set(hooks.get(type(event), []) + hooks.get(AgentEvent, [])))
-            if not applicable_hooks:
-                return
-
             logger.debug(
                 f"Agent '{self.name}' ({session_id}) dispatching '{type(event).__name__}': "
-                f"applicable_hooks={[get_callable_name(h, short=True) for h in applicable_hooks]}"
+                f"hooks={[get_callable_name(h, short=True) for h in self.hooks]}"
             )
 
             # Run all applicable hooks and collect their reactions
             hook_reactions: dict[str, Reaction | None] = {}
-            for hook in applicable_hooks:
+            for hook in self.hooks:
                 hook_name = getattr(
                     hook, "__name__", getattr(hook, "__qualname__", safe_repr(hook))
                 )
@@ -415,6 +394,20 @@ class Agent(Model):
                     )
                     continue
 
+                if isinstance(event, AgentEnd):
+                    warn_at_user_stacklevel(
+                        f"Hook '{hook_name}' returned {reaction} during AgentEnd, but reactions are ignored at this stage.",
+                        AgentWarning,
+                    )
+                    continue
+
+                if isinstance(event, Reacted):
+                    warn_at_user_stacklevel(
+                        f"Hook '{hook_name}' returned {reaction} during Reacted, but reactions are ignored at this stage.",
+                        AgentWarning,
+                    )
+                    continue
+
                 hook_reactions[hook_name] = reaction
 
             if not hook_reactions:
@@ -477,7 +470,9 @@ class Agent(Model):
                 reaction=winning_reaction,
             )
             events.append(reacted_event)
-            yield reacted_event
+
+            async for _event in _dispatch(reacted_event):
+                yield _event
 
             if isinstance(winning_reaction, Continue):
                 messages = winning_reaction.messages
@@ -570,10 +565,12 @@ class Agent(Model):
 
         # Core step loop
 
-        step = 1
+        step = 0
         error: Exception | str | None = None
 
-        while step <= self.max_steps + 1:
+        while step < self.max_steps:
+            step += 1
+
             try:
                 async for event in _dispatch(
                     StepStart(
@@ -589,7 +586,7 @@ class Agent(Model):
 
                 # Generation
 
-                step_chat = await self._generate(messages=messages)
+                step_chat = await self._generate(messages)
                 if step_chat.failed and step_chat.error:
                     async for event in _dispatch(
                         AgentError(
@@ -602,7 +599,9 @@ class Agent(Model):
                         )
                     ):
                         yield event
-                    raise step_chat.error
+
+                    error = t.cast("Exception", step_chat.error)  # Should be Exception in rigging
+                    break
 
                 # Sync extra fields to metadata for storage
                 step_chat.generated[-1].metadata.update(step_chat.extra)
@@ -649,7 +648,8 @@ class Agent(Model):
                     ):
                         yield event
 
-                    continue
+                    # If the agent is stalled and nobody handled it, break out
+                    break
 
                 # Process tool calls
 
@@ -675,8 +675,6 @@ class Agent(Model):
                 if any(cond(events) for cond in stop_conditions):
                     break
 
-                step += 1
-
             except Retry as e:
                 messages = e.messages or messages
                 continue
@@ -689,7 +687,7 @@ class Agent(Model):
                 break
 
         stop_reason: AgentStopReason = "finished"
-        if step > self.max_steps + 1:
+        if step >= self.max_steps:
             error = MaxStepsError(max_steps=self.max_steps)
             stop_reason = "max_steps_reached"
         elif error is not None:
@@ -720,28 +718,32 @@ class Agent(Model):
         else:
             logger.warning(log_message)
 
-        yield AgentEnd(
-            session_id=session_id,
-            agent=self,
-            thread=thread,
-            messages=messages,
-            events=events,
-            stop_reason=stop_reason,
-            result=AgentResult(
+        async for event in _dispatch(
+            AgentEnd(
+                session_id=session_id,
                 agent=self,
+                thread=thread,
                 messages=messages,
-                usage=_total_usage_from_events(events),
-                steps=step,
-                failed=stop_reason != "finished",
-                error=error,
-            ),
-        )
+                events=events,
+                stop_reason=stop_reason,
+                result=AgentResult(
+                    agent=self,
+                    messages=messages,
+                    stop_reason=stop_reason,
+                    usage=_total_usage_from_events(events),
+                    steps=step,
+                    failed=stop_reason != "finished",
+                    error=error,
+                ),
+            )
+        ):
+            yield event
 
     def _log_event_metrics(self, event: AgentEvent) -> None:
         from dreadnode import log_metric
 
         if isinstance(event, AgentEnd):
-            log_metric("steps_taken", min(0, event.result.steps - 1))
+            log_metric("steps_taken", max(0, event.result.steps - 1))
             log_metric(f"stop_{event.stop_reason}", 1)
 
         if not isinstance(event, AgentEventInStep):
@@ -778,7 +780,6 @@ class Agent(Model):
     ) -> t.AsyncGenerator[AgentEvent, None]:
         from dreadnode import log_output, log_outputs, score, task_and_run
 
-        hooks = self._get_hooks()
         messages = [*deepcopy(thread.messages), rg.Message("user", str(user_input))]
 
         configuration = get_config_model(self)()
@@ -822,7 +823,7 @@ class Agent(Model):
             params=trace_params,
         ):
             try:
-                async with aclosing(self._stream(thread, messages, hooks, commit=commit)) as stream:
+                async with aclosing(self._stream(thread, messages, commit=commit)) as stream:
                     async for event in stream:
                         last_event = event
                         self._log_event_metrics(event)
@@ -837,7 +838,7 @@ class Agent(Model):
                 if isinstance(last_event, AgentEnd):
                     log_outputs(
                         to="both",
-                        steps_taken=min(0, last_event.result.steps - 1),
+                        steps_taken=max(0, last_event.result.steps - 1),
                         reason=last_event.stop_reason,
                         failed=last_event.result.failed,
                     )
@@ -874,8 +875,16 @@ class Agent(Model):
         commit: CommitBehavior = "always",
     ) -> t.AsyncIterator[t.AsyncGenerator[AgentEvent, None]]:
         thread = thread or self.thread
-        async with aclosing(self._stream_traced(thread, user_input, commit=commit)) as stream:
-            yield stream
+
+        async with AsyncExitStack() as stack:
+            # Ensure all tools are properly entered if they
+            # are context managers before we start using them
+            for tool_container in self.tools:
+                if hasattr(tool_container, "__aenter__") and hasattr(tool_container, "__aexit__"):
+                    await stack.enter_async_context(tool_container)
+
+            async with aclosing(self._stream_traced(thread, user_input, commit=commit)) as stream:
+                yield stream
 
     async def run(
         self,
@@ -897,7 +906,7 @@ class Agent(Model):
 
 class TaskAgent(Agent):
     """
-    A specialized agent for running tasks with a focus on completion and reporting.
+    A specialized agent mixin for running tasks with a focus on completion and reporting.
     It extends the base Agent class to provide task-specific functionality.
 
     - Automatically includes the `finish_task`, `give_up_on_task`, and `update_todo` tools.
@@ -905,7 +914,12 @@ class TaskAgent(Agent):
     - Uses the `AgentStalled` event to handle stalled tasks by pushing the model to continue or finish the task.
     """
 
-    def model_post_init(self, _: t.Any) -> None:
+    def model_post_init(self, context: t.Any) -> None:
+        super().model_post_init(context)
+
+        # TODO(nick): Would be better to have a pattern here for
+        # add-if-missing for tools, hooks, and stop conditions
+
         if not any(tool for tool in self.tools if tool.name == "finish_task"):
             self.tools.append(finish_task)
 
@@ -916,11 +930,72 @@ class TaskAgent(Agent):
             self.tools.append(update_todo)
 
         # Force the agent to use finish_task
-        self.stop_conditions.append(never())
-        self.hooks.insert(
-            0,
-            retry_with_feedback(
-                event_type=AgentStalled,
-                feedback="Continue the task if possible, use the 'finish_task' tool to complete it, or 'give_up_on_task' if it cannot be completed.",
-            ),
-        )
+        if not any(cond for cond in self.stop_conditions if cond.name == "stop_never"):
+            self.stop_conditions.append(never())
+
+        if not any(
+            hook
+            for hook in self.hooks
+            if get_callable_name(hook, short=True) == "retry_with_feedback"
+        ):
+            self.hooks.append(
+                retry_with_feedback(
+                    event_type=AgentStalled,
+                    feedback="No tool calls were observed. Continue the task if possible, use the 'finish_task' tool to complete it, or 'give_up_on_task' if it cannot be completed.",
+                )
+            )
+
+
+class RegexRefAgent(Agent):
+    """
+    An agent mixin that allows for dynamic references of prior text using regex patterns in tool arguments.
+    This helps prevent repeating large amounts of prior text in tool calls.
+
+    Instructions are automatically added to the agent's instructions to guide usage of the {find:<pattern>} syntax
+    along with a hook that resolves these references during tool calls.
+    """
+
+    @staticmethod
+    async def resolve_regex_ref(event: AgentEvent) -> Reaction | None:
+        if not isinstance(event, ToolStart):
+            return None
+
+        for m in re.finditer(r"\{find:(.*?)\}", event.tool_call.arguments):
+            regex = m.group(1).replace("\\\\", "\\")  # models tend to over-escape
+            logger.info(f"Found find reference: {regex}")
+            all_message_content = "\n\n".join([m.content for m in event.messages])
+            reference_matches = re.findall(regex, all_message_content)
+            if reference_matches:
+                logger.debug(f"Replacing '{m.group(0)}' with '{reference_matches[-1][:50]}...'.")
+                event.tool_call.function.arguments = event.tool_call.arguments.replace(
+                    m.group(0), json.dumps(reference_matches[-1]).strip('"')
+                )
+
+        return None
+
+    def model_post_init(self, context: t.Any) -> None:
+        super().model_post_init(context)
+
+        if not any(
+            hook
+            for hook in self.hooks
+            if get_callable_name(hook, short=True) == "resolve_regex_ref"
+        ):
+            self.hooks.append(RegexRefAgent.resolve_regex_ref)
+
+        instruction_section = dedent("""
+        # Regex Find Instructions
+        To efficiently reuse data from the conversation, you can pass {find:<pattern>} anywhere in tool arguments to dynamically
+        refer to prior text using a regex pattern. This helps prevent costly repetition of prior text.
+
+        You must escape special characters in the regex.
+
+        Example: If the history contains `$krb5tgs$23$*user...<long_hash>`, use:
+        `hashcat(hashes=["{find:\\$krb5tgs\\$.*}"], wordlist="...")`
+        and the system will find the full hash for you and insert it into the tool call.
+        """)
+
+        if self.instructions is None:
+            self.instructions = instruction_section
+        elif self.instructions and instruction_section not in self.instructions:
+            self.instructions += "\n\n" + instruction_section

{dreadnode-1.15.3 → dreadnode-1.17.0}/dreadnode/agent/events.py

@@ -24,13 +24,12 @@ from dreadnode.util import format_dict, shorten_string
 if t.TYPE_CHECKING:
     from dreadnode.agent.agent import Agent
     from dreadnode.agent.reactions import Reaction
-    from dreadnode.agent.result import AgentResult
+    from dreadnode.agent.result import AgentResult, AgentStopReason
     from dreadnode.agent.thread import Thread
     from dreadnode.common_types import AnyDict
 
 
 AgentEventT = t.TypeVar("AgentEventT", bound="AgentEvent")
-AgentStopReason = t.Literal["finished", "max_steps_reached", "error", "stalled"]
 
 
 @dataclass
@@ -292,7 +291,7 @@ class Reacted(AgentEventInStep):
 
 @dataclass
 class AgentEnd(AgentEvent):
-    stop_reason: AgentStopReason
+    stop_reason: "AgentStopReason"
     result: "AgentResult"
 
     def format_as_panel(self, *, truncate: bool = False) -> Panel:  # noqa: ARG002

{dreadnode-1.15.3 → dreadnode-1.17.0}/dreadnode/agent/result.py

@@ -8,10 +8,13 @@ from rigging.message import Message
 if t.TYPE_CHECKING:
     from dreadnode.agent.agent import Agent
 
+AgentStopReason = t.Literal["finished", "max_steps_reached", "error", "stalled"]
+
 
 @dataclass(config=ConfigDict(arbitrary_types_allowed=True))
 class AgentResult:
     agent: "Agent"
+    stop_reason: AgentStopReason
     messages: list[Message]
     usage: Usage
     steps: int
@@ -23,6 +26,7 @@ class AgentResult:
             f"agent={self.agent.name}",
             f"messages={len(self.messages)}",
             f"usage='{self.usage}'",
+            f"stop_reason='{self.stop_reason}'",
             f"steps={self.steps}",
         ]
 

{dreadnode-1.15.3 → dreadnode-1.17.0}/dreadnode/agent/tools/base.py

@@ -1,6 +1,8 @@
+import asyncio
+import functools
 import typing as t
 
-from pydantic import ConfigDict
+from pydantic import ConfigDict, PrivateAttr
 from rigging import tools
 from rigging.tools.base import ToolMethod as RiggingToolMethod
 
@@ -171,18 +173,82 @@ class Toolset(Model):
     - Pydantic's declarative syntax for defining state (fields).
     - Automatic application of the `@configurable` decorator.
     - A `get_tools` method for discovering methods decorated with `@dreadnode.tool_method`.
+    - Support for async context management, with automatic re-entrancy handling.
     """
 
+    model_config = ConfigDict(arbitrary_types_allowed=True, use_attribute_docstrings=True)
+
     variant: str | None = None
     """The variant for filtering tools available in this toolset."""
 
-    model_config = ConfigDict(arbitrary_types_allowed=True, use_attribute_docstrings=True)
+    # Context manager magic
+    _entry_ref_count: int = PrivateAttr(default=0)
+    _context_handle: object = PrivateAttr(default=None)
+    _entry_lock: asyncio.Lock = PrivateAttr(default_factory=asyncio.Lock)
 
     @property
     def name(self) -> str:
         """The name of the toolset, derived from the class name."""
         return self.__class__.__name__
 
+    def __init_subclass__(cls, **kwargs: t.Any) -> None:
+        super().__init_subclass__(**kwargs)
+
+        # This essentially ensures that if the Toolset is any kind of context manager,
+        # it will be re-entrant, and only actually enter/exit once. This means we can
+        # safely build auto-entry/exit logic into our Agent class without worrying about
+        # breaking the code if the user happens to enter a toolset manually before using
+        # it in an agent.
+
+        original_aenter = cls.__dict__.get("__aenter__")
+        original_enter = cls.__dict__.get("__enter__")
+        original_aexit = cls.__dict__.get("__aexit__")
+        original_exit = cls.__dict__.get("__exit__")
+
+        has_enter = callable(original_aenter) or callable(original_enter)
+        has_exit = callable(original_aexit) or callable(original_exit)
+
+        if has_enter and not has_exit:
+            raise TypeError(
+                f"{cls.__name__} defining __aenter__ or __enter__ must also define __aexit__ or __exit__"
+            )
+        if has_exit and not has_enter:
+            raise TypeError(
+                f"{cls.__name__} defining __aexit__ or __exit__ must also define __aenter__ or __enter__"
+            )
+        if original_aenter and original_enter:
+            raise TypeError(f"{cls.__name__} cannot define both __aenter__ and __enter__")
+        if original_aexit and original_exit:
+            raise TypeError(f"{cls.__name__} cannot define both __aexit__ and __exit__")
+
+        @functools.wraps(original_aenter or original_enter)  # type: ignore[arg-type]
+        async def aenter_wrapper(self: "Toolset", *args: t.Any, **kwargs: t.Any) -> t.Any:
+            async with self._entry_lock:
+                if self._entry_ref_count == 0:
+                    handle = None
+                    if original_aenter:
+                        handle = await original_aenter(self, *args, **kwargs)
+                    elif original_enter:
+                        handle = original_enter(self, *args, **kwargs)
+                    self._context_handle = handle if handle is not None else self
+                self._entry_ref_count += 1
+                return self._context_handle
+
+        cls.__aenter__ = aenter_wrapper  # type: ignore[attr-defined]
+
+        @functools.wraps(original_aexit or original_exit)  # type: ignore[arg-type]
+        async def aexit_wrapper(self: "Toolset", *args: t.Any, **kwargs: t.Any) -> t.Any:
+            async with self._entry_lock:
+                self._entry_ref_count -= 1
+                if self._entry_ref_count == 0:
+                    if original_aexit:
+                        await original_aexit(self, *args, **kwargs)
+                    elif original_exit:
+                        original_exit(self, *args, **kwargs)
+                    self._context_handle = None
+
+        cls.__aexit__ = aexit_wrapper  # type: ignore[attr-defined]
+
     def get_tools(self, *, variant: str | None = None) -> list[AnyTool]:
         variant = variant or self.variant