main.py

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