Update agent

Author: BenConstable9

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

@@ -17,7 +17,6 @@
 from aiostream import stream
 from json import JSONDecodeError
 import re
-import os
 from pydantic import BaseModel, Field
 
 
@@ -226,23 +225,12 @@ async def consume_inner_messages_from_agentic_flow(
                 # Create an instance of the InnerAutoGenText2Sql class
                 inner_autogen_text_2_sql = InnerAutoGenText2Sql(**self.kwargs)
 
-                # Add database connection info to injected parameters
-                query_params = injected_parameters.copy() if injected_parameters else {}
-                if "Text2Sql__Tsql__ConnectionString" in os.environ:
-                    query_params["database_connection_string"] = os.environ[
-                        "Text2Sql__Tsql__ConnectionString"
-                    ]
-                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(
                     consume_inner_messages_from_agentic_flow(
                         inner_autogen_text_2_sql.process_user_message(
                             user_message=parallel_message,
-                            injected_parameters=query_params,
+                            injected_parameters=injected_parameters,
                             database_results=filtered_parallel_messages.database_results,
                         ),
                         parallel_message,

text_2_sql/autogen/src/autogen_text_2_sql/inner_autogen_text_2_sql.py

@@ -44,29 +44,6 @@ def __init__(self, **kwargs: dict):
         self.kwargs = kwargs
         self.set_mode()
 
-        # Store original environment variables
-        self.original_db_conn = os.environ.get("Text2Sql__Tsql__ConnectionString")
-        self.original_db_name = os.environ.get("Text2Sql__Tsql__Database")
-
-    def _update_environment(self, injected_parameters: dict = None):
-        """Update environment variables with injected parameters."""
-        if injected_parameters:
-            if "database_connection_string" in injected_parameters:
-                os.environ["Text2Sql__Tsql__ConnectionString"] = injected_parameters[
-                    "database_connection_string"
-                ]
-            if "database_name" in injected_parameters:
-                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__Tsql__Database"] = self.original_db_name
-
     def set_mode(self):
         """Set the mode of the plugin based on the environment variables."""
         self.pre_run_query_cache = (
@@ -195,19 +172,12 @@ def process_user_message(
         """
         logging.info("Processing question: %s", user_message)
 
-        # Update environment with injected parameters
-        self._update_environment(injected_parameters)
-
-        try:
-            agent_input = {
-                "user_message": user_message,
-                "injected_parameters": injected_parameters,
-            }
+        agent_input = {
+            "user_message": user_message,
+            "injected_parameters": injected_parameters,
+        }
 
-            if database_results:
-                agent_input["database_results"] = database_results
+        if database_results:
+            agent_input["database_results"] = database_results
 
-            return self.agentic_flow.run_stream(task=json.dumps(agent_input))
-        finally:
-            # Restore original environment
-            self._restore_environment()
+        return self.agentic_flow.run_stream(task=json.dumps(agent_input))