import streamlit as st
import os
import json
import pandas as pd
import random
from os.path import join
from datetime import datetime
from src import (
preprocess_and_load_df,
get_from_user,
ask_question,
)
from dotenv import load_dotenv
from langchain_groq import ChatGroq
from langchain_google_genai import ChatGoogleGenerativeAI
from streamlit_feedback import streamlit_feedback
from huggingface_hub import HfApi
from datasets import load_dataset, get_dataset_config_info, Dataset
from PIL import Image
import time
import uuid
# Page config with beautiful theme
st.set_page_config(
page_title="VayuChat - AI Air Quality Assistant",
page_icon="V",
layout="wide",
initial_sidebar_state="expanded"
)
# Custom CSS for beautiful styling
st.markdown("""
""", unsafe_allow_html=True)
# JavaScript for interactions
st.markdown("""
""", unsafe_allow_html=True)
# FORCE reload environment variables
load_dotenv(override=True)
# Get API keys
Groq_Token = os.getenv("GROQ_API_KEY")
hf_token = os.getenv("HF_TOKEN")
gemini_token = os.getenv("GEMINI_TOKEN")
models = {
"gpt-oss-20b": "openai/gpt-oss-20b",
"gpt-oss-120b": "openai/gpt-oss-120b",
"llama3.1": "llama-3.1-8b-instant",
"llama3.3": "llama-3.3-70b-versatile",
"deepseek-R1": "deepseek-r1-distill-llama-70b",
"llama4 maverik":"meta-llama/llama-4-maverick-17b-128e-instruct",
"llama4 scout":"meta-llama/llama-4-scout-17b-16e-instruct",
"gemini-pro": "gemini-1.5-pro"
}
self_path = os.path.dirname(os.path.abspath(__file__))
# Initialize session ID for this session
if "session_id" not in st.session_state:
st.session_state.session_id = str(uuid.uuid4())
def upload_feedback(feedback, error, output, last_prompt, code, status):
"""Enhanced feedback upload function with better logging and error handling"""
try:
if not hf_token or hf_token.strip() == "":
st.warning("Cannot upload feedback - HF_TOKEN not available")
return False
# Create comprehensive feedback data
feedback_data = {
"timestamp": datetime.now().isoformat(),
"session_id": st.session_state.session_id,
"feedback_score": feedback.get("score", ""),
"feedback_comment": feedback.get("text", ""),
"user_prompt": last_prompt,
"ai_output": str(output),
"generated_code": code or "",
"error_message": error or "",
"is_image_output": status.get("is_image", False),
"success": not bool(error)
}
# Create unique folder name with timestamp
timestamp_str = datetime.now().strftime("%Y%m%d_%H%M%S")
random_id = str(uuid.uuid4())[:8]
folder_name = f"feedback_{timestamp_str}_{random_id}"
# Create markdown feedback file
markdown_content = f"""# VayuChat Feedback Report
## Session Information
- **Timestamp**: {feedback_data['timestamp']}
- **Session ID**: {feedback_data['session_id']}
## User Interaction
**Prompt**: {feedback_data['user_prompt']}
## AI Response
**Output**: {feedback_data['ai_output']}
## Generated Code
```python
{feedback_data['generated_code']}
```
## Technical Details
- **Error Message**: {feedback_data['error_message']}
- **Is Image Output**: {feedback_data['is_image_output']}
- **Success**: {feedback_data['success']}
## User Feedback
- **Score**: {feedback_data['feedback_score']}
- **Comments**: {feedback_data['feedback_comment']}
"""
# Save markdown file locally
markdown_filename = f"{folder_name}.md"
markdown_local_path = f"/tmp/{markdown_filename}"
with open(markdown_local_path, "w", encoding="utf-8") as f:
f.write(markdown_content)
# Upload to Hugging Face
api = HfApi(token=hf_token)
# Upload markdown feedback
api.upload_file(
path_or_fileobj=markdown_local_path,
path_in_repo=f"data/{markdown_filename}",
repo_id="SustainabilityLabIITGN/VayuChat_Feedback",
repo_type="dataset",
)
# Upload image if it exists and is an image output
if status.get("is_image", False) and isinstance(output, str) and os.path.exists(output):
try:
image_filename = f"{folder_name}_plot.png"
api.upload_file(
path_or_fileobj=output,
path_in_repo=f"data/{image_filename}",
repo_id="SustainabilityLabIITGN/VayuChat_Feedback",
repo_type="dataset",
)
except Exception as img_error:
print(f"Error uploading image: {img_error}")
# Clean up local files
if os.path.exists(markdown_local_path):
os.remove(markdown_local_path)
st.success("Feedback uploaded successfully!")
return True
except Exception as e:
st.error(f"Error uploading feedback: {e}")
print(f"Feedback upload error: {e}")
return False
# Filter available models
available_models = []
model_names = list(models.keys())
groq_models = []
gemini_models = []
for model_name in model_names:
if "gemini" not in model_name:
groq_models.append(model_name)
else:
gemini_models.append(model_name)
if Groq_Token and Groq_Token.strip():
available_models.extend(groq_models)
if gemini_token and gemini_token.strip():
available_models.extend(gemini_models)
if not available_models:
st.error("No API keys available! Please set up your API keys in the .env file")
st.stop()
# Set DeepSeek-R1 as default if available
default_index = 0
if "deepseek-R1" in available_models:
default_index = available_models.index("deepseek-R1")
# Header with logo, title and model selector
header_col1, header_col2 = st.columns([2, 1])
with header_col1:
st.markdown("""
V
VayuChat
Environmental Data Analysis
""", unsafe_allow_html=True)
with header_col2:
st.markdown("AI Model:
", unsafe_allow_html=True)
model_name = st.selectbox(
"Model:",
available_models,
index=default_index,
help="Choose your AI model",
label_visibility="collapsed"
)
st.markdown("
", unsafe_allow_html=True)
# Load data with error handling
try:
df = preprocess_and_load_df(join(self_path, "Data.csv"))
# Data loaded silently - no success message needed
except Exception as e:
st.error(f"Error loading data: {e}")
st.stop()
inference_server = "https://api-inference.huggingface.co/models/mistralai/Mistral-7B-Instruct-v0.2"
image_path = "IITGN_Logo.png"
# Clean sidebar
with st.sidebar:
# Quick Queries Section - moved to top
st.markdown("### Quick Queries")
# Load quick prompts
questions = []
questions_file = join(self_path, "questions.txt")
if os.path.exists(questions_file):
try:
with open(questions_file, 'r', encoding='utf-8') as f:
content = f.read()
questions = [q.strip() for q in content.split("\n") if q.strip()]
except Exception as e:
questions = []
# Add default prompts if file doesn't exist or is empty
if not questions:
questions = [
"Which month had highest pollution?",
"Which city has worst air quality?",
"Show annual PM2.5 average",
"Compare winter vs summer pollution",
"List all cities by pollution level",
"Plot monthly average PM2.5 for 2023"
]
# Quick query buttons in sidebar - compact style
selected_prompt = None
for i, question in enumerate(questions[:6]): # Show only first 6
# Truncate long questions for display
display_text = question[:35] + "..." if len(question) > 35 else question
# Use columns to make buttons more compact
if st.button(display_text, key=f"sidebar_prompt_{i}", help=question, use_container_width=True):
selected_prompt = question
st.markdown("---")
# Dataset Info Section
st.markdown("### Dataset Info")
st.markdown("""
PM2.5 Air Quality Data
Locations: Gujarat cities
Parameters: PM2.5, PM10
""", unsafe_allow_html=True)
# Current Model Info
st.markdown("### Current Model")
st.markdown(f"**{model_name}**")
model_descriptions = {
"llama3.1": "Fast and efficient for general queries",
"llama3.3": "Most advanced LLaMA model for complex reasoning",
"mistral": "Balanced performance and speed",
"gemma": "Google's lightweight model",
"gemini-pro": "Google's most powerful model",
"gpt-oss-20b": "OpenAI's compact open-weight GPT for everyday tasks",
"gpt-oss-120b": "OpenAI's massive open-weight GPT for nuanced responses",
"deepseek-R1": "DeepSeek's distilled LLaMA model for efficient reasoning",
"llama4 maverik": "Meta's LLaMA 4 Maverick — high-performance instruction model",
"llama4 scout": "Meta's LLaMA 4 Scout — optimized for adaptive reasoning"
}
if model_name in model_descriptions:
st.caption(model_descriptions[model_name])
st.markdown("---")
# Clear Chat Button
if st.button("Clear Chat", use_container_width=True):
st.session_state.responses = []
st.session_state.processing = False
st.session_state.session_id = str(uuid.uuid4())
try:
st.rerun()
except AttributeError:
st.experimental_rerun()
# Main content area - removed quick prompts section from here as it's now in sidebar
# Initialize chat history and processing state
if "responses" not in st.session_state:
st.session_state.responses = []
if "processing" not in st.session_state:
st.session_state.processing = False
def show_custom_response(response):
"""Custom response display function with improved styling"""
role = response.get("role", "assistant")
content = response.get("content", "")
if role == "user":
# User message with right alignment - reduced margins
st.markdown(f"""
""", unsafe_allow_html=True)
elif role == "assistant":
# Check if content is an image filename - don't display the filename text
is_image_path = isinstance(content, str) and any(ext in content for ext in ['.png', '.jpg', '.jpeg'])
# Assistant message with left alignment - reduced margins
if not is_image_path:
st.markdown(f"""
VayuChat
{content if isinstance(content, str) else str(content)}
""", unsafe_allow_html=True)
# Show generated code with Streamlit expander
if response.get("gen_code"):
with st.expander("📋 View Generated Code", expanded=False):
st.code(response["gen_code"], language="python")
# Try to display image if content is a file path
try:
if isinstance(content, str) and (content.endswith('.png') or content.endswith('.jpg')):
if os.path.exists(content):
# Display image without showing filename
st.image(content, use_column_width=True)
return {"is_image": True}
# Also handle case where content shows filename but we want to show image
elif isinstance(content, str) and any(ext in content for ext in ['.png', '.jpg']):
# Extract potential filename from content
import re
filename_match = re.search(r'([^/\\]+\.(?:png|jpg|jpeg))', content)
if filename_match:
filename = filename_match.group(1)
if os.path.exists(filename):
st.image(filename, use_column_width=True)
return {"is_image": True}
except:
pass
return {"is_image": False}
def show_processing_indicator(model_name, question):
"""Show processing indicator"""
st.markdown(f"""
VayuChat • Processing with {model_name}
Question: {question}
Generating response...
""", unsafe_allow_html=True)
# Main chat container
chat_container = st.container()
with chat_container:
# Display chat history
for response_id, response in enumerate(st.session_state.responses):
status = show_custom_response(response)
# Show feedback section for assistant responses
if response["role"] == "assistant":
feedback_key = f"feedback_{int(response_id/2)}"
error = response.get("error", "")
output = response.get("content", "")
last_prompt = response.get("last_prompt", "")
code = response.get("gen_code", "")
if "feedback" in st.session_state.responses[response_id]:
feedback_data = st.session_state.responses[response_id]["feedback"]
st.markdown(f"""
Your Feedback: {feedback_data.get('score', '')}
{f"- {feedback_data.get('text', '')}" if feedback_data.get('text') else ""}
""", unsafe_allow_html=True)
else:
# Beautiful feedback section
st.markdown("---")
st.markdown("**How was this response?**")
col1, col2 = st.columns(2)
with col1:
thumbs_up = st.button("👍 Helpful", key=f"{feedback_key}_up", use_container_width=True)
with col2:
thumbs_down = st.button("👎 Not Helpful", key=f"{feedback_key}_down", use_container_width=True)
if thumbs_up or thumbs_down:
thumbs = "👍 Helpful" if thumbs_up else "👎 Not Helpful"
comments = st.text_area(
"Tell us more (optional):",
key=f"{feedback_key}_comments",
placeholder="What could be improved? Any suggestions?",
max_chars=500
)
if st.button("Submit Feedback", key=f"{feedback_key}_submit"):
feedback = {"score": thumbs, "text": comments}
# Upload feedback with enhanced error handling
if upload_feedback(feedback, error, output, last_prompt, code, status or {}):
st.session_state.responses[response_id]["feedback"] = feedback
time.sleep(1) # Give user time to see success message
st.rerun()
else:
st.error("Failed to submit feedback. Please try again.")
# Show processing indicator if processing
if st.session_state.get("processing"):
show_processing_indicator(
st.session_state.get("current_model", "Unknown"),
st.session_state.get("current_question", "Processing...")
)
# Chat input (always visible at bottom)
prompt = st.chat_input("Ask me anything about air quality!", key="main_chat")
# Handle selected prompt from quick prompts
if selected_prompt:
prompt = selected_prompt
# Handle new queries
if prompt and not st.session_state.get("processing"):
# Prevent duplicate processing
if "last_prompt" in st.session_state:
last_prompt = st.session_state["last_prompt"]
last_model_name = st.session_state.get("last_model_name", "")
if (prompt == last_prompt) and (model_name == last_model_name):
prompt = None
if prompt:
# Add user input to chat history
user_response = get_from_user(prompt)
st.session_state.responses.append(user_response)
# Set processing state
st.session_state.processing = True
st.session_state.current_model = model_name
st.session_state.current_question = prompt
# Rerun to show processing indicator
st.rerun()
# Process the question if we're in processing state
if st.session_state.get("processing"):
prompt = st.session_state.get("current_question")
model_name = st.session_state.get("current_model")
try:
response = ask_question(model_name=model_name, question=prompt)
if not isinstance(response, dict):
response = {
"role": "assistant",
"content": "Error: Invalid response format",
"gen_code": "",
"ex_code": "",
"last_prompt": prompt,
"error": "Invalid response format"
}
response.setdefault("role", "assistant")
response.setdefault("content", "No content generated")
response.setdefault("gen_code", "")
response.setdefault("ex_code", "")
response.setdefault("last_prompt", prompt)
response.setdefault("error", None)
except Exception as e:
response = {
"role": "assistant",
"content": f"Sorry, I encountered an error: {str(e)}",
"gen_code": "",
"ex_code": "",
"last_prompt": prompt,
"error": str(e)
}
st.session_state.responses.append(response)
st.session_state["last_prompt"] = prompt
st.session_state["last_model_name"] = model_name
st.session_state.processing = False
# Clear processing state
if "current_model" in st.session_state:
del st.session_state.current_model
if "current_question" in st.session_state:
del st.session_state.current_question
st.rerun()
# Minimal auto-scroll - only scroll when processing
if st.session_state.get("processing"):
st.markdown("", unsafe_allow_html=True)
# Beautiful sidebar footer
# with st.sidebar:
# st.markdown("---")
# st.markdown("""
#
# """, unsafe_allow_html=True)
# Statistics (if logging is enabled)
if hf_token and hf_token.strip():
st.markdown("### 📈 Session Stats")
total_interactions = len([r for r in st.session_state.get("responses", []) if r.get("role") == "assistant"])
st.metric("Interactions", total_interactions)
feedbacks_given = len([r for r in st.session_state.get("responses", []) if r.get("role") == "assistant" and "feedback" in r])
st.metric("Feedbacks Given", feedbacks_given)
# Footer
st.markdown("""
Together for Cleaner Air
VayuChat - Empowering environmental awareness through AI
© 2024 IIT Gandhinagar Sustainability Lab
""", unsafe_allow_html=True)