app.py

import os
import threading
import logging
import uuid
import shutil
import json
import tempfile
from flask import Flask, request as flask_request, make_response
import dash
from dash import dcc, html, Input, Output, State, callback_context
import dash_bootstrap_components as dbc
import openai
import base64
import datetime
from werkzeug.utils import secure_filename

logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(threadName)s %(message)s")
logger = logging.getLogger("AskTricare")

app_flask = Flask(__name__)
SESSION_DATA = {}
SESSION_LOCKS = {}
SESSION_DIR_BASE = os.path.join(tempfile.gettempdir(), "asktricare_sessions")
os.makedirs(SESSION_DIR_BASE, exist_ok=True)

openai.api_key = os.environ.get("OPENAI_API_KEY")

def get_session_id():
    sid = flask_request.cookies.get("asktricare_session_id")
    if not sid:
        sid = str(uuid.uuid4())
    return sid

def get_session_dir(session_id):
    d = os.path.join(SESSION_DIR_BASE, session_id)
    os.makedirs(d, exist_ok=True)
    return d

def get_session_lock(session_id):
    if session_id not in SESSION_LOCKS:
        SESSION_LOCKS[session_id] = threading.Lock()
    return SESSION_LOCKS[session_id]

def get_session_state(session_id):
    if session_id not in SESSION_DATA:
        SESSION_DATA[session_id] = {
            "messages": [],
            "uploads": [],
            "created": datetime.datetime.utcnow().isoformat(),
            "streaming": False,
            "stream_buffer": "",
        }
    return SESSION_DATA[session_id]

def save_session_state(session_id):
    state = get_session_state(session_id)
    d = get_session_dir(session_id)
    with open(os.path.join(d, "state.json"), "w") as f:
        json.dump(state, f)

def load_session_state(session_id):
    d = get_session_dir(session_id)
    path = os.path.join(d, "state.json")
    if os.path.exists(path):
        with open(path, "r") as f:
            SESSION_DATA[session_id] = json.load(f)

def load_system_prompt():
    prompt_path = os.path.join(os.getcwd(), "system_prompt.txt")
    try:
        with open(prompt_path, "r", encoding="utf-8") as f:
            return f.read().strip()
    except Exception as e:
        logger.error(f"Failed to load system prompt: {e}")
        return "You are Ask Tricare, a helpful assistant for TRICARE health benefits. Respond conversationally, and cite relevant sources when possible. If you do not know, say so."

app = dash.Dash(
    __name__,
    server=app_flask,
    suppress_callback_exceptions=True,
    external_stylesheets=[dbc.themes.BOOTSTRAP, "/assets/custom.css"],
    update_title="Ask Tricare"
)

def chat_message_card(msg, is_user):
    align = "end" if is_user else "start"
    color = "primary" if is_user else "secondary"
    avatar = "🧑" if is_user else "🤖"
    return dbc.Card(
        dbc.CardBody([
            html.Div([
                html.Span(avatar, style={"fontSize": "2rem"}),
                html.Span(msg, style={"whiteSpace": "pre-wrap", "marginLeft": "0.75rem"})
            ], style={"display": "flex", "alignItems": "center", "justifyContent": align})
        ]),
        className=f"mb-2 ms-3 me-3",
        color=color,
        inverse=is_user,
        style={"maxWidth": "80%", "alignSelf": f"flex-{align}"}
    )

def uploaded_file_card(filename, is_img):
    ext = os.path.splitext(filename)[1].lower()
    icon = "🖼️" if is_img else "📄"
    return dbc.Card(
        dbc.CardBody([
            html.Span(icon, style={"fontSize": "2rem", "marginRight": "0.5rem"}),
            html.Span(filename)
        ]),
        className="mb-2",
        color="tertiary"
    )

def disclaimer_card():
    return dbc.Card(
        dbc.CardBody([
            html.H5("Disclaimer", className="card-title"),
            html.P("This information is not private. Do not send PII or PHI. For official guidance visit the Tricare website.", style={"fontSize": "0.95rem"})
        ]),
        className="mb-2"
    )

def left_navbar_static():
    return html.Div([
        html.H3("Ask Tricare", className="mb-3 mt-3", style={"fontWeight": "bold"}),
        disclaimer_card(),
        dcc.Upload(
            id="file-upload",
            children=dbc.Button("Upload Document/Image", color="secondary", className="mb-2", style={"width": "100%"}),
            multiple=True,
            style={"width": "100%"}
        ),
        html.Div(id="upload-list"),
        html.Hr(),
        html.H5("Chat History", className="mb-2"),
        html.Ul(id="chat-history-list", style={"listStyle": "none", "paddingLeft": "0"}),
    ], style={"padding": "1rem", "backgroundColor": "#f8f9fa", "height": "100vh", "overflowY": "auto"})

