File size: 47,866 Bytes
c83f30f
 
1befddb
3da458d
c83f30f
a0b841e
4c56041
6a14896
17cd36a
 
c83f30f
da53bed
dc8b4b0
5a4f54c
 
ff62217
c83f30f
72cb6c4
c997974
6a14896
6a76700
1b8b74e
c997974
 
5a3e284
fba1e90
5a4f54c
5a3e284
 
 
 
fba1e90
 
 
5a3e284
579a105
5a3e284
 
a92e5cd
ff62217
fba1e90
 
72cb6c4
 
 
 
 
8279bee
72cb6c4
8279bee
72cb6c4
 
8279bee
 
 
 
6a76700
8279bee
 
 
6a76700
8279bee
 
 
 
 
6a76700
8279bee
 
 
 
6a76700
8279bee
 
 
 
 
6a76700
8279bee
 
 
 
 
 
6a76700
8279bee
6a76700
8279bee
 
fba1e90
 
 
 
 
 
 
 
 
 
c997974
 
 
 
 
 
ff62217
 
 
 
 
 
 
 
5a3e284
c997974
 
 
 
ff62217
6a14896
 
 
 
ff62217
6a14896
 
 
 
 
ff62217
 
6a14896
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
ff62217
6a14896
 
ff62217
 
 
6a14896
 
 
 
 
 
72cb6c4
fba1e90
 
 
 
 
 
 
 
 
 
 
 
 
a0b841e
 
fba1e90
 
a0b841e
 
 
 
 
aa07439
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
fba1e90
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
ff62217
c997974
ff62217
6a14896
 
3ce59a7
ff62217
a0b841e
 
 
 
 
c997974
 
 
 
a0b841e
c997974
 
fba1e90
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
c997974
 
6a14896
fba1e90
 
 
 
 
 
6a14896
fba1e90
c997974
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
a37c4f9
 
 
 
 
ff62217
a37c4f9
 
 
c997974
 
 
fba1e90
 
 
 
ff62217
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
fba1e90
c997974
5a3e284
fba1e90
72cb6c4
 
fba1e90
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
a0b841e
fba1e90
 
a0b841e
17923a5
 
 
 
 
 
 
 
a0b841e
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
fba1e90
 
 
 
 
 
 
 
 
 
 
3e7dd62
 
 
 
 
fba1e90
b36fbf0
 
 
 
 
fba1e90
b36fbf0
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
fba1e90
b36fbf0
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
fba1e90
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
427f510
 
 
 
3e7dd62
 
 
 
 
427f510
3e7dd62
 
 
 
 
 
bcac791
427f510
1dd6c51
 
3e7dd62
 
427f510
bcac791
427f510
 
 
 
 
bcac791
427f510
 
 
 
 
 
 
 
fba1e90
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
ecf9255
 
fba1e90
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
579a105
cfd826b
c997974
 
 
579a105
c997974
 
 
579a105
c997974
cfd826b
1e037b0
 
 
 
 
 
 
 
 
 
 
 
 
c997974
 
 
579a105
c997974
 
 
 
 
579a105
c997974
 
579a105
c997974
6a14896
 
 
1e037b0
6a14896
 
 
 
 
 
d5fd2a5
 
1e037b0
 
 
 
 
 
 
 
 
 
 
d5fd2a5
13e4ba7
1e037b0
13e4ba7
 
 
 
 
 
 
 
 
 
1e037b0
 
13e4ba7
1e037b0
13e4ba7
1e037b0
 
13e4ba7
1e037b0
 
 
 
13e4ba7
 
 
1e037b0
 
 
 
13e4ba7
 
 
 
045aa6f
1e037b0
 
 
 
 
045aa6f
1e037b0
 
 
 
 
045aa6f
1e037b0
045aa6f
13e4ba7
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1e037b0
cfd826b
579a105
cfd826b
a37c4f9
1e037b0
13e4ba7
 
 
 
 
 
1e037b0
13e4ba7
 
 
 
1e037b0
 
 
 
 
cfd826b
 
1e037b0
045aa6f
1e037b0
13e4ba7
045aa6f
cfd826b
 
1e037b0
 
cfd826b
1e037b0
cfd826b
 
1e037b0
 
 
 
cfd826b
 
1e037b0
cfd826b
 
 
13e4ba7
cfd826b
 
 
13e4ba7
 
 
 
 
cfd826b
 
13e4ba7
 
 
 
cfd826b
 
 
 
 
 
 
 
13e4ba7
cfd826b
 
 
 
 
 
 
 
 
 
 
 
 
 
1e037b0
 
cfd826b
1e037b0
 
13e4ba7
1e037b0
 
 
 
 
 
13e4ba7
1e037b0
 
 
 
045aa6f
1e037b0
13e4ba7
045aa6f
cfd826b
1e037b0
cfd826b
1e037b0
 
 
 
 
 
 
 
b36fbf0
1e037b0
 
13e4ba7
 
 
 
 
1e037b0
 
13e4ba7
1e037b0
a37c4f9
13e4ba7
 
 
 
 
 
 
 
 
 
 
 
c997974
fba1e90
 
