import streamlit as st import pandas as pd from datetime import datetime import database as db import sqlite3 import llm_handler # --- Page Configuration --- st.set_page_config(page_title="AI Task Manager", layout="centered", initial_sidebar_state="collapsed") # --- Helper Functions --- def initialize_session_state(): """Initializes session state variables.""" if "user_name" not in st.session_state: st.session_state.user_name = "" if "user_email" not in st.session_state: st.session_state.user_email = "" if "logged_in" not in st.session_state: st.session_state.logged_in = False if "llm" not in st.session_state: try: st.session_state.llm = llm_handler.get_llm() except ValueError as e: st.error(f"LLM Initialization Error: {e}. Please ensure GOOGLE_API_KEY is set in .env") st.stop() if "last_interacted_task_details" not in st.session_state: st.session_state.last_interacted_task_details = None if "chat_history_for_context" not in st.session_state: st.session_state.chat_history_for_context = [] # --- Initialize --- initialize_session_state() db.create_db_and_table() if "table_info" not in st.session_state or st.session_state.get("reload_table_info", False): st.session_state.table_info = db.get_db_info() st.session_state.reload_table_info = False # --- UI Sections --- if not st.session_state.logged_in: st.title("AI Task Manager Login") with st.form("login_form"): name = st.text_input("Your Name", value=st.session_state.get("user_name_input", "")) email = st.text_input("Your Email", value=st.session_state.get("user_email_input", "")) login_button = st.form_submit_button("Login") if login_button: error_messages = [] if not name: error_messages.append("Please enter your name.") if not email: error_messages.append("Please enter your email.") elif not email.strip().lower().endswith(".com"): error_messages.append("Invalid email format.") if not error_messages: st.session_state.user_name = name st.session_state.user_email = email st.session_state.logged_in = True st.session_state.user_name_input = name st.session_state.user_email_input = email st.session_state.last_interacted_task_details = None st.session_state.reload_table_info = True st.rerun() else: for msg in error_messages: st.error(msg) else: st.sidebar.header("User Info") st.sidebar.write(f"👤 **Name:** {st.session_state.user_name}") st.sidebar.write(f"📧 **Email:** {st.session_state.user_email}") if st.sidebar.button("Logout"): st.session_state.logged_in = False st.session_state.user_name = "" st.session_state.user_email = "" st.session_state.last_interacted_task_details = None st.rerun() st.title(f"📝 AI Task Manager for {st.session_state.user_name}") st.markdown("Enter your task command (e.g., 'I have a meeting at 2pm on next Friday' , 'Show pending tasks').") if st.session_state.last_interacted_task_details: task_ctx = st.session_state.last_interacted_task_details # Prepare task context display message task_name = task_ctx.get('task_name', 'N/A') task_id = task_ctx.get('id', 'N/A') due_date = task_ctx.get('due_date', '') due_time = task_ctx.get('due_time', '') ctx_display = f"Last task: '{task_name}' (ID: {task_id})" if due_date: ctx_display += f" - Due: {due_date}" if due_time: ctx_display += f" at {due_time}" st.caption(ctx_display) user_input_query = st.text_input("Your command:", key="task_input", placeholder="e.g., Add task 'Submit report' due next Friday 3pm category Work") if st.button("Process Command", type="primary"): if not user_input_query: st.warning("Please enter a command.") else: llm = st.session_state.llm table_info = st.session_state.table_info context_str = "None" if st.session_state.last_interacted_task_details: details = st.session_state.last_interacted_task_details context_list = [] if details.get('id'): context_list.append(f"id: {details['id']}") if details.get('task_name'): context_list.append(f"name: '{details['task_name']}'") if details.get('due_date'): context_list.append(f"due_date: {details['due_date']}") if details.get('due_time'): context_list.append(f"due_time: {details['due_time']}") context_str = f"Task context - {', '.join(context_list)}" with st.spinner("🤖 Thinking and generating SQL..."): try: generated_sql = llm_handler.generate_sql_query( llm, user_input_query, table_info, st.session_state.user_name, st.session_state.user_email, context_str ) except Exception as e: st.error(f"Error generating SQL: {e}") generated_sql = None if generated_sql: st.write("⚙️ **Generated SQL Query:**") st.code(generated_sql, language="sql") is_select_query = generated_sql.strip().upper().startswith("SELECT") is_insert_query = generated_sql.strip().upper().startswith("INSERT") with st.spinner("💾 Executing query..."): try: summary_context_for_llm = "" if is_select_query: data, columns = db.execute_select_query(generated_sql) if data: df = pd.DataFrame(data, columns=columns) major_display_columns = ['id', 'task_name', 'status', 'category', 'due_date', 'due_time', 'created_at'] display_cols_in_df = [col for col in major_display_columns if col in df.columns] if not display_cols_in_df and df.columns.any(): display_cols_in_df = df.columns.tolist() if display_cols_in_df: st.dataframe(df[display_cols_in_df], use_container_width=True) else: st.info("The query ran but returned no columns to display.") summary_context_for_llm = f"Retrieved {len(data)} task(s)." if len(data) == 1: st.session_state.last_interacted_task_details = dict(zip(columns, data[0])) elif len(data) == 0: summary_context_for_llm = "No tasks found matching your criteria." else: st.info("No tasks found matching your criteria.") summary_context_for_llm = "No tasks found matching your criteria." st.session_state.last_interacted_task_details = None else: result = db.execute_dml_query(generated_sql) action = "processed" if is_insert_query: action = "added" if result: inserted_task_details = db.get_task_by_id(result, st.session_state.user_email) if inserted_task_details: st.session_state.last_interacted_task_details = inserted_task_details summary_context_for_llm = f"Task '{inserted_task_details.get('task_name')}' (ID: {result}) was {action}." else: summary_context_for_llm = f"Task was {action}, but details couldn't be retrieved post-insertion." else: summary_context_for_llm = f"Task addition was attempted but may not have completed as expected (no ID returned)." elif "UPDATE" in generated_sql.upper(): action = "updated" summary_context_for_llm = f"Task(s) {action}. {result} row(s) affected." elif "DELETE" in generated_sql.upper(): action = "deleted" summary_context_for_llm = f"Task(s) {action}. {result} row(s) affected." st.session_state.last_interacted_task_details = None st.success(f"Task command '{action}' processed successfully.") with st.spinner("📜 Generating friendly summary..."): final_summary = llm_handler.summarize_query_result( llm, user_input_query, generated_sql, summary_context_for_llm ) st.markdown(f"**🤖 Summary:**\n {final_summary}") except ValueError as ve: st.error(f"⚠️ Action failed: {ve}") except sqlite3.Error as e: st.error(f"❌ Database Error: {e}") st.error(f" Failed Query: {generated_sql}") except Exception as e: st.error(f"❌ An unexpected error occurred: {e}") st.error(f" Query attempted: {generated_sql}") elif not user_input_query.strip() == "": st.error("Could not generate an SQL query for your request. Please try rephrasing.") st.markdown("---") st.markdown("Example commands:") st.caption(""" - Add: "Schedule a 'Team sync meeting' for next Tuesday at 10:00 category Work" - Add: "I have a 'Doctor's appointment' tomorrow at 3:30 PM" - View: "Show my tasks for this week", "What's pending?" - Update: "Mark 'Team sync meeting' as completed" (if it was the last task discussed) - Update: "I've finished the 'Doctor's appointment'" - Delete: "Cancel the 'Team sync meeting'" - Delete: "Remove task 'Old project idea'" """)