From 2bdeeccc4167b9a0f593e66da00e7a34b21c4d0d Mon Sep 17 00:00:00 2001
From: cthomas <caren@letta.com>
Date: Fri, 23 May 2025 14:38:25 -0700
Subject: [PATCH] feat: add step stream to async sleeptime agent (#2379)

---
 letta/groups/sleeptime_multi_agent_v2.py | 17 +++++++++++++++++
 1 file changed, 17 insertions(+)

diff --git a/letta/groups/sleeptime_multi_agent_v2.py b/letta/groups/sleeptime_multi_agent_v2.py
index f082ca38..4e172bee 100644
--- a/letta/groups/sleeptime_multi_agent_v2.py
+++ b/letta/groups/sleeptime_multi_agent_v2.py
@@ -21,6 +21,7 @@ from letta.services.message_manager import MessageManager
 from letta.services.passage_manager import PassageManager
 from letta.services.step_manager import NoopStepManager, StepManager
 from letta.services.telemetry_manager import NoopTelemetryManager, TelemetryManager
+from letta.tracing import trace_method
 
 
 class SleeptimeMultiAgentV2(BaseAgent):
@@ -55,6 +56,7 @@ class SleeptimeMultiAgentV2(BaseAgent):
         assert group.manager_type == ManagerType.sleeptime, f"Expected group manager type to be 'sleeptime', got {group.manager_type}"
         self.group = group
 
+    @trace_method
     async def step(
         self,
         input_messages: List[MessageCreate],
@@ -119,6 +121,21 @@ class SleeptimeMultiAgentV2(BaseAgent):
         response.usage.run_ids = run_ids
         return response
 
+    @trace_method
+    async def step_stream_no_tokens(
+        self,
+        input_messages: List[MessageCreate],
+        max_steps: int = 10,
+        use_assistant_message: bool = True,
+    ):
+        response = await self.step(input_messages, max_steps, use_assistant_message)
+
+        for message in response.messages:
+            yield f"data: {message.model_dump_json()}\n\n"
+
+        yield f"data: {response.usage.model_dump_json()}\n\n"
+
+    @trace_method
     async def step_stream(
         self,
         input_messages: List[MessageCreate],