fix: comment out done_gen and done_step tokens (#504)

This commit is contained in:
cthomas
2025-01-02 14:12:24 -08:00
committed by GitHub
parent 792e94c12a
commit ef271ca49c
4 changed files with 21 additions and 21 deletions

View File

@@ -30,8 +30,8 @@ class JobStatus(str, Enum):
class MessageStreamStatus(str, Enum):
done_generation = "[DONE_GEN]"
done_step = "[DONE_STEP]"
# done_generation = "[DONE_GEN]"
# done_step = "[DONE_STEP]"
done = "[DONE]"

View File

@@ -292,8 +292,8 @@ class StreamingServerInterface(AgentChunkStreamingInterface):
# if multi_step = True, the stream ends when the agent yields
# if multi_step = False, the stream ends when the step ends
self.multi_step = multi_step
self.multi_step_indicator = MessageStreamStatus.done_step
self.multi_step_gen_indicator = MessageStreamStatus.done_generation
# self.multi_step_indicator = MessageStreamStatus.done_step
# self.multi_step_gen_indicator = MessageStreamStatus.done_generation
# Support for AssistantMessage
self.use_assistant_message = False # TODO: Remove this
@@ -378,8 +378,8 @@ class StreamingServerInterface(AgentChunkStreamingInterface):
"""Clean up the stream by deactivating and clearing chunks."""
self.streaming_chat_completion_mode_function_name = None
if not self.streaming_chat_completion_mode and not self.nonstreaming_legacy_mode:
self._push_to_buffer(self.multi_step_gen_indicator)
# if not self.streaming_chat_completion_mode and not self.nonstreaming_legacy_mode:
# self._push_to_buffer(self.multi_step_gen_indicator)
# Wipe the inner thoughts buffers
self._reset_inner_thoughts_json_reader()
@@ -390,9 +390,9 @@ class StreamingServerInterface(AgentChunkStreamingInterface):
# end the stream
self._active = False
self._event.set() # Unblock the generator if it's waiting to allow it to complete
elif not self.streaming_chat_completion_mode and not self.nonstreaming_legacy_mode:
# signal that a new step has started in the stream
self._push_to_buffer(self.multi_step_indicator)
# elif not self.streaming_chat_completion_mode and not self.nonstreaming_legacy_mode:
# # signal that a new step has started in the stream
# self._push_to_buffer(self.multi_step_indicator)
# Wipe the inner thoughts buffers
self._reset_inner_thoughts_json_reader()

View File

@@ -278,9 +278,9 @@ class ToolExecutionSandbox:
sbx = self.get_running_e2b_sandbox_with_same_state(sbx_config)
if not sbx or self.force_recreate:
if not sbx:
logger.info(f"No running e2b sandbox found with the same state: {sbx_config}")
logger.info(f"No running e2b sandbox found with the same state: {sbx_config}")
else:
logger.info(f"Force recreated e2b sandbox with state: {sbx_config}")
logger.info(f"Force recreated e2b sandbox with state: {sbx_config}")
sbx = self.create_e2b_sandbox_with_metadata_hash(sandbox_config=sbx_config)
# Since this sandbox was used, we extend its lifecycle by the timeout

View File

@@ -249,8 +249,8 @@ def test_streaming_send_message(mock_e2b_api_key_none, client: RESTClient, agent
send_message_ran = False
# 3. Check that we get all the start/stop/end tokens we want
# This includes all of the MessageStreamStatus enums
done_gen = False
done_step = False
# done_gen = False
# done_step = False
done = False
# print(response)
@@ -266,12 +266,12 @@ def test_streaming_send_message(mock_e2b_api_key_none, client: RESTClient, agent
if chunk == MessageStreamStatus.done:
assert not done, "Message stream already done"
done = True
elif chunk == MessageStreamStatus.done_step:
assert not done_step, "Message stream already done step"
done_step = True
elif chunk == MessageStreamStatus.done_generation:
assert not done_gen, "Message stream already done generation"
done_gen = True
# elif chunk == MessageStreamStatus.done_step:
# assert not done_step, "Message stream already done step"
# done_step = True
# elif chunk == MessageStreamStatus.done_generation:
# assert not done_gen, "Message stream already done generation"
# done_gen = True
if isinstance(chunk, LettaUsageStatistics):
# Some rough metrics for a reasonable usage pattern
assert chunk.step_count == 1
@@ -284,8 +284,8 @@ def test_streaming_send_message(mock_e2b_api_key_none, client: RESTClient, agent
assert inner_thoughts_exist, "No inner thoughts found"
assert send_message_ran, "send_message function call not found"
assert done, "Message stream not done"
assert done_step, "Message stream not done step"
assert done_gen, "Message stream not done generation"
# assert done_step, "Message stream not done step"
# assert done_gen, "Message stream not done generation"
def test_humans_personas(client: Union[LocalClient, RESTClient], agent: AgentState):