Harshavard21/FinRAG

1"""2app/pages/1_Chat.py3====================4Single company Q&A page with streaming answers and citation display.5"""6 7import sys8import os9sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))))10from config.settings import settings11 12import streamlit as st13 14st.set_page_config(page_title="Q&A | FinRAG", page_icon="💬", layout="wide")15 16# Shared resource loaders (cached across all pages)17from app.rag_engine import load_retriever, load_reranker, load_groq18 19# Apply same CSS20st.markdown("""21<style>22@import url('https://fonts.googleapis.com/css2?family=Inter:wght@300;400;500;600;700&family=JetBrains+Mono:wght@400;500&display=swap');23:root {24    --bg-card: rgba(255,255,255,0.04);25    --accent-cyan: #06b6d4;26    --accent-blue: #3b82f6;27    --text-primary: #f1f5f9;28    --text-secondary: #94a3b8;29    --border: rgba(255,255,255,0.08);30}31.stApp { background: linear-gradient(135deg, #0a0e1a, #0f1628); font-family: 'Inter', sans-serif; }32#MainMenu, footer, header { visibility: hidden; }33.stButton > button {34    background: linear-gradient(135deg, #3b82f6, #06b6d4) !important;35    color: white !important; border: none !important;36    border-radius: 8px !important; font-weight: 600 !important;37}38.stButton > button:hover { transform: translateY(-1px) !important; }39.citation-chip {40    display: inline-block; background: rgba(59,130,246,0.15);41    border: 1px solid rgba(59,130,246,0.3); border-radius: 20px;42    padding: 0.2rem 0.7rem; font-size: 0.72rem; color: #93c5fd; margin: 0.15rem;43    font-family: 'JetBrains Mono', monospace;44}45.fin-card {46    background: var(--bg-card); border: 1px solid var(--border);47    border-radius: 12px; padding: 1rem 1.2rem; margin-bottom: 0.8rem;48}49</style>50""", unsafe_allow_html=True)51 52 53# ---- Header ----54st.markdown("""55<div style="padding:1.5rem 0 1rem 0;">56    <h1 style="font-size:1.8rem;font-weight:700;margin:0;57               background:linear-gradient(135deg,#3b82f6,#06b6d4);58               -webkit-background-clip:text;-webkit-text-fill-color:transparent;">59        💬 Company Q&A60    </h1>61    <p style="color:#64748b;font-size:0.9rem;margin-top:0.3rem;">62        Ask anything about a company's financials. Answers are grounded in BSE annual reports.63    </p>64</div>65""", unsafe_allow_html=True)66 67# ---- Controls ----68col_company, col_fy = st.columns([2, 1])69 70with col_company:71    companies = list(settings.company_ticker_map.keys())72    selected_company = st.selectbox(73        "Select Company",74        options=companies,75        index=companies.index("TCS") if "TCS" in companies else 0,76        key="qa_company",77    )78 79with col_fy:80    fy_options = ["All Years", "FY2025", "FY2024"]81    selected_fy = st.selectbox("Fiscal Year", fy_options, key="qa_fy")82    fy_filter = None if selected_fy == "All Years" else selected_fy83 84# ---- Sample questions ----85sample_questions = {86    "TCS": [87        "What was TCS's total revenue and net profit for FY2025?",88        "What are TCS's key business segments and their revenue contributions?",89        "What is TCS's headcount and employee attrition rate?",90    ],91    "HDFC": [92        "What is HDFC Bank's gross NPA and net NPA ratio?",93        "What was HDFC Bank's net interest income for FY2025?",94        "What is HDFC Bank's CASA ratio?",95    ],96    "INFOSYS": [97        "What was Infosys's operating margin for FY2025?",98        "What are Infosys's key geographies and their revenue split?",99        "What is Infosys's guidance for the next fiscal year?",100    ],101}102 103if selected_company in sample_questions:104    with st.expander("Sample questions for " + selected_company, expanded=False):105        for q in sample_questions[selected_company]:106            if st.button(q, key=f"sample_{q[:20]}"):107                st.session_state["qa_prefill"] = q108 109# ---- Chat Input ----110default_q = st.session_state.pop("qa_prefill", "")111user_query = st.text_area(112    "Your question",113    value=default_q,114    placeholder="e.g. What was the company's net profit margin for FY2025?",115    height=80,116    key="qa_input",117)118 119col_btn, col_opts = st.columns([1, 3])120with col_btn:121    ask_btn = st.button("Ask Question", type="primary", key="qa_ask")122with col_opts:123    show_chunks = st.checkbox("Show retrieved chunks", value=False, key="qa_show_chunks")124 125# ---- Answer ----126if ask_btn and user_query.strip():127    retriever = load_retriever()128    llm = load_groq()129 130    from src.generation.prompts import build_qa_prompt131 132    # Stage 1: Hybrid retrieval (fast ~1-2s)133    with st.spinner(f"Step 1/3 — Searching {selected_company} documents (hybrid dense+sparse)..."):134        candidates = retriever.retrieve(135            query=user_query,136            top_k=50,137            company_filter=selected_company,138            fiscal_year_filter=fy_filter,139            expand_query=True,140            promote_to_parent=True,141        )142 143    if not candidates:144        st.warning("No relevant documents found. Try a different question or company.")145        st.stop()146 147    # Stage 2: Reranking (first time: downloads ~550MB model, ~2-5 min; subsequent: ~3s)148    reranker = load_reranker()149    with st.spinner("Step 2/3 — Reranking with cross-encoder (first run downloads BGE model ~550MB)..."):150        reranked = reranker.rerank(151            query=user_query,152            candidates=candidates,153            top_n=5,154        )155 156    if not reranked:157        st.warning("No relevant documents found. Try a different question or company.")158        st.stop()159 160    # Build prompt and stream answer161    system_prompt, user_message = build_qa_prompt(162        query=user_query,163        results=reranked,164        company_context=f"{selected_company} ({settings.company_ticker_map.get(selected_company, '')})",165    )166 167    # Stage 3: LLM generation via Groq (streaming)168    st.markdown("---")169    st.markdown(f"""170    <div style="font-size:0.8rem;color:#64748b;margin-bottom:0.5rem;">171        Step 3/3 — Generating answer from {len(reranked)} chunks | {selected_company} {fy_filter or 'all years'}172    </div>173    """, unsafe_allow_html=True)174 175    # Stream the answer176    with st.chat_message("assistant", avatar="📈"):177        answer = st.write_stream(llm.generate_stream(system_prompt, user_message))178 179    # Citations180    st.markdown("**Sources:**")181    citation_html = ""182    for chunk in reranked:183        label = f"{chunk.company} | {chunk.fiscal_year} | Pg.{chunk.page_number}"184        if chunk.content_type == "table":185            label += " [TABLE]"186        citation_html += f'<span class="citation-chip">{label}</span>'187    st.markdown(citation_html, unsafe_allow_html=True)188 189    # Retrieved chunks expander190    if show_chunks:191        with st.expander(f"Retrieved chunks ({len(reranked)})", expanded=False):192            for i, chunk in enumerate(reranked, 1):193                st.markdown(f"""194                <div class="fin-card">195                    <div style="font-size:0.72rem;color:#64748b;margin-bottom:0.4rem;">196                        Chunk {i} | {chunk.company} | {chunk.section} | Page {chunk.page_number} |197                        Score: {chunk.rrf_score:.3f} | Type: {chunk.content_type}198                    </div>199                    <div style="font-size:0.82rem;color:#cbd5e1;white-space:pre-wrap;">{chunk.content[:600]}{"..." if len(chunk.content) > 600 else ""}</div>200                </div>201                """, unsafe_allow_html=True)202 203elif ask_btn:204    st.warning("Please enter a question.")205