Update prompts and agents

Author: BenConstable9

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

@@ -3,14 +3,20 @@
 from typing import AsyncGenerator, List, Sequence
 
 from autogen_agentchat.agents import BaseChatAgent
-from autogen_agentchat.base import Response, TaskResult
-from autogen_agentchat.messages import AgentMessage, ChatMessage, TextMessage
+from autogen_agentchat.base import Response
+from autogen_agentchat.messages import (
+    AgentMessage,
+    ChatMessage,
+    TextMessage,
+    ToolCallResultMessage,
+)
 from autogen_core import CancellationToken
 import json
 import logging
 from autogen_text_2_sql.inner_autogen_text_2_sql import InnerAutoGenText2Sql
 from aiostream import stream
 from json import JSONDecodeError
+import re
 
 
 class ParallelQuerySolvingAgent(BaseChatAgent):
@@ -53,9 +59,6 @@ def parse_inner_message(self, message):
             except JSONDecodeError:
                 pass
 
-            # Try to extract JSON from markdown code blocks
-            import re
-
             json_match = re.search(r"```json\s*(.*?)\s*```", message, re.DOTALL)
             if json_match:
                 try:
@@ -103,12 +106,13 @@ async def consume_inner_messages_from_agentic_flow(
 
                 logging.info(f"Checking Inner Message: {inner_message}")
 
-                if isinstance(inner_message, TaskResult) is False:
-                    try:
+                try:
+                    if isinstance(inner_message, ToolCallResultMessage):
+                        # Check for SQL query results
                         parsed_message = self.parse_inner_message(inner_message.content)
+
                         logging.info(f"Inner Loaded: {parsed_message}")
 
-                        # Search for specific message types and add them to the final output object
                         if isinstance(parsed_message, dict):
                             if (
                                 "type" in parsed_message
@@ -124,6 +128,13 @@ async def consume_inner_messages_from_agentic_flow(
                                     }
                                 )
 
+                    elif isinstance(inner_message, TextMessage):
+                        parsed_message = self.parse_inner_message(inner_message.content)
+
+                        logging.info(f"Inner Loaded: {parsed_message}")
+
+                        # Search for specific message types and add them to the final output object
+                        if isinstance(parsed_message, dict):
                             if ("contains_pre_run_results" in parsed_message) and (
                                 parsed_message["contains_pre_run_results"] is True
                             ):
@@ -139,8 +150,8 @@ async def consume_inner_messages_from_agentic_flow(
                                         }
                                     )
 
-                    except Exception as e:
-                        logging.warning(f"Error processing message: {e}")
+                except Exception as e:
+                    logging.warning(f"Error processing message: {e}")
 
                 yield inner_message
 

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

@@ -6,7 +6,9 @@
 from autogen_agentchat.base import Response
 from autogen_agentchat.messages import AgentMessage, ChatMessage, TextMessage
 from autogen_core import CancellationToken
-from text_2_sql_core.connectors.factory import ConnectorFactory
+from text_2_sql_core.custom_agents.sql_query_cache_agent import (
+    SqlQueryCacheAgentCustomAgent,
+)
 import json
 import logging
 
@@ -18,7 +20,7 @@ def __init__(self):
             "An agent that fetches the queries from the cache based on the user question.",
         )
 
-        self.sql_connector = ConnectorFactory.get_database_connector()
+        self.agent = SqlQueryCacheAgentCustomAgent()
 
     @property
     def produced_message_types(self) -> List[type[ChatMessage]]:
