Update envs

Author: BenConstable9

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

@@ -223,8 +223,8 @@ async def consume_inner_messages_from_agentic_flow(
                 query_params["database_connection_string"] = os.environ[
                     "Text2Sql__Tsql__ConnectionString"
                 ]
-            if "Text2Sql__Database" in os.environ:
-                query_params["database_name"] = os.environ["Text2Sql__Database"]
+            if "Text2Sql__Tsql__Database" in os.environ:
+                query_params["database_name"] = os.environ["Text2Sql__Tsql__Database"]
 
             # Launch tasks for each sub-query
             inner_solving_generators.append(

text_2_sql/autogen/src/autogen_text_2_sql/inner_autogen_text_2_sql.py

@@ -46,7 +46,7 @@ def __init__(self, **kwargs: dict):
 
         # Store original environment variables
         self.original_db_conn = os.environ.get("Text2Sql__Tsql__ConnectionString")
-        self.original_db_name = os.environ.get("Text2Sql__Database")
+        self.original_db_name = os.environ.get("Text2Sql__Tsql__Database")
 
     def _update_environment(self, injected_parameters: dict = None):
         """Update environment variables with injected parameters."""
@@ -56,14 +56,16 @@ def _update_environment(self, injected_parameters: dict = None):
                     "database_connection_string"
                 ]
             if "database_name" in injected_parameters:
-                os.environ["Text2Sql__Database"] = injected_parameters["database_name"]
+                os.environ["Text2Sql__Tsql__Database"] = injected_parameters[
+                    "database_name"
+                ]
 
     def _restore_environment(self):
         """Restore original environment variables."""
         if self.original_db_conn:
             os.environ["Text2Sql__Tsql__ConnectionString"] = self.original_db_conn
         if self.original_db_name:
-            os.environ["Text2Sql__Database"] = self.original_db_name
+            os.environ["Text2Sql__Tsql__Database"] = self.original_db_name
 
     def set_mode(self):
         """Set the mode of the plugin based on the environment variables."""

text_2_sql/text_2_sql_core/src/text_2_sql_core/connectors/snowflake_sql.py

@@ -100,7 +100,7 @@ async def query_execution(
             password=os.environ["Text2Sql__Snowflake__Password"],
             account=os.environ["Text2Sql__Snowflake__Account"],
             warehouse=os.environ["Text2Sql__Snowflake__Warehouse"],
-            database=os.environ["Text2Sql__Database"],
+            database=os.environ["Text2Sql__Snowflake__Database"],
         )
 
         try:

text_2_sql/text_2_sql_core/src/text_2_sql_core/utils/environment.py

@@ -18,7 +18,7 @@ def get_identity_type() -> IdentityType:
     Returns:
         IdentityType: The identity type
     """
-    identity = os.environ.get("IdentityType")
+    identity = os.environ["IdentityType"]
 
     if identity == "user_assigned":
         return IdentityType.USER_ASSIGNED