579a105
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
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
824
825
826
827
828
829
830
831
832
833
834
835
836
837
838
839
840
841
842
843
844
845
846
847
848
849
850
851
852
853
854
855
856
857
858
859
860
861
862
863
864
865
866
867
868
869
870
871
872
873
874
875
876
877
878
879
880
881
882
883
884
885
886
887
888
889
890
891
892
893
894
895
896
897
898
899
900
901
902
903
904
905
906
907
908
909
910
911
912
913
914
915
916
917
918
919
920
921
922
923
924
925
926
927
928
929
930
931
932
933
934
935
936
937
938
939
940
941
942
943
944
945
946
947
948
949
950
951
952
953
954
955
956
957
958
959
960
import gradio as gr
import pandas as pd
import logging
import asyncio
import os
import time
from uuid import uuid4
from datetime import datetime, timedelta
from pathlib import Path
from huggingface_hub import CommitScheduler
from auditqa.sample_questions import QUESTIONS
from auditqa.reports import files, report_list, new_files, new_report_list
from auditqa.process_chunks import load_chunks, getconfig, get_local_qdrant, load_new_chunks
from auditqa.retriever import get_context
from auditqa.reader import nvidia_client, dedicated_endpoint
from auditqa.utils import make_html_source, parse_output_llm_with_sources, save_logs, get_message_template, get_client_location, get_client_ip, get_platform_info
from dotenv import load_dotenv
load_dotenv()
from threading import Lock
from gradio.routes import Request
from qdrant_client import QdrantClient
import json

# # fetch tokens and model config params
SPACES_LOG = os.environ["SPACES_LOG"]
model_config = getconfig("model_params.cfg")

# create the local logs repo
JSON_DATASET_DIR = Path("json_dataset")
JSON_DATASET_DIR.mkdir(parents=True, exist_ok=True)
JSON_DATASET_PATH = JSON_DATASET_DIR / f"logs-{uuid4()}.json"

# the logs are written to dataset repo periodically from local logs
# https://huggingface.co/spaces/Wauplin/space_to_dataset_saver
scheduler = CommitScheduler(
     repo_id=model_config.get('app','repo_id'),
     repo_type="dataset",
     folder_path=JSON_DATASET_DIR,
     path_in_repo="audit_chatbot",
     token=SPACES_LOG)

#####--------------- VECTOR STORE -------------------------------------------------
# reports contain the already created chunks from Markdown version of pdf reports
# document processing was done using : https://github.com/axa-group/Parsr
# We need to create the local vectorstore collection once using load_chunks
# vectorestore colection are stored on persistent storage so this needs to be run only once
# hence, comment out line below when creating for first time
# vectorstores = load_new_chunks()
# once the vectore embeddings  are created we will use qdrant client to access these
# vectorstores = get_local_qdrant()

# Configure cloud Qdrant client #TESTING
def get_cloud_qdrant():
    from langchain_community.embeddings import HuggingFaceEmbeddings
    from langchain_community.vectorstores import Qdrant
    from torch import cuda
    
    # Get config and setup embeddings like in process_chunks.py
    model_config = getconfig("model_params.cfg")
    device = 'cuda' if cuda.is_available() else 'cpu'
    
    embeddings = HuggingFaceEmbeddings(
        model_kwargs = {'device': device},
        encode_kwargs = {'normalize_embeddings': True},
        model_name=model_config.get('retriever','MODEL')
    )
    
    # Get Qdrant API key from environment variable
    qdrant_api_key = os.getenv("QDRANT")
    if not qdrant_api_key:
        raise ValueError("QDRANT API key not found in environment variables")
    
    # Create the Qdrant client
    client = QdrantClient(
        url="https://ff3f0448-0a00-470e-9956-49efa3071db3.europe-west3-0.gcp.cloud.qdrant.io:6333", 
        api_key=qdrant_api_key,
    )
    
    # Wrap the client in Langchain's Qdrant vectorstore
    vectorstore = Qdrant(
        client=client,
        collection_name="allreports",
        embeddings=embeddings,
    )
    
    return {"allreports": vectorstore}

# Replace local Qdrant with cloud Qdrant
vectorstores = get_cloud_qdrant()

#####---------------------CHAT-----------------------------------------------------
def start_chat(query,history):
    history = history + [(query,None)]
    history = [tuple(x) for x in history]
    return (gr.update(interactive = False),gr.update(selected=1),history)

def finish_chat():
    return (gr.update(interactive = True,value = ""))
    
def submit_feedback(feedback, logs_data):
    """Handle feedback submission"""
    try:
        if logs_data is None:
            return gr.update(visible=False), gr.update(visible=True)
            
        session_id = logs_data.get("session_id")
        if session_id:
            # Update session last_activity to now
            session_manager.update_session(session_id)
            # Compute duration from the session manager and update the log.
            logs_data["session_duration_seconds"] = session_manager.get_session_duration(session_id)
            
        # Now save the (feedback) log record
        save_logs(scheduler, JSON_DATASET_PATH, logs_data, feedback)
        return gr.update(visible=False), gr.update(visible=True)
    except Exception as e:
        return gr.update(visible=False), gr.update(visible=True)

# Session Manager added (track session duration, location, and platform)
class SessionManager:
    def __init__(self):
        self.sessions = {}
        
    def create_session(self, client_ip, user_agent):
        session_id = str(uuid4())
        self.sessions[session_id] = {
            'start_time': datetime.now(),
            'last_activity': datetime.now(),
            'client_ip': client_ip,
            'location_info': get_client_location(client_ip),
            'platform_info': get_platform_info(user_agent)
        }
        return session_id
    
    def update_session(self, session_id):
        if session_id in self.sessions:
            self.sessions[session_id]['last_activity'] = datetime.now()
    
    def get_session_duration(self, session_id):
        if session_id in self.sessions:
            start = self.sessions[session_id]['start_time']
            last = self.sessions[session_id]['last_activity']
            return (last - start).total_seconds()
        return 0
    
    def get_session_data(self, session_id):
        return self.sessions.get(session_id)

