app.py

import gradio as gr
import sqlite3
import os
import asyncio
from core import submit_query
from config import logger

DB_PATH = "chat_history.db"

def init_db():
    conn = sqlite3.connect(DB_PATH)
    c = conn.cursor()
    c.execute("""
        CREATE TABLE IF NOT EXISTS users (
            username TEXT PRIMARY KEY,
            password TEXT NOT NULL
        )
    """)
    c.execute("""
        CREATE TABLE IF NOT EXISTS sessions (
            session_id INTEGER PRIMARY KEY AUTOINCREMENT,
            username TEXT,
            created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
        )
    """)
    c.execute("""
        CREATE TABLE IF NOT EXISTS history (
            session_id INTEGER,
            username TEXT,
            providers TEXT,
            message TEXT,
            openai_response TEXT,
            anthropic_response TEXT,
            gemini_response TEXT,
            FOREIGN KEY(session_id) REFERENCES sessions(session_id)
        )
    """)
    conn.commit()
    conn.close()

def register_user(username, password):
    conn = sqlite3.connect(DB_PATH)
    c = conn.cursor()
    try:
        c.execute("INSERT INTO users (username, password) VALUES (?, ?)", (username, password))
        conn.commit()
        return True
    except sqlite3.IntegrityError:
        return False
    finally:
        conn.close()

def login_user(username, password):
    conn = sqlite3.connect(DB_PATH)
    c = conn.cursor()
    c.execute("SELECT * FROM users WHERE username=? AND password=?", (username, password))
    result = c.fetchone()
    conn.close()
    return result is not None

def create_new_session(username):
    conn = sqlite3.connect(DB_PATH)
    c = conn.cursor()
    c.execute("INSERT INTO sessions (username) VALUES (?)", (username,))
    session_id = c.lastrowid
    conn.commit()
    conn.close()
    return session_id

def list_sessions(username):
    conn = sqlite3.connect(DB_PATH)
    c = conn.cursor()
    c.execute("SELECT session_id, created_at FROM sessions WHERE username=? ORDER BY created_at DESC", (username,))
    sessions = c.fetchall()
    conn.close()
    return sessions

def get_chat_history(session_id):
    conn = sqlite3.connect(DB_PATH)
    c = conn.cursor()
    c.execute("SELECT message, openai_response, anthropic_response, gemini_response FROM history WHERE session_id=? ORDER BY rowid", (session_id,))
    rows = c.fetchall()
    conn.close()

    openai_history = []
    anthropic_history = []
    gemini_history = []

    for m, o, a, g in rows:
        user_msg_dict = {"role": "user", "content": m}
        openai_history.append(user_msg_dict)
        anthropic_history.append(user_msg_dict)
        gemini_history.append(user_msg_dict)
        if o:
            openai_history.append({"role": "assistant", "content": o})
        if a:
            anthropic_history.append({"role": "assistant", "content": a})
        if g:
            gemini_history.append({"role": "assistant", "content": g})

    return openai_history, anthropic_history, gemini_history

def save_chat_history(session_id, username, providers, message, openai_resp, anthropic_resp, gemini_resp):
    conn = sqlite3.connect(DB_PATH)
    c = conn.cursor()
    c.execute(
        "INSERT INTO history (session_id, username, providers, message, openai_response, anthropic_response, gemini_response) VALUES (?, ?, ?, ?, ?, ?, ?)",
        (session_id, username, ",".join(providers), message, openai_resp, anthropic_resp, gemini_resp)
    )
    conn.commit()
    conn.close()

async def chatbot_fn(message, openai_history, anthropic_history, gemini_history, username, providers, session_id):
    if not username or not session_id:
        return "", openai_history or [], anthropic_history or [], gemini_history or [], openai_history or [], anthropic_history or [], gemini_history or []
    if not providers:
        user_msg_dict = {"role": "user", "content": message}
        no_provider_msg = {"role": "assistant", "content": "Please select at least one provider."}
        openai_history = (openai_history or []) + [user_msg_dict, no_provider_msg]
        anthropic_history = (anthropic_history or []) + [user_msg_dict, no_provider_msg]
        gemini_history = (gemini_history or []) + [user_msg_dict, no_provider_msg]
        return "", openai_history, anthropic_history, gemini_history, openai_history, anthropic_history, gemini_history

    context_history = openai_history
    final_result = None
    error_message = ""
    async for err, openai_msgs, anthropic_msgs, gemini_msgs, updated_context_history in submit_query(
        message, providers, context_history, username
    ):
        final_result = (openai_msgs, anthropic_msgs, gemini_msgs, updated_context_history)
        error_message = err

    if not final_result:
        logger.error("submit_query returned no result")
        return f"Error: submit_query failed.", openai_history, anthropic_history, gemini_history, openai_history, anthropic_history, gemini_history

    new_openai_history, new_anthropic_history, new_gemini_history, _ = final_result

    def get_last_response_text_from_history(history_list):
        if isinstance(history_list, list) and history_list:
            for m in reversed(history_list):
                if isinstance(m, dict) and m.get("role") == "assistant" and m.get("content"):
                    return m.get("content", "")
        return ""

    openai_text_to_save = get_last_response_text_from_history(new_openai_history)
    anthropic_text_to_save = get_last_response_text_from_history(new_anthropic_history)
    gemini_text_to_save = get_last_response_text_from_history(new_gemini_history)

    save_chat_history(session_id, username, providers, message, openai_text_to_save, anthropic_text_to_save, gemini_text_to_save)

    return error_message, new_openai_history, new_anthropic_history, new_gemini_history, new_openai_history, new_anthropic_history, new_gemini_history

def session_label(session):
    sid, created = session
    return f"{sid} ({created})"

