dashboard/chat.py

"""Chat Q&A section — free-form questions over stored documents for the current ticker.

The chat answers are grounded exclusively in SEC filings, earnings transcripts, and
structured metrics (no live news / web search).  The agent loop is bounded at 5 rounds
and every response must cite its sources via chunk_context headers.
"""
from __future__ import annotations

import streamlit as st

from agent.llm import RunConfig, classify_llm_error
from dashboard.theme import (
    BORDER,
    BG_MUTED,
    TEXT,
    TEXT_MUTED,
    AI_COLOR,
    FS_PAGE,
    FS_META,
    SPACE_4,
    SPACE_6,
)
from storage import metrics_db
from dashboard.i18n import t


# Tool → human label, used both for suggested-question chips (n/a) and the
# "Searched: ..." line shown under each answer.
_TOOL_LABELS: dict[str, str] = {
    "search_filing": "Filings (10-K/10-Q)",
    "search_transcript": "Transcript",
    "get_financial_metrics": "Metrics",
    "get_analyst_expectations": "Analyst data",
}


def _watch_to_question(item: str) -> str:
    """Reword a `what_to_watch` item into a chat question (pure, no LLM call)."""
    item = item.strip().rstrip(".")
    return f'What do the filings and transcript say about: "{item}"?'


def _searched_labels(sources: list[dict]) -> list[str]:
    called = {s.get("tool_name") for s in sources}
    return [label for name, label in _TOOL_LABELS.items() if name in called]


def _suggested_questions(ticker: str, brief: dict | None) -> list[str]:
    questions: list[str] = []
    if brief:
        for item in (brief.get("what_to_watch") or [])[:3]:
            if isinstance(item, str) and item.strip():
                questions.append(_watch_to_question(item))
            elif isinstance(item, dict):
                text = item.get("text") or item.get("watch") or ""
                if text:
                    questions.append(_watch_to_question(text))
    questions.append(t("chat_q_changed"))
    questions.append(t("chat_q_risks"))
    return questions[:4]


# ── Public entry point ────────────────────────────────────────────────────────

def render(ticker: str, brief: dict | None = None, config: RunConfig | None = None) -> None:
    """Render the chat Q&A tab for *ticker*.

    Requires: the ticker must have been ingested (``metrics_db`` has rows for it).
    Does NOT require a brief to have been generated (though *brief*, if passed,
    seeds the suggested-question chips).
    """
    # Header — mirrors quality_tone.py:179-200 pattern.
    st.markdown(
        f"""
        <div style="border-bottom:1px solid {BORDER};padding-bottom:{SPACE_4};
                    margin-bottom:{SPACE_6};">
            <div style="font-size:{FS_PAGE};font-weight:700;color:{TEXT};
                        letter-spacing:-0.02em;">
                Ask about {ticker}
            </div>
            <div style="font-size:{FS_META};color:{TEXT_MUTED};margin-top:4px;">
                {t("chat_subtitle")}
            </div>
        </div>
        """,
        unsafe_allow_html=True,
    )

    if config is None:
        st.warning(t("model_blocked_chat"), icon="⚠️")
        return

    # Guard: ticker must be ingested.
    rows = metrics_db.get_all_metrics(ticker)
    if not rows:
        st.warning(
            t("chat_not_ingested").format(ticker=ticker),
            icon="⚠️",
        )
        return

    # Per-ticker chat history stored in session_state (mirrors reasoning_trace at app.py:233).
    chat_history = st.session_state.setdefault("chat_history", {})
    history: list[dict] = chat_history.setdefault(ticker, [])

    # Render existing messages.
    for msg in history:
        role = msg["role"]
        with st.chat_message(role):
            st.markdown(msg["content"])
            if role == "assistant":
                if msg.get("sources"):
                    _render_sources_expander(msg["sources"])
                if msg.get("searched"):
                    _render_searched_line(msg["searched"])

    # Suggested-question chips — only when the thread is empty, so returning
    # visitors aren't shown stale prompts mid-conversation.
    clicked_chip: str | None = None
    if not history:
        st.caption(t("chat_suggested_label"))
        chips = _suggested_questions(ticker, brief)
        cols = st.columns(len(chips))
        for i, (col, chip_text) in enumerate(zip(cols, chips)):
            with col:
                if st.button(chip_text, key=f"chat_chip_{i}", use_container_width=True):
                    clicked_chip = chip_text

    # Capture new user input; question resolution order: typed input > CTA
    # prefill from another page > a clicked suggested-question chip.
    user_input = st.chat_input(t("chat_input_placeholder").format(ticker=ticker))
    question = user_input or st.session_state.pop("chat_prefill", None) or clicked_chip
    if not question:
        return

    # Display the user message immediately.
    history.append({"role": "user", "content": question})
    with st.chat_message("user"):
        st.markdown(question)

    # Call the chat agent and stream the answer.
    searched: list[str] = []
    with st.chat_message("assistant"):
        with st.spinner(t("chat_searching")):
            try:
                from agent.chat_agent import answer_question
                # Pass history BEFORE the current question (history[-1] is the just-added
                # user message; exclude it since it's passed separately as `question`).
                result = answer_question(ticker, question, history[:-1], config=config)
                answer = result["answer"]
                sources = result["sources"]
                searched = _searched_labels(sources)
            except Exception as exc:
                friendly = classify_llm_error(exc, config.provider)
                answer = f'⚠️ {friendly or t("chat_error").format(error=f"`{exc}`")}'
                sources = []

        st.markdown(answer)
        if sources:
            _render_sources_expander(sources)
        if searched:
            _render_searched_line(searched)

    # Persist the assistant reply (with sources for later re-render).
    history.append({"role": "assistant", "content": answer, "sources": sources, "searched": searched})
    st.rerun()


# ── Private helpers ───────────────────────────────────────────────────────────

def _render_searched_line(searched: list[str]) -> None:
    """Show which document collections were queried for this answer."""
    if not searched:
        return
    st.markdown(
        f"<div style='font-size:{FS_META};color:{AI_COLOR};margin-top:4px;'>"
        f"{t('chat_searched_prefix')} {' · '.join(searched)}</div>",
        unsafe_allow_html=True,
    )


def _render_sources_expander(sources: list[dict]) -> None:
    """Render a collapsible block showing the raw tool outputs used to build the answer."""
    if not sources:
        return
    label = t("chat_sources").format(n=len(sources))
    with st.expander(label, expanded=False):
        for i, src in enumerate(sources, 1):
            st.markdown(
                f"<div style='font-size:{FS_META};font-weight:600;color:{TEXT_MUTED};"
                f"margin-bottom:4px;'>{i}. {src['tool_name']}({_fmt_args(src['args'])})</div>",
                unsafe_allow_html=True,
            )
            st.markdown(
                f"<pre style='background:{BG_MUTED};border:1px solid {BORDER};"
                f"border-radius:6px;padding:8px 10px;font-size:0.72rem;"
                f"white-space:pre-wrap;overflow-x:auto;color:{TEXT_MUTED};"
                f"margin-bottom:10px;'>{src['output']}</pre>",
                unsafe_allow_html=True,
            )


def _fmt_args(args: dict) -> str:
    """Format tool call args for the sources expander header."""
    parts = [f"{k}={v!r}" for k, v in args.items() if v not in (None, "")]
    return ", ".join(parts)