Update agents

Author: BenConstable9

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

@@ -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_message(
-                        message=message_rewrite,
+                    inner_autogen_text_2_sql.process_user_message(
+                        user_message=message_rewrite,
                         injected_parameters=query_params,
                     ),
                     identifier,

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

@@ -44,7 +44,7 @@ async def on_messages_stream(
         try:
             request_details = json.loads(messages[0].content)
             injected_parameters = request_details["injected_parameters"]
-            user_messages = request_details["message"]
+            user_messages = request_details["user_message"]
             logging.info(f"Processing messages: {user_messages}")
             logging.info(f"Input Parameters: {injected_parameters}")
         except json.JSONDecodeError:

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

@@ -18,11 +18,11 @@ async def process_message(
         }
 
         # Process each question sequentially
-        for question in messages:
+        for message in messages:
             # Fetch the queries from the cache based on the question
-            logging.info(f"Fetching queries from cache for question: {question}")
+            logging.info(f"Fetching queries from cache for question: {message}")
             cached_query = await self.sql_connector.fetch_queries_from_cache(
-                question, injected_parameters=injected_parameters
+                message, injected_parameters=injected_parameters
             )
 
             # If any question has pre-run results, set the flag

text_2_sql/text_2_sql_core/src/text_2_sql_core/payloads/interaction_payloads.py

@@ -25,12 +25,6 @@ class InteractionPayloadBase(BaseModel):
 
 
 class PayloadBase(InteractionPayloadBase):
-    prompt_tokens: int | None = Field(
-        None, description="Number of tokens in the prompt", alias="promptTokens"
-    )
-    completion_tokens: int | None = Field(
-        None, description="Number of tokens in the completion", alias="completionTokens"
-    )
     message_id: str = Field(
         ..., default_factory=lambda: str(uuid4()), alias="messageId"
     )

text_2_sql/text_2_sql_core/src/text_2_sql_core/prompts/user_message_rewrite_agent.yaml

@@ -91,7 +91,7 @@ system_message: |
     <output_format>
         Return a JSON object with sub-messages and combination instructions:
         {
-          "decomposed_messages": [
+          "decomposed_user_messages": [
             ["<sub_message_1>"],
             ["<sub_message_2>"],
             ...
@@ -108,7 +108,7 @@ system_message: |
     Input: "Which product categories have shown consistent growth quarter over quarter in 2008, and what were their top selling items?"
     Output:
     {
-      "decomposed_messages": [
+      "decomposed_user_messages": [
         ["Calculate quarterly sales totals by product category for 2008", "For these categories, find their top selling products in 2008"]
       ],
       "combination_logic": "First identify growing categories from quarterly analysis, then find their best-selling products",
@@ -120,7 +120,7 @@ system_message: |
     Input: "How many orders did we have in 2008?"
     Output:
     {
-      "decomposed_messages": [
+      "decomposed_user_messages": [
         ["How many orders did we have in 2008?"]
       ],
       "combination_logic": "Direct count query, no combination needed",
@@ -132,7 +132,7 @@ system_message: |
     Input: "Compare the sales performance of our top 5 products in Europe versus North America, including their market share in each region"
     Output:
     {
-      "decomposed_messages": [
+      "decomposed_user_messages": [
         ["Get total sales by product in European countries"],
         ["Get total sales by product in North American countries"],
         ["Calculate total market size for each region", "Find top 5 products by sales in each region"],
@@ -146,7 +146,7 @@ system_message: |
     Input: "Hello, what can you help me with?"
     Output:
     {
-      "decomposed_messages": [
+      "decomposed_user_messages": [
         ["What are your capabilities?"]
       ],
       "combination_logic": "Simple greeting and capability question",