File size: 2,646 Bytes
5a600b9 2c06037 5a600b9 2c06037 509b1d2 2c06037 5a600b9 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 |
from typing import Dict, Any
import os
from flow_modules.Tachi67.ContentWriterFlowModule import ContentWriterFlow
from flows.base_flows import CircularFlow
class CodeWriterFlow(ContentWriterFlow):
def _on_reach_max_round(self):
self._state_update_dict({
"code": "The maximum amount of rounds was reached before the model generated the code.",
"status": "unfinished"
})
@CircularFlow.output_msg_payload_processor
def detect_finish_or_continue(self, output_payload: Dict[str, Any], src_flow) -> Dict[str, Any]:
command = output_payload["command"]
if command == "finish":
# ~~~ delete the temp code file ~~~
keys_to_fetch_from_state = ["temp_code_file_location", "code"]
fetched_state = self._fetch_state_attributes_by_keys(keys=keys_to_fetch_from_state)
temp_code_file_location = fetched_state["temp_code_file_location"]
code_content = fetched_state["code"]
if os.path.exists(temp_code_file_location):
os.remove(temp_code_file_location)
# ~~~ return the code content ~~~
return {
"EARLY_EXIT": True,
"code": code_content,
"summary": "CodeWriter: " + output_payload["command_args"]["summary"],
"status": "finished"
}
elif command == "manual_finish":
# ~~~ delete the temp code file ~~~
keys_to_fetch_from_state = ["temp_code_file_location"]
fetched_state = self._fetch_state_attributes_by_keys(keys=keys_to_fetch_from_state)
temp_code_file_location = fetched_state["temp_code_file_location"]
if os.path.exists(temp_code_file_location):
os.remove(temp_code_file_location)
# ~~~ return the manual quit status ~~~
return {
"EARLY_EXIT": True,
"code": "no code was generated",
"summary": "CodeWriter: CodeWriter was terminated explicitly by the user, process is unfinished",
"status": "unfinished"
}
elif command == "test":
# ~~~ fetch code string from flow state ~~~
keys_to_fetch_from_state = ["code"]
fetched_state = self._fetch_state_attributes_by_keys(keys=keys_to_fetch_from_state)
# ~~~ add code content to the command args (branch input data) ~~~
code_content = fetched_state["code"]
output_payload["command_args"]["code"] = code_content
return output_payload
else:
return output_payload |