manage_schema.py

import pandas as pd 
import numpy as np 
import json 
from llm_config import generate_llm_response 
from llm_prompts import DETERMINE_DTYPE_PROMPT 
 
SAMPLE_SIZE = 200 
 
 
def determine_column_type(df, column): 
 sample = df[column].sample(n=min(SAMPLE_SIZE, len(df)), random_state=42).tolist() 
 prompt = DETERMINE_DTYPE_PROMPT.format(sample_values=str(sample)) 
 response = generate_llm_response(prompt) 
 
 try: 
 result = json.loads(response) 
 return result['column_type'], result['invalid_indices'] 
 except (json.JSONDecodeError, KeyError): 
 print(f"Error parsing LLM response for column {column}") 
 return 'string', [] 
 
 
def enforce_column_type(df, column, column_type, invalid_indices): 
 if column_type == 'float': 
 df[column] = pd.to_numeric(df[column], errors='coerce') 
 elif column_type == 'integer': 
 df[column] = pd.to_numeric(df[column], errors='coerce').astype('Int64') 
 elif column_type == 'date': 
 df[column] = pd.to_datetime(df[column], errors='coerce') 
 
 # Set invalid values to NaN 
 df.loc[invalid_indices, column] = np.nan 
 
 return df 
 
 
def process_dataframe(df): 
 print("Determining and enforcing column data types...") 
 
 for column in df.columns: 
 print(f"\nProcessing column: {column}") 
 column_type, invalid_indices = determine_column_type(df, column) 
 print(f" Detected type: {column_type}") 
 print(f" Number of invalid values: {len(invalid_indices)}") 
 
 df = enforce_column_type(df, column, column_type, invalid_indices) 
 
 valid_percentage = (df[column].count() / len(df)) * 100 
 print(f" Percentage of valid values after type enforcement: {valid_percentage:.2f}%") 
 
 return df