dreadnode 1.15.1__tar.gz → 1.15.2__tar.gz

{dreadnode-1.15.1 → dreadnode-1.15.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: dreadnode
-Version: 1.15.1
+Version: 1.15.2
 Summary: Dreadnode SDK
 Project-URL: Homepage, https://github.com/dreadnode/sdk
 Project-URL: Repository, https://github.com/dreadnode/sdk

{dreadnode-1.15.1 → dreadnode-1.15.2}/dreadnode/agent/agent.py

@@ -2,10 +2,11 @@ import inspect
 import typing as t
 from contextlib import aclosing, asynccontextmanager
 from copy import deepcopy
+from textwrap import dedent
 
 import rigging as rg
 from loguru import logger
-from pydantic import ConfigDict, Field, PrivateAttr, SkipValidation, field_validator
+from pydantic import AfterValidator, ConfigDict, Field, PrivateAttr, SkipValidation, field_validator
 from rigging.message import inject_system_content
 from ulid import ULID  # can't access via rg
 
@@ -71,14 +72,18 @@ class Agent(Model):
 
     name: str
     """The name of the agent."""
-    description: str = ""
+    description: t.Annotated[str, AfterValidator(dedent)] = ""
     """A brief description of the agent's purpose."""
     tags: list[str] = Config(default_factory=lambda: ["agent"])
     """A list of tags associated with the agent."""
+    label: str | None = Config(default=None)
+    """Specific label for tracing, otherwise derived from the name."""
 
     model: str | None = Config(default=None)
     """Inference model (rigging generator identifier)."""
-    instructions: str | None = Config(default=None)
+    instructions: t.Annotated[str | None, AfterValidator(lambda x: dedent(x) if x else x)] = Config(
+        default=None
+    )
     """The agent's core instructions."""
     max_steps: int = Config(default=10)
     """The maximum number of steps (generation + tool calls)."""
@@ -90,15 +95,15 @@ class Agent(Model):
     tool_mode: ToolMode = Config(default="auto", repr=False)
     """The tool calling mode to use."""
 
-    hooks: list[Hook] = Config(default_factory=list, exclude=True, repr=False)
+    hooks: list[Hook] = Field(default_factory=list, exclude=True, repr=False)
     """Hooks to run at various points in the agent's lifecycle."""
-    stop_conditions: list[StopCondition] = Config(default_factory=list)
+    stop_conditions: list[StopCondition] = Field(default_factory=list)
     """The logical condition for successfully stopping a run."""
     thread: Thread = Field(default_factory=Thread, exclude=True, repr=False)
     """Stateful thread for this agent, for when otherwise not specified during execution."""
-    scorers: ScorersLike[AgentResult] = Config(default_factory=list)
+    scorers: ScorersLike[AgentResult] = Field(default_factory=list)
     """Scorers to evaluate the agent output."""
-    assert_scores: list[str] | t.Literal[True] = Config(default_factory=list)
+    assert_scores: list[str] | t.Literal[True] = Field(default_factory=list)
     """Scores to ensure are truthy, otherwise the agent task is marked as failed."""
 
     _generator: rg.Generator | None = PrivateAttr(None, init=False)
@@ -716,14 +721,25 @@ class Agent(Model):
         )
         trace_params.update(
             {
+                "name": self.name,
                 "model": self.model,
                 "max_steps": self.max_steps,
                 "tool_mode": self.tool_mode,
+                "tool_count": len(self.all_tools),
+                "instructions_length": len(self.instructions or ""),
+                "stop_condition_count": len(self.stop_conditions),
+                "message_count": len(messages),
             }
         )
 
         last_event: AgentEvent | None = None
-        with task_and_run(name=self.name, tags=self.tags, inputs=trace_inputs, params=trace_params):
+        with task_and_run(
+            name=self.name,
+            tags=self.tags,
+            label=self.label,
+            inputs=trace_inputs,
+            params=trace_params,
+        ):
             try:
                 async with aclosing(self._stream(thread, messages, hooks, commit=commit)) as stream:
                     async for event in stream:

{dreadnode-1.15.1 → dreadnode-1.15.2}/dreadnode/agent/hooks/__init__.py

@@ -3,6 +3,7 @@ from dreadnode.agent.hooks.base import (
     Hook,
     retry_with_feedback,
 )
