Spaces:
Sleeping
Sleeping
File size: 3,873 Bytes
a61eb6c | 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 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 |
from dotenv import load_dotenv
import matplotlib.pyplot as plt
import pandas as pd
from smolagents import LiteLLMModel, ToolCallingAgent, tool
import os
# litellm._turn_on_debug() # Disabled for faster responses
API_KEY = os.environ.get("MISTRAL_API_KEY")
TOKEN = os.getenv("token")
model = LiteLLMModel(model_id="mistral/mistral-large-2512",api_key=API_KEY) #mistral/mistral-large-2512
instructions = """You are a helpful assistant that can analyze datasets and create visualizations based on user questions."""
saves = "./conversation_log.txt"
def load_saves(): # Load the conversation history from the file
try:
with open(saves, "r") as f:
chat_history_content = f.read()
print("Previous Conversation:\n", chat_history_content)
return chat_history_content
except FileNotFoundError:
print("No previous conversation found.")
def save_conversation(user_input, response,dataset_path=None): # Save the conversation history to the file
with open(saves, "a") as f:
f.write(f"User: {user_input} with dataset path: {dataset_path}\n")
f.write(f"Agent: {response}\n\n")
print("Conversation saved.")
def clear_conversation_history(): # Clear the conversation history
with open(saves, "w") as f:
f.write("")
print("Conversation history cleared.")
#TOOLS SECTION
@tool
def csv_reader(file_path:str)->str:
"""
Read a CSV file and return a summary and sample rows (not the entire file).
Args:
file_path: The path to the CSV file you want to read.
"""
try:
df = pd.read_csv(file_path)
summary = f"Dataset Shape: {df.shape[0]} rows, {df.shape[1]} columns\n"
summary += f"Columns: {', '.join(df.columns)}\n"
summary += f"Data Types:\n{df.dtypes.to_string()}\n\n"
summary += "First 10 rows:\n"
summary += df.head(10).to_string() #see first 10 rows of the dataset
return summary
except Exception as e:
return f"Error reading CSV: {e}"
@tool
def bar_chart(file_path: str, x: str, y: str) -> str:
"""
Generates and displays a beautiful bar chart from a CSV or Excel file.
Args:
file_path: The path to the data file (CSV or Excel).
x: The name of the column to use for the X-axis (categories).
y: The name of the column to use for the Y-axis (numeric values).
"""
df = pd.read_csv(file_path)
plt.figure(figsize=(10, 6))
plt.bar(df[x], df[y])
plt.xlabel(x)
plt.xticks(rotation=45, ha='right')
plt.ylabel(y)
plt.title(f'{y} by {x}')
plt.tight_layout()
plt.savefig("bar_chart.png")
plt.close()
return "bar_chart.png"
@tool
def delete_chart()-> str:
"""
Don't show old charts - delete the previous one before generating a new one.
Deletes the generated bar chart image file if it exists.
"""
if os.path.exists("bar_chart.png"):
os.remove("bar_chart.png")
return "Bar chart deleted."
else:
return "No bar chart to delete."
agent = ToolCallingAgent(model=model, tools=[csv_reader, bar_chart],instructions=instructions)
def ask_agent(question, dataset_path):
# Only load last few lines of history to speed up response
history = ""
try:
with open(saves, "r") as f:
lines = f.readlines()
# Only use last 20 lines (last ~5 exchanges)
history = "".join(lines[-20:]) if lines else ""
except FileNotFoundError:
history = ""
full_prompt = f"Latest Question: {question}\nDataset Path: {dataset_path}"
if history:
full_prompt = f"Recent context:\n{history}\n\n{full_prompt}"
response = agent.run(full_prompt)
save_conversation(question, response, dataset_path)
return response, "bar_chart.png" if os.path.exists("bar_chart.png") else None
|