"""Streamlit web UI for the RAG Research Chatbot.""" import os import re import uuid import json import zipfile import streamlit as st from datetime import datetime from typing import Dict, List from src.config_loader import load_config, get_api_key from src.ingest import get_chroma_collection, ingest_documents from src.kb_meta import load_kb_meta_brief from src.query_engine import understand_query, categorize_query from src.retriever import clear_collection_cache, retrieve from src.verifier import verify_and_respond from src.llm import list_models from src.stata import explain_ from huggingface_hub import HfApi from huggingface_hub import snapshot_download DATA_DIR = "./data" os.makedirs(DATA_DIR, exist_ok=True) snapshot_download(repo_id="CGIAR/weai-kb", repo_type="dataset", token=os.getenv('HF_TOKEN'), local_dir=DATA_DIR ) # ── Session initialisation ─────────────────────────────────────────────────── def init_session(): """Initialise st.session_state with messages list, cfg, and last_retrieval.""" if "messages" not in st.session_state: st.session_state.messages = [] if "cfg" not in st.session_state: try: st.session_state.cfg = load_config() st.session_state.cfg.setdefault("paths")["knowledge_base"] = f'{DATA_DIR}/{st.session_state.cfg.get("paths", {}).get("knowledge_base", "knowledge_base")}' st.session_state.cfg.setdefault("paths")["vector_db"] = f'{DATA_DIR}/{st.session_state.cfg.get("paths", {}).get("vector_db", "vector_db")}' st.session_state.cfg.setdefault("paths")["sql_db"] = f'{DATA_DIR}/{st.session_state.cfg.get("paths", {}).get("sql_db", "sql_db")}' except Exception as e: st.error(f"Configuration error: {e}\n\nPlease run `python setup.py` first.") st.stop() if "last_retrieval" not in st.session_state: st.session_state.last_retrieval = None if "pending_clarification" not in st.session_state: st.session_state.pending_clarification = None if "unresolved_category" not in st.session_state: st.session_state.unresolved_category = None if "clarification_rounds" not in st.session_state: st.session_state.clarification_rounds = 0 if "resolution_rounds" not in st.session_state: st.session_state.resolution_rounds = 0 if "pending_clarification_question" not in st.session_state: st.session_state.pending_clarification_question = None if "unresolved_category_question" not in st.session_state: st.session_state.unresolved_category_question = None if "stata_options" not in st.session_state: st.session_state.stata_options = None if "stata_file" not in st.session_state: st.session_state.stata_file = None if "chat_id" not in st.session_state: st.session_state.chat_id = datetime.now().strftime("%Y-%m-%d %H:%M:%S") if "feedback_data" not in st.session_state: st.session_state.feedback_data = None @st.dialog("Provide Your Feedback") def feedback_prompt(): st.write("How would you rate your WEAIBot experience?") # Built-in Streamlit feedback widget (thumbs or stars) rating = st.feedback("faces") # or st.feedback("stars") # Optional text details comments = st.text_area("Tell us more (optional):") feedback_data = { "rating": rating, "comments": comments, } col1, col2 = st.columns(2) with col1: if st.button("Submit Feedback", type="primary", use_container_width=True): # Save feedback to session state or database st.session_state["feedback_data"] = feedback_data st.toast("Thank you for your feedback!", icon="🎉") with open(f"{DATA_DIR}/logs/{st.session_state.chat_id}/feedback.json", "w", encoding="utf-8") as f: json.dump(feedback_data, f, indent=4) st.rerun() with col2: if st.button("Skip", use_container_width=True): st.session_state["feedback_data"] = {} st.rerun() # ── Sidebar ────────────────────────────────────────────────────────────────── def render_sidebar(): """Render sidebar with provider/model selectors, web search toggle, and KB stats.""" cfg = st.session_state.cfg with st.sidebar: st.header("Settings") # --- Provider dropdown --- providers = ["openai", "anthropic", "gemini", "hugging-face"] current_provider = cfg.get("llm", {}).get("provider", "openai") provider_index = providers.index(current_provider) if current_provider in providers else 0 provider = st.selectbox( "LLM Provider", providers, index=provider_index, key="sidebar_provider", ) # Update cfg in session when provider changes if provider != cfg.get("llm", {}).get("provider"): cfg.setdefault("llm", {})["provider"] = provider # --- Model dropdown (cached per provider) --- # Invalidate model cache if provider changed prev_provider_key = "prev_provider" if st.session_state.get(prev_provider_key) != provider: for p in providers: st.session_state.pop(f"models_{p}", None) st.session_state[prev_provider_key] = provider models_cache_key = f"models_{provider}" if models_cache_key not in st.session_state: api_key = get_api_key(cfg, provider) if api_key: try: st.session_state[models_cache_key] = list_models(provider, api_key) except Exception: st.session_state[models_cache_key] = [] else: st.session_state[models_cache_key] = [] available_models = st.session_state[models_cache_key] current_model = cfg.get("llm", {}).get("model", "") if available_models: model_index = ( available_models.index(current_model) if current_model in available_models else 0 ) model = st.selectbox( "Model", available_models, index=model_index, key="sidebar_model", ) else: model = st.text_input( "Model", value=current_model, key="sidebar_model_text", ) # Update cfg in session when model changes if model != cfg.get("llm", {}).get("model"): cfg.setdefault("llm", {})["model"] = model # --- Web search toggle --- web_enabled = cfg.get("web_search", {}).get("enabled", False) web_toggle = st.toggle("Web search", value=web_enabled, key="sidebar_web_search") cfg.setdefault("web_search", {})["enabled"] = web_toggle st.divider() # --- Knowledge base stats --- st.subheader("Knowledge Base") try: collection = get_chroma_collection(cfg) chunk_count = collection.count() st.metric("Chunks indexed", chunk_count) except Exception as e: st.warning(f"Could not read knowledge base: {e}") chunk_count = 0 # --- Re-ingest button --- if st.button("Re-ingest documents", use_container_width=True): st.info("Ingestion may take a few minutes for large document collections...") with st.spinner("Ingesting documents..."): try: count = ingest_documents(cfg) clear_collection_cache() st.success(f"Ingested {count} chunks.") # Clear cached data so it refreshes after re-ingest st.session_state.pop("kb_welcome_summary", None) st.rerun() except Exception as e: st.error(f"Ingestion failed: {e}") def save_feedback(index): def feedback_callback(): # Store the feedback value back into the corresponding message in history st.session_state.messages[index]["feedback"] = st.session_state.get(f"feedback_{index}", None) return feedback_callback # ── Chat interface ─────────────────────────────────────────────────────────── def export_outputs(): with st.status("Processing...", expanded=True) as status: api = HfApi() api.upload_folder( folder_path=DATA_DIR, repo_id="CGIAR/weai-kb", repo_type="dataset", # create_repo=True,# Options: "model", "dataset", or "space" ) status.update( label="Done!", state="complete", ) def process_pending(cfg:Dict, pending:str, prompt:str, default:Dict, max_history:int=1, iterate:bool=True, categorize:bool=False, # max_clarification:int=1, ): if pending is not None: # This prompt is the user's clarification answer # Include the clarification question for context if categorize: pending_question = st.session_state.get("unresolved_category_question", "") else: pending_question = st.session_state.get("pending_clarification_question", "") if pending_question: combined = f"{pending} (Clarification: Q: {pending_question} A: {prompt})" else: combined = f"{pending} — {prompt}" # st.session_state.pending_clarification_question = None original_query = pending st.session_state.pending_clarification = None else: pending_question = "" combined = prompt original_query = prompt st.session_state.clarification_rounds = 0 # search_query = combined # display_query = combined # route = "vector" # sql_query = None # max_clarifications = qu_cfg.get("max_clarifications", 1) if iterate: # Exclude the just-appended user message to avoid sending # the current question twice (once in history, once as query) prior_messages = st.session_state.messages[:-1] history = [ {"role": m["role"], "content": m["content"]} for m in prior_messages[-max_history:] ] try: if categorize: result = categorize_query(combined, cfg, history) else: result = understand_query(combined, cfg, history) except Exception: default["search_query"] = combined default["display_query"] = combined default["original_query"] = original_query result = default return result, combined, original_query, pending_question def process_contact_request(original_query:str, qcat_result:Dict, uploaded_files:List): ticket = str(uuid.uuid4()).split("-")[2] email_format = r"^[a-zA-Z0-9._%+-]+@[a-zA-Z0-9.-]+\.[a-zA-Z]{2,}$" os.makedirs(f"{DATA_DIR}/tickets/{ticket}", exist_ok=True) contact_data = { "ticket":ticket, "session_id":st.session_state.chat_id, "email":qcat_result.get("email"), "action":qcat_result.get("action"), "status":"open", "chat_history":st.session_state.messages, } if qcat_result.get("action") == "contact": if not re.match(email_format, qcat_result.get("email")): return_msg = "Your request doesn't seem to contain a valid email address. \ Be sure to include an email through which our team can reach you for further assistance." st.session_state.unresolved_category = original_query st.session_state.unresolved_category_question = return_msg with st.chat_message("assistant"): st.markdown(return_msg) feedback = st.feedback(options="faces", key=f"fb_qu_{st.session_state.clarification_rounds-1}") st.session_state.messages.append({"role": "assistant", "content": return_msg, "feedback":feedback}) return st.markdown("""\ Thank you for reaching out. Someone from the WEAI team will contact you as soon as they are available. For additional assistance, or to check the status of your ticket (**{ticket_id}**), please contact the WEAI Helpdesk at IFPRI-WEAI@cgiar.org.\n\n Please let me know if there's anything else I can assist you with today.""".format(ticket_id=ticket)) #save ticket with open(f"{DATA_DIR}/tickets/{ticket}/contact_data.json", "w", encoding="utf-8") as f: json.dump(contact_data, f, indent=4) if qcat_result.get("action") == "upload": if uploaded_files: # Specify the path where the zip file will be saved zip_filename = "uploaded_files.zip" # Write uploaded files into a single zip archive on disk with zipfile.ZipFile(zip_filename, "w", zipfile.ZIP_DEFLATED) as zipf: for file in uploaded_files: # Write each file's bytes directly into the zip zipf.writestr(f"{DATA_DIR}/tickets/{ticket}/{file.name}", file.getvalue()) with open(f"{DATA_DIR}/tickets/{ticket}/contact_data.json", "w", encoding="utf-8") as f: json.dump(contact_data, f, indent=4) st.success(f"Successfully uploaded {len(uploaded_files)} files!") with st.chat_message("assistant"): st.markdown("Thank you for sharing these resources with us. The WEAI team will work to \ review the files and reach out to you should we need any more information. \ Please let me know if there's anything else I can assist you with today.") export_outputs() def process_stata_query(qcat_result, uploaded_files): save_dir = f"{DATA_DIR}/stata_files" os.makedirs(save_dir, exist_ok=True) file_path = "" if qcat_result.get("action") == "do": options = {"rewrite do-file":"rewrite", "explain the do-file":"explain", "suggest fix(es) for any errors in the do-file":"suggestfix"} # Accept .do files uploaded_file = uploaded_files[0] if uploaded_file is not None: # Define the full file path on disk file_path = os.path.join(save_dir, uploaded_file.name) # Write the uploaded bytes to the local path with open(file_path, "wb") as f: f.write(uploaded_file.getbuffer()) # st.success(f"File successfully uploaded") if qcat_result.get("action") == "code": options = {"rewrite the stata code":"rewrite", "explain the stata code":"explain", "suggest fix(es) for any errors in the code":"suggestfix"} if qcat_result.get("action") == "error": options = {"explain the error":"explain", "suggest fix(es) for the error":"suggestfix"} st.session_state.stata_options = options st.session_state.stata_file = file_path st.session_state.unresolved_category_question = f"Please let me know what you need help with: {' or '.join(options.keys())}. List all options that apply." st.session_state.resolution_rounds += 1 resolution_msg = f"Before I search, please let me know what you need help with: {' or '.join(options.keys())}. List all options that apply." with st.chat_message("assistant"): st.markdown(resolution_msg) feedback = st.feedback(options="faces", key=f"fb_stata_{st.session_state.resolution_rounds-1}") st.session_state.messages.append({"role": "assistant", "content": resolution_msg, "feedback":feedback}) def resolve_stata_query(cfg, qcat_result, combined_cat, original_query_cat): selected_options = [st.session_state.stata_options.get(s) for s in qcat_result.get("code-action", "explain")] opts = { "rewrite": True if "rewrite" in selected_options else False, "explain": True if "explain" in selected_options else False, "suggestfix": True if "suggestfix" in selected_options else False, # "capture": True if "capture" in selected_options else False, # "verbose": True if "verbose" in selected_options else False, "lines": None#lines if "lines" in selected_options else None } explanation = explain_(qcat_result.get("action"), combined_cat, cfg, st.session_state.stata_file, opts) st.session_state.unresolved_category = original_query_cat st.session_state.unresolved_category_question = f"{explanation}\n\n Please let me know if this answer is helpful or if it requires further clarification." st.session_state.resolution_rounds += 1 resolution_msg = f"{explanation}\n\n Please let me know if this answer is helpful or if it requires further clarification." with st.chat_message("assistant"): st.markdown(resolution_msg) feedback = st.feedback(options="faces", key=f"fb_stata__code_{st.session_state.resolution_rounds-1}") st.session_state.messages.append({"role": "assistant", "content": resolution_msg, "feedback":feedback}) def render_chat(): """Render the chat interface with message history and input.""" cfg = st.session_state.cfg queries = 0 # Display chat history for i, msg in enumerate(st.session_state.messages): with st.chat_message(msg["role"]): st.markdown(msg["content"]) # If it's an assistant message, show the feedback widget if msg["role"] == "assistant": st.feedback( options="faces", key=f"feedback_{i}", # on_change=save_feedback(i), # Pre-populate if feedback already exists in history default=msg.get("feedback", None), ) if msg["role"] == "user": queries += 1 if queries>4 and st.session_state.feedback_data is None: feedback_prompt() # Chat input user_input = st.chat_input("Ask a question about your knowledge base...", accept_file="multiple", file_type=["docx", "csv", "xlsx", "xls", "pdf", "rds", "rda", "tsv", "sav", "dta", "txt", "md", "json", "do"] ) if user_input: prompt = user_input.text uploaded_files = user_input.files # Show and store user message st.session_state.messages.append({"role": "user", "content": prompt}) with st.chat_message("user"): st.markdown(prompt) cat_cfg = cfg.get("query_categorization", {}) qu_cfg = cfg.get("query_understanding", {}) # ── Check if this is a clarification response ─────────────────── unresolved = st.session_state.unresolved_category qcat_result, combined_cat, original_query_cat, unresolved_question = process_pending(cfg, unresolved, prompt, {"category": "general", "action": "continue"}, cat_cfg.get("max_history", 6), categorize=True ) print(qcat_result) if qcat_result.get("category") == "contact_and_info": process_contact_request(original_query_cat, qcat_result, uploaded_files) return if qcat_result.get("category") == "stata" and "unresolved" in qcat_result.get("sub-action") and st.session_state.stata_options is None: st.session_state.unresolved_category = original_query_cat process_stata_query(qcat_result, uploaded_files) return if "unresolved" not in qcat_result.get("sub-action", "unresolved"): resolve_stata_query(cfg, qcat_result, combined_cat, original_query_cat) return if "unresolved" in qcat_result.get("sub-action", "unresolved") and "Please let me know if this answer is helpful or if it requires further clarification." in unresolved_question: st.session_state.pending_clarification = prompt st.session_state.unresolved_category = unresolved st.session_state.pending_clarification_question = unresolved_question st.session_state.unresolved_category_question = unresolved_question # qu_enabled = False # st.markdown(result) # ── Check if this is a clarification response ─────────────────── pending = st.session_state.pending_clarification qu_result, combined, original_query, pending_question = process_pending(cfg, pending, prompt, {"action": "search", "route": "vector", "sql_query": None}, qu_cfg.get("max_history", 6) ) if qu_result.get("action") == "clarify" and st.session_state.clarification_rounds < qu_cfg.get("max_clarifications", 1): # Ask clarification — store original query and question for context st.session_state.pending_clarification = original_query st.session_state.pending_clarification_question = qu_result.get('clarification_question', 'Could you be more specific?') st.session_state.clarification_rounds += 1 clarification_msg = f"{qu_result.get('clarification_question', 'Could you be more specific?')}" with st.chat_message("assistant"): st.markdown(clarification_msg) feedback = st.feedback(options="faces", key=f"fb_qu_{st.session_state.clarification_rounds-1}") st.session_state.messages.append({"role": "assistant", "content": clarification_msg, "feedback":feedback}) return # After max clarification rounds, force search (matches CLI behavior) if qu_result.get("action") == "clarify": qu_result["action"] = "search" search_query = qu_result.get("search_query", combined) display_query = qu_result.get("display_query", original_query) route = qu_result.get("route", "vector") sql_query = qu_result.get("sql_query") # Generate assistant response with st.chat_message("assistant"): # print("Searching knowledge base") with st.status("Searching knowledge base...", expanded=True) as status: # Show reformulated query if different if search_query != original_query: status.update(label=f'Searching for: "{search_query}"...') # Retrieval try: retrieval_result = retrieve(search_query, cfg, route=route, sql_query=sql_query) except Exception as e: status.update(label=f"Retrieval error: {e}", state="error") st.error(f"Retrieval failed: {e}") return st.session_state.last_retrieval = retrieval_result n_local = len(retrieval_result.get("db_results", [])) n_web = len(retrieval_result.get("web_results", [])) n_sql = len(retrieval_result.get("sql_results", [])) sql_match = retrieval_result.get("sql_match_type", "") source_label = f"Found {n_local} local" if n_sql: match_label = f" ({sql_match} match)" if sql_match else "" source_label += f" + {n_sql} SQL rows{match_label}" source_label += f" + {n_web} web sources. Generating response..." status.update(label=source_label) # Response generation uses display_query — a clear, # complete question that incorporates any clarification context try: result = verify_and_respond( display_query, retrieval_result, cfg, original_query=original_query, ) except Exception as e: status.update(label=f"Generation error: {e}", state="error") st.error(f"Response generation failed: {e}") return # Update status based on verification outcome if result.get("refused"): status.update(label="No sufficient sources found.", state="error") elif result.get("verification_passed") is True: sql_label = f" + {n_sql} SQL rows" if n_sql else "" status.update( label=f"Verified ({result.get('iterations', 0)} iteration(s)). " f"{n_local} local{sql_label} + {n_web} web sources.", state="complete", ) elif result.get("verification_passed") is False: status.update( label="Response generated (verification did not fully pass).", state="error", ) else: sql_label2 = f" + {n_sql} SQL rows" if n_sql else "" status.update( label=f"Done. {n_local} local{sql_label2} + {n_web} web sources.", state="complete", ) final_response = f"{result.get('response', '')}\n\n For additional assistance, please contact the WEAI Helpdesk at IFPRI-WEAI@cgiar.org.\ Please let me know if there's anything else I can assist you with today." # Display the response st.markdown(final_response) feedback = st.feedback(options="faces", key=f"fb_final_{len(st.session_state.messages)}") # Store assistant message st.session_state.messages.append({"role": "assistant", "content": final_response, "feedback":feedback}) os.makedirs(f"{DATA_DIR}/logs/{st.session_state.chat_id}", exist_ok=True) with open(f"{DATA_DIR}/logs/{st.session_state.chat_id}/chat_history.json", "w", encoding="utf-8") as f: json.dump(st.session_state.messages, f, indent=4) export_outputs() # Cap message history to prevent unbounded memory growth max_messages = 200 # 100 Q&A pairs if len(st.session_state.messages) > max_messages: st.session_state.messages = st.session_state.messages[-max_messages:] if len(st.session_state.messages)%10==0 and st.session_state.feedback_data is None: return # ── Main ───────────────────────────────────────────────────────────────────── def main(): """Orchestrate the Streamlit app.""" st.set_page_config( page_title="ResearchBot", page_icon="🔬", layout="wide", ) init_session() cfg = st.session_state.cfg bot_name = cfg.get("chatbot", {}).get("name", "ResearchBot") st.title(bot_name) # Show KB summary on first visit (LLM-generated welcome summary) if not st.session_state.messages: # if "kb_welcome_summary" not in st.session_state: # st.session_state.kb_welcome_summary = load_kb_meta_brief(cfg) # kb_summary = st.session_state.kb_welcome_summary welcome_msg = """ Welcome to WEAIBot, the AI-powered support assistant for the WEAI Helpdesk.\n\n You can start by typing your question below, and WEAIBot will do its best to provide an answer based on available WEAI resources.\n\n If you would like assistance from a member of the WEAI team, simply send a message through the chat and include the email address where we can reach you. A support ticket will be created, and a team member will contact you as soon as they are available. """ st.markdown(welcome_msg) # if kb_summary: # with st.expander("Knowledge Base Contents", expanded=True): # st.markdown(kb_summary) render_sidebar() render_chat() st.caption( "All answers are sourced from the local knowledge base. " "Web sources are supplementary only. " "Every claim is citation-verified before display." ) if __name__ == "__main__": main()