with gr.Blocks() as demo:
    gr.Markdown("# Multi-Model Chat")

    # Login/Register group
    login_group = gr.Group(visible=True)
    with login_group:
        username = gr.Textbox(label="Username")
        password = gr.Textbox(label="Password", type="password")
        login_btn = gr.Button("Login")
        register_btn = gr.Button("Register")
        login_status = gr.Markdown("")

    # Chat group (hidden by default)
    chat_group = gr.Group(visible=False)
    with chat_group:
        with gr.Row():
            session_selector = gr.Dropdown(choices=[], label="Select Chat Session", interactive=True)
            new_chat_btn = gr.Button("New Chat")
            user_display = gr.Markdown("", elem_id="user_display")
            logout_btn = gr.Button("Logout")
        providers = gr.CheckboxGroup(
            choices=["OpenAI", "Anthropic", "Gemini"],
            label="Select Providers",
            value=["OpenAI"]
        )
        with gr.Row():
            openai_chat = gr.Chatbot(label="OpenAI", type="messages", height=400)
            anthropic_chat = gr.Chatbot(label="Anthropic", type="messages", height=400)
            gemini_chat = gr.Chatbot(label="Gemini", type="messages", height=400)
        msg = gr.Textbox(label="Message")
        send_btn = gr.Button("Send")
        clear_btn = gr.Button("Clear History")
        hidden_username = gr.State("")
        openai_history_state = gr.State([])
        anthropic_history_state = gr.State([])
        gemini_history_state = gr.State([])
        session_id_state = gr.State("")

    # --- Function Definitions ---
    def do_login(u, p):
        if login_user(u, p):
            session_id = create_new_session(u)
            sessions = list_sessions(u)
            session_choices = [session_label(s) for s in sessions]
            session_value = session_label(sessions[0]) if sessions else ""
            return (
                f"Welcome, {u}!",
                gr.update(visible=False),  # Hide login group
                gr.update(visible=True),   # Show chat group
                [], [], [],  # Empty chats for new session
                [], [], [],
                u,
                f"Logged in as: **{u}**",
                session_id,
                gr.update(choices=session_choices, value=session_value)
            )
        else:
            return "Login failed.", gr.update(), gr.update(), [], [], [], [], [], [], "", "", gr.update(choices=[], value="")

    def do_register(u, p):
        if register_user(u, p):
            return "Registration successful! Please login.", gr.update(), gr.update(), [], [], [], [], [], [], "", "", gr.update(choices=[], value="")
        else:
            return "Username already exists.", gr.update(), gr.update(), [], [], [], [], [], [], "", "", gr.update(choices=[], value="")

    def do_logout():
        return (
            "",  # login_status
            gr.update(visible=True),   # Show login group
            gr.update(visible=False),  # Hide chat group
            [], [], [], [], [], [], "", "", gr.update(choices=[], value="")
        )

    def do_new_chat(username):
        if not username:
            return "", [], [], [], gr.update(choices=[], value="")
        session_id = create_new_session(username)
        sessions = list_sessions(username)
        session_choices = [session_label(s) for s in sessions]
        session_value = session_label(sessions[0]) if sessions else ""
        return session_id, [], [], [], gr.update(choices=session_choices, value=session_value)

    def do_select_session(session_label_str, username):
        if not session_label_str or not username:
            return [], [], [], [], [], [], ""
        session_id = int(session_label_str.split(" ")[0])
        openai_hist, anthropic_hist, gemini_hist = get_chat_history(session_id)
        return (
            openai_hist, anthropic_hist, gemini_hist,  # update chat displays
            openai_hist, anthropic_hist, gemini_hist,  # update state variables
            session_id                                 # update session_id_state
        )

    # --- Gradio Wiring ---
    login_btn.click(
        do_login,
        [username, password],
        [
            login_status, login_group, chat_group,
            openai_chat, anthropic_chat, gemini_chat,
            openai_history_state, anthropic_history_state, gemini_history_state,
            hidden_username,
            user_display,
            session_id_state,
            session_selector  # Only one output for dropdown!
        ]
    )

    register_btn.click(
        do_register,
        [username, password],
        [
            login_status, login_group, chat_group,
            openai_chat, anthropic_chat, gemini_chat,
            openai_history_state, anthropic_history_state, gemini_history_state,
            hidden_username,
            user_display,
            session_id_state,
            session_selector
        ]
    )

    send_btn.click(
        chatbot_fn,
        [msg, openai_history_state, anthropic_history_state, gemini_history_state, hidden_username, providers, session_id_state],
        [msg, openai_chat, anthropic_chat, gemini_chat, openai_history_state, anthropic_history_state, gemini_history_state]
    )

    clear_btn.click(
        lambda: ("", [], [], [], [], [], []),
        [],
        [msg, openai_chat, anthropic_chat, gemini_chat, openai_history_state, anthropic_history_state, gemini_history_state]
    )

    logout_btn.click(
        do_logout,
        [],
        [
            login_status, login_group, chat_group,
            openai_chat, anthropic_chat, gemini_chat,
            openai_history_state, anthropic_history_state, gemini_history_state,
            hidden_username,
            user_display,
            session_id_state,
            session_selector
        ]
    )

    new_chat_btn.click(
        do_new_chat,
        [hidden_username],
        [session_id_state, openai_chat, anthropic_chat, gemini_chat, session_selector]
    )

    session_selector.change(
        do_select_session,
        [session_selector, hidden_username],
        [
            openai_chat, anthropic_chat, gemini_chat,           # chat displays
            openai_history_state, anthropic_history_state, gemini_history_state,  # state variables
            session_id_state                                    # session id state
        ]
    )

if __name__ == "__main__":
    init_db()
    demo.launch()