AtharvaThakur/Insights

1import streamlit as st2from litellm import completion3from dotenv import load_dotenv4import os5from Modules.python_interpreter import PythonInterpreter, run_interpreter6import pandas as pd7 8load_dotenv()  # take environment variables from .env.9 10class DataCodeGen:11    def __init__(self):12        pass13 14 15    def get_data_info(self):16        file_path = './data.csv'17        df = pd.read_csv(file_path)18 19        # Get column names20        column_names = ", ".join(df.columns.tolist())21        22        # Get data types23        data_types = ", ".join([f"{col}: {dtype}" for col, dtype in df.dtypes.items()])24        25        # Get number of rows and columns26        num_rows, num_cols = df.shape27        28        # Get unique values and example values for each column29        unique_values_info = []30        example_values_info = []31        for col in df.columns:32            unique_values = df[col].unique()33            unique_values_info.append(f"{col}: {len(unique_values)} unique values")34            example_values = df[col].head(5).tolist()  # Get first 5 values as examples35            example_values_info.append(f"{col}: {example_values}")36 37        # Construct the dataset information string38        info_string = f"Dataset Information:\n"39        info_string += f"Dataset file path: {file_path}\n"40        info_string += f"Columns: {column_names}\n"41        info_string += f"Data Types: {data_types}\n"42        info_string += f"Number of Rows: {num_rows}\n"43        info_string += f"Number of Columns: {num_cols}\n"44        info_string += f"Unique Values per Column: {'; '.join(unique_values_info)}\n"45        # info_string += f"Example Values per Column: {'; '.join(example_values_info)}\n"46 47        return info_string48 49    @st.cache_data(experimental_allow_widgets=True)50    def generate_code(_self,query):51        os.environ['GEMINI_API_KEY'] = os.getenv("GOOGLE_API_KEY")52        data_info= _self.get_data_info()53        output = completion(54            model="gemini/gemini-pro", 55            messages=[56                    {"role": "user", "content": "You are a data analyst 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."},57                    {"role": "assistant", "content": "I am a data analyst 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."},58                    {"role": "user", "content": "You can only use the following python libraries- pandas, numpy, matplotlib.pyplot, seaborn, sklearn"},59                    {"role": "assistant", "content": "I can only use the following python libraries- pandas, numpy, matplotlib.pyplot, seaborn, sklearn"},60                    {"role": "user", "content": "Your job is write the python code the answer for the given query regarding a dataset. The python should find the correct answer the query, also generate a visualization if necessary and store it in `fig.pdf`. Store the answer to query and information regarding the visualization in `data.txt`. Even if the given task is to plot a graph you have to include textual information regarding the graphs like the labels and values in `data.txt`."},61                    {"role": "assistant", "content": "My job is write the python code that will find the answer for the given query regarding a dataset. The python should find the correct answer the query, also generate a visualization if necessary and store it in `fig.pdf`. I have to store the answer to query along with label and value shown in the visualization in `data.txt`. Even if I have to just plot a graph I will include textual information regarding the graphs like the labels and values in `data.txt`."},62                    {"role": "user", "content": f"Here is some information about the dataset.\n {data_info}"},63                    {"role": "user", "content": f"Given query - {query}"},64                ]65        )66 67        response = output.choices[0].message.content68        return response69 70 71    def extract_code(self,response):        72        # else:73        #     print(response.choices[0].message.content)74        #     # Extract plan from the response75        #     plan = response.choices[0].message.content.split("```python")[0]76        #     plan = plan.replace("'", "")77        #     plan = plan.replace('`', "")78        #     print("plan:", plan)79            80        if "```python" in response:81            python_code = response.split("```python")[1].split("```")[0].strip()82            return python_code83        elif "```" in response:84            python_code = response.split("```")[1].split("```")[0].strip()85            print("Code found in the response but not Left out the word python:", python_code)86            return python_code87        elif "```python" in response.choices[0].message.content:88            python_code = response.choices[0].message.content.split(89                "```python")[1].split("```")[0].strip()90            return python_code91        92        93        # if python_code:94        #     interpreter_code_output = run_interpreter(python_code)95        #     print("Python code output:\n", interpreter_code_output)