|
|
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) |
|
|
|