def right_main_static():
    return html.Div([
        dbc.Card([
            dbc.CardBody([
                html.Div(id="chat-window", style={"minHeight": "60vh", "display": "flex", "flexDirection": "column", "justifyContent": "flex-end"}),
                html.Div([
                    dcc.Textarea(
                        id="user-input",
                        placeholder="Type your question...",
                        style={"width": "100%", "height": "60px", "resize": "vertical", "wordWrap": "break-word"},
                        wrap="soft",
                        maxLength=1000
                    ),
                    dbc.Button("Send", id="send-btn", color="primary", className="mt-2", style={"float": "right", "minWidth": "100px"}),
                ], style={"marginTop": "1rem"}),
                html.Div(id="error-message", style={"color": "#bb2124", "marginTop": "0.5rem"}),
            ])
        ], className="mt-3"),
        dcc.Loading(id="loading", type="default", fullscreen=False, style={"position": "absolute", "top": "5%", "left": "50%"}),
        dcc.Interval(id="stream-interval", interval=400, n_intervals=0, disabled=True, max_intervals=1000)
    ], style={"padding": "1rem", "backgroundColor": "#fff", "height": "100vh", "overflowY": "auto"})

app.layout = html.Div([
    dcc.Store(id="session-id", storage_type="local"),
    dcc.Location(id="url"),
    html.Div([
        html.Div(left_navbar_static(), id='left-navbar', style={"width": "30vw", "height": "100vh", "position": "fixed", "left": 0, "top": 0, "zIndex": 2, "overflowY": "auto"}),
        html.Div(right_main_static(), id='right-main', style={"marginLeft": "30vw", "width": "70vw", "overflowY": "auto"})
    ], style={"display": "flex"})
])

def _is_supported_doc(filename):
    ext = os.path.splitext(filename)[1].lower()
    return ext in [".txt", ".pdf", ".md", ".docx"]

@app.callback(
    Output("session-id", "data"),
    Input("url", "href"),
    prevent_initial_call=False
)
def assign_session_id(_):
    sid = get_session_id()
    d = get_session_dir(sid)
    load_session_state(sid)
    logger.info(f"Assigned session id: {sid}")
    return sid

@app.callback(
    Output("upload-list", "children"),
    Output("chat-history-list", "children"),
    Output("chat-window", "children"),
    Output("error-message", "children"),
    Output("stream-interval", "disabled"),
    Output("stream-interval", "n_intervals"),
    Input("session-id", "data"),
    Input("send-btn", "n_clicks"),
    Input("file-upload", "contents"),
    State("file-upload", "filename"),
    State("user-input", "value"),
    State("stream-interval", "n_intervals"),
    prevent_initial_call=False
)
def main_callback(session_id, send_clicks, file_contents, file_names, user_input, stream_n):
    trigger = callback_context.triggered[0]['prop_id'].split('.')[0] if callback_context.triggered else ""
    if not session_id:
        session_id = get_session_id()
    session_lock = get_session_lock(session_id)
    with session_lock:
        load_session_state(session_id)
        state = get_session_state(session_id)
        error = ""
        start_streaming = False

        if trigger == "file-upload" and file_contents and file_names:
            uploads = []
            if not isinstance(file_contents, list):
                file_contents = [file_contents]
                file_names = [file_names]
            for c, n in zip(file_contents, file_names):
                header, data = c.split(',', 1)
                ext = os.path.splitext(n)[1].lower()
                is_img = ext in [".png", ".jpg", ".jpeg", ".gif", ".bmp", ".webp"]
                fname = secure_filename(f"{datetime.datetime.utcnow().strftime('%Y%m%d%H%M%S')}_{n}")
                session_dir = get_session_dir(session_id)
                fp = os.path.join(session_dir, fname)
                with open(fp, "wb") as f:
                    f.write(base64.b64decode(data))
                uploads.append({"name": fname, "is_img": is_img, "path": fp})
            state["uploads"].extend(uploads)
            save_session_state(session_id)
            logger.info(f"Session {session_id}: Uploaded files {[u['name'] for u in uploads]}")

        if trigger == "send-btn" and user_input and user_input.strip():
            state["messages"].append({"role": "user", "content": user_input})
            state["streaming"] = True
            state["stream_buffer"] = ""
            save_session_state(session_id)

            def run_stream(session_id, messages):
                try:
                    system_prompt = load_system_prompt()
                    msg_list = [{"role": "system", "content": system_prompt}]
                    for m in messages:
                        msg_list.append({"role": m["role"], "content": m["content"]})
                    response = openai.ChatCompletion.create(
                        model="gpt-3.5-turbo",
                        messages=msg_list,
                        max_tokens=700,
                        temperature=0.2,
                        stream=True,
                    )
                    reply = ""
                    for chunk in response:
                        delta = chunk["choices"][0]["delta"]
                        content = delta.get("content", "")
                        if content:
                            reply += content
                            # Update buffer in session state
                            session_lock = get_session_lock(session_id)
                            with session_lock:
                                load_session_state(session_id)
                                state = get_session_state(session_id)
                                state["stream_buffer"] = reply
                                save_session_state(session_id)
                    # Finalize message
                    session_lock = get_session_lock(session_id)
                    with session_lock:
                        load_session_state(session_id)
                        state = get_session_state(session_id)
                        state["messages"].append({"role": "assistant", "content": reply})
                        state["stream_buffer"] = ""
                        state["streaming"] = False
                        save_session_state(session_id)
                    logger.info(f"Session {session_id}: User: {user_input} | Assistant: {reply}")
                except Exception as e:
                    session_lock = get_session_lock(session_id)
                    with session_lock:
                        load_session_state(session_id)
                        state = get_session_state(session_id)
                        state["streaming"] = False
                        state["stream_buffer"] = ""
                        save_session_state(session_id)
                    logger.error(f"Session {session_id}: Streaming error: {e}")

            threading.Thread(target=run_stream, args=(session_id, list(state["messages"])), daemon=True).start()
            start_streaming = True

        chat_history = state.get("messages", [])
        uploads = state.get("uploads", [])
        upload_cards = [uploaded_file_card(os.path.basename(f["name"]), f["is_img"]) for f in uploads]
        chat_items = [html.Li(html.Span((msg['role'] + ": " + msg['content'])[:40] + ("..." if len(msg['content']) > 40 else ""), style={"fontSize": "0.92rem"})) for msg in chat_history[-6:]]
        chat_cards = []
        for i, msg in enumerate(chat_history):
            if msg['role'] == "user":
                chat_cards.append(chat_message_card(msg['content'], is_user=True))
            elif msg['role'] == "assistant":
                chat_cards.append(chat_message_card(msg['content'], is_user=False))
        if state.get("streaming", False):
            # Add a partial assistant message at the end
            if state.get("stream_buffer", ""):
                chat_cards.append(chat_message_card(state["stream_buffer"], is_user=False))
            return upload_cards, chat_items, chat_cards, error, False, 0 if not start_streaming else False, 0
        return upload_cards, chat_items, chat_cards, error, (not state.get("streaming", False)), 0

