Spaces:
Sleeping
Sleeping
| """ | |
| FinRAG Gradio UI — web interface for the FinRAG chatbot. | |
| Features: | |
| - Chat interface with conversation history | |
| - Retrieval mode selector (local / global / hybrid) | |
| - Filing filter with live dropdown from Neo4j | |
| - Source citations shown per response | |
| - Sidebar commands: /risks, /financials, /compare | |
| - Reset conversation button | |
| """ | |
| import gradio as gr | |
| from typing import List, Tuple, Optional | |
| # ── Lazy engine singleton (OOM-safe) ─────────────────────────────────── # | |
| _engine = None | |
| _resolver = None | |
| _init_error = None | |
| def _get_engine(): | |
| global _engine, _resolver, _init_error | |
| if _init_error: | |
| raise RuntimeError(_init_error) | |
| if _engine is not None: | |
| return _engine, _resolver | |
| try: | |
| print("Initialising FinRAG engine (first request)...") | |
| from llm.fin_rag_engine import FinRAGEngine | |
| from data.retrieval.filing_resolver import FilingResolver | |
| _engine = FinRAGEngine() | |
| _resolver = FilingResolver() | |
| print("Engine ready.") | |
| return _engine, _resolver | |
| except Exception as e: | |
| _init_error = str(e) | |
| print(f"Engine init failed: {e}") | |
| raise | |
| # ── Filing helpers ───────────────────────────────────────────────────── # | |
| def get_filing_choices() -> List[str]: | |
| """Load all available filings for the dropdown.""" | |
| try: | |
| _, resolver = _get_engine() | |
| docs = resolver.list_all() | |
| choices = ["All filings"] | |
| for d in docs: | |
| label = f"{d.get('company_name', '?')} — {d.get('filing_id', '?')} ({d.get('fiscal_year', '?')})" | |
| choices.append(label) | |
| return choices | |
| except Exception: | |
| return ["All filings"] | |
| def extract_filing_id(choice: str) -> Optional[str]: | |
| """Parse filing_id out of the dropdown label.""" | |
| if not choice or choice == "All filings": | |
| return None | |
| # label format: "NVIDIA — NVDA10K2024.md (2024)" | |
| parts = choice.split("—") | |
| if len(parts) >= 2: | |
| return parts[1].strip().split(" ")[0].strip() | |
| return None | |
| # ── Core chat function ───────────────────────────────────────────────── # | |
| def chat( | |
| message: str, | |
| history: List[dict], | |
| mode: str, | |
| filing_choice: str, | |
| top_k: int, | |
| ) -> Tuple[List[dict], str]: | |
| """ | |
| Main chat handler called by Gradio. | |
| Returns (updated_history, sources_text). | |
| """ | |
| if not message.strip(): | |
| return history, "" | |
| # Load engine lazily | |
| try: | |
| engine, resolver = _get_engine() | |
| except Exception as e: | |
| return history + [ | |
| {"role": "user", "content": message}, | |
| {"role": "assistant", "content": f"Service unavailable: {e}"}, | |
| ], "" | |
| filing_id = extract_filing_id(filing_choice) | |
| # Handle slash commands | |
| if message.strip().startswith("/"): | |
| parts = message.strip().split(maxsplit=1) | |
| cmd = parts[0].lower() | |
| arg = parts[1].strip() if len(parts) > 1 else "" | |
| if cmd == "/reset": | |
| engine.reset_history() | |
| return [], "[OK] Conversation reset." | |
| elif cmd == "/risks": | |
| ref = arg or filing_id | |
| if not ref: | |
| reply = "[X] Please specify a company: `/risks NVIDIA` or select a filing from the dropdown." | |
| history.extend([{"role": "user", "content": message}, {"role": "assistant", "content": reply}]) | |
| return history, "" | |
| resp = engine.summarise_risks(ref) | |
| sources = ", ".join(resp.source_filings) or "N/A" | |
| history.extend([{"role": "user", "content": message}, {"role": "assistant", "content": resp.answer}]) | |
| return history, f" {sources}" | |
| elif cmd == "/financials": | |
| ref = arg or filing_id | |
| if not ref: | |
| reply = "[X] Please specify a company: `/financials MSFT` or select a filing." | |
| history.extend([{"role": "user", "content": message}, {"role": "assistant", "content": reply}]) | |
| return history, "" | |
| resp = engine.extract_financials(ref) | |
| sources = ", ".join(resp.source_filings) or "N/A" | |
| history.extend([{"role": "user", "content": message}, {"role": "assistant", "content": resp.answer}]) | |
| return history, f" {sources}" | |
| elif cmd == "/compare": | |
| if not arg: | |
| reply = "[X] Usage: `/compare NVIDIA, Microsoft | AI revenue`" | |
| history.extend([{"role": "user", "content": message}, {"role": "assistant", "content": reply}]) | |
| return history, "" | |
| # Parse: "NVIDIA, Microsoft | topic" | |
| if "|" in arg: | |
| companies_raw, topic = arg.split("|", 1) | |
| else: | |
| companies_raw = arg | |
| topic = "overall business performance" | |
| companies = [c.strip() for c in companies_raw.split(",") if c.strip()] | |
| resp = engine.compare_companies(companies, topic.strip()) | |
| sources = ", ".join(resp.source_filings) or "N/A" | |
| history.extend([{"role": "user", "content": message}, {"role": "assistant", "content": resp.answer}]) | |
| return history, f" {sources}" | |
| elif cmd == "/help": | |
| help_text = ( | |
| "**Available commands:**\n\n" | |
| "- `/reset` — clear conversation history\n" | |
| "- `/risks <company>` — risk summary (e.g. `/risks NVIDIA`)\n" | |
| "- `/financials <company>` — key financials (e.g. `/financials MSFT`)\n" | |
| "- `/compare <co1>, <co2> | <topic>` — compare companies\n\n" | |
| "Or just ask naturally: *What are Google's main revenue segments?*" | |
| ) | |
| history.extend([{"role": "user", "content": message}, {"role": "assistant", "content": help_text}]) | |
| return history, "" | |
| # Normal question | |
| resp = engine.ask( | |
| question=message, | |
| mode=mode.lower(), | |
| filing_id=filing_id, | |
| top_k=top_k, | |
| rerank_top_k=5, | |
| ) | |
| sources = ", ".join(resp.source_filings) if resp.source_filings else "N/A" | |
| history.extend([{"role": "user", "content": message}, {"role": "assistant", "content": resp.answer}]) | |
| return history, f" Sources: {sources} | 🔍 Mode: {resp.retrieval_mode}" | |
| def reset_chat() -> Tuple[List, str, str]: | |
| if _engine: | |
| _engine.reset_history() | |
| return [], "", "[OK] Conversation cleared." | |
| # ── Gradio UI ────────────────────────────────────────────────────────── # | |
| CSS = """ | |
| @import url('https://fonts.googleapis.com/css2?family=DM+Serif+Display:ital@0;1&family=DM+Mono:wght@400;500&display=swap'); | |
| body, .gradio-container { | |
| font-family: 'DM Mono', monospace !important; | |
| background: #0d0f12 !important; | |
| } | |
| /* Header */ | |
| .fin-header { | |
| background: linear-gradient(135deg, #0d0f12 0%, #111419 100%); | |
| border-bottom: 1px solid #1e2530; | |
| padding: 20px 28px 16px; | |
| margin-bottom: 0; | |
| } | |
| .fin-title { | |
| font-family: 'DM Serif Display', serif !important; | |
| font-size: 28px !important; | |
| color: #e8dfc8 !important; | |
| letter-spacing: -0.5px; | |
| margin: 0 0 4px 0; | |
| } | |
| .fin-subtitle { | |
| font-family: 'DM Mono', monospace !important; | |
| font-size: 11px !important; | |
| color: #4a5568 !important; | |
| letter-spacing: 1.5px; | |
| text-transform: uppercase; | |
| } | |
| /* Chatbot */ | |
| .chatbot-wrap .wrap { | |
| background: #0d0f12 !important; | |
| border: 1px solid #1e2530 !important; | |
| border-radius: 8px !important; | |
| } | |
| .chatbot-wrap .message.user { | |
| background: #141820 !important; | |
| border: 1px solid #1e2530 !important; | |
| color: #c8d0dc !important; | |
| font-family: 'DM Mono', monospace !important; | |
| font-size: 13px !important; | |
| border-radius: 6px 6px 2px 6px !important; | |
| } | |
| .chatbot-wrap .message.bot { | |
| background: #0f1318 !important; | |
| border: 1px solid #1a2235 !important; | |
| color: #d4c9a8 !important; | |
| font-family: 'DM Mono', monospace !important; | |
| font-size: 13px !important; | |
| border-radius: 6px 6px 6px 2px !important; | |
| } | |
| /* Input */ | |
| .input-row textarea { | |
| background: #111419 !important; | |
| border: 1px solid #1e2530 !important; | |
| border-radius: 6px !important; | |
| color: #c8d0dc !important; | |
| font-family: 'DM Mono', monospace !important; | |
| font-size: 13px !important; | |
| resize: none !important; | |
| } | |
| .input-row textarea:focus { | |
| border-color: #c9a84c !important; | |
| box-shadow: 0 0 0 2px rgba(201,168,76,0.1) !important; | |
| } | |
| /* Buttons */ | |
| .send-btn { | |
| background: #c9a84c !important; | |
| color: #0d0f12 !important; | |
| font-family: 'DM Mono', monospace !important; | |
| font-weight: 500 !important; | |
| font-size: 12px !important; | |
| letter-spacing: 0.5px !important; | |
| border: none !important; | |
| border-radius: 6px !important; | |
| min-width: 80px !important; | |
| } | |
| .send-btn:hover { | |
| background: #dbb855 !important; | |
| } | |
| .reset-btn { | |
| background: transparent !important; | |
| color: #4a5568 !important; | |
| font-family: 'DM Mono', monospace !important; | |
| font-size: 11px !important; | |
| border: 1px solid #1e2530 !important; | |
| border-radius: 6px !important; | |
| } | |
| .reset-btn:hover { | |
| border-color: #4a5568 !important; | |
| color: #6b7a8d !important; | |
| } | |
| /* Sidebar controls */ | |
| .sidebar { | |
| background: #0f1318 !important; | |
| border: 1px solid #1e2530 !important; | |
| border-radius: 8px !important; | |
| padding: 16px !important; | |
| } | |
| .sidebar label { | |
| color: #4a5568 !important; | |
| font-size: 10px !important; | |
| letter-spacing: 1.5px !important; | |
| text-transform: uppercase !important; | |
| font-family: 'DM Mono', monospace !important; | |
| } | |
| .sidebar .wrap { | |
| background: #111419 !important; | |
| border: 1px solid #1e2530 !important; | |
| border-radius: 4px !important; | |
| color: #c8d0dc !important; | |
| font-family: 'DM Mono', monospace !important; | |
| font-size: 12px !important; | |
| } | |
| .sidebar input[type=range] { | |
| accent-color: #c9a84c !important; | |
| } | |
| /* Sources bar */ | |
| .sources-bar { | |
| background: #0f1318 !important; | |
| border: 1px solid #1a2235 !important; | |
| border-radius: 4px !important; | |
| color: #4a6080 !important; | |
| font-family: 'DM Mono', monospace !important; | |
| font-size: 11px !important; | |
| padding: 8px 12px !important; | |
| } | |
| /* Section labels */ | |
| .section-label { | |
| color: #2a3545 !important; | |
| font-size: 10px !important; | |
| letter-spacing: 2px !important; | |
| text-transform: uppercase !important; | |
| margin-bottom: 8px !important; | |
| font-family: 'DM Mono', monospace !important; | |
| } | |
| /* Command pills */ | |
| .cmd-pill { | |
| display: inline-block; | |
| background: #141820; | |
| border: 1px solid #1e2530; | |
| color: #c9a84c; | |
| font-size: 11px; | |
| padding: 2px 8px; | |
| border-radius: 3px; | |
| margin: 2px; | |
| font-family: 'DM Mono', monospace; | |
| } | |
| /* Hide Gradio footer */ | |
| footer { display: none !important; } | |
| .built-with { display: none !important; } | |
| /* Scrollbar */ | |
| ::-webkit-scrollbar { width: 4px; } | |
| ::-webkit-scrollbar-track { background: #0d0f12; } | |
| ::-webkit-scrollbar-thumb { background: #1e2530; border-radius: 2px; } | |
| """ | |
| COMMANDS_MD = """ | |
| <div style="font-family:'DM Mono',monospace;font-size:11px;color:#4a5568;line-height:2"> | |
| <span style="color:#2a3545;letter-spacing:1.5px;text-transform:uppercase;font-size:10px">Commands</span><br> | |
| <span style="color:#c9a84c">/risks</span> <company><br> | |
| <span style="color:#c9a84c">/financials</span> <company><br> | |
| <span style="color:#c9a84c">/compare</span> A, B | topic<br> | |
| <span style="color:#c9a84c">/reset</span><br> | |
| <span style="color:#c9a84c">/help</span> | |
| </div> | |
| """ | |
| def build_ui(): | |
| filing_choices = get_filing_choices() | |
| with gr.Blocks( | |
| title="FinRAG", | |
| ) as demo: | |
| # ── Header ──────────────────────────────────────────────────── # | |
| gr.HTML(""" | |
| <div class="fin-header"> | |
| <div class="fin-title">FinRAG</div> | |
| <div class="fin-subtitle">SEC 10-K · HybridRAG · Neo4j · Weaviate · DeepSeek-R1 · Groq</div> | |
| </div> | |
| """) | |
| with gr.Row(equal_height=True): | |
| # ── Main chat panel ──────────────────────────────────────── # | |
| with gr.Column(scale=3): | |
| chatbot = gr.Chatbot( | |
| label="", | |
| height=520, | |
| show_label=False, | |
| elem_classes=["chatbot-wrap"], | |
| render_markdown=True, | |
| avatar_images=(None, None), | |
| placeholder=( | |
| "<div style='text-align:center;color:#1e2530;" | |
| "font-family:DM Mono,monospace;font-size:12px;" | |
| "padding:60px 20px'>" | |
| "Ask about any SEC 10-K filing<br>" | |
| "<span style='font-size:10px;letter-spacing:1px;" | |
| "text-transform:uppercase;color:#151c28'>" | |
| "or type /help for commands</span>" | |
| "</div>" | |
| ), | |
| ) | |
| sources_display = gr.Textbox( | |
| label="", | |
| placeholder="Sources will appear here after each response", | |
| interactive=False, | |
| show_label=False, | |
| elem_classes=["sources-bar"], | |
| lines=1, | |
| ) | |
| with gr.Row(elem_classes=["input-row"]): | |
| msg_input = gr.Textbox( | |
| placeholder="Ask a question or type /help for commands...", | |
| show_label=False, | |
| scale=5, | |
| lines=1, | |
| max_lines=4, | |
| autofocus=True, | |
| ) | |
| send_btn = gr.Button( | |
| "Send", | |
| variant="primary", | |
| scale=1, | |
| elem_classes=["send-btn"], | |
| ) | |
| # ── Sidebar ──────────────────────────────────────────────── # | |
| with gr.Column(scale=1, elem_classes=["sidebar"]): | |
| gr.HTML("<div class='section-label'>Retrieval mode</div>") | |
| mode_radio = gr.Radio( | |
| choices=["Hybrid", "Local", "Global"], | |
| value="Hybrid", | |
| show_label=False, | |
| ) | |
| gr.HTML("<div class='section-label' style='margin-top:16px'>Filing filter</div>") | |
| filing_dropdown = gr.Dropdown( | |
| choices=filing_choices, | |
| value="All filings", | |
| show_label=False, | |
| allow_custom_value=False, | |
| ) | |
| gr.HTML("<div class='section-label' style='margin-top:16px'>Chunks retrieved</div>") | |
| top_k_slider = gr.Slider( | |
| minimum=4, | |
| maximum=20, | |
| value=10, | |
| step=2, | |
| show_label=False, | |
| ) | |
| gr.HTML("<div class='section-label' style='margin-top:16px'>Quick actions</div>") | |
| risks_btn = gr.Button("Risk summary", size="sm", variant="secondary") | |
| financials_btn = gr.Button("Key financials", size="sm", variant="secondary") | |
| gr.HTML("<div style='margin-top:16px'>" + COMMANDS_MD + "</div>") | |
| reset_btn = gr.Button( | |
| "↺ Reset conversation", | |
| size="sm", | |
| elem_classes=["reset-btn"], | |
| ) | |
| # ── Event handlers ───────────────────────────────────────────── # | |
| def submit(message, history, mode, filing, top_k): | |
| new_history, sources = chat(message, history, mode, filing, top_k) | |
| return new_history, sources, "" # clear input after send | |
| # Send on button click | |
| send_btn.click( | |
| fn=submit, | |
| inputs=[msg_input, chatbot, mode_radio, filing_dropdown, top_k_slider], | |
| outputs=[chatbot, sources_display, msg_input], | |
| ) | |
| # Send on Enter (Shift+Enter for newline) | |
| msg_input.submit( | |
| fn=submit, | |
| inputs=[msg_input, chatbot, mode_radio, filing_dropdown, top_k_slider], | |
| outputs=[chatbot, sources_display, msg_input], | |
| ) | |
| # Quick action: risks button | |
| def quick_risks(history, filing): | |
| fid = extract_filing_id(filing) | |
| if not fid: | |
| msg = "Please select a filing from the dropdown first." | |
| return history + [{"role": "user", "content": "/risks"}, {"role": "assistant", "content": msg}], "" | |
| new_h, src = chat(f"/risks {fid}", history, "Local", filing, 10) | |
| return new_h, src | |
| risks_btn.click( | |
| fn=quick_risks, | |
| inputs=[chatbot, filing_dropdown], | |
| outputs=[chatbot, sources_display], | |
| ) | |
| # Quick action: financials button | |
| def quick_financials(history, filing): | |
| fid = extract_filing_id(filing) | |
| if not fid: | |
| msg = "Please select a filing from the dropdown first." | |
| return history + [{"role": "user", "content": "/financials"}, {"role": "assistant", "content": msg}], "" | |
| new_h, src = chat(f"/financials {fid}", history, "Local", filing, 10) | |
| return new_h, src | |
| financials_btn.click( | |
| fn=quick_financials, | |
| inputs=[chatbot, filing_dropdown], | |
| outputs=[chatbot, sources_display], | |
| ) | |
| # Reset | |
| reset_btn.click( | |
| fn=reset_chat, | |
| outputs=[chatbot, sources_display, sources_display], | |
| ) | |
| return demo | |
| # ── Launch ───────────────────────────────────────────────────────────── # | |
| if __name__ == "__main__": | |
| demo = build_ui() | |
| demo.launch( | |
| server_name="0.0.0.0", | |
| server_port=7860, | |
| share=False, # set True to get a public gradio.live URL | |
| show_error=True, | |
| theme=gr.themes.Base( | |
| primary_hue="yellow", | |
| neutral_hue="slate", | |
| font=gr.themes.GoogleFont("DM Mono"), | |
| ), | |
| css=CSS, | |
| ) |