Spaces:
Running
Running
| import re | |
| import sys | |
| from io import BytesIO, StringIO | |
| import matplotlib.pyplot as plt | |
| import streamlit as st | |
| from langchain.callbacks import get_openai_callback | |
| from pandasai import PandasAI | |
| from pandasai.llm.openai import OpenAI | |
| from streamlit_chat import message | |
| class PandasAgent: | |
| def count_tokens_agent(agent, query): | |
| """ | |
| Count the tokens used by the CSV Agent | |
| """ | |
| with get_openai_callback() as cb: | |
| result = agent(query) | |
| st.write(f"Spent a total of {cb.total_tokens} tokens") | |
| return result | |
| def __init__(self): | |
| pass | |
| def get_agent_response(self, uploaded_file_content, query): | |
| llm = OpenAI() | |
| # from modules.llm import ChatGLM | |
| # llm = ChatGLM() | |
| pandas_ai = PandasAI(llm, verbose=True) | |
| old_stdout = sys.stdout | |
| sys.stdout = captured_output = StringIO() | |
| response = pandas_ai.run(data_frame=uploaded_file_content, prompt=query) | |
| fig = plt.gcf() | |
| if fig.get_axes(): | |
| # Adjust the figure size | |
| fig.set_size_inches(12, 6) | |
| # Adjust the layout tightness | |
| plt.tight_layout() | |
| buf = BytesIO() | |
| fig.savefig(buf, format="png") | |
| buf.seek(0) | |
| st.image(buf, caption="Generated Plot") | |
| sys.stdout = old_stdout | |
| return response, captured_output | |
| def process_agent_thoughts(self, captured_output): | |
| thoughts = captured_output.getvalue() | |
| cleaned_thoughts = re.sub(r"\x1b\[[0-9;]*[a-zA-Z]", "", thoughts) | |
| cleaned_thoughts = re.sub(r"\[1m>", "", cleaned_thoughts) | |
| return cleaned_thoughts | |
| def display_agent_thoughts(self, cleaned_thoughts): | |
| with st.expander("Display the agent's thoughts"): | |
| st.write(cleaned_thoughts) | |
| def update_chat_history(self, query, result): | |
| st.session_state.chat_history.append(("user", query)) | |
| st.session_state.chat_history.append(("agent", result)) | |
| def display_chat_history(self): | |
| for i, (sender, message_text) in enumerate(st.session_state.chat_history): | |
| if sender == "user": | |
| message(message_text, is_user=True, key=f"{i}_user") | |
| else: | |
| message(message_text, key=f"{i}") | |