@app.callback(
    Output("chat-window", "children"),
    Output("stream-interval", "disabled"),
    Input("stream-interval", "n_intervals"),
    State("session-id", "data"),
    prevent_initial_call=True
)
def poll_stream(n_intervals, session_id):
    session_lock = get_session_lock(session_id)
    with session_lock:
        load_session_state(session_id)
        state = get_session_state(session_id)
        chat_history = state.get("messages", [])
        chat_cards = []
        for i, msg in enumerate(chat_history):
            if msg['role'] == "user":
                chat_cards.append(chat_message_card(msg['content'], is_user=True))
            elif msg['role'] == "assistant":
                chat_cards.append(chat_message_card(msg['content'], is_user=False))
        if state.get("streaming", False):
            if state.get("stream_buffer", ""):
                chat_cards.append(chat_message_card(state["stream_buffer"], is_user=False))
            return chat_cards, False
        return chat_cards, True

@app_flask.after_request
def set_session_cookie(resp):
    sid = flask_request.cookies.get("asktricare_session_id")
    if not sid:
        sid = str(uuid.uuid4())
        resp.set_cookie("asktricare_session_id", sid, max_age=60*60*24*7, path="/")
    return resp

def cleanup_sessions(max_age_hours=48):
    now = datetime.datetime.utcnow()
    for sid in os.listdir(SESSION_DIR_BASE):
        d = os.path.join(SESSION_DIR_BASE, sid)
        try:
            state_path = os.path.join(d, "state.json")
            if os.path.exists(state_path):
                with open(state_path, "r") as f:
                    st = json.load(f)
                created = st.get("created")
                if created and (now - datetime.datetime.fromisoformat(created)).total_seconds() > max_age_hours * 3600:
                    shutil.rmtree(d)
                    logger.info(f"Cleaned up session {sid}")
        except Exception as e:
            logger.error(f"Cleanup error for {sid}: {e}")

try:
    import torch
    if torch.cuda.is_available():
        torch.set_default_tensor_type(torch.cuda.FloatTensor)
        logger.info("CUDA GPU detected and configured.")
except Exception as e:
    logger.warning(f"CUDA config failed: {e}")

if __name__ == '__main__':
    print("Starting the Dash application...")
    app.run(debug=True, host='0.0.0.0', port=7860, threaded=True)
    print("Dash application has finished running.")