[deploy] Merge pull request #150 from microsoft/dev

Dev

Author: Chenglong-MS

py-src/data_formulator/agent_routes.py

@@ -181,8 +181,10 @@ def process_data_on_load_request():
         client = get_client(content['model'])
 
         logger.info(f" model: {content['model']}")
+
+        conn = db_manager.get_connection(session['session_id'])
+        agent = DataLoadAgent(client=client, conn=conn)
 
-        agent = DataLoadAgent(client=client)
         candidates = agent.run(content["input_data"])
 
         candidates = [c['content'] for c in candidates if c['status'] == 'ok']

py-src/data_formulator/agents/agent_data_load.py

@@ -4,6 +4,8 @@
 import json
 
 from data_formulator.agents.agent_utils import extract_json_objects, generate_data_summary
+from data_formulator.agents.agent_sql_data_transform import  sanitize_table_name, get_sql_table_statistics_str
+
 import logging
 
 logger = logging.getLogger(__name__)
@@ -124,12 +126,18 @@
 
 class DataLoadAgent(object):
 
-    def __init__(self, client):
+    def __init__(self, client, conn):
         self.client = client
+        self.conn = conn
 
     def run(self, input_data, n=1):
 
-        data_summary = generate_data_summary([input_data], include_data_samples=True, field_sample_size=30)
+        if input_data['virtual']:
+            table_name = sanitize_table_name(input_data['name'])
+            table_summary_str = get_sql_table_statistics_str(self.conn, table_name, row_sample_size=5, field_sample_size=30)
+            data_summary = f"[TABLE {table_name}]\n\n{table_summary_str}"
+        else:
+            data_summary = generate_data_summary([input_data], include_data_samples=True, field_sample_size=30)
 
         user_query = f"[DATA]\n\n{data_summary}\n\n[OUTPUT]"
 

py-src/data_formulator/agents/agent_sql_data_transform.py