@@ -49,31 +51,9 @@ async def on_messages_stream(
             # If not JSON array, process as single question
             raise ValueError("Could not load message")
 
-        # Initialize results dictionary
-        cached_results = {
-            "cached_questions_and_schemas": [],
-            "contains_pre_run_results": False,
-        }
-
-        # Process each question sequentially
-        for question in user_questions:
-            # 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(
-                question, injected_parameters=injected_parameters
-            )
-
-            # If any question has pre-run results, set the flag
-            if cached_query.get("contains_pre_run_results", False):
-                cached_results["contains_pre_run_results"] = True
-
-            # Add the cached results for this question
-            if cached_query.get("cached_questions_and_schemas"):
-                cached_results["cached_questions_and_schemas"].extend(
-                    cached_query["cached_questions_and_schemas"]
-                )
-
-        logging.info(f"Final cached results: {cached_results}")
+        cached_results = await self.agent.process_message(
+            user_questions, injected_parameters
+        )
         yield Response(
             chat_message=TextMessage(
                 content=json.dumps(cached_results), source=self.name

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

@@ -6,12 +6,11 @@
 from autogen_agentchat.base import Response
 from autogen_agentchat.messages import AgentMessage, ChatMessage, TextMessage
 from autogen_core import CancellationToken
-from text_2_sql_core.connectors.factory import ConnectorFactory
 import json
 import logging
-from text_2_sql_core.prompts.load import load
-from jinja2 import Template
-import asyncio
+from text_2_sql_core.custom_agents.sql_schema_selection_agent import (
+    SqlSchemaSelectionAgentCustomAgent,
+)
 
 
 class SqlSchemaSelectionAgent(BaseChatAgent):
@@ -21,15 +20,7 @@ def __init__(self, **kwargs):
             "An agent that fetches the schemas from the cache based on the user question.",
         )
 
-        self.ai_search_connector = ConnectorFactory.get_ai_search_connector()
-
-        self.open_ai_connector = ConnectorFactory.get_open_ai_connector()
-
-        self.sql_connector = ConnectorFactory.get_database_connector()
-
-        system_prompt = load("sql_schema_selection_agent")["system_message"]
-
-        self.system_prompt = Template(system_prompt).render(kwargs)
+        self.agent = SqlSchemaSelectionAgentCustomAgent(**kwargs)
 
     @property
     def produced_message_types(self) -> List[type[ChatMessage]]:
@@ -49,64 +40,15 @@ async def on_messages(
     async def on_messages_stream(
         self, messages: Sequence[ChatMessage], cancellation_token: CancellationToken
     ) -> AsyncGenerator[AgentMessage | Response, None]:
-        last_response = messages[-1].content
-
-        # load the json of the last message and get the user question's
-
-        user_questions = json.loads(last_response)
-
-        logging.info(f"User questions: {user_questions}")
-
-        entity_tasks = []
-
-        for user_question in user_questions:
-            messages = [
-                {"role": "system", "content": self.system_prompt},
-                {"role": "user", "content": user_question},
-            ]
-            entity_tasks.append(self.open_ai_connector.run_completion_request(messages))
-
-        entity_results = await asyncio.gather(*entity_tasks)
-
-        entity_search_tasks = []
-        column_search_tasks = []
-
-        for entity_result in entity_results:
-            loaded_entity_result = json.loads(entity_result)
-
-            logging.info(f"Loaded entity result: {loaded_entity_result}")
-
-            for entity_group in loaded_entity_result["entities"]:
-                entity_search_tasks.append(
-                    self.sql_connector.get_entity_schemas(
-                        " ".join(entity_group), as_json=False
-                    )
-                )
-
-            for filter_condition in loaded_entity_result["filter_conditions"]:
-                column_search_tasks.append(
-                    self.ai_search_connector.get_column_values(
-                        filter_condition, as_json=False
-                    )
-                )
-
-        schemas_results = await asyncio.gather(*entity_search_tasks)
-        column_value_results = await asyncio.gather(*column_search_tasks)
-
-        # deduplicate schemas
-        final_schemas = []
-
-        for schema_result in schemas_results:
-            for schema in schema_result:
-                if schema not in final_schemas:
-                    final_schemas.append(schema)
-
-        final_results = {
-            "COLUMN_OPTIONS_AND_VALUES_FOR_FILTERS": column_value_results,
-            "SCHEMA_OPTIONS": final_schemas,
-        }
-
-        logging.info(f"Final results: {final_results}")
+        try:
+            request_details = json.loads(messages[0].content)
+            user_questions = request_details["question"]
+            logging.info(f"Processing questions: {user_questions}")
+        except json.JSONDecodeError:
+            # If not JSON array, process as single question
+            raise ValueError("Could not load message")
+
+        final_results = await self.agent.process_message(user_questions)
 
         yield Response(
             chat_message=TextMessage(

text_2_sql/text_2_sql_core/src/text_2_sql_core/connectors/open_ai.py

@@ -28,7 +28,12 @@ def get_authentication_properties(cls) -> dict:
         return token_provider, api_key
 
     async def run_completion_request(
-        self, messages: list[dict], temperature=0, max_tokens=2000, model="4o-mini"
+        self,
+        messages: list[dict],
+        temperature=0,
+        max_tokens=2000,
+        model="4o-mini",
+        response_format=None,
     ) -> str:
         if model == "4o-mini":
             model_deployment = os.environ["OpenAI__MiniCompletionDeployment"]
@@ -45,13 +50,29 @@ async def run_completion_request(
             azure_ad_token_provider=token_provider,
             api_key=api_key,
         ) as open_ai_client:
-            response = await open_ai_client.chat.completions.create(
-                model=model_deployment,
-                messages=messages,
-                temperature=temperature,
-                max_tokens=max_tokens,
-            )
-        return response.choices[0].message.content
+            if response_format is not None:
+                response = await open_ai_client.beta.chat.completions.parse(
+                    model=model_deployment,
+                    messages=messages,
+                    temperature=temperature,
+                    max_tokens=max_tokens,
+                    response_format=response_format,
+                )
+            else:
+                response = await open_ai_client.chat.completions.create(
+                    model=model_deployment,
+                    messages=messages,
+                    temperature=temperature,
+                    max_tokens=max_tokens,
+                )
+
+        message = response.choices[0].message
+        if response_format is not None and message.parsed is not None:
+            return message.parsed
+        elif response_format is not None:
+            return message.refusal
+        else:
+            return message.content
 
     async def run_embedding_request(self, batch: list[str]):
         token_provider, api_key = self.get_authentication_properties()

text_2_sql/text_2_sql_core/src/text_2_sql_core/connectors/tsql_sql.py

@@ -127,6 +127,7 @@ async def get_entity_schemas(
 
             del schema["Entity"]
             del schema["Schema"]
+            del schema["Database"]
 
         if as_json:
             return json.dumps(schemas, default=str)

text_2_sql/text_2_sql_core/src/text_2_sql_core/custom_agents/__init__.py

@@ -0,0 +1,39 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+from text_2_sql_core.connectors.factory import ConnectorFactory
+import logging
+
+
+class SqlQueryCacheAgentCustomAgent:
+    def __init__(self):
+        self.sql_connector = ConnectorFactory.get_database_connector()
+
+    async def process_message(
+        self, user_questions: list[str], injected_parameters: dict
+    ) -> dict:
+        # Initialize results dictionary
+        cached_results = {
+            "cached_questions_and_schemas": [],
+            "contains_pre_run_results": False,
+        }
+
+        # Process each question sequentially
+        for question in user_questions:
+            # 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(
+                question, injected_parameters=injected_parameters
+            )
+
+            # If any question has pre-run results, set the flag
+            if cached_query.get("contains_pre_run_results", False):
+                cached_results["contains_pre_run_results"] = True
+
+            # Add the cached results for this question
+            if cached_query.get("cached_questions_and_schemas"):
+                cached_results["cached_questions_and_schemas"].extend(
+                    cached_query["cached_questions_and_schemas"]
+                )
+
+        logging.info(f"Final cached results: {cached_results}")
+        return cached_results