# Initialize session manager
session_manager = SessionManager()

async def chat(query, history, sources, reports, subtype, year, client_ip=None, session_id=None, request: gr.Request = None):
    """Update chat function to handle session data"""
    
    if not session_id:
        user_agent = request.headers.get('User-Agent', '') if request else ''
        session_id = session_manager.create_session(client_ip, user_agent)
    else:
        session_manager.update_session(session_id)
    
    # Get session data
    session_data = session_manager.get_session_data(session_id)
    session_duration = session_manager.get_session_duration(session_id)

    print(f">> NEW QUESTION : {query}")
    print(f"history:{history}")
    print(f"sources:{sources}")
    print(f"reports:{reports}")
    print(f"subtype:{subtype}")
    print(f"year:{year}")
    docs_html = ""
    output_query = ""

    ##------------------------fetch collection from vectorstore------------------------------
    vectorstore = vectorstores["allreports"]

    ##------------------------------get context---------------------------------------------- 
    ### adding for assessing computation time
    start_time = time.time()
    context_retrieved = get_context(vectorstore=vectorstore,query=query,reports=reports,
                                                sources=sources,subtype=subtype,year=year)
    end_time = time.time()
    print("Time for retriever:",end_time - start_time)

    # WARNING FOR NO CONTEXT: Check if any paragraphs were retrieved, add warning if none found
    # We use this in the Gradio UI below (displays in the chat dialogue box)
    if not context_retrieved or len(context_retrieved) == 0:
        warning_message = "⚠️ **No relevant information was found in the audit reports pertaining your query.** Please try rephrasing your question or selecting different report filters."
        history[-1] = (query, warning_message)
        # Update logs with the warning instead of answer
        logs_data = {
            "record_id": str(uuid4()),
            "session_id": session_id,
            "session_duration_seconds": session_duration,
            "client_location": session_data['location_info'],
            "platform": session_data['platform_info'],
            "year": year,
            "question": query,
            "retriever": model_config.get('retriever','MODEL'),
            "endpoint_type": model_config.get('reader','TYPE'),
            "reader": model_config.get('reader','NVIDIA_MODEL'),
            "answer": warning_message,
            "no_results": True  # Flag to indicate no results were found
        }
        yield [tuple(x) for x in history], "", logs_data, session_id
        # Save log for the warning response
        save_logs(scheduler, JSON_DATASET_PATH, logs_data)
        return
    
    context_retrieved_formatted = "||".join(doc.page_content for doc in context_retrieved)
    context_retrieved_lst = [doc.page_content for doc in context_retrieved]
    
    ##------------------- -------------Define Prompt-------------------------------------------
    SYSTEM_PROMPT = """
        You are AuditQ&A, an AI Assistant created by Auditors and Data Scientist. \
            You are given a question and extracted passages of the consolidated/departmental/thematic focus audit reports.\
            Provide a clear and structured answer based on the passages/context provided and the guidelines.
        Guidelines:
        - Passeges are provided as comma separated list of strings
        - If the passages have useful facts or numbers, use them in your answer.
        - When you use information from a passage, mention where it came from by using [Doc i] at the end of the sentence. i stands for the number of the document.
        - Do not use the sentence 'Doc i says ...' to say where information came from.
        - If the same thing is said in more than one document, you can mention all of them like this: [Doc i, Doc j, Doc k]
        - Do not just summarize each passage one by one. Group your summaries to highlight the key parts in the explanation.
        - If it makes sense, use bullet points and lists to make your answers easier to understand.
        - You do not need to use every passage. Only use the ones that help answer the question.
        - If the documents do not have the information needed to answer the question, just say you do not have enough information.
        """
    
    USER_PROMPT = """Passages:
        {context}
        -----------------------
        Question: {question}  - Explained to audit expert
        Answer in english with the passages citations:
        """.format(context = context_retrieved_lst, question=query)
    
    ##-------------------- apply message template ------------------------------
    messages = get_message_template(model_config.get('reader','TYPE'),SYSTEM_PROMPT,USER_PROMPT)

    ## -----------------Prepare HTML for displaying source documents --------------
    docs_html = []
    for i, d in enumerate(context_retrieved, 1):
        docs_html.append(make_html_source(d, i))
    docs_html = "".join(docs_html)

    ##-----------------------get answer from endpoints------------------------------
    answer_yet = ""
    # Logs strcuture updated for session data (feedback and timestamp added separately via save_logs)
    logs_data = {
        "record_id": str(uuid4()),  # Add unique record ID
        "session_id": session_id,
        "session_duration_seconds": session_duration,
        "client_location": session_data['location_info'],
        "platform": session_data['platform_info'], 
        "system_prompt": SYSTEM_PROMPT, 
        "sources": sources, 
        "reports": reports, 
        "subtype": subtype, 
        #"year": year,
        "question": query,
        "retriever": model_config.get('retriever','MODEL'),
        "endpoint_type": model_config.get('reader','TYPE'),
        "reader": model_config.get('reader','NVIDIA_MODEL'),
        "docs": [doc.page_content for doc in context_retrieved],
    }

    if model_config.get('reader','TYPE') == 'NVIDIA':
        chat_model = nvidia_client()
        async def process_stream():
            nonlocal answer_yet # Use the outer scope's answer_yet variable
            # Without nonlocal, Python would create a new local variable answer_yet inside process_stream(), 
            # instead of modifying the one from the outer scope.
            # Iterate over the streaming response chunks
            response = chat_model.chat_completion(
                                model=model_config.get("reader","NVIDIA_MODEL"),
                                messages=messages,
                                stream=True,
                                max_tokens=int(model_config.get('reader','MAX_TOKENS')),
                            )
            for message in response:
                token = message.choices[0].delta.content
                if token:
                    answer_yet += token
                    parsed_answer = parse_output_llm_with_sources(answer_yet)
                    history[-1] = (query, parsed_answer)
                    # Update logs_data with current answer
                    logs_data["answer"] = parsed_answer
                    yield [tuple(x) for x in history], docs_html, logs_data, session_id

        # Stream the response updates
        async for update in process_stream():
            yield update

    else:
        chat_model = dedicated_endpoint() # TESTING: ADAPTED FOR HF INFERENCE API (needs to be reverted for production version)
        async def process_stream():
            nonlocal answer_yet
            try:
                formatted_messages = [
                    {
                        "role": msg.type if hasattr(msg, 'type') else msg.role,
                        "content": msg.content
                    }
                    for msg in messages
                ]
                
                response = chat_model.chat_completion(
                    messages=formatted_messages,
                    max_tokens=int(model_config.get('reader', 'MAX_TOKENS'))
                )
                
                response_text = response.choices[0].message.content
                words = response_text.split()
                for word in words:
                    answer_yet += word + " "
                    parsed_answer = parse_output_llm_with_sources(answer_yet)
                    history[-1] = (query, parsed_answer)
                    # Update logs_data with current answer (and get a new timestamp)
                    logs_data["answer"] = parsed_answer
                    yield [tuple(x) for x in history], docs_html, logs_data, session_id
                    await asyncio.sleep(0.05)
                
            except Exception as e:
                raise

        async for update in process_stream():
            yield update

        # chat_model = dedicated_endpoint()
        # async def process_stream():
        # # Without nonlocal, Python would create a new local variable answer_yet inside process_stream(), 
        # # instead of modifying the one from the outer scope.
        #     nonlocal answer_yet # Use the outer scope's answer_yet variable
        #     # Iterate over the streaming response chunks
        #     async for chunk in chat_model.astream(messages):
        #         token = chunk.content
        #         answer_yet += token
        #         parsed_answer = parse_output_llm_with_sources(answer_yet)
        #         history[-1] = (query, parsed_answer)
        #         yield [tuple(x) for x in history], docs_html

        # # Stream the response updates
        # async for update in process_stream():
        #     yield update

    try:
        # Save log after streaming is complete
        save_logs(scheduler, JSON_DATASET_PATH, logs_data)
    except Exception as e:
        raise





