File size: 33,781 Bytes
0017280 f5c65c3 63152e4 84d4499 6c1a39c 0017280 84d4499 a705d5d 0017280 ca519bd 6c1a39c f95cdb7 0017280 6c1a39c f6f4639 0017280 84d4499 c28e4d0 9718fcb c28e4d0 9718fcb c28e4d0 a5a33f5 15ee6a2 6c1a39c 0017280 1330e5c 28c564f 9718fcb 1330e5c 0017280 ebd834c 7d228f9 2308eeb ae2262c 7d228f9 143ed5d e65bc7c bbe348b 84d4499 7dcc533 b5e4252 8c5b8e9 edc7e93 b5e4252 0017280 4ac52be 6c1a39c 4ac52be 876dcd1 4ac52be 6c1a39c 4ac52be 6c1a39c 4ac52be 6c1a39c 4ac52be 6c1a39c 4ac52be 0017280 4ac52be 6c1a39c 4ac52be 6c1a39c 4ac52be 6c1a39c 4ac52be 6c1a39c 4ac52be 6c1a39c 4ac52be 6c1a39c 4ac52be 6c1a39c 4ac52be 15ee6a2 4b4dbd6 0017280 84d4499 af313ea 84d4499 0017280 84d4499 a705d5d 144784b a705d5d f58d785 a705d5d 63152e4 a705d5d 273a9f8 a705d5d 144784b a705d5d 932bfa2 a96bc94 a705d5d 63152e4 1c80fab bbe348b a705d5d 1c80fab a705d5d 1c80fab a705d5d bbe348b a705d5d bbe348b 63152e4 bbe348b 63152e4 bbe348b a705d5d bbe348b b5e4252 a705d5d bbe348b a705d5d bbe348b a705d5d bbe348b a705d5d bbe348b 7aee7b3 a705d5d 7d228f9 a705d5d 63152e4 a705d5d edc7e93 a705d5d bbe348b a705d5d 7d228f9 8c5b8e9 a705d5d 7d228f9 a705d5d edc7e93 bbe348b a705d5d 0017280 b5e4252 0017280 4b4dbd6 0017280 4b4dbd6 c9ec354 4b4dbd6 8275dd3 4b4dbd6 c9ec354 4b4dbd6 b5e4252 4b4dbd6 fb6dec0 b5e4252 0017280 2a9cbda d0f3e65 a229e23 d0f3e65 0017280 e469e35 2a9cbda 0017280 6c1a39c d89ddb6 4cf36f7 3feae8a a705d5d 95778c5 7dcc533 f6f4639 7725730 f6f4639 7dcc533 84d4499 63152e4 6c1a39c bfcb47b 14677f1 7d228f9 a705d5d d89ddb6 0fd3783 f6f4639 88061d3 62b61b0 0fd3783 4cf36f7 ebd834c 3feae8a a705d5d 932bfa2 a705d5d 7d228f9 a705d5d 63152e4 a705d5d bbe348b a705d5d ebd834c 2308eeb a705d5d 2308eeb a705d5d ebd834c bbe348b 0017280 7d228f9 0017280 ebd834c 9815e49 ebd834c 0017280 9815e49 20d594f 9815e49 0017280 20d594f 6b689a5 20d594f 6b689a5 20d594f 0017280 9815e49 c23d859 9815e49 0017280 bbe348b fbd18c8 0017280 9815e49 fbd18c8 0017280 bbe348b 0017280 fbd18c8 0017280 fbd18c8 0017280 e3907f4 a131ce6 bbe348b 0017280 a131ce6 4b4dbd6 0017280 bbe348b 0017280 df99b8b 84d4499 4cee389 84d4499 214f864 f95cdb7 5c33fa2 0017280 1330e5c 0017280 1330e5c 0017280 7e9f52f 89e2c52 745d768 7aee7b3 745d768 322784b 745d768 6c1a39c 7725730 0017280 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 324 325 326 327 328 329 330 331 332 333 334 335 336 337 338 339 340 341 342 343 344 345 346 347 348 349 350 351 352 353 354 355 356 357 358 359 360 361 362 363 364 365 366 367 368 369 370 371 372 373 374 375 376 377 378 379 380 381 382 383 384 385 386 387 388 389 390 391 392 393 394 395 396 397 398 399 400 401 402 403 404 405 406 407 408 409 410 411 412 413 414 415 416 417 418 419 420 421 422 423 424 425 426 427 428 429 430 431 432 433 434 435 436 437 438 439 440 441 442 443 444 445 446 447 448 449 450 451 452 453 454 455 456 457 458 459 460 461 462 463 464 465 466 467 468 469 470 471 472 473 474 475 476 477 478 479 480 481 482 483 484 485 486 487 488 489 490 491 492 493 494 495 496 497 498 499 500 501 502 503 504 505 506 507 508 509 510 511 512 513 514 515 516 517 518 519 520 521 522 523 524 525 526 527 528 529 530 531 532 533 534 535 536 537 538 539 540 541 542 543 544 545 546 547 548 549 550 551 552 553 554 555 556 557 558 559 560 561 562 563 564 565 566 567 568 569 570 571 572 573 574 575 576 577 578 579 580 581 582 583 584 585 586 587 588 589 590 591 592 593 594 595 596 597 598 599 600 601 602 603 604 605 606 607 608 609 610 611 612 613 614 615 616 617 618 619 620 621 622 623 624 625 626 627 628 629 630 631 632 633 634 635 636 637 638 639 640 641 642 643 644 645 646 647 648 649 650 651 652 653 654 655 656 657 658 659 660 661 662 663 664 665 666 667 668 669 670 671 672 673 674 675 676 677 678 679 680 681 682 683 684 685 686 687 688 689 690 691 692 693 694 695 696 697 698 699 700 701 702 703 704 705 706 707 708 709 710 711 712 713 714 715 716 717 718 719 720 | """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 summarize_kb_for_welcome
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 src.stata.explain import read_file
from src.prompts import build_stata_prompt
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 "verification_status" not in st.session_state:
st.session_state.verification_status = 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,
}
#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") == "contact":
response = """\
Thank you for reaching out. I am not able to arrange meetings or facilitate introductions directly.
For additional assistance, 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."""
# 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))
if qcat_result.get("action") == "upload":
response = """\
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 share with us your email in case you haven't done so already. \n\n
Please let me know if there's anything else I can assist you with today."""
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!")
if len(qcat_result.get("response", "").strip())>0:
response = qcat_result.get("response")
with st.chat_message("assistant"):
st.markdown(response)
feedback = st.feedback(options="faces", key=f"fb_contact_{st.session_state.resolution_rounds-1}")
st.session_state.messages.append({"role": "assistant",
"content": response,
"feedback":feedback,
"verification_passed":None})
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,
"verification_passed":None})
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,
"verification_passed":None})
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
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
)
st.session_state.messages.append({"role": "user",
"content": prompt,
"category":qcat_result.get("category"),
})
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 qcat_result.get("category") == "stata" and "unresolved" not in qcat_result.get("sub-action", "unresolved"):
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
}
dofx = read_file(st.session_state.stata_fil, opts['lines'])
context = build_stata_prompt(combined_cat, qcat_result.get("action"), dofx, opts)
st.session_state.pending_clarification = context
st.session_state.pending_clarification_question = prompt
# 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,
"verification_passed":None})
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")
status = True
# 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:
st.session_state.verification_status = 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",
)
st.session_state.verification_status = False
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,
"verification_passed":st.session_state.verification_status})
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 = summarize_kb_for_welcome(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.
# """
welcome_msg = """
Welcome!
The WEAIbot is an 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.
"""
st.markdown(welcome_msg)
# if kb_summary:
# with st.expander("Welcome!", 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()
|