supporting saving intermediate tables as virtual tables

Author: Chenglong-MS

py-src/data_formulator/agents/agent_data_filter.py

@@ -187,7 +187,9 @@ def process_gpt_response(self, input_table, output_field, response, messages):
 
                     if result['status'] == 'ok':
                         new_data = json.loads(result['content'])
-                        result['content'] = new_data
+                        result['content'] = {
+                            'rows': new_data,
+                        }
                     else:
                         logger.info(result['content'])
                     result['code'] = code_str

py-src/data_formulator/agents/agent_generic_py_concept.py

@@ -181,7 +181,9 @@ def derive({arg_string}):
 
                     if result['status'] == 'ok':
                         new_data = json.loads(result['content'])
-                        result['content'] = new_data
+                        result['content'] = {
+                            'rows': new_data,
+                        }
                     else:
                         print(result['content'])
                     result['code'] = code_str

py-src/data_formulator/agents/agent_py_concept_derive.py

@@ -164,7 +164,9 @@ def process_gpt_response(self, input_tables, messages, response):
                     result['code'] = code_str
 
                     if result['status'] == 'ok':
-                        result['content'] = json.loads(result['content'])
+                        result['content'] = {
+                            'rows': json.loads(result['content']),
+                        }
                     else:
                         logger.info(result['content'])
                 except Exception as e:

py-src/data_formulator/agents/agent_py_data_rec.py

@@ -221,7 +221,9 @@ def process_gpt_response(self, input_tables, messages, response):
 
                     if result['status'] == 'ok':
                         # parse the content
-                        result['content'] = json.loads(result['content'])
+                        result['content'] = {
+                            'rows': json.loads(result['content']),
+                        }
                     else:
                         logger.info(result['content'])
                 except Exception as e:

py-src/data_formulator/agents/agent_py_data_transform.py

@@ -6,8 +6,6 @@
 from data_formulator.agents.agent_utils import extract_json_objects, extract_code_from_gpt_response
 from data_formulator.agents.agent_sql_data_transform import get_sql_table_statistics_str
 
-import data_formulator.py_sandbox as py_sandbox
-
 import traceback
 
 
@@ -160,7 +158,9 @@ def process_gpt_response(self, input_tables, messages, response):
                     result = {
                         "status": "ok",
                         "code": code_str,
-                        "content": query_output.to_dict('records'),
+                        "content": {
+                            'rows': query_output.to_dict('records'),
+                        },
                     }
                 except Exception as e:
                     logger.warning('other error:')