Update app.py

app.py

@@ -1,33 +1,40 @@
 import streamlit as st
 import pandas as pd
-
 from langchain.agents.agent_types import AgentType
 from langchain_experimental.agents.agent_toolkits import create_pandas_dataframe_agent
-from langchain_openai import ChatOpenAI
+from langchain_community.llms import Ollama
+
+# Streamlit UI
+st.title("Excel ChatBot (Free - Local Model)")
+st.subheader("Stack: LangChain Agent, Streamlit, Ollama (Mistral)")
 
-st.title("Excel ChatBot")
-st.subheader("Stack used: LangChain Agent, Streamlit, OpenAI LLM - by https://github.com/jaglinux", divider='rainbow')
+uploaded_file = st.file_uploader("Upload CSV or Excel", type=['csv','xlsx'])
 
-uploaded_file = st.file_uploader("Choose a file", type=['csv','xlsx'])
+# Load dataframe
 if uploaded_file is None:
     df = pd.read_csv("titanic.csv")
-    st.write("Default file uploaded, titanic.csv")
+    st.write("Using default Titanic dataset.")
 else:
-    # Can be used wherever a "file-like" object is accepted:
     if uploaded_file.name.endswith(".csv"):
         df = pd.read_csv(uploaded_file)
     elif uploaded_file.name.endswith(".xlsx"):
         df = pd.read_excel(uploaded_file)
-st.dataframe(df, height=5)
 
+st.dataframe(df, height=300)
+
+# Load local LLM using Ollama (make sure `ollama run mistral` is running)
+llm = Ollama(model="mistral")
+
+# Create LangChain agent
 agent = create_pandas_dataframe_agent(
-    ChatOpenAI(temperature=0),
+    llm,
     df,
     verbose=True,
-    agent_type=AgentType.OPENAI_FUNCTIONS,
+    agent_type=AgentType.OPENAI_FUNCTIONS  # Still works, just the name
 )
-if question := st.chat_input("Ask Question to the csv/xlsx provided"):
+
+# Chat input
+if question := st.chat_input("Ask a question about the data"):
     response = agent.invoke(question)
-    print(response['output'])
     st.chat_message("user").markdown(question)
-    st.chat_message("assistant").markdown(response['output'])
+    st.chat_message("assistant").markdown(response["output"])