Add parallel agent flow

Author: BenConstable9

text_2_sql/autogen/agentic_text_2_sql.py

@@ -13,7 +13,10 @@
     target_engine="Microsoft SQL Server",
     engine_specific_rules="Use TOP X to limit the number of rows returned instead of LIMIT X. NEVER USE LIMIT X as it produces a syntax error.",
 )
-SQL_SCHEMA_SELECTION_AGENT = LLMAgentCreator.create("sql_schema_selection_agent")
+SQL_SCHEMA_SELECTION_AGENT = LLMAgentCreator.create(
+    "sql_schema_selection_agent",
+    use_case="Sales data for a company that specializes in selling products online.",
+)
 SQL_QUERY_CORRECTION_AGENT = LLMAgentCreator.create(
     "sql_query_correction_agent",
     target_engine="Microsoft SQL Server",
@@ -42,7 +45,12 @@ def text_2_sql_generator_selector_func(messages):
             decision = "sql_schema_selection_agent"
 
     elif messages[-1].source == "question_decomposition_agent":
-        decision = "sql_schema_selection_agent"
+        decomposition_result = json.loads(messages[-1].content)
+
+        if len(decomposition_result["entities"]) == 1:
+            decision = "sql_schema_selection_agent"
+        else:
+            decision = "parallel_sql_flow_agent"
 
     elif messages[-1].source == "sql_schema_selection_agent":
         decision = "sql_query_generation_agent"

…om_agents/sql_schema_extraction_agent.py …custom_agents/parallel_sql_flow_agent.pytext_2_sql/autogen/agents/custom_agents/sql_schema_extraction_agent.py renamed to text_2_sql/autogen/agents/custom_agents/parallel_sql_flow_agent.py text_2_sql/autogen/agents/custom_agents/sql_schema_extraction_agent.py renamed to text_2_sql/autogen/agents/custom_agents/parallel_sql_flow_agent.py

@@ -6,12 +6,12 @@
 from autogen_agentchat.base import Response
 from autogen_agentchat.messages import AgentMessage, ChatMessage, TextMessage
 from autogen_core.base import CancellationToken
-from text_2_sql.autogen.utils.sql import get_entity_schemas
-from keybert import KeyBERT
+from utils.sql import fetch_queries_from_cache
+import json
 import logging
 
 
-class SqlSchemaExtractionAgent(BaseChatAgent):
+class ParallelSqlFlowAgent(BaseChatAgent):
     def __init__(self):
         super().__init__(
             "sql_query_cache_agent",
@@ -41,24 +41,12 @@ async def on_messages_stream(
         # Fetch the queries from the cache based on the user question.
         logging.info("Fetching queries from cache based on the user question...")
 
-        kw_model = KeyBERT()
-
-        top_keywords = kw_model.extract_keywords(
-            user_question, keyphrase_ngram_range=(1, 3), top_n=5
-        )
-
-        # Extract just the key phrases (ignoring the score)
-        key_phrases = [keyword[0] for keyword in top_keywords]
-
-        # Join them into a string list
-        key_phrases_str = ", ".join(key_phrases)
-
-        entity_schemas = await get_entity_schemas(key_phrases_str)
-
-        logging.info(entity_schemas)
+        cached_queries = await fetch_queries_from_cache(user_question)
 
         yield Response(
-            chat_message=TextMessage(content=entity_schemas, source=self.name)
+            chat_message=TextMessage(
+                content=json.dumps(cached_queries), source=self.name
+            )
         )
 
     async def on_reset(self, cancellation_token: CancellationToken) -> None:

text_2_sql/autogen/agents/llm_agents/sql_schema_selection_agent.yaml

@@ -7,6 +7,8 @@ description:
 system_message:
   "You are a helpful AI Assistant that specialises in selecting relevant SQL schemas to answer a given user's question.
 
+  The user's question will be related to {{ use_case }}.
+
   Perform the following steps to select the correct schema:
 
     1. Extract key terms and entities from the user's question.

text_2_sql/autogen/requirements.txt

@@ -8,3 +8,4 @@ azure-identity
 python-dotenv
 openai
 jinja2
+pyyaml

text_2_sql/semantic_kernel/requirements.txt

@@ -9,3 +9,4 @@ numpy
 seaborn
 pydantic
 openai
+pyyaml