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