Update user input

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.question_rewrite_agent = LLMAgentCreator.create(
-            "question_rewrite_agent", current_datetime=current_datetime
+        self.user_input_rewrite_agent = LLMAgentCreator.create(
+            "user_input_rewrite_agent", current_datetime=current_datetime
         )
 
         self.parallel_query_solving_agent = ParallelQuerySolvingAgent(**self.kwargs)
 
         self.answer_agent = LLMAgentCreator.create("answer_agent")
 
         agents = [
-            self.question_rewrite_agent,
+            self.user_input_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 question_rewrite_agent
+        # If this is the first message start with user_input_rewrite_agent
         if current_agent == "user":
-            decision = "question_rewrite_agent"
+            decision = "user_input_rewrite_agent"
         # Handle transition after query rewriting
-        elif current_agent == "question_rewrite_agent":
+        elif current_agent == "user_input_rewrite_agent":
             decision = "parallel_query_solving_agent"
         # Handle transition after parallel query solving
         elif current_agent == "parallel_query_solving_agent":
@@ -145,22 +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_question_results = {}
+            sub_user_input_results = {}
         else:
-            # Only load sub-question results if we have a database result
-            sub_question_results = self.parse_message_content(messages[1].content)
-            logging.info("Sub-Question Results: %s", sub_question_results)
+            # 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)
 
         try:
-            sub_questions = [
-                sub_question
-                for sub_question_group in sub_question_results.get("sub_questions", [])
-                for sub_question in sub_question_group
+            sub_user_inputs = [
+                sub_user_input
+                for sub_user_input_group in sub_user_input_results.get(
+                    "sub_user_inputs", []
+                )
+                for sub_user_input in sub_user_input_group
             ]
 
             logging.info("SQL Query Results: %s", sql_query_results)
             payload = AnswerWithSourcesPayload(
-                answer=answer, sub_questions=sub_questions
+                answer=answer, sub_user_inputs=sub_user_inputs
             )
 
             if not isinstance(sql_query_results, dict):
@@ -171,9 +173,11 @@ def extract_answer_payload(self, messages: list) -> AnswerWithSourcesPayload:
                 logging.error("No 'results' key in sql_query_results")
                 return payload
 
-            for question, sql_query_result_list in sql_query_results["results"].items():
+            for user_input, 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 question: {question}")
+                    logging.warning(f"No results for user_input: {user_input}")
                     continue
 
                 for sql_query_result in sql_query_result_list:
@@ -209,30 +213,30 @@ def extract_answer_payload(self, messages: list) -> AnswerWithSourcesPayload:
                 answer=f"{answer}\nError processing results: {str(e)}"
             )
 
-    async def process_question(
+    async def process_user_input(
         self,
-        question_payload: UserInputPayload,
+        user_input_payload: UserInputPayload,
         chat_history: list[InteractionPayload] = None,
     ) -> AsyncGenerator[InteractionPayload, None]:
-        """Process the complete question through the unified system.
+        """Process the complete user_input through the unified system.
 
         Args:
         ----
-            task (str): The user question to process.
+            task (str): The user user_input 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 question: %s", question_payload.body.question)
+        logging.info("Processing user_input: %s", user_input_payload.body.user_input)
         logging.info("Chat history: %s", chat_history)
 
         agent_input = {
-            "question": question_payload.body.question,
+            "user_input": user_input_payload.body.user_input,
             "chat_history": {},
-            "injected_parameters": question_payload.body.injected_parameters,
+            "injected_parameters": user_input_payload.body.injected_parameters,
         }
 
         if chat_history is not None:
@@ -241,15 +245,15 @@ async def process_question(
                 if chat.root.payload_type == PayloadType.USER_INPUT:
                     # For now only consider the user query
                     chat_history_key = f"chat_{idx}"
-                    agent_input[chat_history_key] = chat.root.body.question
+                    agent_input[chat_history_key] = chat.root.body.user_input
 
         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 == "question_rewrite_agent":
+                if message.source == "user_input_rewrite_agent":
                     payload = ProcessingUpdatePayload(
                         message="Rewriting the query...",
                     )
@@ -272,10 +276,10 @@ async def process_question(
                 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 == "question_rewrite_agent":
+                elif message.messages[-1].source == "user_input_rewrite_agent":
                     # Load into empty response
                     payload = AnswerWithSourcesPayload(
-                        answer="Apologies, I cannot answer that question as it is not relevant. Please try another question or rephrase your current question."
+                        answer="Apologies, I cannot answer that user_input as it is not relevant. Please try another user_input or rephrase your current user_input."
                     )
             else:
                 logging.error("Unexpected TaskResult: %s", message)

text_2_sql/autogen/src/autogen_text_2_sql/creators/llm_agent_creator.py

@@ -43,7 +43,7 @@ def get_tool(cls, sql_helper, tool_name: str):
         elif tool_name == "sql_get_entity_schemas_tool":
             return FunctionToolAlias(
                 sql_helper.get_entity_schemas,
-                description="Gets the schema of a view or table in the SQL Database by selecting the most relevant entity based on the search term. Extract key terms from the user question and use these as the search term. Several entities may be returned. Only use when the provided schemas in the message history are not sufficient to answer the question.",
+                description="Gets the schema of a view or table in the SQL Database by selecting the most relevant entity based on the search term. Extract key terms from the user input and use these as the search term. Several entities may be returned. Only use when the provided schemas in the message history are not sufficient to answer the question.",
             )
         elif tool_name == "sql_get_column_values_tool":
             return FunctionToolAlias(

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
-        question_rewrites = json.loads(last_response)
+        user_input_rewrites = json.loads(last_response)
 
-        logging.info(f"Query Rewrites: {question_rewrites}")
+        logging.info(f"Query Rewrites: {user_input_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_questions_and_schemas"
+                                    "cached_user_inputs_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(
-            question_rewrites.get("all_non_database_query", "false")
+            user_input_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 question_rewrite in question_rewrites["sub_questions"]:
-            logging.info(f"Processing sub-query: {question_rewrite}")
+        for user_input_rewrite in user_input_rewrites["sub_user_inputs"]:
+            logging.info(f"Processing sub-query: {user_input_rewrite}")
             # Create an instance of the InnerAutoGenText2Sql class
             inner_autogen_text_2_sql = InnerAutoGenText2Sql(**self.kwargs)
 
-            identifier = ", ".join(question_rewrite)
+            identifier = ", ".join(user_input_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_question(
-                        question=question_rewrite,
+                    inner_autogen_text_2_sql.process_user_input(
+                        user_input=user_input_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 question.",
+            "An agent that fetches the queries from the cache based on the user user_input.",
         )
 
         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 questions from the question_rewrite_agent
+        # Get the decomposed user_inputs from the user_input_rewrite_agent
         try:
             request_details = json.loads(messages[0].content)
             injected_parameters = request_details["injected_parameters"]
-            user_questions = request_details["question"]
-            logging.info(f"Processing questions: {user_questions}")
+            user_user_inputs = request_details["user_input"]
+            logging.info(f"Processing user_inputs: {user_user_inputs}")
             logging.info(f"Input Parameters: {injected_parameters}")
         except json.JSONDecodeError:
-            # If not JSON array, process as single question
+            # If not JSON array, process as single user_input
             raise ValueError("Could not load message")
 
         cached_results = await self.agent.process_message(
-            user_questions, injected_parameters
+            user_user_inputs, injected_parameters
         )
         yield Response(
             chat_message=TextMessage(

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

@@ -17,7 +17,7 @@ class SqlSchemaSelectionAgent(BaseChatAgent):
     def __init__(self, **kwargs):
         super().__init__(
             "sql_schema_selection_agent",
-            "An agent that fetches the schemas from the cache based on the user question.",
+            "An agent that fetches the schemas from the cache based on the user input.",
         )
 
         self.agent = SqlSchemaSelectionAgentCustomAgent(**kwargs)
@@ -43,19 +43,19 @@ async def on_messages_stream(
         # Try to parse as JSON first
         try:
             request_details = json.loads(messages[0].content)
-            user_questions = request_details["question"]
+            user_inputs = request_details["question"]
         except (json.JSONDecodeError, KeyError):
             # If not JSON or missing question key, use content directly
-            user_questions = messages[0].content
+            user_inputs = messages[0].content
 
-        if isinstance(user_questions, str):
-            user_questions = [user_questions]
-        elif not isinstance(user_questions, list):
-            user_questions = [str(user_questions)]
+        if isinstance(user_inputs, str):
+            user_inputs = [user_inputs]
+        elif not isinstance(user_inputs, list):
+            user_inputs = [str(user_inputs)]
 
-        logging.info(f"Processing questions: {user_questions}")
+        logging.info(f"Processing questions: {user_inputs}")
 
-        final_results = await self.agent.process_message(user_questions)
+        final_results = await self.agent.process_message(user_inputs)
 
         yield Response(
             chat_message=TextMessage(

text_2_sql/autogen/src/autogen_text_2_sql/inner_autogen_text_2_sql.py

@@ -178,7 +178,7 @@ def process_question(
 
         Args:
         ----
-            task (str): The user question to process.
+            task (str): The user input to process.
             injected_parameters (dict, optional): Parameters to pass to agents. Defaults to None.
 
         Returns:

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_questions: list[str], injected_parameters: dict
+        self, user_inputs: 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_questions:
+        for question in user_inputs:
             # 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_questions: list[str]) -> dict:
-        logging.info(f"User questions: {user_questions}")
+    async def process_message(self, user_inputs: list[str]) -> dict:
+        logging.info(f"user inputs: {user_inputs}")
 
         entity_tasks = []
 
-        for user_question in user_questions:
+        for user_input in user_inputs:
             messages = [
                 {"role": "system", "content": self.system_prompt},
-                {"role": "user", "content": user_question},
+                {"role": "user", "content": user_input},
             ]
             entity_tasks.append(
                 self.open_ai_connector.run_completion_request(
@@ -47,7 +47,7 @@ async def process_message(self, user_questions: list[str]) -> dict:
             logging.info(f"Entity result: {entity_result}")
 
             for entity_group in entity_result.entities:
-                logging.info(f"Searching for schemas for entity group: {entity_group}")
+                logging.info("Searching for schemas for entity group: %s", entity_group)
                 entity_search_tasks.append(
                     self.sql_connector.get_entity_schemas(
                         " ".join(entity_group), as_json=False
@@ -56,7 +56,7 @@ async def process_message(self, user_questions: list[str]) -> dict:
 
             for filter_condition in entity_result.filter_conditions:
                 logging.info(
-                    f"Searching for column values for filter: {filter_condition}"
+                    "Searching for column values for filter: %s", filter_condition
                 )
                 column_search_tasks.append(
                     self.sql_connector.get_column_values(

…core/prompts/question_rewrite_agent.yaml …re/prompts/user_input_rewrite_agent.yamltext_2_sql/text_2_sql_core/src/text_2_sql_core/prompts/question_rewrite_agent.yaml renamed to text_2_sql/text_2_sql_core/src/text_2_sql_core/prompts/user_input_rewrite_agent.yaml text_2_sql/text_2_sql_core/src/text_2_sql_core/prompts/question_rewrite_agent.yaml renamed to text_2_sql/text_2_sql_core/src/text_2_sql_core/prompts/user_input_rewrite_agent.yaml

@@ -1,5 +1,5 @@
 model: "4o-mini"
-description: "An agent that preprocesses user questions by decomposing complex queries into simpler sub-queries that can be processed independently and then combined."
+description: "An agent that preprocesses user inputs by decomposing complex queries into simpler sub-queries that can be processed independently and then combined."
 system_message: |
   <role_and_objective>
       You are a helpful AI Assistant specializing in breaking down complex questions into simpler sub-queries that can be processed independently and then combined for the final answer. You should identify when a question can be solved through simpler sub-queries and provide clear instructions for combining their results.