AtharvaThakur/Insights

1import streamlit as st2from litellm import completion3from dotenv import load_dotenv4import os5import pandas as pd6def get_data_info():7        file_path = './data.csv'8        df = pd.read_csv(file_path)9 10        # Get column names11        column_names = ", ".join(df.columns.tolist())12        13        # Get data types14        data_types = ", ".join([f"{col}: {dtype}" for col, dtype in df.dtypes.items()])15        16        # Get number of rows and columns17        num_rows, num_cols = df.shape18        19        # Get unique values and example values for each column20        unique_values_info = []21        example_values_info = []22        for col in df.columns:23            unique_values = df[col].unique()24            unique_values_info.append(f"{col}: {len(unique_values)} unique values")25            example_values = df[col].head(5).tolist()  # Get first 5 values as examples26            example_values_info.append(f"{col}: {example_values}")27 28        # Construct the dataset information string29        info_string = f"Dataset Information:\n"30        info_string += f"Dataset file path: {file_path}\n"31        info_string += f"Columns: {column_names}\n"32        info_string += f"Data Types: {data_types}\n"33        info_string += f"Number of Rows: {num_rows}\n"34        info_string += f"Number of Columns: {num_cols}\n"35        info_string += f"Unique Values per Column: {'; '.join(unique_values_info)}\n"36        # info_string += f"Example Values per Column: {'; '.join(example_values_info)}\n"37 38        return info_string39 40def generate_answer(query):41    os.environ['GEMINI_API_KEY'] = os.getenv("GOOGLE_API_KEY")42    output = completion(43        model="gemini/gemini-pro", 44        messages=[45                {"role": "user", "content": "You are a data analyst who's job is to give an answer to the asked query based on the given information about the dataset."},46                {"role": "assistant", "content": "I am a data analyst who would give an answer to the asked query based on the given information about the dataset."},47                {"role": "user", "content": f"Here is information about the dataset.\n {get_data_info()}"},48                {"role": "user", "content": f"Given query - {query}"},49            ]50    )51 52    response = output.choices[0].message.content53    return response