financial-rag / src /ui.py
anasxs's picture
Fix Gradio CSS launch for Space
cdb66aa
Raw
History Blame Contribute Delete
2.97 kB
from __future__ import annotations
from functools import lru_cache
from typing import Any
import gradio as gr
from pipeline import build_default_pipeline, QueryFilters
COMPANY_OPTIONS = [
"All companies",
"Apple",
"Amazon",
"Alphabet",
"Berkshire Hathaway",
"Johnson and Johnson",
]
FILING_OPTIONS = ["All filings", "10-K", "10-Q"]
CUSTOM_CSS = """
.gradio-container { max-width: 1100px !important; margin: auto !important; }
#title { text-align: center; margin-bottom: 8px; }
#subtitle { text-align: center; color: #475569; margin-bottom: 24px; }
"""
@lru_cache(maxsize=1)
def get_pipeline() -> Any :
return build_default_pipeline()
def normalize_filter(value: str) -> str | None:
return None if value.startswith("All") else value
def format_sources(sources: list[Any]) -> str:
if not sources:
return "No sources returned."
lines: list[str] = []
for source in sources:
metadata = source.metadata
label = build_source_label(source.source_id,metadata)
lines.append(f"**{label}**\n\n{source.text[:900]}...")
return "\n\n---\n\n".join(lines)
def build_source_label(source_id: int, metadata: dict[str, Any]) -> str:
company = metadata.get("company_name", "Unknown company")
filing = metadata.get("filing_type", "Unknown filing")
year = metadata.get("fiscal_year", "Unknown year")
page = metadata.get("page_number", "Unknown page")
return f"Source {source_id}: {company} | {filing} | {year} | page {page}"
def answer_question(question:str, company:str, filing_type:str) -> tuple[str, str]:
if not question.strip():
return "Please enter a question.", ""
filters = QueryFilters(normalize_filter(company), normalize_filter(filing_type))
try:
response = get_pipeline().answer_question(question, filters)
except Exception as exc:
return f"Query failed: {exc}", ""
return response.answer, format_sources(response.sources)
def build_interface() -> gr.Blocks:
"""Build a clean Gradio interface for portfolio demos."""
with gr.Blocks(title="Financial RAG") as demo:
gr.Markdown("# Financial RAG", elem_id="title")
gr.Markdown("Ask SEC filing questions with hybrid retrieval and cited sources.", elem_id="subtitle")
with gr.Row():
company = gr.Dropdown(COMPANY_OPTIONS, value="All companies", label="Company")
filing = gr.Dropdown(FILING_OPTIONS, value="All filings", label="Filing type")
question = gr.Textbox(label="Question", lines=3, placeholder="What are Amazon's major business segments?")
submit = gr.Button("Ask", variant="primary")
answer = gr.Textbox(label="Answer", lines=8)
sources = gr.Markdown(label="Sources")
submit.click(answer_question, inputs=[question, company, filing], outputs=[answer, sources])
return demo
if __name__ == "__main__":
get_pipeline()
build_interface().launch(css=CUSTOM_CSS)