import streamlit as st import csv import pandas as pd import logging import sys import os import re import shutil from pathlib import Path sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) from src.agents.CleanReponseAgent import ai_format_response from src.data.FinancialKBIngestor import FinancialKBIngestor from src.core.FinAgentEngine import FinAgentEngine from src.agents.ModerationAgent import ModerationAgent from src.core.settings import get_settings from src.core.errors import add_error from src.core.logging_config import configure_logging logger = logging.getLogger(__name__) settings = get_settings() configure_logging() def reset_chroma_if_configured() -> None: """ If FIN_ASSISTANT_RESET_CHROMA=1, wipe the local Chroma persist directory. This clears *both* the Knowledge Base vectors and the semantic cache because they share `src/data/.chroma`. """ flag = str(os.getenv("FIN_ASSISTANT_RESET_CHROMA", "") or "").strip().lower() if flag not in {"1", "true", "yes", "y"}: return # Avoid repeated resets on Streamlit reruns within the same session/process. try: if st.session_state.get("_chroma_reset_done"): return except Exception: pass # If any Chroma-backed objects were cached via Streamlit, clear them first so we # don't keep stale in-memory clients pointing at a deleted directory. try: st.cache_resource.clear() except Exception: pass try: st.cache_data.clear() except Exception: pass chroma_dir = (Path(__file__).resolve().parents[0] / "data" / ".chroma").resolve() if chroma_dir.exists(): try: shutil.rmtree(chroma_dir) logger.info("Reset Chroma persist directory: %s", str(chroma_dir)) except Exception as e: logger.exception("Failed to reset Chroma persist directory: %s", e) return try: chroma_dir.mkdir(parents=True, exist_ok=True) except Exception: pass try: st.session_state["_chroma_reset_done"] = True except Exception: pass def _strip_html(text: str) -> str: """ Best-effort HTML removal for converting rendered assistant messages back into plain text. """ if not text: return "" # Remove tags and collapse whitespace. no_tags = re.sub(r"<[^>]+>", " ", text) return re.sub(r"\s+", " ", no_tags).strip() def build_llm_conversation_history(messages: list[dict]) -> list[dict]: """ Convert UI chat messages into a clean, plain-text conversation history for the LLM. - Prefer `msg["raw"]` when present (authoritative plain text). - Fall back to stripping HTML from `msg["content"]` (for older sessions). """ out: list[dict] = [] for m in messages or []: if not isinstance(m, dict): continue role = m.get("role") content = m.get("raw") if content is None: content = _strip_html(str(m.get("content", ""))) out.append({"role": role, "content": str(content or "")}) return out def get_router(): return FinAgentEngine() def get_moderation_agent(): return ModerationAgent() @st.cache_resource def get_tax_kb(): """ Initialize and ingest the local JSON knowledge base(s) into Chroma only once. """ logger.info("Initializing local JSON Knowledge Bases (Chroma ingest)...") ingestor = FinancialKBIngestor( json_file=[ "tax_kb.json", "us_insurance_kb.json", "us_market_basics.json", "us_portfolio_kb.json", "us_stock_crypto_kb.json", ], use_qa_format=True ) ingestor.run() logger.info("Local JSON Knowledge Bases initialized.") BAD_LANGUAGE_RESPONSE = ( "Please try with different query, like : What is the current price of Apple Stock ? " "Explain 401k in simple terms ?" ) BAD_WORDS_CSV_PATH = str( (Path(__file__).resolve().parents[1] / "data" / "bad_words.csv").resolve() ) def load_bad_words(csv_path: str) -> set[str]: """ Load bad words from a CSV file. """ bad_words: set[str] = set() path = Path(csv_path) if not path.exists(): return bad_words with path.open("r", encoding="utf-8", newline="") as f: reader = csv.reader(f) for row in reader: if not row: continue w = (row[0] or "").strip().lower() if not w or w.startswith("#") or w == "word": continue bad_words.add(w) return bad_words def _contains_bad_language(text: str) -> bool: bad_words = load_bad_words(BAD_WORDS_CSV_PATH) normalized = (text or "").lower() tokens = re.findall(r"[a-z]+", normalized) return any(t in bad_words for t in tokens) def sanitize_user_input(user_input: str) -> tuple[bool, str]: """ Returns (is_allowed, response_text_if_blocked). """ if _contains_bad_language(user_input): return False, BAD_LANGUAGE_RESPONSE moderation_agent = get_moderation_agent() moderation_result = moderation_agent.classify(user_input) if moderation_result.get("flagged"): return False, BAD_LANGUAGE_RESPONSE return True, "" reset_chroma_if_configured() get_tax_kb() # --- Page Config --- st.set_page_config(page_title="AI Finance Assistant", page_icon="๐Ÿ’ฐ", layout="wide") # --- Custom CSS --- st.markdown(""" """, unsafe_allow_html=True) # --- Session State --- if "messages" not in st.session_state: st.session_state.messages = [] if "history" not in st.session_state: st.session_state.history = [] if "user_profile" not in st.session_state: st.session_state.user_profile = { "risk": settings.default_user_profile.risk, "experience": settings.default_user_profile.experience, } if "portfolio" not in st.session_state: st.session_state.portfolio = [] if "last_errors" not in st.session_state: st.session_state.last_errors = [] # --- Sidebar --- st.sidebar.title("โš™๏ธ Settings") st.sidebar.markdown("Customize your experience") # Load S&P 500 data def load_sp500_data(): try: file_path = "data/sp500_top100.csv" if not os.path.isabs(file_path): base_dir = os.path.dirname(os.path.abspath(__file__)) file_path = os.path.join(base_dir, file_path) df = pd.read_csv(file_path) return df except Exception as e: logger.error(f"Error loading S&P 500 data: {e}") return pd.DataFrame(columns=["Symbol", "Company Name"]) sp500_df = load_sp500_data() risk = st.sidebar.selectbox( "Risk Profile", ["Low", "Moderate", "High"], index=["Low", "Moderate", "High"].index(st.session_state.user_profile["risk"].capitalize()) ) experience = st.sidebar.selectbox( "Experience", ["Beginner", "Intermediate", "Advanced"], index=["Beginner", "Intermediate", "Advanced"].index(st.session_state.user_profile["experience"].capitalize()) ) st.session_state.user_profile["risk"] = risk.lower() st.session_state.user_profile["experience"] = experience.lower() st.sidebar.markdown("---") st.sidebar.subheader("๐Ÿ‘ค Profile") # Dropdown for portfolio selection selected_stocks = st.sidebar.multiselect( "Select Stocks for Portfolio", options=sp500_df["Symbol"].tolist(), default=[p["symbol"] for p in st.session_state.portfolio if "symbol" in p], max_selections=5 ) # New section for entering quantities new_portfolio = [] if selected_stocks: st.sidebar.markdown("##### Set Quantities") for symbol in selected_stocks: # Try to find existing quantity in session state existing_qty = next((p["quantity"] for p in st.session_state.portfolio if p.get("symbol") == symbol), 1) qty = st.sidebar.number_input(f"Shares of {symbol}", min_value=1, value=existing_qty, key=f"qty_{symbol}") new_portfolio.append({"symbol": symbol, "quantity": qty}) # Update session state portfolio immediately st.session_state.portfolio = new_portfolio # --- Header --- st.markdown('
๐Ÿ’ฐ AI Finance Assistant
', unsafe_allow_html=True) st.markdown('
Ask about investing, portfolio, market trends, or taxes. Stock prices are sourced from Alpha Vantage and Finnhub.
', unsafe_allow_html=True) # Profile info display p_risk = st.session_state.user_profile.get("risk") p_exp = st.session_state.user_profile.get("experience") p_portfolio_display = [f"{p['symbol']} ({p['quantity']})" for p in st.session_state.portfolio if "symbol" in p] if p_risk and p_exp and p_portfolio_display: st.markdown(f"""
Risk Profile: {p_risk.capitalize()} | Experience: {p_exp.capitalize()} | Portfolio: {', '.join(p_portfolio_display)}
""", unsafe_allow_html=True) else: st.warning("Please select Risk Profile, Experience and Portfolio from settings") # --- Status banner (best-effort telemetry) --- last_errors = st.session_state.get("last_errors") or [] if isinstance(last_errors, list) and last_errors: st.warning("Some tools/data sources returned errors. Results may be incomplete.") with st.expander("Show details"): for e in last_errors[:12]: try: agent = e.get("agent", "unknown") msg = e.get("message", "") code = e.get("code", "") st.write(f"- {agent} [{code}]: {msg}") except Exception: continue if st.button("Dismiss errors"): st.session_state.last_errors = [] st.rerun() # --- Layout --- main_col, history_col = st.columns([3, 1]) # --- Chat UI --- with main_col: for msg in st.session_state.messages: if msg["role"] == "user": st.markdown(f'
{msg["content"]}
', unsafe_allow_html=True) else: st.markdown(f'
{msg["content"]}
', unsafe_allow_html=True) user_input = st.chat_input("Ask something like 'Explain ETFs'") # --- Backend --- def call_backend(query, state): router = get_router() try: result_state = router.invoke(query, initial_state=state) return result_state except Exception as e: result_state = {"response": f"An error occurred during processing: {e}", "errors": []} add_error( result_state, # type: ignore[arg-type] code="backend_error", message=str(e), agent="streamlit_app", ) return result_state # --- Handle Input --- if user_input: allowed, blocked_response = sanitize_user_input(user_input) if not allowed: st.session_state.messages.append({"role": "user", "content": user_input, "raw": user_input}) st.markdown(f'
{user_input}
', unsafe_allow_html=True) st.session_state.messages.append({"role": "assistant", "content": blocked_response, "raw": blocked_response}) st.markdown(f'
{blocked_response}
', unsafe_allow_html=True) st.rerun() st.session_state.messages.append({"role": "user", "content": user_input, "raw": user_input}) with st.spinner("Working on it..."): graph_state = { "user_query": user_input, "conversation_history": build_llm_conversation_history(st.session_state.messages), "user_profile": st.session_state.user_profile, "portfolio": st.session_state.portfolio, "errors": [], } result_state = call_backend(user_input, graph_state) response_raw = result_state.get("response", "No response was generated by the agents.") st.session_state.last_errors = result_state.get("errors") or [] # Update persistent portfolio from graph state (if present) if "portfolio" in result_state: st.session_state.portfolio = result_state["portfolio"] response_rendered = ai_format_response(response_raw) st.session_state.messages.append( {"role": "assistant", "content": response_rendered, "raw": response_raw} ) if user_input not in [h["query"] for h in st.session_state.history]: item = {"query": user_input, "pinned": False} st.session_state.history.insert(0, item) st.session_state.history = st.session_state.history[:50] st.rerun() # --- History Panel --- with history_col: st.markdown("### ๐Ÿ•˜ Search History") col1, col2 = st.columns(2) with col1: if st.button("๐Ÿงน Clear All"): st.session_state.history = [] st.session_state.messages = [] st.rerun() with col2: show_pinned = st.checkbox("โญ Pinned only") history_sorted = sorted(st.session_state.history, key=lambda x: (not x.get("pinned", False))) if show_pinned: history_sorted = [h for h in history_sorted if h.get("pinned")] if history_sorted: for idx, item in enumerate(history_sorted): q = item["query"] pinned = item.get("pinned", False) c1, c2, c3 = st.columns([6,1,1]) with c1: if st.button(q, key=f"run_{idx}"): st.session_state.messages.append({"role": "user", "content": q, "raw": q}) with st.spinner("Working on it..."): graph_state = { "user_query": q, "conversation_history": build_llm_conversation_history(st.session_state.messages), "user_profile": st.session_state.user_profile, "portfolio": st.session_state.portfolio, "errors": [], } result_state = call_backend(q, graph_state) st.session_state.last_errors = result_state.get("errors") or [] resp_raw = result_state.get("response", "No response was generated by the agents.") resp_rendered = ai_format_response(resp_raw) st.session_state.messages.append( {"role": "assistant", "content": resp_rendered, "raw": resp_raw} ) st.rerun() with c2: icon = "โญ" if pinned else "โ˜†" if st.button(icon, key=f"pin_{idx}"): for i, h in enumerate(st.session_state.history): if h["query"] == q: st.session_state.history[i]["pinned"] = not h.get("pinned", False) break st.rerun() with c3: if st.button("๐Ÿ—‘๏ธ", key=f"del_{idx}"): # remove from history st.session_state.history = [h for h in st.session_state.history if h["query"] != q] # remove related messages (user + assistant) new_messages = [] skip_next = False for i, msg in enumerate(st.session_state.messages): if msg["role"] == "user" and msg["content"] == q: skip_next = True continue if skip_next: skip_next = False continue new_messages.append(msg) st.session_state.messages = new_messages st.rerun() else: st.write("No history yet")