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],