various data-loader templates

Author: Chenglong-MS

py-src/data_formulator/agents/agent_py_data_rec.py

@@ -165,7 +165,7 @@ def process_gpt_response(self, input_tables, messages, response):
                     if result['status'] == 'ok':
                         result_df = result['content']
                         result['content'] = {
-                            'rows': result_df.to_dict(orient='records'),
+                            'rows': json.loads(result_df.to_json(orient='records')),
                         }
                     else:
                         logger.info(result['content'])

py-src/data_formulator/agents/agent_py_data_transform.py

@@ -221,13 +221,11 @@ def process_gpt_response(self, input_tables, messages, response):
                     result = py_sandbox.run_transform_in_sandbox2020(code_str, [pd.DataFrame.from_records(t['rows']) for t in input_tables], self.exec_python_in_subprocess)
                     result['code'] = code_str
 
-                    print(f"result: {result}")
-
                     if result['status'] == 'ok':
                         # parse the content
                         result_df = result['content']
                         result['content'] = {
-                            'rows': result_df.to_dict(orient='records'),
+                            'rows': json.loads(result_df.to_json(orient='records')),
                         }
                     else:
                         logger.info(result['content'])

py-src/data_formulator/db_manager.py

@@ -1,8 +1,6 @@
 import duckdb
 import pandas as pd
-from typing import Optional, Dict, List, ContextManager, Any, Tuple
-import time
-from flask import session
+from typing import Dict
 import tempfile
 import os
 from contextlib import contextmanager
@@ -15,7 +13,7 @@ def __init__(self, local_db_dir: str):
         self._local_db_dir: str = local_db_dir
 
     @contextmanager
-    def connection(self, session_id: str) -> ContextManager[duckdb.DuckDBPyConnection]:
+    def connection(self, session_id: str):
         """Get a DuckDB connection as a context manager that will be closed when exiting the context"""
         conn = None
         try: