backend/gradio_ui/chat_tab.py

# chat_tab.py
"""Streaming QA chat panel."""

from __future__ import annotations

import asyncio
import tempfile
import threading
from typing import Any, Dict, List, Optional, Tuple

import gradio as gr

from core.bootstrap import get_chat_store, get_db, get_qa_service
from gradio_ui.renderers import render_chat_transcript
from gradio_ui.state import SearchMode, footer_hint
from utils.response_cleaner import clean_model_response

ChatMessage = Dict[str, Any]
ChatHistory = List[ChatMessage]

_THINKING = "⏳ *Thinking…*"


def _friendly_error(exc: BaseException) -> str:
    msg = str(exc).strip()
    lower = msg.lower()
    if "504" in msg or "timeout" in lower or "serviceerror" in lower:
        return (
            "The AI service timed out while loading or generating a reply. "
            "This often happens on a cold start — wait a minute and try again."
        )
    if "preemption" in lower or "preempt" in lower:
        return "The AI worker was restarted. Please send your question again."
    if not msg:
        return "Something went wrong while generating a reply. Please try again."
    if len(msg) > 220:
        return msg[:217] + "…"
    return msg


def _history_to_api(history: ChatHistory) -> List[Dict[str, str]]:
    api: List[Dict[str, str]] = []
    for msg in history or []:
        role = msg.get("role")
        content = msg.get("content") or ""
        if role in ("user", "assistant") and content and content != _THINKING:
            cleaned = clean_model_response(content) if role == "assistant" else content
            if cleaned:
                api.append({"role": role, "content": cleaned})
    return api


def _transcript(history: ChatHistory) -> str:
    return render_chat_transcript(history or [])


def _start_exchange(history: ChatHistory, message: str) -> ChatHistory:
    updated = list(history or [])
    updated.append({"role": "user", "content": message})
    updated.append({"role": "assistant", "content": _THINKING})
    return updated


def _patch_assistant(
    history: ChatHistory,
    *,
    content: str | None = None,
    sources: List[Dict[str, Any]] | None = None,
    confidence: float | None = None,
) -> ChatHistory:
    updated = list(history or [])
    if not updated or updated[-1].get("role") != "assistant":
        updated.append({"role": "assistant", "content": content or _THINKING})
    msg = dict(updated[-1])
    if content is not None:
        msg["content"] = content
    if sources is not None:
        msg["sources"] = sources
    if confidence is not None:
        msg["confidence"] = float(confidence)
    updated[-1] = msg
    return updated


def _streaming_content(raw: str) -> str:
    cleaned = clean_model_response(raw)
    if cleaned:
        return cleaned
    return _THINKING


def _persist_message(
    session_id: str,
    query: str,
    answer: str,
    sources: List[Dict[str, Any]],
    confidence: Optional[float],
) -> None:
    store = get_chat_store()
    store.add_message(session_id, "user", query)
    store.add_message(
        session_id,
        "assistant",
        answer,
        sources=sources,
        confidence=confidence,
    )


async def stream_chat(
    message: str,
    history: ChatHistory,
    search_mode: SearchMode,
    selected_doc_ids: List[str],
    session_id: Optional[str],
):
    if not message or not message.strip():
        yield history, _transcript(history)
        return

    qa = get_qa_service()
    history = _start_exchange(history, message)
    yield history, _transcript(history)

    prior = _history_to_api(history[:-2])
    doc_ids = selected_doc_ids if selected_doc_ids else None
    sources: List[Dict[str, Any]] = []
    confidence: Optional[float] = None
    tokens: List[str] = []

    loop = asyncio.get_running_loop()
    item_queue: asyncio.Queue[Any] = asyncio.Queue()
    stream_error: List[BaseException] = []

    def _produce_stream() -> None:
        try:
            for stream_item in qa.stream_answer(
                message,
                prior,
                document_ids=doc_ids,
                mode=search_mode,
            ):
                loop.call_soon_threadsafe(item_queue.put_nowait, stream_item)
        except BaseException as exc:
            stream_error.append(exc)
        finally:
            loop.call_soon_threadsafe(item_queue.put_nowait, None)

    threading.Thread(target=_produce_stream, daemon=True).start()

    while True:
        item = await item_queue.get()
        if item is None:
            break
        if isinstance(item, dict):
            if item.get("type") == "sources":
                sources = item.get("data") or []
                history = _patch_assistant(history, sources=sources)
                yield history, _transcript(history)
            elif item.get("type") == "confidence":
                confidence = item.get("data")
                history = _patch_assistant(
                    history,
                    confidence=float(confidence) if confidence is not None else None,
                )
                yield history, _transcript(history)
        else:
            tokens.append(str(item))
            history = _patch_assistant(
                history,
                content=_streaming_content("".join(tokens)),
            )
            yield history, _transcript(history)

    if stream_error:
        history = _patch_assistant(
            history,
            content=f"**Error:** {_friendly_error(stream_error[0])}",
            sources=sources,
            confidence=float(confidence) if confidence is not None else None,
        )
        yield history, _transcript(history)
        return

    answer = clean_model_response("".join(tokens))
    if not answer:
        answer = "Sorry, I couldn't generate a response. Please try again."
    history = _patch_assistant(
        history,
        content=answer,
        sources=sources,
        confidence=float(confidence) if confidence is not None else None,
    )

    if session_id:
        _persist_message(session_id, message, answer, sources, confidence)

    yield history, _transcript(history)


