File size: 3,383 Bytes
d5506a4 49a0727 d5506a4 49a0727 d5506a4 49a0727 d5506a4 49a0727 d5506a4 49a0727 6c1b6e4 f36c42b |
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 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 |
import json
from copy import deepcopy
from typing import Any, Dict
from flow_modules.aiflows.ChatFlowModule import ChatAtomicFlow
class CodeGeneratorAtomicFlow(ChatAtomicFlow):
"""Generates one function with docstrings to finish the given goal (from the controller).
"""
def __init__(self, **kwargs):
super().__init__(**kwargs)
self.system_message_prompt_template = self.system_message_prompt_template.partial(
code_library_file_location="no location yet",
code_library="no code yet"
)
self.hint_for_model = """
Make sure your response is in the following format:
Response Format:
{
"language_of_code": "language of the code",
"code": "String of the code and docstrings corresponding to the goal",
}
"""
@classmethod
def instantiate_from_config(cls, config):
flow_config = deepcopy(config)
kwargs = {"flow_config": flow_config}
# ~~~ Set up prompts ~~~
kwargs.update(cls._set_up_prompts(flow_config))
# ~~~ Set up backend ~~~
kwargs.update(cls._set_up_backend(flow_config))
# ~~~ Instantiate flow ~~~
return cls(**kwargs)
def _get_code_library_file(self, input_data: Dict[str, Any]):
assert "memory_files" in input_data, "memory_files not passed to CodeGeneratorAtomicFlow"
assert "code_library" in input_data['memory_files'], "code_library not in memory_files"
code_library_file_location = input_data['memory_files']['code_library']
return code_library_file_location
def _get_code_library_content(self, input_data: Dict[str, Any]):
assert "code_library" in input_data, "code_library not passed to CodeGeneratorAtomicFlow"
code_library = input_data['code_library']
if len(code_library) == 0:
code_library = "No code yet"
return code_library
def _update_prompts_and_input(self, input_data: Dict[str, Any]):
if 'goal' in input_data:
input_data['goal'] += self.hint_for_model
code_library_file_location = self._get_code_library_file(input_data)
code_library = self._get_code_library_content(input_data)
self.system_message_prompt_template = self.system_message_prompt_template.partial(
code_library_file_location=code_library_file_location,
code_library=code_library
)
def run(self, input_data: Dict[str, Any]) -> Dict[str, Any]:
self._update_prompts_and_input(input_data)
while True:
api_output = super().run(input_data)["api_output"].strip()
try:
response = json.loads(api_output)
return response
except (json.decoder.JSONDecodeError, json.JSONDecodeError):
new_goal = "The previous response cannot be parsed with json.loads, it cannot be parsed with json.loads, it could be the backslashes usesd for escaping single quotes in the string arguments of the code are not properly escaped themselves within the JSON context. Next time, do not provide any comments or code blocks. Make sure your next response is purely json parsable."
new_input_data = input_data.copy()
new_input_data['goal'] = new_goal
input_data = new_input_data
|