reorganize code

ChatAtomicFlow.py

@@ -357,15 +357,18 @@ class ChatAtomicFlow(AtomicFlow):
 
         self._state_update_add_chat_message(role=self.flow_config["user_name"],
                                             content=user_message_content)
-
-    def run(self,input_message: FlowMessage):
-        """ This method runs the flow. It processes the input, calls the backend and updates the state of the flow.
         
-        :param input_message: The input data of the flow.
-        :type input_message: aiflows.messages.FlowMessage
+    def query_llm(self,input_data: Dict[str, Any]):
+        """ This method queries the LLM. It processes the input, calls the backend and returns the response.
+        
+        :param input_data: The input data of the flow.
+        :type input_data: Dict[str, Any]
+        :return: The response of the LLM to the input.
+        :rtype: Union[str, List[str]]
         """
-        input_data = input_message.data
-        # ~~~ Process input ~~~
+        
+        
+         # ~~~ Process input ~~~
         self._process_input(input_data)
 
         # ~~~ Call ~~~
@@ -378,6 +381,18 @@ class ChatAtomicFlow(AtomicFlow):
                 content=answer
             )
         response = response if len(response) > 1 or len(response) == 0 else response[0]
+        return response
+
+    def run(self,input_message: FlowMessage):
+        """ This method runs the flow. It processes the input, calls the backend and updates the state of the flow.
+        
+        :param input_message: The input data of the flow.
+        :type input_message: aiflows.messages.FlowMessage
+        """
+        input_data = input_message.data
+        
+        response = self.query_llm(input_data=input_data)
+       
         
         reply_message = self._package_output_message(input_message, response = {"api_output": response})