app.py

import gradio as gr
import pandas as pd
import logging
import asyncio
import os
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
from dotenv import load_dotenv
from threading import Lock
# import json
# from functools import partial
# import time
from gradio.routes import Request
from qdrant_client import QdrantClient

# TESTING DEBUG LOG
from auditqa.logging_config import setup_logging
setup_logging()
import logging
logger = logging.getLogger(__name__)
logger.setLevel(logging.DEBUG)  # Ensure debug logging is enabled

load_dotenv()

# # fetch tokens and model config params
SPACES_LOG = os.environ["SPACES_LOG"]
SPACES_LOG = os.getenv('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="mtyrrell/spaces_log",
     repo_type="dataset",
     folder_path=JSON_DATASET_DIR,
     path_in_repo=".",
     token=SPACES_LOG,
     every=2)

#####--------------- VECTOR STORE -------------------------------------------------
# Configure cloud Qdrant client
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:
            logger.error("No logs data available for feedback")
            return gr.update(visible=False), gr.update(visible=True)
            
        save_logs(scheduler, JSON_DATASET_PATH, logs_data, feedback)
        return gr.update(visible=False), gr.update(visible=True)
    except Exception as e:
        logger.error(f"Error saving feedback: {e}")
        # Still need to return the expected outputs even on error
        return gr.update(visible=False), gr.update(visible=True)

# Session Manager added (track session duration & location)
class SessionManager:
    def __init__(self):
        self.sessions = {}
        
    def create_session(self, client_ip):
        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)
        }
        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):
    """Update chat function to handle session data"""
    # TESTING: DEBUG LOG
    logger.debug(f"Chat function called with query: {query}")
    logger.debug(f"Client IP: {client_ip}")
    logger.debug(f"Session ID: {session_id}")
    
    if not session_id: # Session managment
        session_id = session_manager.create_session(client_ip)
        logger.debug(f"Created new session: {session_id}")
    else:
        session_manager.update_session(session_id)
        logger.debug(f"Updated existing session: {session_id}")
    
    # Get session data
    session_data = session_manager.get_session_data(session_id)
    session_duration = session_manager.get_session_duration(session_id)
    logger.debug(f"Session duration: {session_duration}")
    
    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---------------------------------------------- 

    context_retrieved = get_context(vectorstore=vectorstore,query=query,reports=reports,
                                                sources=sources,subtype=subtype,year=year)
    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 feedback + session data (moved up here because: feedback)
    timestamp = str(datetime.now().timestamp())
    logs_data = {
        "session_id": session_id,
        "client_ip": client_ip,
        "client_location": session_data['location_info'],
        "session_duration_seconds": session_duration,
        # "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],
        "answer": "",
        "time": timestamp,
    }

    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
                    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:
                logger.error(f"Error in process_stream: {str(e)}")
                raise

        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:
        logging.error(e)




#####-------------------------- 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"),
                )
                



                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):
                        gr.Markdown("Reminder: To get better results select the specific report/reports")

                        
                        #----- 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="Consolidated",
                            interactive=True,
                        )

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

                        #----------- update the secodn level filter abse don values from first level ----------------
                        def rs_change(rs):
                            return gr.update(choices=new_files[rs], value=new_files[rs])
                        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=['2023'],
                            interactive=True,
                        )
                        gr.Markdown("-------------------------------------------------------------------------")
                        #---------------- Another way to select reports across category and sub-type ------------
                        dropdown_reports = gr.Dropdown(
                        new_report_list,
                        label="Or select specific reports",
                        multiselect=True,
                        value=[],
                        interactive=True,)

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

    # 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.
                """)
            
                
                

    #-------------------- Feedback UI elements + state management -------------------------
    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()  # Add state to store logs data
    debug_ip = gr.Markdown(visible=False)  # Add debug display for IP

    def show_feedback(logs):
        """Show feedback buttons and store logs in state"""
        return gr.update(visible=True), gr.update(visible=False), logs

    def submit_feedback_okay(logs_data):
        """Handle 'okay' feedback submission"""
        ip = logs_data.get('client_ip', 'No IP found')
        location = logs_data.get('client_location', 'No location found')
        return submit_feedback("okay", logs_data) + (gr.update(visible=True, value=f"TESTING - Client IP: {ip}, Location: {location}"),)

    def submit_feedback_not_okay(logs_data):
        """Handle 'not okay' feedback submission"""
        ip = logs_data.get('client_ip', 'No IP found')
        location = logs_data.get('client_location', 'No location found')
        return submit_feedback("not_okay", logs_data) + (gr.update(visible=True, value=f"TESTING - Client IP: {ip}, Location: {location}"),)
    
    okay_btn.click(
        submit_feedback_okay,
        [feedback_state],
        [feedback_row, feedback_thanks, debug_ip]
    )
    
    not_okay_btn.click(
        submit_feedback_not_okay,
        [feedback_state],
        [feedback_row, feedback_thanks, debug_ip]
    )

    #-------------------- Gradio voodoo continued -------------------------

    # 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)
    
    # Update the event handlers
    (textbox
        .submit(get_client_ip_handler, [textbox], [client_ip], api_name="get_ip_textbox")
        .then(start_chat, [textbox, chatbot], [textbox, tabs, chatbot], queue=False, api_name="start_chat_textbox")
        .then(chat, 
            [textbox, chatbot, dropdown_sources, dropdown_reports, dropdown_category, dropdown_year, client_ip, session_id], 
            [chatbot, sources_textbox, feedback_state, session_id], 
            queue=True, concurrency_limit=8, api_name="chat_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"))

    (examples_hidden
        .change(start_chat, [examples_hidden, chatbot], [textbox, tabs, chatbot], queue=False, api_name="start_chat_examples")
        .then(get_client_ip_handler, [examples_hidden], [client_ip], api_name="get_ip_examples")
        .then(chat, 
            [examples_hidden, chatbot, dropdown_sources, dropdown_reports, dropdown_category, dropdown_year, client_ip, session_id], 
            [chatbot, sources_textbox, feedback_state, session_id], 
            concurrency_limit=8, api_name="chat_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"))

    demo.queue()

demo.launch()
logger.info("App launched")