dreadnode 1.15.3__tar.gz → 1.16.0__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: dreadnode
-Version: 1.15.3
+Version: 1.16.0
 Summary: Dreadnode SDK
 Project-URL: Homepage, https://github.com/dreadnode/sdk
 Project-URL: Repository, https://github.com/dreadnode/sdk
@@ -222,7 +222,7 @@ 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'

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

@@ -1,4 +1,6 @@
 import inspect
+import json
+import re
 import typing as t
 from contextlib import aclosing, asynccontextmanager
 from copy import deepcopy
@@ -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,
                     )
@@ -897,7 +898,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 +906,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 +922,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.16.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.16.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.16.0}/dreadnode/api/client.py

@@ -6,6 +6,7 @@ import typing as t
 from datetime import datetime, timezone
 from pathlib import Path
 from urllib.parse import urlparse
+from uuid import UUID
 
 import httpx
 from loguru import logger
@@ -19,6 +20,8 @@ from dreadnode.api.models import (
     ExportFormat,
     GithubTokenResponse,
     MetricAggregationType,
+    Organization,
+    PaginatedWorkspaces,
     Project,
     RawRun,
     RawTask,
@@ -34,6 +37,8 @@ from dreadnode.api.models import (
     TraceTree,
     UserDataCredentials,
     UserResponse,
+    Workspace,
+    WorkspaceFilter,
 )
 from dreadnode.api.util import (
     convert_flat_tasks_to_tree,
@@ -276,16 +281,47 @@ class ApiClient:
         response = self.request("GET", "/strikes/projects")
         return [Project(**project) for project in response.json()]
 
-    def get_project(self, project: str) -> Project:
+    def get_project(self, project_identifier: str | UUID, workspace_id: UUID) -> Project:
         """Retrieves details of a specific project.
 
         Args:
-            project (str): The project identifier.
+            project_identifier (str | UUID): The project identifier. ID, name, or slug.
 
         Returns:
             Project: The Project object.
         """
-        response = self.request("GET", f"/strikes/projects/{project!s}")
+        response = self.request(
+            "GET",
+            f"/strikes/projects/{project_identifier!s}",
+            params={"workspace_id": workspace_id},
+        )
+        return Project(**response.json())
+
+    def create_project(
+        self,
+        name: str | UUID | None = None,
+        workspace_id: UUID | None = None,
+        organization_id: UUID | None = None,
+    ) -> Project:
+        """Creates a new project.
+
+        Args:
+            name: The name of the project. If None, a default name will be used.
+            workspace_id: The workspace ID to create the project in. If None, the default workspace will be used.
+            organization_id: The organization ID to create the project in. If None, the default organization will be used.
+
+        Returns:
+            Project: The created Project object.
+        """
+        payload: dict[str, t.Any] = {}
+        if name is not None:
+            payload["name"] = name
+        if workspace_id is not None:
+            payload["workspace_id"] = str(workspace_id)
+        if organization_id is not None:
+            payload["org_id"] = str(organization_id)
+
+        response = self.request("POST", "/strikes/projects", json_data=payload)
         return Project(**response.json())
 
     def list_runs(self, project: str) -> list[RunSummary]:
@@ -757,3 +793,113 @@ class ApiClient:
         response = self.request("GET", "/platform/templates/all", params=params)
         zip_content: bytes = response.content
         return zip_content
+
+    # RBAC
+    def list_organizations(self) -> list[Organization]:
+        """
+        Retrieves a list of organizations the user belongs to.
+
+        Returns:
+            A list of organization names.
+        """
+        response = self.request("GET", "/organizations")
+        return [Organization(**org) for org in response.json()]
+
+    def get_organization(self, org_id_or_key: UUID | str) -> Organization:
+        """
+        Retrieves details of a specific organization.
+
+        Args:
+            org_id_or_key (str | UUID): The organization identifier.
+
+        Returns:
+            Organization: The Organization object.
+        """
+        response = self.request("GET", f"/organizations/{org_id_or_key!s}")
+        return Organization(**response.json())
+
+    def list_workspaces(self, filters: WorkspaceFilter | None = None) -> list[Workspace]:
+        """
+        Retrieves a list of workspaces the user has access to.
+
+        Returns:
+            A list of workspace names.
+        """
+        response = self.request(
+            "GET", "/workspaces", params=filters.model_dump() if filters else None
+        )
+        paginated_workspaces = PaginatedWorkspaces(**response.json())
+        # handle the pagination
+        all_workspaces: list[Workspace] = paginated_workspaces.workspaces.copy()
+        while paginated_workspaces.has_next:
+            response = self.request(
+                "GET",
+                "/workspaces",
+                params={
+                    "page": paginated_workspaces.page + 1,
+                    "limit": paginated_workspaces.limit,
+                    **(filters.model_dump() if filters else {}),
+                },
+            )
+            next_page = PaginatedWorkspaces(**response.json())
+            all_workspaces.extend(next_page.workspaces)
+            paginated_workspaces.page = next_page.page
+            paginated_workspaces.has_next = next_page.has_next
+
+        return all_workspaces
+
+    def get_workspace(
+        self, workspace_id_or_key: UUID | str, org_id: UUID | None = None
+    ) -> Workspace:
+        """
+        Retrieves details of a specific workspace.
+
+        Args:
+            workspace_id_or_key (str | UUID): The workspace identifier.
+
+        Returns:
+            Workspace: The Workspace object.
+        """
+        params: dict[str, str] = {}
+        if org_id:
+            params = {"org_id": str(org_id)}
+        response = self.request("GET", f"/workspaces/{workspace_id_or_key!s}", params=params)
+        return Workspace(**response.json())
+
+    def create_workspace(
+        self,
+        name: str,
+        key: str,
+        organization_id: UUID,
+        description: str | None = None,
+    ) -> Workspace:
+        """
+        Creates a new workspace.
+
+        Args:
+            name (str): The name of the workspace.
+            organization_id (str | UUID): The organization ID to create the workspace in.
+
+        Returns:
+            Workspace: The created Workspace object.
+        """
+
+        payload = {
+            "name": name,
+            "key": key,
+            "description": description,
+            "org_id": str(organization_id),
+        }
+
+        response = self.request("POST", "/workspaces", json_data=payload)
+        return Workspace(**response.json())
+
+    def delete_workspace(self, workspace_id: str | UUID) -> None:
+        """
+        Deletes a specific workspace.
+
+        Args:
+            workspace_id (str | UUID): The workspace key.
+        """
+
+        self.request("DELETE", f"/workspaces/{workspace_id!s}")