2. Dreu - implement ai functions

app/__init__.py

@@ -30,6 +30,7 @@ def format(self, record):
 _werkzeug_handler = logging.StreamHandler()
 _werkzeug_handler.setFormatter(_WhiteFormatter('%(message)s'))
 logging.getLogger('werkzeug').handlers = [_werkzeug_handler]
+logger = logging.getLogger('werkzeug')
 
 from psycopg2.extensions import ISOLATION_LEVEL_AUTOCOMMIT
 
@@ -107,8 +108,6 @@ def load_database_info():
 # Create the database if it does not exist
 create_database_if_not_exists(connection, db_name)
 
-
-
 """
  then we use SQLAlchemy (create_engine) for everything else
  this is the engine that gets passed around to fetch_sql, execute_sql, 

app/ablation_study/ablation_study.py

@@ -0,0 +1,127 @@
+import os
+
+from dotenv import load_dotenv
+import datetime
+
+
+from app import app as app_module
+from app import db_operations, engine
+
+import json
+
+from app.db_utils.ai_utils import parse_json_response
+from app.server_utils.logger_utils import initialize_action_log
+
+datasets_paths = ['provided_datasets/mari_dataset.csv']
+#models = [{"model": "qwen/qwen3.6-27b", "provider": "groq"}]
+models = [{"model": "openai/gpt-oss-20b", "provider": "groq"}]
+
+def variant(name, **overrides):
+    baseline = {
+        "name": "baseline",
+        "include_error_log": True,
+        "include_data_profile": True,
+
+        "include_action_log": True,
+        "action_log_limit": 10,
+        "include_full_dataset": True
+    }
+
+    # gets the baseline and overrides the key(s) in the overrides variable
+    config = {**baseline, **overrides}
+    config["name"] = name
+    return config
+
+
+def is_stop_action(action_name):
+    if action_name == "stop":
+        return True
+    else:
+        return False
+
+
+if __name__ == "__main__":
+    load_dotenv()
+
+    # TODO: double check this
+    ablation_configs = [
+        variant("baseline"),
+        variant("no_error_log", include_error_log=False),
+        variant("no_data_profile", include_data_profile=False),
+        variant("no_action_log", include_action_log=False),
+        variant("no_full_dataset", include_full_dataset=False),
+        variant("no_action_log_limit", action_log_limit=None) # Includes full action log
+    ]
+
+    client = app_module.test_client()
+    app_module.testing = False
+
+    results = []
+
+    for model_dict in models:
+
+        model = model_dict["model"]
+        model_provider = model_dict["provider"]
+
+        update_settings_table_result = client.post('/api/ai_helper/update_settings_table', json={"model_name": model, "provider": model_provider})
+        data = update_settings_table_result.get_json()
+        assert data["success"] == True
+
+        for config in ablation_configs:
+            # reset globals living in the app package namespace
+            app_module.wrangle_occurred = False
+
+            # reset attributes on the Flask object itself
+            app_module.pgraph_for_session = None
+
+            # reset your stateful class instance
+            db_operations.reset()
+            for dataset in datasets_paths:
+                print(f"Running config {config} wth model {model} and dataset {dataset}")
+                initialize_action_log(engine, reset_log=True)
+
+                with open(dataset, 'rb') as f:
+                    upload_result = client.post('/api/upload',  data={'file': (f, dataset)},
+        content_type='multipart/form-data')
+                data = upload_result.get_json()
+                assert data["success"] == True
+
+                result_dict = {}  # Dict added to the result json
+                result_dict["model"] = model_dict
+                result_dict["config_name"] = config["name"]
+                result_dict["dataset"] = dataset
+                result_dict["actions"] = []
+
+                action_plan_batch = 0
+
+                stop_action_found = False
+
+                while not stop_action_found:
+                    action_plan_batch += 1
+
+                    actions = None
+
+                    # Get actions
+                    get_action_plan_result = client.post('/api/ai_helper/get_action_plan')
+                    action_plan_result_json = get_action_plan_result.get_json()
+                    assert action_plan_result_json["success"] == True
+                    action_plan_json = action_plan_result_json["json_action_plan"]
+
+                    # Go through all actions and perform each of them
+                    for action_dict in action_plan_json:
+
+                        if is_stop_action(action_dict["action_name"]):
+                            stop_action_found = True
+
+                        action_result = client.post('/api/ai_helper/perform_llm_action', json=action_dict)
+                        action_dict["action_plan_batch"] = action_plan_batch
+                        action_dict["timestamp"] = datetime.datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
+                        action_dict["success"] = action_result.get_json()
+
+                        results.append(action_dict)
+
+
+    with open(f'ablation_results_{datetime.datetime.now().strftime("%Y-%m-%d_%H-%M-%S")}.json', 'w') as outfile:
+        json.dump(results, outfile)
+
+

app/db_utils/ai_utils.py

@@ -0,0 +1,97 @@
+import os
+from app.db_utils.execute_sql import copy_table_to_csv
+from app import logger
+import json
+import ast
+
+import random
+import time
+
+
+# Overwriting the csv every time the LLM needs it to be updated; we don't really need to save the old ones
+
+def update_csvs_for_llm(error_table_name, data_profile_name, action_log_name, full_dataset_name):
+    print(f"UPDATE CSVS FOR LLM PATHS error_table_name: {error_table_name}, data_profile_name: {data_profile_name} action_log_name: {action_log_name}")
+
+    action_log_csv_path = "action_log.csv"
+    error_log_csv_path = "error_log.csv"
+    data_profile_csv_path = "data_profile.csv"
+    full_dataset_csv_path = "full_dataset.csv"
+
+    _THIS_DIR = os.path.dirname(os.path.abspath(__file__))
+    FILES_FOR_LLM_PATH = os.path.abspath(os.path.join(_THIS_DIR, '..', 'files_for_llm'))
+
+    action_log_csv_path = FILES_FOR_LLM_PATH + '/' + f'{action_log_csv_path}'
+    error_log_csv_path = FILES_FOR_LLM_PATH + '/' + f'{error_log_csv_path}'
+    data_profile_csv_path = FILES_FOR_LLM_PATH + '/' + f'{data_profile_csv_path}'
+    full_dataset_csv_path = FILES_FOR_LLM_PATH + '/' + f'{full_dataset_csv_path}'
+
+
+    table_name_tuple_list =  [(action_log_name, action_log_csv_path),
+                              (error_table_name, error_log_csv_path),
+                              (data_profile_name, data_profile_csv_path),
+                              (full_dataset_name, full_dataset_csv_path)]
+
+    write_tables_to_csv(table_name_tuple_list)
+
+    return (error_log_csv_path, data_profile_csv_path, action_log_csv_path, full_dataset_csv_path)
+
+
+
+def write_tables_to_csv(table_name_tuple_list):
+    from app import engine
+    for (table_name, csv_path) in table_name_tuple_list:
+
+        # Clear the existing CSV file if it exists
+        if os.path.exists(csv_path):
+            os.remove(csv_path)
+
+        copy_table_to_csv(table_name, csv_path, engine)
+
+
+def parse_json_response(llm_json_response):
+    try:
+        return json.loads(llm_json_response)
+    except json.JSONDecodeError:
+        try:
+            return ast.literal_eval(llm_json_response)
+        except (ValueError, SyntaxError) as e:
+            logger.exception(f"Could not parse response as dict or JSON")
+
+
+def call_with_retry(function, func_args, max_tries=5):
+    for attempt in range(max_tries):
+        try:
+            result = function(*func_args)
+
+            return result
+        except Exception:
+            logger.exception("Error occurred while calling LLM function")
+            if attempt == max_tries - 1:
+                raise
+
+            # TODO: is this okay
+            delay = 10
+            delay *= random.uniform(0.5, 1.5)  # jitter
+            time.sleep(delay)
+
+
+
+def get_api_key(provider):
+    key_map = {
+        "openai": "OPENAI_API_KEY",
+        "anthropic": "ANTHROPIC_API_KEY",
+        "groq": "GROQ_API_KEY"
+    }
+
+    key = os.environ.get(key_map[provider])
+    if key is None:
+        raise ValueError(f"Could not find API key for {provider}")
+    return key
+
+
+
+
+
+
+