updating db support, if you are reading here, hello

Author: Chenglong-MS

py-src/data_formulator/agents/__init__.py

@@ -3,19 +3,21 @@
 
 from data_formulator.agents.agent_concept_derive import ConceptDeriveAgent
 from data_formulator.agents.agent_py_concept_derive import PyConceptDeriveAgent
-from data_formulator.agents.agent_data_transformation import DataTransformationAgent
-from data_formulator.agents.agent_data_transform_v2 import DataTransformationAgentV2
+from data_formulator.agents.agent_py_data_transform import PythonDataTransformationAgent
+from data_formulator.agents.agent_sql_data_transform import SQLDataTransformationAgent
 from data_formulator.agents.agent_data_load import DataLoadAgent
 from data_formulator.agents.agent_sort_data import SortDataAgent
 from data_formulator.agents.agent_data_clean import DataCleanAgent
-from data_formulator.agents.agent_data_rec import DataRecAgent
+from data_formulator.agents.agent_py_data_rec import PythonDataRecAgent
+from data_formulator.agents.agent_sql_data_rec import SQLDataRecAgent
 
 __all__ = [
     "ConceptDeriveAgent",
     "PyConceptDeriveAgent",
-    "DataTransformationAgent",
-    "DataTransformationAgentV2",
-    "DataRecAgent",
+    "PythonDataTransformationAgent",
+    "SQLDataTransformationAgent",
+    "PythonDataRecAgent",
+    "SQLDataRecAgent",
     "DataLoadAgent",
     "SortDataAgent",
     "DataCleanAgent"

py-src/data_formulator/agents/agent_data_transformation.py

@@ -4,7 +4,6 @@
 import json
 
 from data_formulator.agents.agent_utils import extract_json_objects, generate_data_summary, extract_code_from_gpt_response
-from data_formulator.agents.agent_data_transform_v2 import completion_response_wrapper
 
 import data_formulator.py_sandbox as py_sandbox
 
@@ -128,7 +127,7 @@ def transform_data(df):
 ```
 """
 
-class DataRecAgent(object):
+class PythonDataRecAgent(object):
 
     def __init__(self, client, system_prompt=None):
         self.client = client
@@ -177,7 +176,7 @@ def process_gpt_response(self, input_tables, messages, response):
                 result = {'status': 'error', 'code': "", 'content': "No code block found in the response. The model is unable to generate code to complete the task."}
 
             result['dialog'] = [*messages, {"role": choice.message.role, "content": choice.message.content}]
-            result['agent'] = 'DataRecAgent'
+            result['agent'] = 'PythonDataRecAgent'
             result['refined_goal'] = refined_goal
             candidates.append(result)
 
@@ -203,7 +202,7 @@ def run(self, input_tables, description, n=1):
         messages = [{"role":"system", "content": self.system_prompt},
                     {"role":"user","content": user_query}]
 
-        response = completion_response_wrapper(self.client, messages, n)
+        response = self.client.get_completion(messages = messages)
 
         return self.process_gpt_response(input_tables, messages, response)
 
@@ -215,6 +214,6 @@ def followup(self, input_tables, dialog, new_instruction: str, n=1):
 
         messages = [*dialog, {"role":"user", "content": f"Update: \n\n{new_instruction}"}]
 
-        response = completion_response_wrapper(self.client, messages, n)
+        response = self.client.get_completion(messages = messages)
 
         return self.process_gpt_response(input_tables, messages, response)

py-src/data_formulator/agents/agent_data_rec.py renamed to py-src/data_formulator/agents/agent_py_data_rec.py

@@ -7,8 +7,6 @@
 from data_formulator.agents.agent_utils import extract_json_objects, generate_data_summary, extract_code_from_gpt_response
 import data_formulator.py_sandbox as py_sandbox
 
-import traceback
-
 import logging
 
 # Replace/update the logger configuration
@@ -184,17 +182,7 @@ def transform_data(df):
 ```
 '''
 
-def completion_response_wrapper(client, messages, n):
-    ### wrapper for completion response, especially handling errors
-    try:
-        response = client.get_completion(messages = messages)
-    except Exception as e:
-        response = e
-
-    return response
-
-
-class DataTransformationAgentV2(object):
+class PythonDataTransformationAgent(object):
 
     def __init__(self, client, system_prompt=None):
         self.client = client
@@ -245,7 +233,7 @@ def process_gpt_response(self, input_tables, messages, response):
                 result = {'status': 'error', 'code': "", 'content': "No code block found in the response. The model is unable to generate code to complete the task."}
 
             result['dialog'] = [*messages, {"role": choice.message.role, "content": choice.message.content}]
-            result['agent'] = 'DataTransformationAgent'
+            result['agent'] = 'PythonDataTransformationAgent'
             result['refined_goal'] = refined_goal
             candidates.append(result)
 
@@ -286,7 +274,7 @@ def run(self, input_tables, description, expected_fields: list[str], prev_messag
                     *prev_messages,
                     {"role":"user","content": user_query}]
 
-        response = completion_response_wrapper(self.client, messages, n)
+        response = self.client.get_completion(messages = messages)
 
         return self.process_gpt_response(input_tables, messages, response)
 
@@ -308,6 +296,6 @@ def followup(self, input_tables, dialog, output_fields: list[str], new_instructi
         messages = [*updated_dialog, {"role":"user", 
                               "content": f"Update the code above based on the following instruction:\n\n{json.dumps(goal, indent=4)}"}]
 
-        response = completion_response_wrapper(self.client, messages, n)
+        response = self.client.get_completion(messages = messages)
 
         return self.process_gpt_response(input_tables, messages, response)