halfway through working out new concept derivation function

Author: Chenglong-MS

py-src/data_formulator/agents/agent_sql_data_rec.py

@@ -6,6 +6,9 @@
 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 random
+import string
+
 import traceback
 
 
@@ -153,13 +156,32 @@ def process_gpt_response(self, input_tables, messages, response):
                 code_str = code_blocks[-1]
 
                 try:
-                    query_output = self.conn.execute(code_str).fetch_df()
+                    random_suffix = ''.join(random.choices(string.ascii_lowercase, k=4))
+                    table_name = f"view_{random_suffix}"
+                    
+                    create_query = f"CREATE VIEW IF NOT EXISTS {table_name} AS {code_str}"
+                    self.conn.execute(create_query)
+                    self.conn.commit()
+
+                    # Check how many rows are in the table
+                    row_count = self.conn.execute(f"SELECT COUNT(*) FROM {table_name}").fetchone()[0]
+                    
+                    # Only limit to 5000 if there are more rows
+                    if row_count > 5000:
+                        query_output = self.conn.execute(f"SELECT * FROM {table_name} LIMIT 5000").fetch_df()
+                    else:
+                        query_output = self.conn.execute(f"SELECT * FROM {table_name}").fetch_df()
+                        self.conn.execute(f"DROP VIEW {table_name}")
 
                     result = {
                         "status": "ok",
                         "code": code_str,
                         "content": {
                             'rows': query_output.to_dict('records'),
+                            'virtual': {
+                                'table_name': table_name,
+                                'row_count': row_count
+                            } if row_count > 5000 else None
                         },
                     }
                 except Exception as e:

py-src/data_formulator/agents/agent_sql_data_transform.py

@@ -2,7 +2,6 @@
 # Licensed under the MIT License.
 
 import json
-import time
 import random
 import string
 
@@ -201,9 +200,8 @@ def process_gpt_sql_response(self, response, messages):
 
                 try:
                     # Generate unique table name directly with timestamp and random suffix
-                    timestamp = int(time.time())
-                    random_suffix = ''.join(random.choices(string.ascii_lowercase, k=5))
-                    table_name = f"result_{timestamp}_{random_suffix}"
+                    random_suffix = ''.join(random.choices(string.ascii_lowercase, k=4))
+                    table_name = f"view_{random_suffix}"
 
                     create_query = f"CREATE VIEW IF NOT EXISTS {table_name} AS {query_str}"
                     self.conn.execute(create_query)

py-src/data_formulator/app.py

@@ -27,6 +27,8 @@
 from vega_datasets import data as vega_data
 
 from data_formulator.agents.agent_concept_derive import ConceptDeriveAgent
+from data_formulator.agents.agent_py_concept_derive import PyConceptDeriveAgent
+
 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_py_data_rec import PythonDataRecAgent
@@ -364,6 +366,31 @@ def derive_concept_request():
     return response
 
 
+@app.route('/api/derive-py-concept', methods=['GET', 'POST'])
+def derive_py_concept():
+
+    if request.is_json:
+        app.logger.info("# code query: ")
+        content = request.get_json()
+        token = content["token"]
+
+        client = get_client(content['model'])
+
+        app.logger.info(f" model: {content['model']}")
+        agent = PyConceptDeriveAgent(client=client)
+
+        #print(content["input_data"])
+
+        results = agent.run(content["input_data"], [f['name'] for f in content["input_fields"]], 
+                                       content["output_name"], content["description"])
+        
+        response = flask.jsonify({ "status": "ok", "token": token, "results": results })
+    else:
+        response = flask.jsonify({ "token": -1, "status": "error", "results": [] })
+
+    response.headers.add('Access-Control-Allow-Origin', '*')
+    return response
+
 @app.route('/api/clean-data', methods=['GET', 'POST'])
 def clean_data_request():
 
@@ -612,22 +639,20 @@ def list_tables():
                 sample_rows = db.execute(f"SELECT * FROM {table_name} LIMIT 1000").fetchall()
 
                 # Check if this is a view or a table
-                is_view = False
                 try:
