Spaces:
Sleeping
Sleeping
| import shutil | |
| import os | |
| import streamlit as st | |
| from huggingface_hub import snapshot_download | |
| from src.config_loader import load_config, get_api_key | |
| from src.ingest import get_chroma_collection, ingest_documents | |
| from src.retriever import clear_collection_cache, retrieve | |
| from src.llm import list_models | |
| from src.peri import APPROVED_OPTIONS | |
| from src.query_engine import understand_query, categorize_query, init_query, analyze_committment_to_vc | |
| from src.prompts import dict_to_string | |
| snapshot_download(repo_id="CGIAR/peri-kb", | |
| repo_type="dataset", | |
| allow_patterns="ug/*", | |
| token=os.getenv('HF_TOKEN'), | |
| local_dir="./" | |
| ) | |
| shutil.copytree("./ug", "./", dirs_exist_ok=True) | |
| 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", "meta-llama"] | |
| 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 render_chat(): | |
| """Render the chat interface with message history and input.""" | |
| cfg = st.session_state.cfg | |
| # Display chat history | |
| for msg in st.session_state.messages: | |
| with st.chat_message(msg["role"]): | |
| st.markdown(msg["content"]) | |
| # 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) | |
| # ── Categorize user query ─────────────────── | |
| save_dir = "uploaded_do_files" | |
| os.makedirs(save_dir, exist_ok=True) | |
| # Exclude the just-appended user message to avoid sending | |
| # the current question twice (once in history, once as query) | |
| cat_cfg = cfg.get("query_categorization", {}) | |
| max_history = cat_cfg.get("max_history", 6) | |
| # ── Query understanding ───────────────────────────────────────── | |
| qu_cfg = cfg.get("query_understanding", {}) | |
| qu_enabled = qu_cfg.get("enabled", True) | |
| max_history = qu_cfg.get("max_history", 6) | |
| # ── Check if this is a clarification response ─────────────────── | |
| unresolved = st.session_state.unresolved_category | |
| if unresolved is not None: | |
| # This prompt is the user's clarification answer | |
| # Include the clarification question for context | |
| unresolved_question = st.session_state.get("unresolved_category_question", "") | |
| if unresolved_question: | |
| combined_cat = f"{unresolved} (Clarification: Q: {unresolved_question} A: {prompt})" | |
| else: | |
| combined_cat = f"{unresolved} — {prompt}" | |
| st.session_state.unresolved_category_question = None | |
| original_query_cat = unresolved | |
| st.session_state.unresolved_category = None | |
| else: | |
| combined_cat = prompt | |
| original_query_cat = prompt | |
| st.session_state.resolution_rounds = 0 | |
| unresolved_question = [] | |
| prior_messages = st.session_state.messages[:-1] | |
| history = [ | |
| {"role": m["role"], "content": m["content"]} | |
| for m in prior_messages[-max_history:] | |
| ] | |
| try: | |
| qinit_result = init_query(combined_cat, cfg, history) | |
| print("result: ", qinit_result) | |
| if qinit_result.get("country", None) is None: | |
| with st.chat_message("assistant"): | |
| st.markdown("Please specify the country for which you'd like to run a PERI analysis") | |
| return | |
| if qinit_result.get("value_chain", None) is None and qinit_result.get("investment", None) is None: | |
| with st.chat_message("assistant"): | |
| st.markdown(f"Please specify the value chain or investment area in {qinit_result.get('country', None)} for which you'd like to run a PERI analysis") | |
| return | |
| qcat_result = categorize_query(combined_cat, cfg, history) | |
| print("result: ", qcat_result) | |
| except Exception as e: | |
| print(e) | |
| qinit_result = {"country": None, "value_chain": None, "investment": None} | |
| qcat_result = {"category": "pillar_1", "action": "unresolved"} | |
| if not isinstance(qinit_result.get("country", None), list) and qinit_result.get("country", None)==None: | |
| resolution_msg = f"It seems there is no country specified for the analysis. Currently the PERI framework supports analysis for the countries listed below:\n\n {', '.join([c.capitalize() for c in APPROVED_OPTIONS.get('countries')])}\n\n We are also continuously \ | |
| working to expand the framework and you can submit a form if the country you would like to run the analysis on is not included. In the meantime please let me know if you would like to run the anlysis for one of the included countries." | |
| if qinit_result.get("country", None)[0].lower() not in [c.lower() for c in APPROVED_OPTIONS.get('countries')]: | |
| resolution_msg = f"It seems you are trying to run a PERI analysis for {qinit_result.get('country', None)[0].capitalize()}! Currently the PERI framework only supports analysis for the countries listed below:\n\n {', '.join([c.capitalize() for c in APPROVED_OPTIONS.get('countries')])}\n\n We are continuously \ | |
| working to expand the framework and you can submit a form to request. In the meantime please let me know if you would like to run the anlysis for one of the included countries." | |
| if qinit_result.get("value_chain", None)[0]==None and qinit_result.get('investment', None)[0]==None: | |
| resolution_msg = f"It seems there is no value chain or investment area specified for the analysis. Please select one of the value chains or investment areas included in the current PERI framework." | |
| if qinit_result.get("value_chain", None)[0].lower() not in [c.lower() for v in APPROVED_OPTIONS.get('value_chains').values() for c in v]: | |
| resolution_msg = f"It seems you are trying to run a PERI analysis for {qinit_result.get('value_chain', None)[0]} in {qinit_result.get('country', None)[0].capitalize()}! Currently the PERI framework only supports analysis for the value chains listed below:\n\n {dict_to_string(APPROVED_OPTIONS.get('value_chains'), 2)}\n\n We are continuously \ | |
| working to expand the framework and you can submit a form to request. In the meantime please let me know if you would like to run the anlysis for one of the included countries." | |
| if qinit_result.get("investment", None)[0].lower() not in [i.lower() for i in APPROVED_OPTIONS.get('investments')]: | |
| resolution_msg = f"It seems you are trying to run a PERI analysis for {qinit_result.get('investment', None)[0]} investemnt in {qinit_result.get('country', None)[0].capitalize()}! Currently the PERI framework only supports analysis for the investment areas listed below:\n\n {', '.join([c.capitalize() for c in APPROVED_OPTIONS.get('investments')])}\n\n We are continuously \ | |
| working to expand the framework and you can submit a form to request. In the meantime please let me know if you would like to run the anlysis for one of the included countries." | |
| if qinit_result.get("country", None)[0] in APPROVED_OPTIONS.get('countries') and (qinit_result.get('value_chain', None)[0] in APPROVED_OPTIONS.get('value_chains') or qinit_result.get('investment', None)[0] in APPROVED_OPTIONS.get('investments')): | |
| pillar_dict = { | |
| "pillar_1":f"would like to understand whether {','.join(qinit_result.get('value_chain'))} aligns with the government's political incentives.", | |
| "pillar_2":f"would like to understand to what degree decisions are impacted by the lobbying of particular groups or by elite influence", | |
| "pillar_3":f"would like to understand if {','.join(qinit_result.get('value_chain'))} and/or investing in {','.join(qinit_result.get('investment'))} can be feasibly implemented given the broader institutional and policy environment", | |
| } | |
| resolution_msg = f"**Before I search, could you clarify?** Please let me know if you {' and '.join([pillar_dict[c] for c in qcat_result.get('category')])}." | |
| # st.session_state.unresolved_category_question = f"Please let me know if you {' and '.join([pillar_dict[c] for c in qcat_result.get('category')])}." | |
| st.session_state.unresolved_category_question = resolution_msg | |
| st.session_state.unresolved_category = original_query_cat | |
| st.session_state.resolution_rounds += 1 | |
| st.session_state.messages.append({"role": "assistant", "content": resolution_msg}) | |
| with st.chat_message("assistant"): | |
| st.markdown(resolution_msg) | |
| return | |