+from dreadnode.agent.hooks.metrics import tool_metrics
 from dreadnode.agent.hooks.summarize import summarize_when_long
 
 __all__ = [
@@ -11,4 +12,5 @@ __all__ = [
     "backoff_on_ratelimit",
     "retry_with_feedback",
     "summarize_when_long",
+    "tool_metrics",
 ]

dreadnode-1.15.2/dreadnode/agent/hooks/metrics.py

@@ -0,0 +1,84 @@
+import typing as t
+
+from dreadnode.agent.events import AgentEvent, ToolEnd, ToolStart
+from dreadnode.agent.hooks import Hook
+from dreadnode.meta import Config, component
+
+if t.TYPE_CHECKING:
+    from datetime import datetime
+
+
+def tool_metrics(*, detailed: bool = False) -> Hook:
+    """
+    Creates an agent hook to log metrics about tool usage, execution time, and success rates.
+
+    Args:
+        detailed: If True, logs metrics for each specific tool in addition to general stats.
+                  If False, only logs aggregate statistics across all tools.
+
+    Returns:
+        An async hook function that can be registered with an agent.
+    """
+    _start_times: dict[str, datetime] = {}
+
+    @component
+    async def tool_metrics(
+        event: AgentEvent,
+        *,
+        detailed: bool = Config(
+            default=detailed,
+            help="If True, logs metrics for each specific tool in addition to general stats.",
+        ),
+    ) -> None:
+        """The actual hook implementation that processes agent events."""
+        from dreadnode import log_metric
+
+        if isinstance(event, ToolStart):
+            log_metric("tool/total_count", 1, step=event.step, mode="count")
+            _start_times[event.tool_call.id] = event.timestamp
+
+            if detailed:
+                tool_name = event.tool_call.name
+                log_metric(f"tool/count.{tool_name}", 1, step=event.step, mode="count")
+
+        elif isinstance(event, ToolEnd):
+            tool_name = event.tool_call.name
+            start_time = _start_times.pop(event.tool_call.id, event.timestamp)
+            duration_seconds = (event.timestamp - start_time).total_seconds()
+            errored = "error" in event.message.metadata
+
+            log_metric("tool/total_time", duration_seconds, step=event.step, mode="sum")
+            log_metric("tool/success_rate", 0 if errored else 1, step=event.step, mode="avg")
+
+            if errored:
+                log_metric("tool/failed_count", 1, step=event.step, mode="count")
+
+            if detailed:
+                log_metric(
+                    f"tool/time.{tool_name}",
+                    duration_seconds,
+                    step=event.step,
+                    mode="sum",
+                )
+                log_metric(
+                    f"tool/avg_time.{tool_name}",
+                    duration_seconds,
+                    step=event.step,
+                    mode="avg",
+                )
+                log_metric(
+                    f"tool/success_rate.{tool_name}",
+                    0 if errored else 1,
+                    step=event.step,
+                    mode="avg",
+                )
+
+                if errored:
+                    log_metric(
+                        f"tool/failed_count.{tool_name}",
+                        1,
+                        step=event.step,
+                        mode="count",
+                    )
+
+    return tool_metrics

{dreadnode-1.15.1 → dreadnode-1.15.2}/dreadnode/agent/tools/base.py

@@ -4,7 +4,7 @@ from pydantic import ConfigDict
 from rigging import tools
 from rigging.tools.base import ToolMethod as RiggingToolMethod
 
-from dreadnode.meta import Component, Config, Model
+from dreadnode.meta import Component, Model
 
 Tool = tools.Tool
 ToolMode = tools.ToolMode
@@ -103,18 +103,18 @@ def tool_method(
     description: str | None = None,
     catch: bool | t.Iterable[type[Exception]] | None = None,
     truncate: int | None = None,
-) -> t.Callable[[t.Callable[P, R]], RiggingToolMethod[P, R]]: ...
+) -> t.Callable[[t.Callable[t.Concatenate[t.Any, P], R]], RiggingToolMethod[P, R]]: ...
 
 
 @t.overload
 def tool_method(
-    func: t.Callable[P, R],
+    func: t.Callable[t.Concatenate[t.Any, P], R],
     /,
 ) -> RiggingToolMethod[P, R]: ...
 
 
 def tool_method(
-    func: t.Callable[P, R] | None = None,
+    func: t.Callable[t.Concatenate[t.Any, P], R] | None = None,
     /,
     *,
     variants: list[str] | None = None,
@@ -122,7 +122,10 @@ def tool_method(
     description: str | None = None,
     catch: bool | t.Iterable[type[Exception]] | None = None,
     truncate: int | None = None,
-) -> t.Callable[[t.Callable[P, R]], RiggingToolMethod[P, R]] | RiggingToolMethod[P, R]:
+) -> (
+    t.Callable[[t.Callable[t.Concatenate[t.Any, P], R]], RiggingToolMethod[P, R]]
+    | RiggingToolMethod[P, R]
+):
     """
     Marks a method on a Toolset as a tool, adding it to specified variants.
 
@@ -143,7 +146,9 @@ def tool_method(
         truncate: The maximum number of characters for the tool's output.
     """
 
