Update CodeGeneratorAtomicFlow.py

CodeGeneratorAtomicFlow.py

@@ -9,6 +9,14 @@ from dataclasses import dataclass
 class CodeGeneratorAtomicFlow(OpenAIChatAtomicFlow):
     def __init__(self, **kwargs):
         super().__init__(**kwargs)
+        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 corresponding to the goal",
+              }
+        """
 
     @classmethod
     def instantiate_from_config(cls, config):
@@ -22,18 +30,13 @@ class CodeGeneratorAtomicFlow(OpenAIChatAtomicFlow):
 
         # ~~~ Instantiate flow ~~~
         return cls(**kwargs)
+    
+    def _update_input(self, input_data: Dict[str, Any]):
+        if 'goal' in input_data:
+            input_data['goal'] += self.hint_for_model
 
     def run(self, input_data: Dict[str, Any]) -> Dict[str, Any]:
-        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 corresponding to the goal",
-              }
-        """
-        if 'goal' in input_data:
-            input_data['goal'] += hint_for_model
+        self._update_input(input_data)
         api_output = super().run(input_data)["api_output"].strip()
         try:
             response = json.loads(api_output)