Skip to content

Commit 1570059

Browse files
KillianLucasKillianLucas
committed
More GeneratorExit Protection
1 parent 6bdd6be commit 1570059

File tree

1 file changed

+81
-78
lines changed

1 file changed

+81
-78
lines changed

‎interpreter/core/core.py

Lines changed: 81 additions & 78 deletions
Original file line numberDiff line numberDiff line change
@@ -307,86 +307,89 @@ def is_active_line_chunk(chunk):
307307

308308
last_flag_base = None
309309

310-
for chunk in respond(self):
311-
# For async usage
312-
if hasattr(self, "stop_event") and self.stop_event.is_set():
313-
break
314-
315-
if chunk["content"] == "":
316-
continue
317-
318-
# Handle the special "confirmation" chunk, which neither triggers a flag or creates a message
319-
if chunk["type"] == "confirmation":
320-
# Emit a end flag for the last message type, and reset last_flag_base
321-
if last_flag_base:
322-
yield {**last_flag_base, "end": True}
323-
last_flag_base = None
324-
325-
if self.auto_run == False:
326-
yield chunk
327-
328-
# We want to append this now, so even if content is never filled, we know that the execution didn't produce output.
329-
# ... rethink this though.
330-
self.messages.append(
331-
{
332-
"role": "computer",
333-
"type": "console",
334-
"format": "output",
335-
"content": "",
336-
}
337-
)
338-
continue
339-
340-
# Check if the chunk's role, type, and format (if present) match the last_flag_base
341-
if (
342-
last_flag_base
343-
and "role" in chunk
344-
and "type" in chunk
345-
and last_flag_base["role"] == chunk["role"]
346-
and last_flag_base["type"] == chunk["type"]
347-
and (
348-
"format" not in last_flag_base
349-
or (
350-
"format" in chunk
351-
and chunk["format"] == last_flag_base["format"]
310+
try:
311+
for chunk in respond(self):
312+
# For async usage
313+
if hasattr(self, "stop_event") and self.stop_event.is_set():
314+
break
315+
316+
if chunk["content"] == "":
317+
continue
318+
319+
# Handle the special "confirmation" chunk, which neither triggers a flag or creates a message
320+
if chunk["type"] == "confirmation":
321+
# Emit a end flag for the last message type, and reset last_flag_base
322+
if last_flag_base:
323+
yield {**last_flag_base, "end": True}
324+
last_flag_base = None
325+
326+
if self.auto_run == False:
327+
yield chunk
328+
329+
# We want to append this now, so even if content is never filled, we know that the execution didn't produce output.
330+
# ... rethink this though.
331+
self.messages.append(
332+
{
333+
"role": "computer",
334+
"type": "console",
335+
"format": "output",
336+
"content": "",
337+
}
338+
)
339+
continue
340+
341+
# Check if the chunk's role, type, and format (if present) match the last_flag_base
342+
if (
343+
last_flag_base
344+
and "role" in chunk
345+
and "type" in chunk
346+
and last_flag_base["role"] == chunk["role"]
347+
and last_flag_base["type"] == chunk["type"]
348+
and (
349+
"format" not in last_flag_base
350+
or (
351+
"format" in chunk
352+
and chunk["format"] == last_flag_base["format"]
353+
)
354+
)
355+
):
356+
# If they match, append the chunk's content to the current message's content
357+
# (Except active_line, which shouldn't be stored)
358+
if not is_active_line_chunk(chunk):
359+
self.messages[-1]["content"] += chunk["content"]
360+
else:
361+
# If they don't match, yield a end message for the last message type and a start message for the new one
362+
if last_flag_base:
363+
yield {**last_flag_base, "end": True}
364+
365+
last_flag_base = {"role": chunk["role"], "type": chunk["type"]}
366+
367+
# Don't add format to type: "console" flags, to accommodate active_line AND output formats
368+
if "format" in chunk and chunk["type"] != "console":
369+
last_flag_base["format"] = chunk["format"]
370+
371+
yield {**last_flag_base, "start": True}
372+
373+
# Add the chunk as a new message
374+
if not is_active_line_chunk(chunk):
375+
self.messages.append(chunk)
376+
377+
# Yield the chunk itself
378+
yield chunk
379+
380+
# Truncate output if it's console output
381+
if chunk["type"] == "console" and chunk["format"] == "output":
382+
self.messages[-1]["content"] = truncate_output(
383+
self.messages[-1]["content"],
384+
self.max_output,
385+
add_scrollbars=self.computer.import_computer_api, # I consider scrollbars to be a computer API thing
352386
)
353-
)
354-
):
355-
# If they match, append the chunk's content to the current message's content
356-
# (Except active_line, which shouldn't be stored)
357-
if not is_active_line_chunk(chunk):
358-
self.messages[-1]["content"] += chunk["content"]
359-
else:
360-
# If they don't match, yield a end message for the last message type and a start message for the new one
361-
if last_flag_base:
362-
yield {**last_flag_base, "end": True}
363-
364-
last_flag_base = {"role": chunk["role"], "type": chunk["type"]}
365-
366-
# Don't add format to type: "console" flags, to accommodate active_line AND output formats
367-
if "format" in chunk and chunk["type"] != "console":
368-
last_flag_base["format"] = chunk["format"]
369-
370-
yield {**last_flag_base, "start": True}
371-
372-
# Add the chunk as a new message
373-
if not is_active_line_chunk(chunk):
374-
self.messages.append(chunk)
375-
376-
# Yield the chunk itself
377-
yield chunk
378-
379-
# Truncate output if it's console output
380-
if chunk["type"] == "console" and chunk["format"] == "output":
381-
self.messages[-1]["content"] = truncate_output(
382-
self.messages[-1]["content"],
383-
self.max_output,
384-
add_scrollbars=self.computer.import_computer_api, # I consider scrollbars to be a computer API thing
385-
)
386387

387-
# Yield a final end flag
388-
if last_flag_base:
389-
yield {**last_flag_base, "end": True}
388+
# Yield a final end flag
389+
if last_flag_base:
390+
yield {**last_flag_base, "end": True}
391+
except GeneratorExit:
392+
pass # It's fine
390393

391394
def reset(self):
392395
self.computer.terminate() # Terminates all languages

0 commit comments

Comments
 (0)