YashB1/VayuBuddy_DevelopmentApp

1import os2import pandas as pd3from pandasai import Agent, SmartDataframe4from typing import Tuple5from PIL import Image6from pandasai.llm import HuggingFaceTextGen7from dotenv import load_dotenv8from langchain_groq.chat_models import ChatGroq9 10load_dotenv("Groq.txt")11Groq_Token = os.environ["GROQ_API_KEY"]12models = {"mixtral": "mixtral-8x7b-32768", "llama": "llama2-70b-4096", "gemma": "gemma-7b-it"}13 14hf_token = os.getenv("HF_READ")15 16def preprocess_and_load_df(path: str) -> pd.DataFrame:17    df = pd.read_csv(path)18    df["Timestamp"] = pd.to_datetime(df["Timestamp"])19    return df20 21def load_agent(df: pd.DataFrame, context: str, inference_server: str, name="mixtral") -> Agent:22    # llm = HuggingFaceTextGen(23    # inference_server_url=inference_server,24    # max_new_tokens=250,25    # temperature=0.1,26    # repetition_penalty=1.2,27    # top_k=5,28    # )29    # llm.client.headers = {"Authorization": f"Bearer {hf_token}"}30    llm = ChatGroq(model=models[name], api_key=os.getenv("GROQ_API"), temperature=0.1)31    32    agent = Agent(df, config={"llm": llm, "enable_cache": False, "options": {"wait_for_model": True}})33    agent.add_message(context)34    return agent35 36def load_smart_df(df: pd.DataFrame, inference_server: str, name="mixtral") -> SmartDataframe:37    # llm = HuggingFaceTextGen(38    # inference_server_url=inference_server,39    # )40    # llm.client.headers = {"Authorization": f"Bearer {hf_token}"}41    llm = ChatGroq(model=models[name], api_key=os.getenv("GROQ_API"), temperature=0.1)42    df = SmartDataframe(df, config={"llm": llm, "max_retries": 5, "enable_cache": False})43    return df44 45def get_from_user(prompt):46    return {"role": "user", "content": prompt}47 48def ask_agent(agent: Agent, prompt: str) -> Tuple[str, str, str]:49    response = agent.chat(prompt)50    gen_code = agent.last_code_generated51    ex_code = agent.last_code_executed52    last_prompt = agent.last_prompt53    return {"role": "assistant", "content": response, "gen_code": gen_code, "ex_code": ex_code, "last_prompt": last_prompt}54 55def decorate_with_code(response: dict) -> str:56    return f"""<details>57<summary>Generated Code</summary>58    59```python60{response["gen_code"]}61```62</details>63 64<details>65<summary>Prompt</summary>66 67{response["last_prompt"]}68"""69 70def show_response(st, response):71    with st.chat_message(response["role"]):72        try:73            image = Image.open(response["content"])74            if "gen_code" in response:75                st.markdown(decorate_with_code(response), unsafe_allow_html=True)76            st.image(image)77        except Exception as e:78            if "gen_code" in response:79                display_content = decorate_with_code(response) + f"""</details>80 81{response["content"]}"""82            else:83                display_content = response["content"]84            st.markdown(display_content, unsafe_allow_html=True)85 86def ask_question(model_name, question):87    llm = ChatGroq(model=models[model_name], api_key=os.getenv("GROQ_API"), temperature=0.1)88 89    df_check = pd.read_csv("Data.csv")90    df_check["Timestamp"] = pd.to_datetime(df_check["Timestamp"])91    df_check = df_check.head(5)92 93    new_line = "\n"94 95    template = f"""```python96import pandas as pd97import matplotlib.pyplot as plt98 99df = pd.read_csv("Data.csv")100df["Timestamp"] = pd.to_datetime(df["Timestamp"])101 102    # df.dtypes103    {new_line.join(map(lambda x: '# '+x, str(df_check.dtypes).split(new_line)))}104 105    # {question.strip()}106    # <your code here>107    ```108    """109 110    query = f"""I have a pandas dataframe data of PM2.5 and PM10.111    * Frequency of data is daily. 112    * `pollution` generally means `PM2.5`.113    * Save result in a variable `answer` and make it global.114    * If result is a plot, save it and save path in `answer`. Example: `answer='plot.png'`115    * If result is not a plot, save it as a string in `answer`. Example: `answer='The city is Mumbai'`116 117    Complete the following code.118 119    {template}120 121    """122 123    answer = llm.invoke(query)124    code = f"""125    {template.split("```python")[1].split("```")[0]}126    {answer.content.split("```python")[1].split("```")[0]}127    """128    # update variable `answer` when code is executed129    exec(code)130 131    return {"role": "assistant", "content": answer.content, "gen_code": code, "ex_code": code, "last_prompt": question}