apps/shared/components/chat.py

"""Chat UI components — message rendering and streaming."""

from __future__ import annotations

from typing import Iterator

import streamlit as st

from src.serving.citation import cite, CitationResult


def render_chat_history(messages: list[dict]):
    """Render all messages in the chat history."""
    for msg in messages:
        with st.chat_message(msg["role"]):
            st.markdown(msg["content"])


def stream_response(generator: Iterator[str]) -> str:
    """Stream LLM output to the chat UI, return the full response text."""
    with st.chat_message("assistant"):
        response = st.write_stream(generator)
    return response


def render_agent_response(result) -> None:
    """Render agent reasoning trace + final answer in a chat bubble."""
    with st.chat_message("assistant"):
        for i, step in enumerate(result.steps):
            if step.thought:
                st.caption("Step {} — {}".format(i + 1, step.thought))
            if step.tool_call:
                args = ", ".join(
                    "{}={}".format(k, v)
                    for k, v in step.tool_call.arguments.items()
                )
                st.info("🔧 {}({})".format(step.tool_call.tool_name, args))
            if step.observation:
                with st.expander("Tool result (step {})".format(i + 1)):
                    st.code(step.observation[:800], language=None)
        st.markdown(result.answer)


def render_cited_response(
    answer: str,
    chunks: list[tuple[str, float]],
    threshold: float = 0.10,
) -> CitationResult:
    """Render an answer with citation markers and a references section.

    Displays the annotated answer with ``[N]`` markers inline, followed
    by an expandable references section showing the source chunk for
    each citation.

    Returns the CitationResult so callers can store it in session state.
    """
    result = cite(answer, chunks, threshold=threshold)

    with st.chat_message("assistant"):
        st.markdown(result.annotated_answer)
        if result.citations:
            with st.expander(
                "Sources ({} citation{})".format(
                    len(result.citations),
                    "s" if len(result.citations) != 1 else "",
                ),
                expanded=False,
            ):
                for ref_line in result.references.split("\n"):
                    if ref_line.strip():
                        st.caption(ref_line)

    return result