| 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} |
|
|
| |
| kwargs.update(cls._set_up_prompts(flow_config)) |
|
|
| |
| kwargs.update(cls._set_up_backend(flow_config)) |
|
|
| |
| 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) |
| api_output = super().run(input_data)["api_output"].strip() |
| try: |
| response = json.loads(api_output) |
| return response |
| except json.decoder.JSONDecodeError: |
| new_input_data = input_data.copy() |
| new_input_data['goal'] += ("The previous respond cannot be parsed with json.loads, it could be the " |
| "backslashes used for escaping single quotes in the string arguments of the " |
| "Python code are not properly escaped themselves within the JSON context. " |
| "Make sure your next response is in JSON format.") |
| new_api_output = super().run(new_input_data)["api_output"].strip() |
| return json.loads(new_api_output) |
|
|