-    def make_tool_method(func: t.Callable[P, R]) -> RiggingToolMethod[P, R]:
+    def make_tool_method(
+        func: t.Callable[t.Concatenate[t.Any, P], R],
+    ) -> RiggingToolMethod[P, R]:
         tool_method_descriptor: RiggingToolMethod[P, R] = tools.tool_method(
             name=name,
             description=description,
@@ -168,7 +173,7 @@ class Toolset(Model):
     - A `get_tools` method for discovering methods decorated with `@dreadnode.tool_method`.
     """
 
-    variant: str = Config("all")
+    variant: str | None = None
     """The variant for filtering tools available in this toolset."""
 
     model_config = ConfigDict(arbitrary_types_allowed=True, use_attribute_docstrings=True)
@@ -190,7 +195,7 @@ class Toolset(Model):
                     continue
 
                 variants = getattr(class_member, TOOL_VARIANTS_ATTR, [])
-                if variant in variants:
+                if not variant or not variants or variant in variants:
                     bound_tool = t.cast("AnyTool", getattr(self, name))
                     tools.append(bound_tool)
                     seen_names.add(name)

{dreadnode-1.15.1 → dreadnode-1.15.2}/dreadnode/cli/agent/cli.py

@@ -139,8 +139,6 @@ async def run(  # noqa: PLR0912, PLR0915
     agent_cli.__annotations__["config"] = config_annotation
 
     help_text = f"Run the '{agent_name}' agent."
-    if agent_blueprint.__doc__:
-        help_text += "\n\n" + agent_blueprint.__doc__
     if agent_blueprint.description:
         help_text += "\n\n" + agent_blueprint.description
 

{dreadnode-1.15.1 → dreadnode-1.15.2}/dreadnode/cli/attack/cli.py

@@ -128,8 +128,6 @@ async def run(  # noqa: PLR0912, PLR0915
     attack_cli.__annotations__["config"] = config_annotation
 
     help_text = f"Run the '{attack_name}' attack."
-    if attack_blueprint.__doc__:
-        help_text += "\n\n" + attack_blueprint.__doc__
     if attack_blueprint.description:
         help_text += "\n\n" + attack_blueprint.description
 

{dreadnode-1.15.1 → dreadnode-1.15.2}/dreadnode/cli/eval/cli.py

@@ -129,8 +129,6 @@ async def run(  # noqa: PLR0912, PLR0915
     eval_cli.__annotations__["config"] = config_annotation
 
     help_text = f"Run the '{eval_name}' eval."
-    if eval_blueprint.__doc__:
-        help_text += "\n\n" + eval_blueprint.__doc__
     if eval_blueprint.description:
         help_text += "\n\n" + eval_blueprint.description
 

{dreadnode-1.15.1 → dreadnode-1.15.2}/dreadnode/cli/study/cli.py

@@ -130,8 +130,6 @@ async def run(  # noqa: PLR0912, PLR0915
     study_cli.__annotations__["config"] = config_annotation
 
     help_text = f"Run the '{study_name}' study."
-    if study_blueprint.__doc__:
-        help_text += "\n\n" + study_blueprint.__doc__
     if study_blueprint.description:
         help_text += "\n\n" + study_blueprint.description
 

{dreadnode-1.15.1 → dreadnode-1.15.2}/dreadnode/cli/task/cli.py

@@ -4,6 +4,7 @@ import itertools
 import typing as t
 from inspect import isawaitable
 from pathlib import Path
+from textwrap import dedent
 
 import cyclopts
 import rich
@@ -131,7 +132,7 @@ async def run(  # noqa: PLR0912, PLR0915
 
     help_text = f"Run the '{task_name}' task."
     if task_blueprint.__doc__:
-        help_text += "\n\n" + task_blueprint.__doc__
+        help_text += "\n\n" + dedent(task_blueprint.__doc__)
 
     task_app = cyclopts.App(
         name=task_name,

{dreadnode-1.15.1 → dreadnode-1.15.2}/dreadnode/eval/eval.py

@@ -359,6 +359,7 @@ class Eval(Model, t.Generic[In, Out]):
                     tags=self.tags,
                     inputs=trace_inputs,
                     params={**trace_params, **scenario_params},
+                    label=self.label,
                 )
                 if self.trace
                 else contextlib.nullcontext()

{dreadnode-1.15.1 → dreadnode-1.15.2}/dreadnode/optimization/study.py

@@ -58,6 +58,8 @@ class Study(Model, t.Generic[CandidateT, OutputT]):
     """A brief description of the study's purpose."""
     tags: list[str] = Config(default_factory=lambda: ["study"])
     """A list of tags associated with the study for logging."""
+    label: str | None = Config(default=None)
+    """Specific label for tracing, otherwise derived from the name."""
 
     search_strategy: SkipValidation[Search[CandidateT]]
     """The search strategy to use for suggesting new trials."""
@@ -575,7 +577,13 @@ class Study(Model, t.Generic[CandidateT, OutputT]):
             log_outputs(to=log_to, **outputs)
 
         with (
-            task_and_run(name=self.name, tags=self.tags, inputs=trace_inputs, params=trace_params),
+            task_and_run(
+                name=self.name,
+                tags=self.tags,
+                inputs=trace_inputs,
+                params=trace_params,
+                label=self.label,
+            ),
             contextlib.ExitStack() as stack,
         ):
             stack.callback(log_study, last_event)

{dreadnode-1.15.1 → dreadnode-1.15.2}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "dreadnode"
-version = "1.15.1"
+version = "1.15.2"
 description = "Dreadnode SDK"
 authors = [{ name = "Nick Landers", email = "monoxgas@gmail.com" }]
 readme = "README.md"