def create_new_session(
    search_mode: SearchMode,
    selected_doc_ids: List[str],
) -> Tuple[Optional[str], ChatHistory, str, str, gr.update]:
    store = get_chat_store()
    doc_ids = selected_doc_ids if search_mode == "single" else None
    session = store.create_session("New Chat", doc_ids)
    session_id = session["id"]
    choices = _session_choices()
    doc_count = len(get_db().list_documents())
    return (
        session_id,
        [],
        _transcript([]),
        footer_hint(search_mode, selected_doc_ids, doc_count),
        gr.update(choices=choices, value=session_id),
    )


def load_session(session_id: Optional[str]) -> Tuple[ChatHistory, str]:
    if not session_id:
        return [], _transcript([])
    store = get_chat_store()
    messages = store.get_messages(session_id)
    history: ChatHistory = []

    for msg in messages:
        role = msg.get("role")
        content = msg.get("content") or ""
        if role not in ("user", "assistant"):
            continue
        entry: ChatMessage = {"role": role, "content": content}
        if role == "assistant":
            entry["sources"] = msg.get("sources") or []
            conf = msg.get("confidence")
            if conf is not None:
                entry["confidence"] = float(conf)
        history.append(entry)

    return history, _transcript(history)


def export_session_file(session_id: Optional[str]) -> Optional[str]:
    if not session_id:
        return None
    store = get_chat_store()
    content = store.export_session(session_id)
    if not content:
        return None
    tmp = tempfile.NamedTemporaryFile(
        mode="w",
        suffix=".txt",
        prefix=f"finsight_session_{session_id[:8]}_",
        delete=False,
        encoding="utf-8",
    )
    tmp.write(content)
    tmp.close()
    return tmp.name


def _session_choices() -> List[Tuple[str, str]]:
    store = get_chat_store()
    sessions = store.list_sessions()
    return [(f"{s['title']} ({s['id'][:8]})", s["id"]) for s in sessions]


def on_load_sessions() -> Tuple[Optional[str], gr.update, ChatHistory, str]:
    choices = _session_choices()
    if choices:
        session_id = choices[0][1]
        history, transcript = load_session(session_id)
        return session_id, gr.update(choices=choices, value=session_id), history, transcript
    session = get_chat_store().create_session("New Chat")
    sid = session["id"]
    choices = _session_choices()
    return sid, gr.update(choices=choices, value=sid), [], _transcript([])


def build_chat_panel() -> dict:
    chat_transcript = gr.HTML(
        render_chat_transcript([]),
        padding=False,
        container=False,
        autoscroll=False,
        elem_id="fs-chat-transcript",
        elem_classes=["fs-chat-transcript"],
    )

    export_file = gr.File(label="Exported session", visible=False, elem_classes=["hidden-control"])

    return {
        "chat_transcript": chat_transcript,
        "export_file": export_file,
        "on_load_sessions": on_load_sessions,
        "create_new_session": create_new_session,
        "load_session": load_session,
        "export_session_file": export_session_file,
        "stream_chat": stream_chat,
    }