From 8dcd7d082d1c9ad7703fb7412cbf15e976b87c0c Mon Sep 17 00:00:00 2001
From: Vincent Abruzzo <vincent@dreadnode.io>
Date: Wed, 17 Dec 2025 17:11:33 -0500
Subject: [PATCH] completion task structure

---
 dreadnode/agent/agent.py | 9 +++++++++
 1 file changed, 9 insertions(+)

diff --git a/dreadnode/agent/agent.py b/dreadnode/agent/agent.py
index a680b877..80bf9f56 100644
--- a/dreadnode/agent/agent.py
+++ b/dreadnode/agent/agent.py
@@ -299,6 +299,8 @@ async def _generate(
             post_transforms.append(post_transform)
 
         try:
+            from dreadnode import log_output, task_and_run
+
             messages = rg.caching.apply_cache_mode_to_messages(self.caching, [messages])[0]
 
             logger.trace(f"Generating with model '{self.generator.model}'. Messages: {messages!r}")
@@ -307,6 +309,13 @@ async def _generate(
             if isinstance(generated, BaseException):
                 raise generated  # noqa: TRY301
 
+            with task_and_run(
+                name="Chat completion",
+                tags=["completion"],
+                label="Chat completion",
+            ):
+                log_output("messages", [messages[-1], generated.message])
+
             chat = rg.Chat(
                 messages,
                 [generated.message],