@@ -300,20 +300,24 @@ def followup(self, input_tables, dialog, output_fields: list[str], new_instructi
         return self.process_gpt_sql_response(response, messages)
 
 
-def get_sql_table_statistics_str(conn, table_name: str) -> str:
+def get_sql_table_statistics_str(conn, table_name: str, 
+        row_sample_size: int = 5, # number of rows to be sampled in the sample data part
+        field_sample_size: int = 7, # number of example values for each field to be sampled
+        max_val_chars: int = 140 # max number of characters to be shown for each example value
+    ) -> str:
     """Get a string representation of the table statistics"""
 
     table_name = sanitize_table_name(table_name)
 
     # Get column information
     columns = conn.execute(f"DESCRIBE {table_name}").fetchall()
-    sample_data = conn.execute(f"SELECT * FROM {table_name} LIMIT 5").fetchall()
+    sample_data = conn.execute(f"SELECT * FROM {table_name} LIMIT {row_sample_size}").fetchall()
 
     # Format sample data as pipe-separated string
     col_names = [col[0] for col in columns]
     formatted_sample_data = "| " + " | ".join(col_names) + " |\n"
     for i, row in enumerate(sample_data):
-        formatted_sample_data += f"{i}| " + " | ".join(str(val) for val in row) + " |\n"
+        formatted_sample_data += f"{i}| " + " | ".join(str(val)[:max_val_chars]+ "..." if len(str(val)) > max_val_chars else str(val) for val in row) + " |\n"
 
     col_metadata_list = []
     for col in columns:
@@ -364,12 +368,12 @@ def get_sql_table_statistics_str(conn, table_name: str) -> str:
             (SELECT DISTINCT {quoted_col_name}
                 FROM {table_name} 
                 WHERE {quoted_col_name} IS NOT NULL 
-                LIMIT 5)
+                LIMIT {field_sample_size})
             """
 
             sample_values = conn.execute(query_for_sample_values).fetchall()
 
-            stats_dict['sample_values'] = sample_values
+            stats_dict['sample_values'] = [str(val)[:max_val_chars]+ "..." if len(str(val)) > max_val_chars else str(val) for val in sample_values]
 
         col_metadata_list.append({
             "column": col_name,

py-src/data_formulator/agents/agent_utils.py

@@ -181,7 +181,7 @@ def dedup_data_transform_candidates(candidates):
     return [items[0] for _, items in candidate_groups.items()]
 
 
-def get_field_summary(field_name, df, field_sample_size):
+def get_field_summary(field_name, df, field_sample_size, max_val_chars=100):
     try:
         values = sorted([x for x in list(set(df[field_name].values)) if x != None])
     except:
@@ -196,11 +196,22 @@ def get_field_summary(field_name, df, field_sample_size):
     else:
         val_sample = values[:int(sample_size / 2)] + ["..."] + values[-(sample_size - int(sample_size / 2)):]
 
-    val_str = ', '.join([str(s) if ',' not in str(s) else f'"{str(s)}"' for s in val_sample])
+    def sample_val_cap(val):
+        if len(str(val)) > max_val_chars:
+            s = str(val)[:max_val_chars] + "..."
+        else:
+            s = str(val)
+
+        if ',' in s:
+            s = f'"{s}"'
+
+        return s
+
+    val_str = ', '.join([sample_val_cap(str(s)) for s in val_sample])
 
     return f"{field_name} -- type: {df[field_name].dtype}, values: {val_str}"
 
-def generate_data_summary(input_tables, include_data_samples=True, field_sample_size=7):
+def generate_data_summary(input_tables, include_data_samples=True, field_sample_size=7, max_val_chars=140):
 
     input_table_names = [f'{string_to_py_varname(t["name"])}' for t in input_tables]
 
@@ -209,7 +220,7 @@ def generate_data_summary(input_tables, include_data_samples=True, field_sample_
     field_summaries = []
     for input_data in input_tables:
         df = pd.DataFrame(input_data['rows'])
-        s = '\n\t'.join([get_field_summary(fname, df, field_sample_size)  for fname in list(df.columns.values)])
+        s = '\n\t'.join([get_field_summary(fname, df, field_sample_size, max_val_chars)  for fname in list(df.columns.values)])
         field_summaries.append(s)
 
     table_field_summaries = [f'table_{i} ({input_table_names[i]}) fields:\n\t{s}' for i, s in enumerate(field_summaries)]

py-src/data_formulator/agents/client_utils.py

@@ -64,28 +64,12 @@ def get_completion(self, messages):
         # Configure LiteLLM 
 
         if self.endpoint == "openai":
-
-            print("--------------------------------")
-            print(f"self.params: {self.params}")
-            print(f"self.model: {self.model}")
-            print(f"self.endpoint: {self.endpoint}")
-            print(f"self.params['api_key']: {self.params.get('api_key', 'None')}")
-            print(f"self.params['api_base']: {self.params.get('api_base', 'None')}")
-            print(f"self.params['api_version']: {self.params.get('api_version', 'None')}")
-            print("--------------------------------")
-
-
             client = openai.OpenAI(
-                base_url=self.params.get("api_base", 'placeholder'),
-                api_key=self.params.get("api_key", 'placeholder'),
+                base_url=self.params.get("api_base", None),
+                api_key=self.params.get("api_key", ""),
                 timeout=120
             )
 
-
-            print("--------------------------------")
-            print(f"client: {client}")
-            print("--------------------------------")
-
             completion_params = {
                 "model": self.model,
                 "messages": messages,

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "data_formulator"
-version = "0.2.1"
+version = "0.2.1.1"
 
 requires-python = ">=3.9"
 authors = [

src/app/dfSlice.tsx

@@ -163,7 +163,7 @@ export const fetchFieldSemanticType = createAsyncThunk(
             headers: { 'Content-Type': 'application/json', },
             body: JSON.stringify({
                 token: Date.now(),
-                input_data: {name: table.id, rows: table.rows},
+                input_data: {name: table.id, rows: table.rows, virtual: table.virtual ? true : false},
                 model: dfSelectors.getActiveModel(state)
             }),
         };

src/app/utils.tsx

@@ -319,10 +319,14 @@ export const assembleVegaChart = (
                     let sortedValues = JSON.parse(encoding.sortBy)['values'];
                     encodingObj['sort'] = sortOrder == "ascending" ? sortedValues : sortedValues.reverse();
 
-                    // // special hack: ensure stack bar and stacked area charts are ordered correctly
-                    // if (channel == 'color' && (vgObj['mark'] == 'bar' || vgObj['mark'] == 'area')) {
-                    //     vgObj['encoding']['order'] = {'values': sortedValues};
-                    // }
+                    // special hack: ensure stack bar and stacked area charts are ordered correctly
+                    if (channel == 'color' && (vgObj['mark'] == 'bar' || vgObj['mark'] == 'area')) {
+                        // this is a very interesting hack, it leverages the hidden derived field name used in compiled Vega script to 
+                        // handle order of stack bar and stacked area charts
+                        vgObj['encoding']['order'] = {
+                            "field": `color_${field?.name}_sort_index`,
+                        }
+                    }
                 } catch {
                     console.warn(`sort error > ${encoding.sortBy}`)
                 }

src/views/DBTableManager.tsx

@@ -53,6 +53,8 @@ import { useDispatch, useSelector } from 'react-redux';
 import { dfActions } from '../app/dfSlice';
 import { alpha } from '@mui/material';
 import { DataFormulatorState } from '../app/dfSlice';
+import { fetchFieldSemanticType } from '../app/dfSlice';
+import { AppDispatch } from '../app/store';
 
 export const handleDBDownload = async (sessionId: string) => {
     try {
@@ -253,7 +255,7 @@ export const DBTableManager: React.FC = () => {
 
 export const DBTableSelectionDialog: React.FC<{ buttonElement: any }> = function DBTableSelectionDialog({ buttonElement }) {
 
-    const dispatch = useDispatch();
+    const dispatch = useDispatch<AppDispatch>();
     const sessionId = useSelector((state: DataFormulatorState) => state.sessionId);
 
     const [tableDialogOpen, setTableDialogOpen] = useState<boolean>(false);
@@ -471,6 +473,7 @@ export const DBTableSelectionDialog: React.FC<{ buttonElement: any }> = function
             anchored: true, // by default, db tables are anchored
         }
        dispatch(dfActions.loadTable(table));
+       dispatch(fetchFieldSemanticType(table));
        setTableDialogOpen(false);
     }
 

src/views/VisualizationView.tsx

@@ -463,10 +463,6 @@ export const ChartEditorFC: FC<{  cachedCandidates: DictTable[],
             return renderTableChart(chart, conceptShelfItems, visTableRows);
         }
 
-        console.log('assembled chart');
-        console.log(chart.chartType);
-        console.log(chart.encodingMap);
-        console.log(visTableRows.slice(0, 10));
 
         let element = <></>;
         if (!chart || !checkChartAvailabilityOnPreparedData(chart, conceptShelfItems, visTableRows)) {
@@ -478,8 +474,6 @@ export const ChartEditorFC: FC<{  cachedCandidates: DictTable[],
         element = <Box id={id} key={`focused-chart`} ></Box>    
 
         let assembledChart = assembleVegaChart(chart.chartType, chart.encodingMap, conceptShelfItems, visTableRows, 48, true);
-        console.log('assembled chart');
-        console.log(assembledChart);
 
         assembledChart['resize'] = true;
         assembledChart['config'] = {