app_web.py

"""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()