Version to 1.15.0. Small logging/formatting improvements.

Author: monoxgas

dreadnode/agent/agent.py

@@ -4,6 +4,7 @@
 from copy import deepcopy
 
 import rigging as rg
+from loguru import logger
 from pydantic import ConfigDict, Field, PrivateAttr, SkipValidation, field_validator
 from rigging.message import inject_system_content
 from ulid import ULID  # can't access via rg
@@ -237,6 +238,8 @@ async def _generate(
         try:
             messages = rg.caching.apply_cache_mode_to_messages(self.caching, [messages])[0]
 
+            logger.trace(f"Generating with model '{generator.model}'. Messages: {messages!r}")
+
             generated = (await generator.generate_messages([messages], [params]))[0]
             if isinstance(generated, BaseException):
                 raise generated  # noqa: TRY301
@@ -252,6 +255,7 @@ async def _generate(
             )
 
         except Exception as error:  # noqa: BLE001
+            logger.error(f"Error during generation: {error}", exc_info=True)
             chat = rg.Chat(
                 messages,
                 [],
@@ -278,6 +282,13 @@ async def _stream(  # noqa: PLR0912, PLR0915
         stop_conditions = self.stop_conditions
         session_id = ULID()
 
+        logger.info(
+            f"Starting Agent '{self.name}' ({session_id}): "
+            f"model='{self.model}', "
+            f"max_steps={self.max_steps}, "
+            f"tools={[tool.name for tool in self.all_tools]}"
+        )
+
         # Event dispatcher
 
         async def _dispatch(event: AgentEvent) -> t.AsyncIterator[AgentEvent]:
@@ -292,6 +303,11 @@ async def _dispatch(event: AgentEvent) -> t.AsyncIterator[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]}"
+            )
+
             # Run all applicable hooks and collect their reactions
             hook_reactions: dict[str, Reaction | None] = {}
             for hook in applicable_hooks:
@@ -310,6 +326,10 @@ async def _dispatch(event: AgentEvent) -> t.AsyncIterator[AgentEvent]:
                 if reaction is None:
                     continue
 
+                logger.debug(
+                    f"Agent '{self.name}' ({session_id}) hook '{hook_name}' returned reaction: {reaction!r}"
+                )
+
                 if not isinstance(reaction, Reaction):
                     warn_at_user_stacklevel(
                         f"Hook '{hook_name}' returned {reaction}, but expected a Reaction.",
@@ -386,6 +406,9 @@ async def _dispatch(event: AgentEvent) -> t.AsyncIterator[AgentEvent]:
                 return
 
             if isinstance(winning_reaction, RetryWithFeedback):
+                logger.debug(
+                    f"Agent '{self.name}' ({session_id}) injecting feedback for retry: '{winning_reaction.feedback}'"
+                )
                 messages.append(rg.Message("user", winning_reaction.feedback))
                 raise Retry(messages=messages) from winning_reaction
 
@@ -408,6 +431,10 @@ async def _process_tool_call(
             ):
                 yield event
 
+            logger.debug(
+                f"Executing tool '{tool_call.name}' with args: {tool_call.function.arguments}"
+            )
+
             message: rg.Message
             stop = False
             tool = next((t for t in self.all_tools if t.name == tool_call.name), None)
@@ -431,6 +458,7 @@ async def _process_tool_call(
                         yield event
                     raise
             else:
+                logger.warning(f"Tool '{tool_call.name}' not found.")
                 message = rg.Message.from_model(
                     rg.model.SystemErrorModel(content=f"Tool '{tool_call.name}' not found.")
                 )
@@ -516,6 +544,7 @@ async def _process_tool_call(
                 # Check for stop conditions
 
                 if any(cond(events) for cond in stop_conditions):
+                    logger.info("A stop condition was met. Ending run.")
                     break
 
                 # Check if stalled
@@ -524,6 +553,10 @@ async def _process_tool_call(
                     if not stop_conditions:
                         break
 
+                    logger.warning(
+                        f"Agent '{self.name}' ({session_id}) stalled: No tool calls and no stop conditions met."
+                    )
+
                     async for event in _dispatch(
                         AgentStalled(
                             session_id=session_id,
@@ -589,6 +622,23 @@ async def _process_tool_call(
             thread.messages = messages
             thread.events.extend(events)
 
+        total_usage = _total_usage_from_events(events)
+        log_message = (
+            f"Agent '{self.name}' finished: "
+            f"reason='{stop_reason}', "
+            f"steps={step - 1}, "
+            f"total_tokens={total_usage.total_tokens}, "
+            f"in_tokens={total_usage.input_tokens}, "
+            f"out_tokens={total_usage.output_tokens}"
+        )
+
+        if stop_reason == "finished":
+            logger.success(log_message)
+        elif stop_reason == "error":
+            logger.error(f"{log_message}, error='{error!r}'")
+        else:
+            logger.warning(log_message)
+
         yield AgentEnd(
             session_id=session_id,
             agent=self,

dreadnode/cli/agent/cli.py

@@ -1,4 +1,5 @@
 import contextlib
+import inspect
 import itertools
 import typing as t
 from inspect import isawaitable
@@ -55,7 +56,8 @@ async def run(  # noqa: PLR0912, PLR0915
     agent: str,
     *tokens: t.Annotated[str, cyclopts.Parameter(show=False, allow_leading_hyphen=True)],
     config: Path | None = None,
-    dreadnode_config: DreadnodeConfig | None = None,
+    raw: t.Annotated[bool, cyclopts.Parameter(["-r", "--raw"], negative=False)] = False,
+    dn_config: DreadnodeConfig | None = None,
 ) -> None:
     """
     Run an agent by name, file, or module.
@@ -70,6 +72,7 @@ async def run(  # noqa: PLR0912, PLR0915
     Args:
         agent: The agent to run, e.g., 'my_agents.py:basic' or 'basic'.
         config: Optional path to a TOML/YAML/JSON configuration file for the agent.
+        raw: If set, only display raw logging output without additional formatting.
     """
     from dreadnode.agent import Agent
 
@@ -107,38 +110,43 @@ async def run(  # noqa: PLR0912, PLR0915
     agent_blueprint = agents_by_lower_name[agent_name.lower()]
 
     config_model = get_config_model(agent_blueprint)
-    config_parameter = cyclopts.Parameter(name="*", group="Agent Config")(config_model)
-
-    config_default = None
+    config_annotation = cyclopts.Parameter(name="*", group="Agent Config")(config_model)
+    config_default: t.Any = inspect.Parameter.empty
     with contextlib.suppress(Exception):
         config_default = config_model()
-        config_parameter = config_parameter | None  # type: ignore [assignment]
 
     async def agent_cli(
         input: t.Annotated[str, cyclopts.Parameter(help="Input to the agent")],
         *,
         config: t.Any = config_default,
-        dreadnode_config: DreadnodeConfig | None = dreadnode_config,
+        dn_config: DreadnodeConfig | None = dn_config,
     ) -> None:
-        (dreadnode_config or DreadnodeConfig()).apply()
-
+        dn_config = dn_config or DreadnodeConfig()
+        if raw and dn_config.log_level is None:
+            dn_config.log_level = "info"
+        dn_config.apply()
         agent = hydrate(agent_blueprint, config)
-        flat_config = flatten_model(config)
 
         rich.print(f"Running agent: [bold]{agent.name}[/bold] with config:")
-        for key, value in flat_config.items():
+        for key, value in flatten_model(config).items():
             rich.print(f" |- {key}: {value}")
         rich.print()
 
         async with agent.stream(input) as stream:
             async for event in stream:
                 rich.print(event)
 
-    agent_cli.__annotations__["config"] = config_parameter
+    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
 
     agent_app = cyclopts.App(
         name=agent_name,
-        help=f"Run the '{agent_name}' agent.",
+        help=help_text,
         help_on_error=True,
         help_flags=("help"),
         version_flags=(),

dreadnode/cli/task/cli.py

@@ -1,4 +1,5 @@
 import contextlib
+import inspect
 import itertools
 import typing as t
 from inspect import isawaitable
@@ -54,7 +55,7 @@ async def run(  # noqa: PLR0912, PLR0915
     task: str,
     *tokens: t.Annotated[str, cyclopts.Parameter(show=False, allow_leading_hyphen=True)],
     config: Path | None = None,
-    dreadnode_config: DreadnodeConfig | None = None,
+    dn_config: DreadnodeConfig | None = None,
 ) -> None:
     """
     Run a task by name, file, or module.
@@ -104,19 +105,17 @@ async def run(  # noqa: PLR0912, PLR0915
     task_blueprint = tasks_by_lower_name[task_name.lower()]
 
     config_model = get_config_model(task_blueprint)
-    config_parameter = cyclopts.Parameter(name="*", group="Task Config")(config_model)
-
-    config_default = None
+    config_annotation = cyclopts.Parameter(name="*", group="Task Config")(config_model)
+    config_default: t.Any = inspect.Parameter.empty
     with contextlib.suppress(Exception):
         config_default = config_model()
-        config_parameter = config_parameter | None  # type: ignore[assignment]
 
     async def task_cli(
         *,
         config: t.Any = config_default,
-        dreadnode_config: DreadnodeConfig | None = dreadnode_config,
+        dn_config: DreadnodeConfig | None = dn_config,
     ) -> None:
-        (dreadnode_config or DreadnodeConfig()).apply()
+        (dn_config or DreadnodeConfig()).apply()
 
         hydrated_task = hydrate(task_blueprint, config)
         flat_config = flatten_model(config)
@@ -128,11 +127,15 @@ async def task_cli(
 
         await hydrated_task()
 
-    task_cli.__annotations__["config"] = config_parameter
+    task_cli.__annotations__["config"] = config_annotation
+
+    help_text = f"Run the '{task_name}' task."
+    if task_blueprint.__doc__:
+        help_text += "\n\n" + task_blueprint.__doc__
 
     task_app = cyclopts.App(
         name=task_name,
-        help=task_blueprint.__doc__ or f"Run the '{task_name}' task.",
+        help=help_text,
         help_on_error=True,
         help_flags=("help"),
         version_flags=(),

dreadnode/optimization/format.py

@@ -42,6 +42,8 @@ def format_study(study: "Study") -> RenderableType:
     """
     Format a single Study object into a detailed rich Panel.
     """
+    from dreadnode.airt import Attack
+
     details = Table(
         box=box.MINIMAL,
         show_header=False,
@@ -51,7 +53,15 @@ def format_study(study: "Study") -> RenderableType:
     details.add_column("Value", style="white")
 
     details.add_row(Text("Description", justify="right"), study.description or "-")
-    details.add_row(Text("Task Factory", justify="right"), get_callable_name(study.task_factory))
+
+    # TODO(nick): Move attack formatting out of here
+    if isinstance(study, Attack):
+        details.add_row(Text("Target", justify="right"), repr(study.target))
+    else:
+        details.add_row(
+            Text("Task Factory", justify="right"), get_callable_name(study.task_factory)
+        )
+
     details.add_row(Text("Search Strategy", justify="right"), study.search_strategy.name)
 
     if study.dataset is not None:
@@ -60,10 +70,13 @@ def format_study(study: "Study") -> RenderableType:
         )
 
     if study.objectives:
-        objective_names = ", ".join(f"[cyan]{name}[/]" for name in study.objective_names)
-        details.add_row(Text("Objectives", justify="right"), objective_names)
-        directions = ", ".join(f"[yellow]{direction}[/]" for direction in study.directions)
-        details.add_row(Text("Directions", justify="right"), directions)
+        objectives = " | ".join(
+            f"[cyan]{name} :arrow_upper_right:[/]"
+            if direction == "maximize"
+            else f"[magenta]{name} :arrow_lower_right:[/]"
+            for name, direction in zip(study.objective_names, study.directions, strict=True)
+        )
+        details.add_row(Text("Objectives", justify="right"), objectives)
 
     if study.constraints:
         constraint_names = ", ".join(

dreadnode/util.py

@@ -33,15 +33,12 @@
 from loguru import logger
 
 get_user_frame_and_stacklevel = _get_user_frame_and_stacklevel
-warn_at_user_stacklevel = _warn_at_user_stacklevel
 get_filepath_attribute = _get_filepath_attribute
 is_user_code = _is_user_code
 
 
 import dreadnode  # noqa: E402
 
-warn_at_user_stacklevel = _warn_at_user_stacklevel
-
 SysExcInfo = (
     tuple[type[BaseException], BaseException, TracebackType | None] | tuple[None, None, None]
 )
@@ -742,6 +739,18 @@ def _internal_error_exc_info() -> SysExcInfo:
         return original_exc_info
 
 
+def warn_at_user_stacklevel(msg: str, category: type[Warning]) -> None:
+    """
+    Issue a warning at the user code stack level and log it.
+
+    Args:
+        msg: The warning message.
+        category: The warning category.
+    """
+    logger.warning(msg)
+    _warn_at_user_stacklevel(msg, category)
+
+
 @contextmanager
 def handle_internal_errors() -> t.Iterator[None]:
     """

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "dreadnode"
-version = "1.14.1"
+version = "1.15.0"
 description = "Dreadnode SDK"
 authors = [{ name = "Nick Landers", email = "monoxgas@gmail.com" }]
 readme = "README.md"
@@ -44,7 +44,6 @@ scoring = [
     "presidio-analyzer>=2.2.359,<3.0.0",
     "scikit-learn>=1.7.1,<2.0.0",
     "confusables>=1.2.0,<2.0.0",
-
     "nltk>=3.9.1,<4.0.0",
     "textblob>=0.19.0,<1.0.0",
     "textstat>=0.7.10,<1.0.0",
@@ -80,10 +79,6 @@ dn = "dreadnode.__main__:run"
 [tool.poetry.plugins."pipx.run"]
 dreadnode = 'dreadnode.__main__:run'
 
-
-[tool.poetry.group.dev.dependencies]
-
-
 [project.urls]
 Homepage = "https://github.com/dreadnode/sdk"
 Repository = "https://github.com/dreadnode/sdk"