#####-------------------------- Gradio App--------------------------------------####

# Set up Gradio Theme
theme = gr.themes.Base(
    primary_hue="blue",
    secondary_hue="red",
    font=[gr.themes.GoogleFont("Poppins"), "ui-sans-serif", "system-ui", "sans-serif"],
    text_size = gr.themes.utils.sizes.text_sm,
)

init_prompt =  """
Hello, I am Audit Q&A, a conversational assistant designed to help you understand audit Reports. I will answer your questions by using **Audit reports publishsed by Auditor General Office**.
💡 How to use (tabs on right)
- **Reports**: You can choose to address your question to either specific report or a collection of report like District or Ministry focused reports. \
If you dont select any then the Consolidated report is relied upon to answer your question.
- **Examples**: We have curated some example questions,select a particular question from category of questions.
- **Sources**: This tab will display the relied upon paragraphs from the report, to help you in assessing or fact checking if the answer provided by Audit Q&A assitant is correct or not.
⚠️ For limitations of the tool please check **Disclaimer** tab.
"""


with gr.Blocks(title="Audit Q&A", css= "style.css", theme=theme,elem_id = "main-component") as demo:
    #----------------------------------------------------------------------------------------------
    # main tab where chat interaction happens
    # ---------------------------------------------------------------------------------------------
    with gr.Tab("AuditQ&A"):
        
        with gr.Row(elem_id="chatbot-row"):
            # chatbot output screen
            with gr.Column(scale=2):
                chatbot = gr.Chatbot(
                    value=[(None,init_prompt)],
                    show_copy_button=True,show_label = False,elem_id="chatbot",layout = "panel",
                    avatar_images = (None,"data-collection.png")
                )
                
                #---------------- FEEDBACK ----------------------
                with gr.Column(elem_id="feedback-container"):
                    with gr.Row(visible=False) as feedback_row:
                        gr.Markdown("Was this response helpful?")
                        with gr.Row():
                            okay_btn = gr.Button("👍 Okay", elem_classes="feedback-button")
                            not_okay_btn = gr.Button("👎 Not to expectations", elem_classes="feedback-button")
                    feedback_thanks = gr.Markdown("Thanks for the feedback!", visible=False)
                    feedback_state = gr.State()
                #---------------- WARNINGS ----------------------
                # No filters selected warning
                with gr.Row(visible=False, elem_id="warning-row", elem_classes="warning-message") as warning_row:
                    with gr.Column():
                        gr.Markdown("<span class='warning-icon'>⚠️</span> **No report filter selected. Are you sure you want to proceed?**")
                        with gr.Row(elem_classes="warning-buttons"):
                            proceed_btn = gr.Button("Proceed", elem_classes="proceed")
                            cancel_btn = gr.Button("Cancel", elem_classes="cancel")

                # Short query warning (< 4 words)
                with gr.Row(visible=False, elem_id="warning-row", elem_classes="warning-message") as short_query_warning_row:
                    with gr.Column():
                        gr.Markdown("<span class='warning-icon'>⚠️</span> **Your query is too short. Please lengthen your query to ensure the app has adequate context.**")
                        with gr.Row(elem_classes="warning-buttons"):
                            short_query_proceed_btn = gr.Button("OK", elem_classes="proceed")

                #---------------- QUERY INPUT ----------------------
                with gr.Row(elem_id = "input-message"):
                    textbox=gr.Textbox(placeholder="Ask me anything here!",show_label=False,scale=7,
                                       lines = 1,interactive = True,elem_id="input-textbox")

            # second column with playground area for user to select values
            with gr.Column(scale=1, variant="panel",elem_id = "right-panel"):
                # creating tabs on right panel
                with gr.Tabs() as tabs:
                    #---------------- tab for REPORTS SELECTION ----------------------
                    
                    with gr.Tab("Reports",elem_id = "tab-config",id = 2):
                        with gr.Row():
                            gr.Markdown("Reminder: To get better results select the specific report/reports")
                            gr.Markdown("""<div class="question-tooltip">?
                                <div class="tooltip-content">Select the audit reports that you want to analyse directly or browse through categories and select reports</div>
                            </div>""", elem_id="reports-tooltip")
                        
                        #---------------- SELECTION METHOD - RADIO BUTTON ------------
                        search_method = gr.Radio(
                            choices=["Search by Report Name", "Search by Filters"],
                            label="Choose search method",
                            value="Search by Report Name",
                        )
                        
                        #---------------- SELECT BY REPORT NAME SECTION ------------
                        with gr.Group(visible=True) as report_name_section:
                            # Get default report value from config if present
                            default_report = model_config.get('app', 'dropdown_default', fallback=None)
                            # Check if it actually exists in the master list
                            default_report_value = [default_report] if default_report in new_report_list else None
                            
                            dropdown_reports = gr.Dropdown(
                                new_report_list,
                                label="Select one or more reports (scroll or type to search)",
                                multiselect=True,
                                value=default_report_value,
                                interactive=True,
                            )
                        
                        #---------------- SELECT BY FILTERS SECTION ------------
                        with gr.Group(visible=False) as filters_section:
                            #----- First level filter for selecting Report source/category ----------
                            dropdown_sources = gr.Dropdown(
                                ["Consolidated","Ministry, Department, Agency and Projects","Local Government","Value for Money","Thematic"],
                                label="Select Report Category",
                                value=None,
                                interactive=True,
                            )

                            #------ second level filter for selecting subtype within the report category selected above
                            dropdown_category = gr.Dropdown(
                                [],  # Start with empty choices
                                value=None,
                                label = "Filter for Sub-Type",
                                interactive=True)

                            #----------- update the second level filter based on values from first level ----------------
                            def rs_change(rs):
                                if rs:  # Only update choices if a value is selected
                                    return gr.update(choices=new_files[rs], value=None)  # Set value to None (no preselection)
                                else:
                                    return gr.update(choices=[], value=None)  # Empty choices if nothing selected
                            dropdown_sources.change(fn=rs_change, inputs=[dropdown_sources], outputs=[dropdown_category])

                            #--------- Select the years for reports -------------------------------------
                            dropdown_year = gr.Dropdown(
                                ['2018','2019','2020','2021','2022','2023'],
                                label="Filter for year",
                                multiselect=True,
                                value=None,
                                interactive=True,
                            )
                        
                        # Toggle visibility based on search method
                        def toggle_search_method(method):
                            """Note - this function removes the default value from report search when toggled"""
                            if method == "Search by Report Name":
                                # Show report selection, hide filters, and clear filter values
                                return (
                                    gr.update(visible=True),  # report_name_section
                                    gr.update(visible=False),  # filters_section
                                    gr.update(value=None),    # dropdown_sources
                                    gr.update(value=None),    # dropdown_category
                                    gr.update(value=None),    # dropdown_year
                                    gr.update()               # dropdown_reports
                                )
                            else:  # "Search by Filters"
                                # Show filters, hide report selection, and clear report values
                                return (
                                    gr.update(visible=False), # report_name_section 
                                    gr.update(visible=True),  # filters_section
                                    gr.update(),              # dropdown_sources
                                    gr.update(),              # dropdown_category
                                    gr.update(),              # dropdown_year
                                    gr.update(value=[])       # dropdown_reports
                                )
                        
                        # Pass to the event handler
                        search_method.change(
                            fn=toggle_search_method,
                            inputs=[search_method],
                            outputs=[
                                report_name_section, 
                                filters_section,
                                dropdown_sources,
                                dropdown_category, 
                                dropdown_year,
                                dropdown_reports
                            ]
                        )

                    ############### tab for Question selection ###############
                    with gr.TabItem("Examples",elem_id = "tab-examples",id = 0):
                        examples_hidden = gr.Textbox(visible = False)

                        # getting defualt key value to display
                        first_key = list(QUESTIONS.keys())[0]
                        # create the question category dropdown
                        dropdown_samples = gr.Dropdown(QUESTIONS.keys(),value = first_key,
                                                       interactive = True,show_label = True,
                                                       label = "Select a category of sample questions",
                                                       elem_id = "dropdown-samples")
                        
                        
                        # iterate through the questions list
                        samples = []
                        for i,key in enumerate(QUESTIONS.keys()):
                            examples_visible = True if i == 0 else False
                            with gr.Row(visible = examples_visible) as group_examples:
                                examples_questions = gr.Examples(
                                    QUESTIONS[key],
                                    [examples_hidden],
                                    examples_per_page=8,
                                    run_on_click=False,
                                    elem_id=f"examples{i}",
                                    api_name=f"examples{i}",
                                    # label = "Click on the example question or enter your own",
                                    # cache_examples=True,
                                )
                            
                            samples.append(group_examples)
                    ##------------------- tab for Sources reporting ##------------------
                    with gr.Tab("Sources",elem_id = "tab-citations",id = 1):
                        sources_textbox = gr.HTML(show_label=False, elem_id="sources-textbox")
                        docs_textbox = gr.State("")

    def change_sample_questions(key):
        # update the questions list based on key selected
        index = list(QUESTIONS.keys()).index(key)
        visible_bools = [False] * len(samples)
        visible_bools[index] = True
        return [gr.update(visible=visible_bools[i]) for i in range(len(samples))]

    dropdown_samples.change(change_sample_questions,dropdown_samples,samples)
                        
        # ---- New Guidelines Tab ----
    with gr.Tab("Guidelines", elem_classes="max-height other-tabs"):
        gr.Markdown("""

        #### Welcome to Audit Q&A, your AI-powered assistant for exploring and understanding Uganda's audit reports. This tool leverages advanced language models to help you get clear and structured answers based on audit publications. To get you started, here a few tips on how to use the tool:

        ## 💬 Crafting Effective Prompts

        Clear, specific questions will give you the best results. Here are some examples:

        | ❌ Less Effective | ✅ More Effective |
        |------------------|-------------------|
        | "What are the findings?" | "What were the main issues identified in procurement practices in the Ministry of Health in 2022?" |
        | "Tell me about revenue collection" | "What specific challenges were identified in revenue collection at the local government level in 2021-2022?" |
        | "Is there corruption?" | "What audit findings related to misappropriation of funds were reported in the education sector between 2020-2023?" |

        ## ⭐ Best Practices

        - **Be Clear and Specific**: Frame your questions clearly and focus on what you want to learn.
        - **One Topic at a Time**: Break complex queries into simpler, focused questions.
        - **Provide Context**: Mentioning specific ministries, years, or projects helps narrow the focus.
        - **Follow Up**: Ask follow-up questions to explore a topic more deeply.
        
        ## 🔍 Utilizing Filters

        - **Report Category & Subtype**: Use the "Reports" tab to choose your preferred report category and refine your query by selecting a specific sub-type. This will help narrow down the context for your question.
        - **Year Selection**: Choose one or more years from the "Year" filter to target your query to specific time periods.
        - **Specific Reports**: Optionally, select specific reports using the dropdown to focus on a particular document or set of documents.

        ## 📚 Useful Resources
                            
        - <ins>[**Short Course: Generative AI for Everyone** (3 hours)](https://www.deeplearning.ai/courses/generative-ai-for-everyone/)</ins>
        - <ins>[**Short Course: Advanced Prompting** (1 hour)](https://www.deeplearning.ai/courses/ai-for-everyone/)</ins>
        - <ins>[**Short Course: Introduction to AI with IBM** (13 hours)](https://www.coursera.org/learn/introduction-to-ai)</ins>

        Enjoy using Audit Q&A and happy prompting!
                """)
        
    # static tab 'about us'
    with gr.Tab("About",elem_classes = "max-height other-tabs"):
        with gr.Row():
            with gr.Column(scale=1):
                    gr.Markdown("""The <ins>[**Office of the Auditor General (OAG)**](https://www.oag.go.ug/welcome)</ins> in Uganda, \
                consistent with the mandate of Supreme Audit Institutions (SAIs),\
                remains integral in ensuring transparency and fiscal responsibility.\
                Regularly, the OAG submits comprehensive audit reports to Parliament, \
                which serve as instrumental references for both policymakers and the public, \
                facilitating informed decisions regarding public expenditure. 
                
                However, the prevalent underutilization of these audit reports, \
                leading to numerous unimplemented recommendations, has posed significant challenges\
                to the effectiveness and impact of the OAG's operations. The audit reports made available \
                to the public have not been effectively used by them and other relevant stakeholders. \
                The current format of the audit reports is considered a challenge to the \
                stakeholders' accessibility and usability. This in one way constrains transparency \
                and accountability in the utilization of public funds and effective service delivery. 
                
                In the face of this, modern advancements in Artificial Intelligence (AI),\
                particularly Retrieval Augmented Generation (RAG) technology, \
                emerge as a promising solution. By harnessing the capabilities of such AI tools, \
                there is an opportunity not only to improve the accessibility and understanding \
                of these audit reports but also to ensure that their insights are effectively \
                translated into actionable outcomes, thereby reinforcing public transparency \
                and service delivery in Uganda. 
                
                To address these issues, the OAG has initiated several projects, \
                such as the Audit Recommendation Tracking (ART) System and the Citizens Feedback Platform (CFP). \
                These systems are designed to increase the transparency and relevance of audit activities. \
                However, despite these efforts, engagement and awareness of the audit findings remain low, \
                and the complexity of the information often hinders effective public utilization. Recognizing the need for further\
                enhancement in how audit reports are processed and understood, \
                the **Civil Society and Budget Advocacy Group (CSBAG)** in partnership with the **GIZ**, \
                has recognizing the need for further enhancement in how audit reports are processed and understood.   
                
                This prototype tool leveraging AI (Artificial Intelligence) aims at offering critical capabilities such as '
                summarizing complex texts, extracting thematic insights, and enabling interactive, \
                user-friendly analysis through a chatbot interface. By making the audit reports more accessible,\
                this aims to increase readership and utilization among stakeholders, \
                which can lead to better accountability and improve service delivery
                
                """)


    
    # static tab for disclaimer
    with gr.Tab("Disclaimer",elem_classes = "max-height other-tabs"):
        with gr.Row():
            with gr.Column(scale=1):
                gr.Markdown("""
                - This chatbot is intended for specific use of answering the questions based on audit reports published by OAG, for any use beyond this scope we have no liability to response provided by chatbot.
                - We do not guarantee the accuracy, reliability, or completeness of any information provided by the chatbot and disclaim any liability or responsibility for actions taken based on its responses.
                - The chatbot may occasionally provide inaccurate or inappropriate responses, and it is important to exercise judgment and critical thinking when interpreting its output.
                - The chatbot responses should not be considered professional or authoritative advice and are generated based on patterns in the data it has been trained on.
                - The chatbot's responses do not reflect the opinions or policies of our organization or its affiliates.
                - Any personal or sensitive information shared with the chatbot is at the user's own risk, and we cannot guarantee complete privacy or confidentiality.
                - the chatbot is not deterministic, so there might be change in answer to same question when asked by different users or multiple times.
                - By using this chatbot, you agree to these terms and acknowledge that you are solely responsible for any reliance on or actions taken based on its responses.
                - **This is just a prototype and being tested and worked upon, so its not perfect and may sometimes give irrelevant answers**. If you are not satisfied with the answer, please ask a more specific question or report your feedback to help us improve the system.
                """)
            
                
                

    #-------------------- New UI elements for Feedback -------------------------


    def submit_feedback_okay(logs_data):
        """Handle 'okay' feedback submission"""
        return submit_feedback("okay", logs_data)

    def submit_feedback_not_okay(logs_data):
        """Handle 'not okay' feedback submission"""
        return submit_feedback("not_okay", logs_data)
    
    def show_feedback(logs_data):
        """Handle feedback display with proper output format"""
        if logs_data is None:
            return (
                gr.update(visible=False),  # feedback_row
                gr.update(visible=False),  # feedback_thanks
                None  # feedback_state
            )
        return (
            gr.update(visible=True),  # feedback_row
            gr.update(visible=False),  # feedback_thanks
            logs_data  # feedback_state
        )
    
    okay_btn.click(
        submit_feedback_okay,
        [feedback_state],
        [feedback_row, feedback_thanks]
    )
    
    not_okay_btn.click(
        submit_feedback_not_okay,
        [feedback_state],
        [feedback_row, feedback_thanks]
    )

   #-------------------- Session Management + Geolocation -------------------------

    # Add these state components at the top level of the Blocks
    session_id = gr.State(None)
    client_ip = gr.State(None)

    
    @demo.load(api_name="get_client_ip")
    def get_client_ip_handler(dummy_input="", request: gr.Request = None):
        """Handler for getting client IP in Gradio context"""
        return get_client_ip(request)
    
    #-------------------- No Filters Set Warning -------------------------
    # Warn users when no filters are selected
    warning_state = gr.State(False)
    pending_query = gr.State(None)

    def show_warning():
        """Show warning popup when no filters selected"""
        return gr.update(visible=True)

    def hide_warning():
        """Hide warning popup"""
        return gr.update(visible=False)

    # Logic needs to be changed to accomodate default filter values (currently I have them all set to None)
    def check_filters(check_status, textbox_value, sources, reports, subtype, year):
        """Check if any filters are selected"""
        # If a previous check failed, don't continue with this check
        if check_status is not None:
            return (
                check_status,  # keep current check status
                False,  # keep warning state unchanged
                gr.update(visible=False),  # keep warning row visibility unchanged
                textbox_value,  # keep the textbox value
                None  # no need to store query
            )
        
        no_filters = (not reports) and (not sources) and (not subtype) and (not year)
        if no_filters:
            # If no filters, show warning and set status
            return (
                "filter",  # check status - no filters selected
                True,  # warning state
                gr.update(visible=True),  # warning row visibility
                gr.update(value=""),  # clear textbox
                textbox_value  # store the query
            )
        # If filters exist, proceed normally
        return (
            None,  # no check failed
            False,  # normal state
            gr.update(visible=False),  # hide warning
            textbox_value,  # keep the original value
            None  # no need to store query
        )
    
    async def handle_chat_flow(check_status, warning_active, short_query_warning_active, query, chatbot, sources, reports, subtype, year, client_ip, session_id):
        """Handle chat flow with explicit check for status"""
        # Don't proceed if any check failed or query is None
        if check_status is not None or warning_active or short_query_warning_active or query is None or query == "":
            yield (
                chatbot,  # unchanged chatbot
                "",  # empty sources
                None,  # no feedback state
                session_id  # keep session
            )
            return  # Exit the generator
        
        # Include start_chat functionality here
        history = chatbot + [(query, None)]
        history = [tuple(x) for x in history]
        
        # Proceed with chat and yield each update
        async for update in chat(query, history, sources, reports, subtype, year, client_ip, session_id):
            yield update

    #-------------------- Short Query Warning -------------------------
    # Warn users when query is too short (less than 4 words)
    short_query_warning_state = gr.State(False)
    check_status = gr.State(None)

    def check_query_length(textbox_value):
        """Check if query has at least 4 words"""
        if textbox_value and len(textbox_value.split()) < 4:
            # If query is too short, show warning and set status
            return (
                "short",  # check status - this query is too short
                True,  # short query warning state
                gr.update(visible=True),  # short query warning row visibility
                gr.update(value=""),  # clear textbox 
                textbox_value  # store the query
            )
        # If query is long enough, proceed normally
        return (
            None,  # no check failed
            False,  # normal state
            gr.update(visible=False),  # hide warning
            gr.update(value=textbox_value),  # keep the textbox value
            None  # no need to store query
        )

    
    #-------------------- Gradio Handlers -------------------------
    
    # Hanlders: Text input from Textbox
    (textbox
        .submit(
            check_query_length,
            [textbox],
            [check_status, short_query_warning_state, short_query_warning_row, textbox, pending_query],
            api_name="check_query_length_textbox"
        )
        .then(
            check_filters,
            [check_status, textbox, dropdown_sources, dropdown_reports, dropdown_category, dropdown_year],
            [check_status, warning_state, warning_row, textbox, pending_query],
            api_name="submit_textbox",
            show_progress=False
        )
        .then(
            get_client_ip_handler,
            [textbox],
            [client_ip],
            show_progress=False,
            api_name="get_client_ip_textbox"
        )
        .then(
            handle_chat_flow,
            [check_status, warning_state, short_query_warning_state, textbox, chatbot, dropdown_sources, dropdown_reports, dropdown_category, dropdown_year, client_ip, session_id],
            [chatbot, sources_textbox, feedback_state, session_id],
            queue=True,
            api_name="handle_chat_flow_textbox"
        )
        .then(
            show_feedback,
            [feedback_state],
            [feedback_row, feedback_thanks, feedback_state],
            api_name="show_feedback_textbox"
        )
        .then(
            finish_chat,
            None,
            [textbox],
            api_name="finish_chat_textbox"
        ))
    
    # Hanlders: Text input from Examples (same chain as textbox)
    examples_hidden.change(
        lambda x: x,
        inputs=examples_hidden,
        outputs=textbox,
        api_name="submit_examples"
        ).then(
            check_query_length,
            [textbox],
            [check_status, short_query_warning_state, short_query_warning_row, textbox, pending_query],
            api_name="check_query_length_examples"
        ).then(
            check_filters,
            [check_status, textbox, dropdown_sources, dropdown_reports, dropdown_category, dropdown_year],
            [check_status, warning_state, warning_row, textbox, pending_query],
            api_name="check_filters_examples",
            show_progress=False
        ).then(
            get_client_ip_handler,
            [textbox],
            [client_ip],
            show_progress=False,
            api_name="get_client_ip_examples"
        ).then( 
            handle_chat_flow,
            [check_status, warning_state, short_query_warning_state, textbox, chatbot, dropdown_sources, dropdown_reports, dropdown_category, dropdown_year, client_ip, session_id],
            [chatbot, sources_textbox, feedback_state, session_id],
            queue=True,
            api_name="handle_chat_flow_examples"
        ).then(
            show_feedback,
            [feedback_state],
            [feedback_row, feedback_thanks, feedback_state],
            api_name="show_feedback_examples"
        ).then(
            finish_chat,
            None,
            [textbox],
            api_name="finish_chat_examples"
        )
        

    # Handlers for the warning buttons
    proceed_btn.click(
        lambda query: (
            None,  # reset check status
            False,  # warning state
            gr.update(visible=False),  # warning row
            gr.update(value=query if query else "", interactive=True),  # restore query
            None  # clear pending query
        ),
        pending_query,
        [check_status, warning_state, warning_row, textbox, pending_query]
    ).then(
        get_client_ip_handler,
        [textbox],
        [client_ip]
    ).then(
        handle_chat_flow,
        [check_status, warning_state, short_query_warning_state, textbox, chatbot, dropdown_sources, dropdown_reports, dropdown_category, dropdown_year, client_ip, session_id],
        [chatbot, sources_textbox, feedback_state, session_id],
        queue=True  
    ).then(
        show_feedback,
        [feedback_state],
        [feedback_row, feedback_thanks, feedback_state]
    ).then(
        finish_chat,
        None,
        [textbox]
    )

    # Cancel button for no filters
    cancel_btn.click(
        lambda: (
            None,  # reset check status
            False,  # warning state
            gr.update(visible=False),  # warning row
            gr.update(value="", interactive=True),  # clear textbox
            None  # clear pending query
        ),
        None,
        [check_status, warning_state, warning_row, textbox, pending_query]
    )

    # short query warning OK button
    short_query_proceed_btn.click(
        lambda query: (
            None,  # reset check status
            False,  # short query warning state
            gr.update(visible=False),  # short query warning row
            gr.update(value=query if query else "", interactive=True),  # restore query
            None  # clear pending query
        ),
        pending_query,
        [check_status, short_query_warning_state, short_query_warning_row, textbox, pending_query]
    )

    demo.queue()

demo.launch()