""" Shared UI components used by both assistant apps. Keeps the individual app files thin and DRY. """ from __future__ import annotations import json import time from typing import Callable, Generator, List, Optional import streamlit as st from core.memory import ConversationMemory from core.observability import ObservabilityManager from core.safety import SafetyFilter, SafetyResult # --------------------------------------------------------------------------- # Page config helper # --------------------------------------------------------------------------- def configure_page(title: str, icon: str = "🤖") -> None: st.set_page_config( page_title=title, page_icon=icon, layout="wide", initial_sidebar_state="expanded", ) # --------------------------------------------------------------------------- # Sidebar components # --------------------------------------------------------------------------- def render_sidebar_info( model_name: str, backend: str, memory: ConversationMemory, obs_manager: ObservabilityManager, safety_filter: SafetyFilter, ) -> dict: """ Render the sidebar with model info, memory controls, and observability stats. Returns a dict of user-configured settings. """ with st.sidebar: st.markdown("## ⚙️ Configuration") st.markdown(f"**Model:** `{model_name}`") st.markdown(f"**Backend:** `{backend}`") st.divider() # Memory controls st.markdown("### 🧠 Memory") st.markdown(f"Conversation turns: **{memory.turn_count}**") if st.button("🗑️ Clear Memory", use_container_width=True): memory.clear() st.success("Memory cleared.") st.rerun() st.divider() # Safety mode st.markdown("### 🛡️ Safety") strict_mode = st.toggle("Strict safety mode", value=True) st.markdown(f"Refusals this session: **{safety_filter.refusal_count}**") st.divider() # Observability summary st.markdown("### 📊 Session Stats") summary = obs_manager.get_summary() if summary: st.metric("Total Requests", summary.get("total_requests", 0)) st.metric("Avg Latency", f"{summary.get('avg_latency_ms', 0):.0f}ms") st.metric("Refusal Rate", f"{summary.get('refusal_rate', 0)*100:.1f}%") else: st.caption("No requests yet.") st.divider() # Export st.markdown("### 💾 Export") if st.button("📥 Download Chat Log", use_container_width=True): logs = obs_manager.get_all_logs() st.download_button( label="Save logs.json", data=json.dumps(logs, indent=2), file_name="chat_logs.json", mime="application/json", ) # Chat history export history = memory.history if history: chat_export = "\n\n".join( f"**{m.role.upper()}**: {m.content}" for m in history ) st.download_button( label="Save Chat History", data=chat_export, file_name="chat_history.txt", mime="text/plain", use_container_width=True, ) return {"strict_mode": strict_mode} # --------------------------------------------------------------------------- # Chat message rendering # --------------------------------------------------------------------------- def render_chat_history(memory: ConversationMemory) -> None: """Render all messages in the conversation history.""" for msg in memory.history: with st.chat_message(msg.role): st.markdown(msg.content) def render_metrics_row( latency_ms: float, prompt_tokens: int, response_tokens: int, was_refused: bool, ) -> None: """Render a compact metrics row below a response.""" col1, col2, col3, col4 = st.columns(4) col1.metric("⏱️ Latency", f"{latency_ms:.0f}ms") col2.metric("📥 Prompt Tokens", f"~{prompt_tokens}") col3.metric("📤 Response Tokens", f"~{response_tokens}") col4.metric("🚫 Refused", "Yes" if was_refused else "No") # --------------------------------------------------------------------------- # Core chat handler # --------------------------------------------------------------------------- def handle_chat_turn( user_input: str, memory: ConversationMemory, prompt_builder, assistant, safety_filter: SafetyFilter, obs_manager: ObservabilityManager, backend: str, stream: bool = True, ) -> None: """ Execute one full chat turn: 1. Input safety check 2. Build prompt 3. Generate response (streaming or not) 4. Output safety check 5. Update memory 6. Log to observability """ from core.observability import estimate_tokens, Timer # --- Input safety check --- input_check: SafetyResult = safety_filter.check_input(user_input) if not input_check.is_safe: with st.chat_message("assistant"): st.warning(f"🛡️ **Safety Filter Triggered** ({input_check.violation_type.value})") st.markdown(input_check.refusal_message) obs_manager.log_request( model_used=assistant.model_name, backend=backend, prompt=user_input, response=input_check.refusal_message, latency_ms=0.0, was_refused=True, violation_type=input_check.violation_type.value, ) return # --- Build prompt --- messages = prompt_builder.build(memory, user_input) prompt_str = " ".join(m["content"] for m in messages) # --- Generate response --- response_text = "" with st.chat_message("assistant"): if stream: response_placeholder = st.empty() with Timer() as timer: for chunk in assistant.stream_response(messages): response_text += chunk response_placeholder.markdown(response_text + "▌") response_placeholder.markdown(response_text) else: with Timer() as timer: response_text = assistant.generate_response(messages) st.markdown(response_text) latency_ms = timer.elapsed_ms # --- Output safety check --- output_check: SafetyResult = safety_filter.check_output(response_text) if not output_check.is_safe: response_text = output_check.refusal_message with st.chat_message("assistant"): st.warning(f"🛡️ **Output filtered** ({output_check.violation_type.value})") st.markdown(response_text) # --- Update memory --- memory.add_user_message(user_input) memory.add_assistant_message(response_text) # --- Observability --- log = obs_manager.log_request( model_used=assistant.model_name, backend=backend, prompt=prompt_str, response=response_text, latency_ms=latency_ms, was_refused=not output_check.is_safe, violation_type=output_check.violation_type.value if not output_check.is_safe else None, ) # Show metrics with st.expander("📊 Request metrics", expanded=False): render_metrics_row( latency_ms=log.latency_ms, prompt_tokens=log.estimated_prompt_tokens, response_tokens=log.estimated_response_tokens, was_refused=log.was_refused, ) # --------------------------------------------------------------------------- # Observability dashboard # --------------------------------------------------------------------------- def render_observability_dashboard(obs_manager: ObservabilityManager) -> None: """Render a full observability dashboard in a Streamlit expander.""" import pandas as pd summary = obs_manager.get_summary() logs = obs_manager.get_all_logs() st.markdown("## 📊 Observability Dashboard") if not logs: st.info("No requests logged yet. Start chatting to see metrics.") return # Summary metrics col1, col2, col3, col4 = st.columns(4) col1.metric("Total Requests", summary.get("total_requests", 0)) col2.metric("Avg Latency", f"{summary.get('avg_latency_ms', 0):.0f}ms") col3.metric("Refusal Rate", f"{summary.get('refusal_rate', 0)*100:.1f}%") col4.metric("Total Tokens (est.)", f"{summary.get('total_estimated_tokens', 0):,}") # Per-backend breakdown by_backend = summary.get("by_backend", {}) if by_backend: st.markdown("### By Backend") df_backend = pd.DataFrame(by_backend).T.reset_index() df_backend.columns = ["Backend", "Requests", "Avg Latency (ms)", "Refusals"] st.dataframe(df_backend, use_container_width=True) # Recent logs table st.markdown("### Recent Requests") df = pd.DataFrame(logs[-20:]) # Last 20 if not df.empty: display_cols = [ "timestamp", "backend", "model_used", "latency_ms", "estimated_prompt_tokens", "estimated_response_tokens", "was_refused", ] display_cols = [c for c in display_cols if c in df.columns] st.dataframe(df[display_cols], use_container_width=True)