-                    # In most SQL databases, views are listed in a system table
-                    # For DuckDB, we can check if it's a view by querying the system tables
-                    view_check = db.execute(f"SELECT * FROM duckdb_views() WHERE view_name = '{table_name}'").fetchone()
-                    is_view = view_check is not None
-                except Exception:
+                    # Get both view existence and source in one query
+                    view_info = db.execute(f"SELECT view_name, sql FROM duckdb_views() WHERE view_name = '{table_name}'").fetchone()
+                    view_source = view_info[1] if view_info else None
+                except Exception as e:
                     # If the query fails, assume it's a regular table
-                    pass
+                    view_source = None
 
                 result.append({
                     "name": table_name,
                     "columns": [{"name": col[0], "type": col[1]} for col in columns],
                     "row_count": row_count,
                     "sample_rows": [dict(zip([col[0] for col in columns], row)) for row in sample_rows],
-                    "is_view": is_view,
+                    "view_source": view_source
                 })
 
         return jsonify({
@@ -652,6 +677,8 @@ def sample_table():
         projection_fields = data.get('projection_fields', []) # if empty, we want to include all fields
         method = data.get('method', 'random') # one of 'random', 'head', 'bottom'
         order_by_fields = data.get('order_by_fields', [])
+
+        print(f"sample_table: {table_id}, {sample_size}, {projection_fields}, {method}, {order_by_fields}")
 
         # Validate field names against table columns to prevent SQL injection
         with db_manager.connection(session['session_id']) as db:
@@ -817,7 +844,7 @@ def create_table():
 
 @app.route('/api/tables/delete-table', methods=['POST'])
 def drop_table():
-    """Drop a table"""
+    """Drop a table or view"""
     try:
         data = request.get_json()
         table_name = data.get('table_name')
@@ -826,11 +853,14 @@ def drop_table():
             return jsonify({"status": "error", "message": "No table name provided"}), 400
 
         with db_manager.connection(session['session_id']) as db:
+            # First try to drop it as a view
+            db.execute(f"DROP VIEW IF EXISTS {table_name}")
+            # Then try to drop it as a table
             db.execute(f"DROP TABLE IF EXISTS {table_name}")
 
             return jsonify({
                 "status": "success",
-                "message": f"Table {table_name} dropped"
+                "message": f"Table/view {table_name} dropped"
             })
 
     except Exception as e:

src/app/dfSlice.tsx

@@ -43,6 +43,11 @@ export interface DataFormulatorState {
     testedModels: {id: string, status: 'ok' | 'error' | 'testing' | 'unknown', message: string}[];
 
     tables : DictTable[];
+    extTables: {
+        baseTableRef: string,
+        rows: any[],
+        virtualTableRef?: string, // whether this is a virtual table
+    }[]; // extensions to base tables with derived tables
     charts: Chart[];
 
     activeChallenges: {tableId: string, challenges: { text: string; difficulty: 'easy' | 'medium' | 'hard'; }[]}[];
@@ -79,6 +84,7 @@ const initialState: DataFormulatorState = {
     testedModels: [],
 
     tables: [],
+    extTables: [],
     charts: [],
 
     activeChallenges: [],
@@ -261,6 +267,7 @@ export const dataFormulatorSlice = createSlice({
             state.testedModels = [];
 
             state.tables = [];
+            state.extTables = [];
             state.charts = [];
             state.activeChallenges = [];
 
@@ -291,6 +298,7 @@ export const dataFormulatorSlice = createSlice({
 
             //state.table = undefined;
             state.tables = savedState.tables || [];
+            state.extTables = savedState.extTables || [];
             state.charts = savedState.charts || [];
 
             state.activeChallenges = savedState.activeChallenges || [];
@@ -370,6 +378,37 @@ export const dataFormulatorSlice = createSlice({
         addChallenges: (state, action: PayloadAction<{tableId: string, challenges: { text: string; difficulty: 'easy' | 'medium' | 'hard'; }[]}>) => {
             state.activeChallenges = [...state.activeChallenges, action.payload];
         },
+        setExtTables: (state, action: PayloadAction<{baseTableRef: string, rows: any[], virtualTableRef?: string}>) => {
+            let extTable = action.payload;
+
+            let baseTable = state.tables.find(t => t.id == extTable.baseTableRef) as DictTable;
+            let existingExtTable = state.extTables.find(t => t.baseTableRef == extTable.baseTableRef);
+
+            // we want to extend base rows with the new rows
+            let baseRows = existingExtTable ? existingExtTable.rows : baseTable.rows;
+
+            if (existingExtTable) {
+                for (let i = 0; i < extTable.rows.length; i++) {
+                    extTable.rows[i] = {...extTable.rows[i], ...baseRows[i]};
+                }
+            } 
+
+            if (state.extTables.some(t => t.baseTableRef == extTable.baseTableRef)) {
+                state.extTables = state.extTables.map(t => t.baseTableRef == extTable.baseTableRef ? {...t, rows: extTable.rows} : t);
+            } else {
+                state.extTables = [...state.extTables, extTable];
+            }
+        },
+        removeExtTable: (state, action: PayloadAction<string>) => {
+            let baseTableRef = action.payload;
+            state.extTables = state.extTables.filter(t => t.baseTableRef != baseTableRef);
+        },
+        updateExtTableFieldName: (state, action: PayloadAction<{baseTableRef: string, oldName: string, newName: string}>) => {
+            let baseTableRef = action.payload.baseTableRef;
+            let oldName = action.payload.oldName;
+            let newName = action.payload.newName;
+            state.extTables = state.extTables.map(t => t.baseTableRef == baseTableRef ? {...t, rows: t.rows.map(r => r[oldName] = newName)} : t);
+        },
         createNewChart: (state, action: PayloadAction<{chartType?: string, tableId?: string}>) => {
             let chartType = action.payload.chartType;
             let tableId = action.payload.tableId || state.tables[0].id;
@@ -542,7 +581,14 @@ export const dataFormulatorSlice = createSlice({
                 && Object.entries(chart.encodingMap).some(([channel, encoding]) => encoding.fieldID && conceptID == encoding.fieldID))) {
                 console.log("cannot delete!")
             } else {
-                state.conceptShelfItems = state.conceptShelfItems.filter(field => field.id != conceptID);
+
+                let field = state.conceptShelfItems.find(f => f.id == conceptID);
+                if (field?.source == "derived") {
+                    // delete generated column from the derived table
+                    state.extTables = state.extTables.map(t => t.baseTableRef == field.tableRef ? {...t, rows: t.rows.map(r => delete r[field.name])} : t);
+                }
+                state.conceptShelfItems = state.conceptShelfItems.filter(f => f.id != conceptID);
+
                 for (let chart of state.charts)  {
                     for (let [channel, encoding] of Object.entries(chart.encodingMap)) {
                         if (encoding.fieldID && conceptID == encoding.fieldID) {

src/app/utils.tsx

@@ -34,6 +34,8 @@ export function getUrls() {
 
         // these functions involves openai models
         DERIVE_CONCEPT_URL: `/api/derive-concept-request`,
+        DERIVE_PY_CONCEPT: `/api/derive-py-concept`,
+
         SORT_DATA_URL: `/api/codex-sort-request`,
         CLEAN_DATA_URL: `/api/clean-data`,
 
@@ -127,7 +129,45 @@ export function runCodeOnInputListsInVM(
     return ioPairs;
 }
 
-export function baseTableToExtTable(table: any[], derivedFields: FieldItem[], allFields: FieldItem[]) {
+export function prepVisTable(table: any[], derivedFields: FieldItem[], allFields: FieldItem[], encodingMap: EncodingMap) {
+
+    let binningFields = []
+    let aggregateFields = []
+    let groupByFields = []
+
+    for (const [channel, encoding] of Object.entries(encodingMap)) {
+        const field = encoding.fieldID ? _.find(allFields, (f) => f.id === encoding.fieldID) : undefined;
+        if (field) {
+            if (encoding.bin) {
+                binningFields.push(field.name);
+            }
+        }
+        if (encoding.aggregate) {
+            aggregateFields.push([field?.name, encoding.aggregate]);
+        } else {
+            if (field) {
+                groupByFields.push(field.name);
+            }
+        }
+    }
+
+    ///// TODOTODO: 
+
+    return [];
+}
+
+
+
+export function baseTableToExtTable(table: DictTable, extTable?: {rows: any[], baseTableRef: string, virtualTableRef?: string}) {
+    // derive fields from derivedFields from the original table
+    if (extTable) {
+        return structuredClone(extTable.rows);
+    } else {
+        return structuredClone(table.rows);
+    }
+}
+
+export function baseTableToExtTableOld(table: any[], derivedFields: FieldItem[], allFields: FieldItem[]) {
     // derive fields from derivedFields from the original table
 
     if (table.length == 0) {