various fixes

Author: Chenglong-MS

py-src/data_formulator/agents/agent_code_explanation.py

@@ -1,8 +1,7 @@
 # Copyright (c) Microsoft Corporation.
 # Licensed under the MIT License.
 
-import pandas as pd
-from data_formulator.agents.agent_utils import generate_data_summary, extract_code_from_gpt_response
+from data_formulator.agents.agent_utils import generate_data_summary
 
 import logging
 

py-src/data_formulator/agents/agent_py_concept_derive.py

@@ -1,7 +1,6 @@
 # Copyright (c) Microsoft Corporation.
 # Licensed under the MIT License.
 
-import json
 import time
 
 from data_formulator.agents.agent_utils import generate_data_summary, extract_code_from_gpt_response
@@ -10,7 +9,6 @@
 import traceback
 
 import logging
-import datetime
 
 logger = logging.getLogger(__name__)
 

py-src/data_formulator/agents/agent_py_data_transform.py

@@ -2,7 +2,6 @@
 # Licensed under the MIT License.
 
 import json
-import sys
 
 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

py-src/data_formulator/agents/agent_query_completion.py

@@ -1,10 +1,9 @@
 # Copyright (c) Microsoft Corporation.
 # Licensed under the MIT License.
 
-import pandas as pd
 import json
 
-from data_formulator.agents.agent_utils import extract_code_from_gpt_response, extract_json_objects
+from data_formulator.agents.agent_utils import extract_json_objects
 import re
 import logging
 

py-src/data_formulator/agents/agent_utils.py

@@ -6,10 +6,6 @@
 import pandas as pd
 import numpy as np
 
-import base64
-
-from pprint import pprint
-
 import re
 
 def string_to_py_varname(var_str): 

py-src/data_formulator/agents/client_utils.py

@@ -1,4 +1,3 @@
-import os
 import litellm
 import openai
 from azure.identity import DefaultAzureCredential, get_bearer_token_provider

py-src/data_formulator/data_loader/kusto_data_loader.py

@@ -51,41 +51,6 @@ def query(self, kql: str) -> pd.DataFrame:
         return dataframe_from_result_table(result.primary_results[0])
 
     def list_tables(self) -> List[Dict[str, Any]]:
-        # first list functions (views)
-        query = ".show functions"
-        function_result_df = self.query(query)
-
-        functions = []
-        for func in function_result_df.to_dict(orient="records"):
-            func_name = func['Name']
-            result = self.query(f".show function ['{func_name}'] schema as json").to_dict(orient="records")
-            schema = json.loads(result[0]['Schema'])
-            parameters = schema['InputParameters']
-            columns = [{
-                'name': r["Name"],
-                'type': r["Type"]
-            } for r in schema['OutputColumns']]
-
-            # skip functions with parameters at the moment
-            if len(parameters) > 0:
-                continue
-
-            sample_query = f"['{func_name}'] | take {10}"
-            sample_result = self.query(sample_query).to_dict(orient="records")
-        
-            function_metadata = {
-                "row_count": 0,
-                "columns": columns,
-                "parameters": parameters,
-                "sample_rows": sample_result
-            }
-            functions.append({
-                "type": "function",
-                "name": func_name,
-                "metadata": function_metadata
-            })
-
-        # then list tables
         query = ".show tables"
         tables_df = self.query(query)
 
@@ -101,8 +66,8 @@ def list_tables(self) -> List[Dict[str, Any]]:
             row_count_result = self.query(f".show table ['{table_name}'] details").to_dict(orient="records")
             row_count = row_count_result[0]["TotalRowCount"]
 
-            sample_query = f"['{table_name}'] | take {10}"
-            sample_result = self.query(sample_query).to_dict(orient="records")
+            sample_query = f"['{table_name}'] | take {5}"
+            sample_result = json.loads(self.query(sample_query).to_json(orient="records"))
 
             table_metadata = {
                 "row_count": row_count,
@@ -116,7 +81,7 @@ def list_tables(self) -> List[Dict[str, Any]]:
                 "metadata": table_metadata
             })
 
-        return functions + tables
+        return tables
 
     def ingest_data(self, table_name: str, name_as: str = None, size: int = 5000000) -> pd.DataFrame:
         if name_as is None:
@@ -167,7 +132,7 @@ def ingest_data(self, table_name: str, name_as: str = None, size: int = 5000000)
             total_rows_ingested += len(chunk_df)
 
     def view_query_sample(self, query: str) -> str:
-        return self.query(query).head(10).to_dict(orient="records")
+        return json.loads(self.query(query).head(10).to_json(orient="records"))
 
     def ingest_data_from_query(self, query: str, name_as: str) -> pd.DataFrame:
         # Sanitize the table name for SQL compatibility

py-src/data_formulator/data_loader/mysql_data_loader.py

@@ -93,7 +93,7 @@ def ingest_data(self, table_name: str, name_as: str | None = None, size: int = 1
         """)
 
     def view_query_sample(self, query: str) -> str:
-        return self.duck_db_conn.execute(query).df().head(10).to_dict(orient="records")
+        return json.loads(self.duck_db_conn.execute(query).df().head(10).to_json(orient="records"))
 
     def ingest_data_from_query(self, query: str, name_as: str) -> pd.DataFrame:
         # Execute the query and get results as a DataFrame

pyproject.toml

@@ -21,7 +21,6 @@ classifiers = [
 ]
 
 dependencies = [  
-    "autopep8",  
     "jupyter",  
     "pandas",  
     "docker",  
@@ -31,6 +30,7 @@ dependencies = [
     "flask-cors",  
     "openai",  
     "azure-identity",  
+    "azure-kusto-data",
     "azure-keyvault-secrets",  
     "python-dotenv",  
     "vega_datasets",

requirements.txt

@@ -1,4 +1,3 @@
-autopep8
 jupyter
 pandas
 docker
@@ -7,6 +6,7 @@ matplotlib
 flask
 openai
 azure-identity
+azure-kusto-data
 azure-keyvault-secrets
 python-dotenv
 vega_datasets