Update payload

Author: BenConstable9

text_2_sql/autogen/src/autogen_text_2_sql/autogen_text_2_sql.py

@@ -40,16 +40,16 @@ def get_all_agents(self):
         # Get current datetime for the Query Rewrite Agent
         current_datetime = datetime.now()
 
-        self.user_input_rewrite_agent = LLMAgentCreator.create(
-            "user_input_rewrite_agent", current_datetime=current_datetime
+        self.message_rewrite_agent = LLMAgentCreator.create(
+            "message_rewrite_agent", current_datetime=current_datetime
         )
 
         self.parallel_query_solving_agent = ParallelQuerySolvingAgent(**self.kwargs)
 
         self.answer_agent = LLMAgentCreator.create("answer_agent")
 
         agents = [
-            self.user_input_rewrite_agent,
+            self.message_rewrite_agent,
             self.parallel_query_solving_agent,
             self.answer_agent,
         ]
@@ -73,11 +73,11 @@ def unified_selector(self, messages):
         current_agent = messages[-1].source if messages else "user"
         decision = None
 
-        # If this is the first message start with user_input_rewrite_agent
+        # If this is the first message start with message_rewrite_agent
         if current_agent == "user":
-            decision = "user_input_rewrite_agent"
+            decision = "message_rewrite_agent"
         # Handle transition after query rewriting
-        elif current_agent == "user_input_rewrite_agent":
+        elif current_agent == "message_rewrite_agent":
             decision = "parallel_query_solving_agent"
         # Handle transition after parallel query solving
         elif current_agent == "parallel_query_solving_agent":
@@ -145,24 +145,24 @@ def extract_answer_payload(self, messages: list) -> AnswerWithSourcesPayload:
         except json.JSONDecodeError:
             logging.warning("Unable to read SQL query results: %s", sql_query_results)
             sql_query_results = {}
-            sub_user_input_results = {}
+            sub_message_results = {}
         else:
-            # Only load sub-user_input results if we have a database result
-            sub_user_input_results = self.parse_message_content(messages[1].content)
-            logging.info("Sub-user_input Results: %s", sub_user_input_results)
+            # Only load sub-message results if we have a database result
+            sub_message_results = self.parse_message_content(messages[1].content)
+            logging.info("Sub-message Results: %s", sub_message_results)
 
         try:
-            sub_user_inputs = [
-                sub_user_input
-                for sub_user_input_group in sub_user_input_results.get(
-                    "sub_user_inputs", []
+            decomposed_messages = [
+                sub_message
+                for sub_message_group in sub_message_results.get(
+                    "decomposed_messages", []
                 )
-                for sub_user_input in sub_user_input_group
+                for sub_message in sub_message_group
             ]
 
             logging.info("SQL Query Results: %s", sql_query_results)
             payload = AnswerWithSourcesPayload(
-                answer=answer, sub_user_inputs=sub_user_inputs
+                answer=answer, decomposed_messages=decomposed_messages
             )
 
             if not isinstance(sql_query_results, dict):
@@ -173,11 +173,9 @@ def extract_answer_payload(self, messages: list) -> AnswerWithSourcesPayload:
                 logging.error("No 'results' key in sql_query_results")
                 return payload
 
-            for user_input, sql_query_result_list in sql_query_results[
-                "results"
-            ].items():
+            for message, sql_query_result_list in sql_query_results["results"].items():
                 if not sql_query_result_list:  # Check if list is empty
-                    logging.warning(f"No results for user_input: {user_input}")
+                    logging.warning(f"No results for message: {message}")
                     continue
 
                 for sql_query_result in sql_query_result_list:
@@ -213,47 +211,47 @@ def extract_answer_payload(self, messages: list) -> AnswerWithSourcesPayload:
                 answer=f"{answer}\nError processing results: {str(e)}"
             )
 
-    async def process_user_input(
+    async def process_message(
         self,
-        user_input_payload: UserInputPayload,
+        message_payload: UserInputPayload,
         chat_history: list[InteractionPayload] = None,
     ) -> AsyncGenerator[InteractionPayload, None]:
-        """Process the complete user_input through the unified system.
+        """Process the complete message through the unified system.
 
         Args:
         ----
-            task (str): The user user_input to process.
+            task (str): The user message to process.
             chat_history (list[str], optional): The chat history. Defaults to None.
             injected_parameters (dict, optional): Parameters to pass to agents. Defaults to None.
 
         Returns:
         -------
             dict: The response from the system.
         """
-        logging.info("Processing user_input: %s", user_input_payload.body.user_input)
+        logging.info("Processing message: %s", message_payload.body.message)
         logging.info("Chat history: %s", chat_history)
 
         agent_input = {
-            "user_input": user_input_payload.body.user_input,
+            "message": message_payload.body.message,
             "chat_history": {},
-            "injected_parameters": user_input_payload.body.injected_parameters,
+            "injected_parameters": message_payload.body.injected_parameters,
         }
 
         if chat_history is not None:
             # Update input
             for idx, chat in enumerate(chat_history):
-                if chat.root.payload_type == PayloadType.USER_INPUT:
+                if chat.root.payload_type == PayloadType.message:
                     # For now only consider the user query
                     chat_history_key = f"chat_{idx}"
-                    agent_input[chat_history_key] = chat.root.body.user_input
+                    agent_input[chat_history_key] = chat.root.body.message
 
         async for message in self.agentic_flow.run_stream(task=json.dumps(agent_input)):
             logging.debug("Message: %s", message)
 
             payload = None
 
             if isinstance(message, TextMessage):
-                if message.source == "user_input_rewrite_agent":
+                if message.source == "message_rewrite_agent":
                     payload = ProcessingUpdatePayload(
                         message="Rewriting the query...",
                     )
@@ -276,10 +274,10 @@ async def process_user_input(
                 elif message.messages[-1].source == "parallel_query_solving_agent":
                     # Load into disambiguation request
                     payload = self.extract_disambiguation_request(message.messages)
-                elif message.messages[-1].source == "user_input_rewrite_agent":
+                elif message.messages[-1].source == "message_rewrite_agent":
                     # Load into empty response
                     payload = AnswerWithSourcesPayload(
-                        answer="Apologies, I cannot answer that user_input as it is not relevant. Please try another user_input or rephrase your current user_input."
+                        answer="Apologies, I cannot answer that message as it is not relevant. Please try another message or rephrase your current message."
                     )
             else:
                 logging.error("Unexpected TaskResult: %s", message)

text_2_sql/autogen/src/autogen_text_2_sql/custom_agents/parallel_query_solving_agent.py

@@ -84,9 +84,9 @@ async def on_messages_stream(
             injected_parameters = {}
 
         # Load the json of the last message to populate the final output object
-        user_input_rewrites = json.loads(last_response)
+        message_rewrites = json.loads(last_response)
 
-        logging.info(f"Query Rewrites: {user_input_rewrites}")
+        logging.info(f"Query Rewrites: {message_rewrites}")
 
         async def consume_inner_messages_from_agentic_flow(
             agentic_flow, identifier, database_results
@@ -143,7 +143,7 @@ async def consume_inner_messages_from_agentic_flow(
                             ):
                                 logging.info("Contains pre-run results")
                                 for pre_run_sql_query, pre_run_result in parsed_message[
-                                    "cached_user_inputs_and_schemas"
+                                    "cached_messages_and_schemas"
                                 ].items():
                                     database_results[identifier].append(
                                         {
@@ -164,7 +164,7 @@ async def consume_inner_messages_from_agentic_flow(
 
         # Convert all_non_database_query to lowercase string and compare
         all_non_database_query = str(
-            user_input_rewrites.get("all_non_database_query", "false")
+            message_rewrites.get("all_non_database_query", "false")
         ).lower()
 
         if all_non_database_query == "true":
@@ -177,12 +177,12 @@ async def consume_inner_messages_from_agentic_flow(
             return
 
         # Start processing sub-queries
-        for user_input_rewrite in user_input_rewrites["sub_user_inputs"]:
-            logging.info(f"Processing sub-query: {user_input_rewrite}")
+        for message_rewrite in message_rewrites["decomposed_messages"]:
+            logging.info(f"Processing sub-query: {message_rewrite}")
             # Create an instance of the InnerAutoGenText2Sql class
             inner_autogen_text_2_sql = InnerAutoGenText2Sql(**self.kwargs)
 
-            identifier = ", ".join(user_input_rewrite)
+            identifier = ", ".join(message_rewrite)
 
             # Add database connection info to injected parameters
             query_params = injected_parameters.copy() if injected_parameters else {}
@@ -196,8 +196,8 @@ async def consume_inner_messages_from_agentic_flow(
             # Launch tasks for each sub-query
             inner_solving_generators.append(
                 consume_inner_messages_from_agentic_flow(
-                    inner_autogen_text_2_sql.process_user_input(
-                        user_input=user_input_rewrite,
+                    inner_autogen_text_2_sql.process_message(
+                        message=message_rewrite,
                         injected_parameters=query_params,
                     ),
                     identifier,

text_2_sql/autogen/src/autogen_text_2_sql/custom_agents/sql_query_cache_agent.py

@@ -17,7 +17,7 @@ class SqlQueryCacheAgent(BaseChatAgent):
     def __init__(self):
         super().__init__(
             "sql_query_cache_agent",
-            "An agent that fetches the queries from the cache based on the user user_input.",
+            "An agent that fetches the queries from the cache based on the user message.",
         )
 
         self.agent = SqlQueryCacheAgentCustomAgent()
@@ -40,19 +40,19 @@ async def on_messages(
     async def on_messages_stream(
         self, messages: Sequence[ChatMessage], cancellation_token: CancellationToken
     ) -> AsyncGenerator[AgentMessage | Response, None]:
-        # Get the decomposed user_inputs from the user_input_rewrite_agent
+        # Get the decomposed messages from the message_rewrite_agent
         try:
             request_details = json.loads(messages[0].content)
             injected_parameters = request_details["injected_parameters"]
-            user_user_inputs = request_details["user_input"]
-            logging.info(f"Processing user_inputs: {user_user_inputs}")
+            user_messages = request_details["message"]
+            logging.info(f"Processing messages: {user_messages}")
             logging.info(f"Input Parameters: {injected_parameters}")
         except json.JSONDecodeError:
-            # If not JSON array, process as single user_input
+            # If not JSON array, process as single message
             raise ValueError("Could not load message")
 
         cached_results = await self.agent.process_message(
-            user_user_inputs, injected_parameters
+            user_messages, injected_parameters
         )
         yield Response(
             chat_message=TextMessage(

text_2_sql/autogen/src/autogen_text_2_sql/custom_agents/sql_schema_selection_agent.py

@@ -43,19 +43,19 @@ async def on_messages_stream(
         # Try to parse as JSON first
         try:
             request_details = json.loads(messages[0].content)
-            user_inputs = request_details["question"]
+            messages = request_details["question"]
         except (json.JSONDecodeError, KeyError):
             # If not JSON or missing question key, use content directly
-            user_inputs = messages[0].content
+            messages = messages[0].content
 
-        if isinstance(user_inputs, str):
-            user_inputs = [user_inputs]
-        elif not isinstance(user_inputs, list):
-            user_inputs = [str(user_inputs)]
+        if isinstance(messages, str):
+            messages = [messages]
+        elif not isinstance(messages, list):
+            messages = [str(messages)]
 
-        logging.info(f"Processing questions: {user_inputs}")
+        logging.info(f"Processing questions: {messages}")
 
-        final_results = await self.agent.process_message(user_inputs)
+        final_results = await self.agent.process_message(messages)
 
         yield Response(
             chat_message=TextMessage(

text_2_sql/text_2_sql_core/src/text_2_sql_core/custom_agents/sql_query_cache_agent.py

@@ -9,7 +9,7 @@ def __init__(self):
         self.sql_connector = ConnectorFactory.get_database_connector()
 
     async def process_message(
-        self, user_inputs: list[str], injected_parameters: dict
+        self, messages: list[str], injected_parameters: dict
     ) -> dict:
         # Initialize results dictionary
         cached_results = {
@@ -18,7 +18,7 @@ async def process_message(
         }
 
         # Process each question sequentially
-        for question in user_inputs:
+        for question in messages:
             # Fetch the queries from the cache based on the question
             logging.info(f"Fetching queries from cache for question: {question}")
             cached_query = await self.sql_connector.fetch_queries_from_cache(

text_2_sql/text_2_sql_core/src/text_2_sql_core/custom_agents/sql_schema_selection_agent.py

@@ -22,15 +22,15 @@ def __init__(self, **kwargs):
 
         self.system_prompt = Template(system_prompt).render(kwargs)
 
-    async def process_message(self, user_inputs: list[str]) -> dict:
-        logging.info(f"user inputs: {user_inputs}")
+    async def process_message(self, messages: list[str]) -> dict:
+        logging.info(f"user inputs: {messages}")
 
         entity_tasks = []
 
-        for user_input in user_inputs:
+        for message in messages:
             messages = [
                 {"role": "system", "content": self.system_prompt},
-                {"role": "user", "content": user_input},
+                {"role": "user", "content": message},
             ]
             entity_tasks.append(
                 self.open_ai_connector.run_completion_request(