Update engine specific rules logix

Author: BenConstable9

text_2_sql/autogen/src/autogen_text_2_sql/autogen_text_2_sql.py

@@ -31,9 +31,8 @@
 
 
 class AutoGenText2Sql:
-    def __init__(self, engine_specific_rules: str, **kwargs: dict):
+    def __init__(self, **kwargs: dict):
         self.target_engine = os.environ["Text2Sql__DatabaseEngine"].upper()
-        self.engine_specific_rules = engine_specific_rules
         self.kwargs = kwargs
 
     def get_all_agents(self):
@@ -45,9 +44,7 @@ def get_all_agents(self):
             "question_rewrite_agent", current_datetime=current_datetime
         )
 
-        self.parallel_query_solving_agent = ParallelQuerySolvingAgent(
-            engine_specific_rules=self.engine_specific_rules, **self.kwargs
-        )
+        self.parallel_query_solving_agent = ParallelQuerySolvingAgent(**self.kwargs)
 
         self.answer_agent = LLMAgentCreator.create("answer_agent")
 

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

@@ -6,6 +6,7 @@
 from text_2_sql_core.prompts.load import load
 from autogen_text_2_sql.creators.llm_model_creator import LLMModelCreator
 from jinja2 import Template
+import logging
 
 
 class LLMAgentCreator:
@@ -89,6 +90,17 @@ def create(cls, name: str, **kwargs) -> AssistantAgent:
 
         sql_helper = ConnectorFactory.get_database_connector()
 
+        # Handle engine specific rules
+        if "engine_specific_rules" not in kwargs:
+            if sql_helper.engine_specific_rules is not None:
+                kwargs["engine_specific_fields"] = sql_helper.engine_specific_rules
+                logging.info(
+                    "Engine specific fields pulled from in-built: %s",
+                    kwargs["engine_specific_fields"],
+                )
+            else:
+                kwargs["engine_specific_fields"] = ""
+
         tools = []
         if "tools" in agent_file and len(agent_file["tools"]) > 0:
             for tool in agent_file["tools"]:

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

@@ -20,13 +20,12 @@
 
 
 class ParallelQuerySolvingAgent(BaseChatAgent):
-    def __init__(self, engine_specific_rules: str, **kwargs: dict):
+    def __init__(self, **kwargs: dict):
         super().__init__(
             "parallel_query_solving_agent",
             "An agent that solves each query in parallel.",
         )
 
-        self.engine_specific_rules = engine_specific_rules
         self.kwargs = kwargs
 
     @property
@@ -177,9 +176,7 @@ async def consume_inner_messages_from_agentic_flow(
         for question_rewrite in question_rewrites["sub_questions"]:
             logging.info(f"Processing sub-query: {question_rewrite}")
             # Create an instance of the InnerAutoGenText2Sql class
-            inner_autogen_text_2_sql = InnerAutoGenText2Sql(
-                self.engine_specific_rules, **self.kwargs
-            )
+            inner_autogen_text_2_sql = InnerAutoGenText2Sql(**self.kwargs)
 
             identifier = ", ".join(question_rewrite)
 

text_2_sql/autogen/src/autogen_text_2_sql/inner_autogen_text_2_sql.py

@@ -38,10 +38,9 @@ async def on_messages_stream(self, messages, sender=None, config=None):
 
 
 class InnerAutoGenText2Sql:
-    def __init__(self, engine_specific_rules: str, **kwargs: dict):
+    def __init__(self, **kwargs: dict):
         self.pre_run_query_cache = False
         self.target_engine = os.environ["Text2Sql__DatabaseEngine"].upper()
-        self.engine_specific_rules = engine_specific_rules
         self.kwargs = kwargs
         self.set_mode()
 
@@ -73,21 +72,18 @@ def get_all_agents(self):
 
         self.sql_schema_selection_agent = SqlSchemaSelectionAgent(
             target_engine=self.target_engine,
-            engine_specific_rules=self.engine_specific_rules,
             **self.kwargs,
         )
 
         self.sql_query_correction_agent = LLMAgentCreator.create(
             "sql_query_correction_agent",
             target_engine=self.target_engine,
-            engine_specific_rules=self.engine_specific_rules,
             **self.kwargs,
         )
 
         self.disambiguation_and_sql_query_generation_agent = LLMAgentCreator.create(
             "disambiguation_and_sql_query_generation_agent",
             target_engine=self.target_engine,
-            engine_specific_rules=self.engine_specific_rules,
             **self.kwargs,
         )
         agents = [