AtharvaThakur/Insights

1from litellm import completion2from dotenv import load_dotenv3import os4import pandas as pd5 6load_dotenv()  # take environment variables from .env.7 8def get_data_info():9        file_path = './data.csv'10        df = pd.read_csv(file_path)11 12        # Get column names13        column_names = ", ".join(df.columns.tolist())14 15        # Get data types16        data_types = ", ".join([f"{col}: {dtype}" for col, dtype in df.dtypes.items()])17 18        # Get number of rows and columns19        num_rows, num_cols = df.shape20 21        # Get unique values and example values for each column22        unique_values_info = []23        example_values_info = []24        for col in df.columns:25                unique_values = df[col].unique()26                unique_values_info.append(f"{col}: {len(unique_values)} unique values")27                example_values = df[col].head(5).tolist()  # Get first 5 values as examples28                example_values_info.append(f"{col}: {example_values}")29 30        # Construct the dataset information string31        info_string = f"Dataset Information:\n"32        info_string += f"Dataset file path: {file_path}\n"33        info_string += f"Columns: {column_names}\n"34        info_string += f"Data Types: {data_types}\n"35        info_string += f"Number of Rows: {num_rows}\n"36        info_string += f"Number of Columns: {num_cols}\n"37        info_string += f"Unique Values per Column: {'; '.join(unique_values_info)}\n"38        # info_string += f"Example Values per Column: {'; '.join(example_values_info)}\n"39 40        return info_string41 42def code_debugger(python_code,error_message):43        os.environ['GEMINI_API_KEY'] = os.getenv("GOOGLE_API_KEY")44        output = completion(45            model="gemini/gemini-pro", 46            messages=[47                    {"role": "user", "content": "You are a computer with the ability to run any code you want when you are given a prompt and return a response with a plan of what code you want to run. You should start your response with the python program, The commands you provide should be in a single code block encapsulated in '''python and ''' for Python and should be valid Python programs."},48                    {"role": "assistant", "content": "I am a computer with the ability to run any code I want when I am given a prompt and return a response with a python program. I will start my response with python program. The commands I provide should be in a single code block encapulated in ```python and ``` and should be a valid Python program."},49                    {"role": "user", "content": "Your are given a python code that has an error. you have to solve that error"},50                    {"role": "assistant", "content": "my job is write the correct python code to solve the error."},51                    {"role": "user", "content": f"Here is the info regarding dataset, python code and the associated error\n dataset_info:-{get_data_info()} python code:-{python_code} \n error message:- {error_message}"},52                ]53        )54 55        response = output.choices[0].message.content56        57        return response