Spaces:
Running
Running
File size: 2,309 Bytes
bf12aca |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 |
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:
@staticmethod
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}")
|