app.py

# Import necessary libraries
import os
import json
import time
import re
import base64
from datetime import datetime
from io import BytesIO

import numpy as np
import pandas as pd
import gradio as gr
from gradio import ChatMessage
from gradio_modal import Modal
from sentence_transformers import CrossEncoder
from azure.storage.fileshare import ShareServiceClient

# Import custom modules
from climateqa.engine.embeddings import get_embeddings_function
from climateqa.engine.llm import get_llm
from climateqa.engine.vectorstore import get_pinecone_vectorstore
from climateqa.engine.reranker import get_reranker
from climateqa.sample_questions import QUESTIONS
from climateqa.constants import POSSIBLE_REPORTS
from climateqa.utils import get_image_from_azure_blob_storage
from climateqa.engine.graph import make_graph_agent
from climateqa.engine.chains.retrieve_papers import find_papers
from front.utils import serialize_docs, process_figures
from climateqa.event_handler import (
    init_audience,
    handle_retrieved_documents,
    stream_answer,
    handle_retrieved_owid_graphs,
    convert_to_docs_to_html
)
from utils import create_user_id
from front.utils import make_html_source
import logging

logging.basicConfig(level=logging.WARNING)
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'  # Suppresses INFO and WARNING logs
logging.getLogger().setLevel(logging.WARNING)


# Load environment variables in local mode
try:
    from dotenv import load_dotenv
    load_dotenv()
except Exception as e:
    pass


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

# Initialize prompt and system template
init_prompt = """
Hello, I am ClimateQ&A, a conversational assistant designed to help you understand climate change and biodiversity loss. I will answer your questions by **sifting through the IPCC and IPBES scientific reports**.

❓ How to use
- **Language**: You can ask me your questions in any language. 
- **Audience**: You can specify your audience (children, general public, experts) to get a more adapted answer.
- **Sources**: You can choose to search in the IPCC or IPBES reports, or both.
- **Relevant content sources**: You can choose to search for figures, papers, or graphs that can be relevant for your question.

⚠️ Limitations
*Please note that the AI is 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.*

🛈 Information
Please note that we log your questions for meta-analysis purposes, so avoid sharing any sensitive or personal information.

What do you want to learn ?
"""

# Azure Blob Storage credentials
account_key = os.environ["BLOB_ACCOUNT_KEY"]
if len(account_key) == 86:
    account_key += "=="

credential = {
    "account_key": account_key,
    "account_name": os.environ["BLOB_ACCOUNT_NAME"],
}

account_url = os.environ["BLOB_ACCOUNT_URL"]
file_share_name = "climateqa"
service = ShareServiceClient(account_url=account_url, credential=credential)
share_client = service.get_share_client(file_share_name)

user_id = create_user_id()

# Citation information
CITATION_LABEL = "BibTeX citation for ClimateQ&A"
CITATION_TEXT = r"""@misc{climateqa,
    author={Théo Alves Da Costa, Timothée Bohe},
    title={ClimateQ&A, AI-powered conversational assistant for climate change and biodiversity loss},
    year={2024},
    howpublished= {\url{https://climateqa.com}},
}
@software{climateqa,
    author = {Théo Alves Da Costa, Timothée Bohe},
    publisher = {ClimateQ&A},
    title = {ClimateQ&A, AI-powered conversational assistant for climate change and biodiversity loss},
}
"""

# Create vectorstore and retriever
embeddings_function = get_embeddings_function()
vectorstore = get_pinecone_vectorstore(embeddings_function, index_name=os.getenv("PINECONE_API_INDEX"))
vectorstore_graphs = get_pinecone_vectorstore(embeddings_function, index_name=os.getenv("PINECONE_API_INDEX_OWID"), text_key="description")
vectorstore_region = get_pinecone_vectorstore(embeddings_function, index_name=os.getenv("PINECONE_API_INDEX_REGION"))

llm = get_llm(provider="openai",max_tokens = 1024,temperature = 0.0)
reranker = get_reranker("nano")

agent = make_graph_agent(llm=llm, vectorstore_ipcc=vectorstore, vectorstore_graphs=vectorstore_graphs, vectorstore_region = vectorstore_region, reranker=reranker, threshold_docs=0)#TODO put back default 0.2

# Function to update modal visibility
def update_config_modal_visibility(config_open):
    new_config_visibility_status = not config_open
    return gr.update(visible=new_config_visibility_status), new_config_visibility_status

# Main chat function
async def chat(
    query: str, 
    history: list[ChatMessage], 
    audience: str, 
    sources: list[str], 
    reports: list[str], 
    relevant_content_sources_selection: list[str], 
    search_only: bool
) -> tuple[list, str, str, str, list, str]:
    """Process a chat query and return response with relevant sources and visualizations.
    
    Args:
        query (str): The user's question
        history (list): Chat message history
        audience (str): Target audience type
        sources (list): Knowledge base sources to search
        reports (list): Specific reports to search within sources
        relevant_content_sources_selection (list): Types of content to retrieve (figures, papers, etc)
        search_only (bool): Whether to only search without generating answer
        
    Yields:
        tuple: Contains:
            - history: Updated chat history
            - docs_html: HTML of retrieved documents
            - output_query: Processed query
            - output_language: Detected language
            - related_contents: Related content
            - graphs_html: HTML of relevant graphs
    """
    # Log incoming question
    date_now = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    print(f">> NEW QUESTION ({date_now}) : {query}")

    audience_prompt = init_audience(audience)
    sources = sources or ["IPCC", "IPBES"]
    reports = reports or []

    # Prepare inputs for agent
    inputs = {
        "user_input": query,
        "audience": audience_prompt,
        "sources_input": sources,
        "relevant_content_sources_selection": relevant_content_sources_selection,
        "search_only": search_only,
        "reports": reports
    }

    # Get streaming events from agent
    result = agent.astream_events(inputs, version="v1")

    # Initialize state variables
    docs = []
    related_contents = []
    docs_html = ""
    new_docs_html = ""
    output_query = ""
    output_language = ""
    output_keywords = ""
    start_streaming = False
    graphs_html = ""    
    used_documents = []
    answer_message_content = ""

    # Define processing steps
    steps_display = {
        "categorize_intent": ("🔄️ Analyzing user message", True),
        "transform_query": ("🔄️ Thinking step by step to answer the question", True),
        "retrieve_documents": ("🔄️ Searching in the knowledge base", False),
        "retrieve_local_data": ("🔄️ Searching in the knowledge base", False),
    }

    try:
        # Process streaming events
        async for event in result:

            if "langgraph_node" in event["metadata"]:
                node = event["metadata"]["langgraph_node"]

                # Handle document retrieval
                if event["event"] == "on_chain_end" and event["name"] in ["retrieve_documents","retrieve_local_data"] and event["data"]["output"] != None:
                    history, used_documents = handle_retrieved_documents(
                        event, history, used_documents
                    )
                if event["event"] == "on_chain_end" and event["name"] == "answer_search" :
                    docs = event["data"]["input"]["documents"]
                    docs_html = convert_to_docs_to_html(docs)                    
                    related_contents = event["data"]["input"]["related_contents"]
        
                # Handle intent categorization
                elif (event["event"] == "on_chain_end" and 
                      node == "categorize_intent" and 
                      event["name"] == "_write"):
                    intent = event["data"]["output"]["intent"]
                    output_language = event["data"]["output"].get("language", "English")
                    history[-1].content = f"Language identified: {output_language}\nIntent identified: {intent}"

                # Handle processing steps display
                elif event["name"] in steps_display and event["event"] == "on_chain_start":
                    event_description, display_output = steps_display[node]
                    if (not hasattr(history[-1], 'metadata') or 
                        history[-1].metadata["title"] != event_description):
                        history.append(ChatMessage(
                            role="assistant",
                            content="",
                            metadata={'title': event_description}
                        ))

                # Handle answer streaming
                elif (event["name"] != "transform_query" and 
                      event["event"] == "on_chat_model_stream" and
                      node in ["answer_rag","answer_rag_no_docs", "answer_search", "answer_chitchat"]):
                    history, start_streaming, answer_message_content = stream_answer(
                        history, event, start_streaming, answer_message_content
                    )

                # Handle graph retrieval
                elif event["name"] in ["retrieve_graphs", "retrieve_graphs_ai"] and event["event"] == "on_chain_end":
                    graphs_html = handle_retrieved_owid_graphs(event, graphs_html)

                # Handle query transformation
                if event["name"] == "transform_query" and event["event"] == "on_chain_end":
                    if hasattr(history[-1], "content"):
                        sub_questions = [q["question"] for q in event["data"]["output"]["questions_list"]]
                        history[-1].content += "Decompose question into sub-questions:\n\n - " + "\n - ".join(sub_questions)

            yield history, docs_html, output_query, output_language, related_contents, graphs_html 

    except Exception as e:
        print(f"Event {event} has failed")
        raise gr.Error(str(e))

    try:
        # Log interaction to Azure if not in local environment
        if os.getenv("GRADIO_ENV") != "local":
            timestamp = str(datetime.now().timestamp())
            prompt = history[1]["content"]
            logs = {
                "user_id": str(user_id),
                "prompt": prompt,
                "query": prompt,
                "question": output_query,
                "sources": sources,
                "docs": serialize_docs(docs),
                "answer": history[-1].content,
                "time": timestamp,
            }
            log_on_azure(f"{timestamp}.json", logs, share_client)
    except Exception as e:
        print(f"Error logging on Azure Blob Storage: {e}")
        error_msg = f"ClimateQ&A Error: {str(e)[:100]} - The error has been noted, try another question and if the error remains, you can contact us :)"
        raise gr.Error(error_msg)

    yield history, docs_html, output_query, output_language, related_contents, graphs_html 

# Function to save feedback
def save_feedback(feed: str, user_id):
    if len(feed) > 1:
        timestamp = str(datetime.now().timestamp())
        file = user_id + timestamp + ".json"
        logs = {
            "user_id": user_id,
            "feedback": feed,
            "time": timestamp,
        }
        log_on_azure(file, logs, share_client)
        return "Feedback submitted, thank you!"

# Function to log data on Azure
def log_on_azure(file, logs, share_client):
    logs = json.dumps(logs)
    file_client = share_client.get_file_client(file)
    file_client.upload_file(logs)





# --------------------------------------------------------------------
# Gradio
# --------------------------------------------------------------------



def vote(data: gr.LikeData):
    if data.liked:
        print(data.value)
    else:
        print(data)

def save_graph(saved_graphs_state, embedding, category):
    print(f"\nCategory:\n{saved_graphs_state}\n")
    if category not in saved_graphs_state:
        saved_graphs_state[category] = []
    if embedding not in saved_graphs_state[category]:
        saved_graphs_state[category].append(embedding)
    return saved_graphs_state, gr.Button("Graph Saved")


# Functions to toggle visibility
def toggle_summary_visibility():
    global summary_visible
    summary_visible = not summary_visible
    return gr.update(visible=summary_visible)

def toggle_relevant_visibility():
    global relevant_visible
    relevant_visible = not relevant_visible
    return gr.update(visible=relevant_visible)

def change_completion_status(current_state):
    current_state = 1 - current_state
    return current_state

def update_sources_number_display(sources_textbox, figures_cards, current_graphs, papers_html):
    sources_number = sources_textbox.count("<h2>")
    figures_number = figures_cards.count("<h2>")
    graphs_number = current_graphs.count("<iframe")
    papers_number = papers_html.count("<h2>")
    sources_notif_label = f"Sources ({sources_number})"
    figures_notif_label = f"Figures ({figures_number})"
    graphs_notif_label = f"Graphs ({graphs_number})"
    papers_notif_label = f"Papers ({papers_number})"
    recommended_content_notif_label = f"Recommended content ({figures_number + graphs_number + papers_number})"

    return gr.update(label=recommended_content_notif_label), gr.update(label=sources_notif_label), gr.update(label=figures_notif_label), gr.update(label=graphs_notif_label), gr.update(label=papers_notif_label)

def change_sample_questions(key):
    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))]



# Chat functions
def start_chat(query, history, search_only):
    history = history + [ChatMessage(role="user", content=query)]
    if not search_only:
        return (gr.update(interactive=False), gr.update(selected=1), history, [])
    else:
        return (gr.update(interactive=False), gr.update(selected=2), history, [])

def finish_chat():
    return gr.update(interactive=True, value="")

# Initialize visibility states
summary_visible = False
relevant_visible = False

# UI Layout Components
def create_chat_interface():
    chatbot = gr.Chatbot(
        value=[ChatMessage(role="assistant", content=init_prompt)],
        type="messages", 
        show_copy_button=True,
        show_label=False,
        elem_id="chatbot",
        layout="panel",
        avatar_images=(None, "https://i.ibb.co/YNyd5W2/logo4.png"),
        max_height="80vh",
        height="100vh"
    )
    
    with gr.Row(elem_id="input-message"):
        
        textbox = gr.Textbox(
            placeholder="Ask me anything here!",
            show_label=False,
            scale=12,
            lines=1,
            interactive=True,
            elem_id="input-textbox"
        )
        
        config_button = gr.Button("", elem_id="config-button")
    
    return chatbot, textbox, config_button

def create_examples_tab():
    examples_hidden = gr.Textbox(visible=False)
    first_key = list(QUESTIONS.keys())[0]
    dropdown_samples = gr.Dropdown(
        choices=QUESTIONS.keys(),
        value=first_key,
        interactive=True,
        label="Select a category of sample questions",
        elem_id="dropdown-samples"
    )

    samples = []
    for i, key in enumerate(QUESTIONS.keys()):
        examples_visible = (i == 0)
        with gr.Row(visible=examples_visible) as group_examples:
            examples_questions = gr.Examples(
                examples=QUESTIONS[key],
                inputs=[examples_hidden],
                examples_per_page=8,
                run_on_click=False,
                elem_id=f"examples{i}",
                api_name=f"examples{i}"
            )
        samples.append(group_examples)
        
    return examples_hidden, dropdown_samples, samples

def create_figures_tab():
    sources_raw = gr.State()
    new_figures = gr.State([])
    used_figures = gr.State([])
    
    with Modal(visible=False, elem_id="modal_figure_galery") as figure_modal:
        gallery_component = gr.Gallery(
            object_fit='scale-down',
            elem_id="gallery-component",
            height="80vh"
        )
        
    show_full_size_figures = gr.Button(
        "Show figures in full size",
        elem_id="show-figures",
        interactive=True
    )
    show_full_size_figures.click(
        lambda: Modal(visible=True),
        None,
        figure_modal
    )

    figures_cards = gr.HTML(show_label=False, elem_id="sources-figures")
    
    return sources_raw, new_figures, used_figures, gallery_component, figures_cards, figure_modal

def create_papers_tab():
    with gr.Accordion(
        visible=True,
        elem_id="papers-summary-popup",
        label="See summary of relevant papers",
        open=False
    ) as summary_popup:
        papers_summary = gr.Markdown("", visible=True, elem_id="papers-summary")

    with gr.Accordion(
        visible=True,
        elem_id="papers-relevant-popup",
        label="See relevant papers",
        open=False
    ) as relevant_popup:
        papers_html = gr.HTML(show_label=False, elem_id="papers-textbox")

    btn_citations_network = gr.Button("Explore papers citations network")
    with Modal(visible=False) as papers_modal:
        citations_network = gr.HTML(
            "<h3>Citations Network Graph</h3>",
            visible=True,
            elem_id="papers-citations-network"
        )
    btn_citations_network.click(
        lambda: Modal(visible=True),
        None,
        papers_modal
    )
    
    return papers_summary, papers_html, citations_network, papers_modal

def create_config_modal(config_open):
    with Modal(visible=False, elem_id="modal-config") as config_modal:
        gr.Markdown("Reminders: You can talk in any language, ClimateQ&A is multi-lingual!")

        dropdown_sources = gr.CheckboxGroup(
            choices=["IPCC", "IPBES", "IPOS"],
            label="Select source (by default search in all sources)",
            value=["IPCC"],
            interactive=True
        )

        dropdown_reports = gr.Dropdown(
            choices=POSSIBLE_REPORTS,
            label="Or select specific reports",
            multiselect=True,
            value=None,
            interactive=True
        )

        dropdown_external_sources = gr.CheckboxGroup(
            choices=["Figures (IPCC/IPBES)", "Papers (OpenAlex)", "Graphs (OurWorldInData)","POC region"],
            label="Select database to search for relevant content",
            value=["Figures (IPCC/IPBES)","POC region"],
            interactive=True
        )

        search_only = gr.Checkbox(
            label="Search only for recommended content without chating",
            value=False,
            interactive=True,
            elem_id="checkbox-chat"
        )

        dropdown_audience = gr.Dropdown(
            choices=["Children", "General public", "Experts"],
            label="Select audience", 
            value="Experts",
            interactive=True
        )

        after = gr.Slider(
            minimum=1950,
            maximum=2023,
            step=1,
            value=1960,
            label="Publication date",
            show_label=True,
            interactive=True,
            elem_id="date-papers",
            visible=False
        )

        output_query = gr.Textbox(
            label="Query used for retrieval",
            show_label=True,
            elem_id="reformulated-query",
            lines=2,
            interactive=False,
            visible=False
        )

        output_language = gr.Textbox(
            label="Language",
            show_label=True,
            elem_id="language",
            lines=1,
            interactive=False,
            visible=False
        )

        dropdown_external_sources.change(
            lambda x: gr.update(visible="Papers (OpenAlex)" in x),
            inputs=[dropdown_external_sources],
            outputs=[after]
        )

        close_config_modal = gr.Button("Validate and Close", elem_id="close-config-modal")
        close_config_modal.click(
            fn=update_config_modal_visibility,
            inputs=[config_open],
            outputs=[config_modal, config_open]
        )
        
        return {
            "config_open" : config_open,
            "config_modal": config_modal,
            "dropdown_sources": dropdown_sources,
            "dropdown_reports": dropdown_reports,
            "dropdown_external_sources": dropdown_external_sources,
            "search_only": search_only,
            "dropdown_audience": dropdown_audience,
            "after": after,
            "output_query": output_query,
            "output_language": output_language,
        }

def cqa_tab(tab_name):
    # State variables
    current_graphs = gr.State([])
    with gr.Tab(tab_name):
        with gr.Row(elem_id="chatbot-row"):
            # Left column - Chat interface
            with gr.Column(scale=2):
                chatbot, textbox, config_button = create_chat_interface()

            # Right column - Content panels
            with gr.Column(scale=2, variant="panel", elem_id="right-panel"):
                with gr.Tabs(elem_id="right_panel_tab") as tabs:
                    # Examples tab
                    with gr.TabItem("Examples", elem_id="tab-examples", id=0):
                        examples_hidden, dropdown_samples, samples = create_examples_tab()

                    # Sources tab
                    with gr.Tab("Sources", elem_id="tab-sources", id=1) as tab_sources:
                        sources_textbox = gr.HTML(show_label=False, elem_id="sources-textbox")


                    # Recommended content tab
                    with gr.Tab("Recommended content", elem_id="tab-recommended_content", id=2) as tab_recommended_content:
                        with gr.Tabs(elem_id="group-subtabs") as tabs_recommended_content:
                            # Figures subtab
                            with gr.Tab("Figures", elem_id="tab-figures", id=3) as tab_figures:
                                sources_raw, new_figures, used_figures, gallery_component, figures_cards, figure_modal = create_figures_tab()

                            # Papers subtab
                            with gr.Tab("Papers", elem_id="tab-citations", id=4) as tab_papers:
                                papers_summary, papers_html, citations_network, papers_modal = create_papers_tab()

                            # Graphs subtab
                            with gr.Tab("Graphs", elem_id="tab-graphs", id=5) as tab_graphs:
                                graphs_container = gr.HTML(
                                    "<h2>There are no graphs to be displayed at the moment. Try asking another question.</h2>",
                                    elem_id="graphs-container"
                                )
    return {
        "chatbot": chatbot,
        "textbox": textbox,
        "tabs": tabs,
        "sources_raw": sources_raw,
        "new_figures": new_figures,
        "current_graphs": current_graphs,
        "examples_hidden": examples_hidden,
        "dropdown_samples": dropdown_samples,
        "samples": samples,
        "sources_textbox": sources_textbox,
        "figures_cards": figures_cards,
        "gallery_component": gallery_component,
        "config_button": config_button,
        "papers_html": papers_html,
        "citations_network": citations_network,
        "papers_summary": papers_summary,
        "tab_recommended_content": tab_recommended_content,
        "tab_sources": tab_sources,
        "tab_figures": tab_figures,
        "tab_graphs": tab_graphs,
        "tab_papers": tab_papers,
        "graph_container": graphs_container
    }
                                
def about_tab():
    with gr.Tab("About", elem_classes="max-height other-tabs"):
        with gr.Row():
            with gr.Column(scale=1):
                gr.Markdown(
                    """
                    ### More info
                    - See more info at [https://climateqa.com](https://climateqa.com/docs/intro/)
                    - Feedbacks on this [form](https://forms.office.com/e/1Yzgxm6jbp)
                                                
                    ### Citation
                    """
                )
                with gr.Accordion(CITATION_LABEL, elem_id="citation", open=False):
                    gr.Textbox(
                        value=CITATION_TEXT,
                        label="",
                        interactive=False,
                        show_copy_button=True,
                        lines=len(CITATION_TEXT.split('\n')),
                    )
    
def event_handling(
    main_tab_components,
    config_components,
    tab_name="ClimateQ&A"
):
    chatbot = main_tab_components["chatbot"]
    textbox = main_tab_components["textbox"]
    tabs = main_tab_components["tabs"]
    sources_raw = main_tab_components["sources_raw"]
    new_figures = main_tab_components["new_figures"]
    current_graphs = main_tab_components["current_graphs"]
    examples_hidden = main_tab_components["examples_hidden"]
    dropdown_samples = main_tab_components["dropdown_samples"]
    samples = main_tab_components["samples"]
    sources_textbox = main_tab_components["sources_textbox"]
    figures_cards = main_tab_components["figures_cards"]
    gallery_component = main_tab_components["gallery_component"]
    config_button = main_tab_components["config_button"]
    papers_html = main_tab_components["papers_html"]
    citations_network = main_tab_components["citations_network"]
    papers_summary = main_tab_components["papers_summary"]
    tab_recommended_content = main_tab_components["tab_recommended_content"]
    tab_sources = main_tab_components["tab_sources"]
    tab_figures = main_tab_components["tab_figures"]
    tab_graphs = main_tab_components["tab_graphs"]
    tab_papers = main_tab_components["tab_papers"]
    graphs_container = main_tab_components["graph_container"]
    
    config_open = config_components["config_open"]
    config_modal = config_components["config_modal"]
    dropdown_sources = config_components["dropdown_sources"]
    dropdown_reports = config_components["dropdown_reports"]
    dropdown_external_sources = config_components["dropdown_external_sources"]
    search_only = config_components["search_only"]
    dropdown_audience = config_components["dropdown_audience"]
    after = config_components["after"]
    output_query = config_components["output_query"]
    output_language = config_components["output_language"]
    
    new_sources_hmtl = gr.State([])
    

     
    config_button.click(
        fn=update_config_modal_visibility,
        inputs=[config_open],
        outputs=[config_modal, config_open]
    )
    
    
    (textbox
        .submit(start_chat, [textbox, chatbot, search_only], [textbox, tabs, chatbot, sources_raw], queue=False, api_name=f"start_chat_textbox_{tab_name}")
        .then(chat, [textbox, chatbot, dropdown_audience, dropdown_sources, dropdown_reports, dropdown_external_sources, search_only], [chatbot, new_sources_hmtl, output_query, output_language, new_figures, current_graphs], concurrency_limit=8, api_name="chat_textbox")
        .then(finish_chat, None, [textbox], api_name=f"finish_chat_textbox_{tab_name}")
    )



    (examples_hidden
        .change(start_chat, [examples_hidden, chatbot, search_only], [textbox, tabs, chatbot, sources_raw], queue=False, api_name=f"start_chat_examples_{tab_name}")
        .then(chat, [examples_hidden, chatbot, dropdown_audience, dropdown_sources, dropdown_reports, dropdown_external_sources, search_only], [chatbot, new_sources_hmtl, output_query, output_language, new_figures, current_graphs], concurrency_limit=8, api_name="chat_textbox")
        .then(finish_chat, None, [textbox], api_name=f"finish_chat_examples_{tab_name}")
    )

    new_sources_hmtl.change(lambda x : x, inputs = [new_sources_hmtl], outputs = [sources_textbox])
    current_graphs.change(lambda x: x, inputs=[current_graphs], outputs=[graphs_container])
    new_figures.change(process_figures, inputs=[sources_raw, new_figures], outputs=[sources_raw, figures_cards, gallery_component])

    # Update sources numbers
    sources_textbox.change(update_sources_number_display, [sources_textbox, figures_cards, current_graphs, papers_html], [tab_recommended_content, tab_sources, tab_figures, tab_graphs, tab_papers])
    figures_cards.change(update_sources_number_display, [sources_textbox, figures_cards, current_graphs, papers_html], [tab_recommended_content, tab_sources, tab_figures, tab_graphs, tab_papers])
    current_graphs.change(update_sources_number_display, [sources_textbox, figures_cards, current_graphs, papers_html], [tab_recommended_content, tab_sources, tab_figures, tab_graphs, tab_papers])
    papers_html.change(update_sources_number_display, [sources_textbox, figures_cards, current_graphs, papers_html], [tab_recommended_content, tab_sources, tab_figures, tab_graphs, tab_papers])

    # Other questions examples
    dropdown_samples.change(change_sample_questions, dropdown_samples, samples)

    # Search for papers
    textbox.submit(find_papers, [textbox, after, dropdown_external_sources], [papers_html, citations_network, papers_summary])
    examples_hidden.change(find_papers, [examples_hidden, after, dropdown_external_sources], [papers_html, citations_network, papers_summary])
    

def main_ui():
    config_open = gr.State(False)
    with gr.Blocks(title="Climate Q&A", css_paths=os.getcwd()+ "/style.css", theme=theme, elem_id="main-component") as demo: 
        config_components = create_config_modal(config_open)       
        with gr.Tabs():
            cqa_components = cqa_tab(tab_name = "ClimateQ&A")
            local_cqa_components = cqa_tab(tab_name = "Beta - POC Adapt'Action")
        
            about_tab()
        
        event_handling(cqa_components, config_components, tab_name = 'ClimateQ&A')
        event_handling(local_cqa_components, config_components, tab_name = 'Beta - POC Adapt\'Action')
        demo.queue()
        
    return demo


    

    

# # Main UI Assembly
# with gr.Blocks(title="Climate Q&A", css_paths=os.getcwd()+ "/style.css", theme=theme, elem_id="main-component") as demo:
    
#     # State variables
#     # chat_completed_state = gr.State(0)
#     current_graphs = gr.State([])
#     saved_graphs = gr.State({})
#     new_sources_hmtl = gr.State([])

#     config_open = gr.State(False)

#     with gr.Tab("ClimateQ&A"):
#         with gr.Row(elem_id="chatbot-row"):
#             # Left column - Chat interface
#             with gr.Column(scale=2):
#                 chatbot, textbox, config_button = create_chat_interface()

#             # Right column - Content panels
#             with gr.Column(scale=2, variant="panel", elem_id="right-panel"):
#                 with gr.Tabs(elem_id="right_panel_tab") as tabs:
#                     # Examples tab
#                     with gr.TabItem("Examples", elem_id="tab-examples", id=0):
#                         examples_hidden, dropdown_samples, samples = create_examples_tab()

#                     # Sources tab
#                     with gr.Tab("Sources", elem_id="tab-sources", id=1) as tab_sources:
#                         sources_textbox = gr.HTML(show_label=False, elem_id="sources-textbox")


#                     # Recommended content tab
#                     with gr.Tab("Recommended content", elem_id="tab-recommended_content", id=2) as tab_recommended_content:
#                         with gr.Tabs(elem_id="group-subtabs") as tabs_recommended_content:
#                             # Figures subtab
#                             with gr.Tab("Figures", elem_id="tab-figures", id=3) as tab_figures:
#                                 sources_raw, new_figures, used_figures, gallery_component, figures_cards, figure_modal = create_figures_tab()

#                             # Papers subtab
#                             with gr.Tab("Papers", elem_id="tab-citations", id=4) as tab_papers:
#                                 papers_summary, papers_html, citations_network, papers_modal = create_papers_tab()

#                             # Graphs subtab
#                             with gr.Tab("Graphs", elem_id="tab-graphs", id=5) as tab_graphs:
#                                 graphs_container = gr.HTML(
#                                     "<h2>There are no graphs to be displayed at the moment. Try asking another question.</h2>",
#                                     elem_id="graphs-container"
#                                 )


                    
#     with gr.Tab("About", elem_classes="max-height other-tabs"):
#         with gr.Row():
#             with gr.Column(scale=1):
#                 gr.Markdown(
#                     """
#                     ### More info
#                     - See more info at [https://climateqa.com](https://climateqa.com/docs/intro/)
#                     - Feedbacks on this [form](https://forms.office.com/e/1Yzgxm6jbp)
                                                
#                     ### Citation
#                     """
#                 )
#                 with gr.Accordion(CITATION_LABEL, elem_id="citation", open=False):
#                     gr.Textbox(
#                         value=CITATION_TEXT,
#                         label="",
#                         interactive=False,
#                         show_copy_button=True,
#                         lines=len(CITATION_TEXT.split('\n')),
#                     )
#     # Configuration pannel
#     config_modal, dropdown_sources, dropdown_reports, dropdown_external_sources, search_only, dropdown_audience, after, output_query, output_language = create_config_modal(config_open)
    
#     # Event handlers    
#     config_button.click(
#         fn=update_config_modal_visibility,
#         inputs=[config_open],
#         outputs=[config_modal, config_open]
#     )
    
    
#     (textbox
#         .submit(start_chat, [textbox, chatbot, search_only], [textbox, tabs, chatbot, sources_raw], queue=False, api_name="start_chat_textbox")
#         .then(chat, [textbox, chatbot, dropdown_audience, dropdown_sources, dropdown_reports, dropdown_external_sources, search_only], [chatbot, new_sources_hmtl, output_query, output_language, new_figures, current_graphs], concurrency_limit=8, api_name="chat_textbox")
#         .then(finish_chat, None, [textbox], api_name="finish_chat_textbox")
#     )



#     (examples_hidden
#         .change(start_chat, [examples_hidden, chatbot, search_only], [textbox, tabs, chatbot, sources_raw], queue=False, api_name="start_chat_examples")
#         .then(chat, [examples_hidden, chatbot, dropdown_audience, dropdown_sources, dropdown_reports, dropdown_external_sources, search_only], [chatbot, new_sources_hmtl, output_query, output_language, new_figures, current_graphs], concurrency_limit=8, api_name="chat_textbox")
#         .then(finish_chat, None, [textbox], api_name="finish_chat_examples")
#     )

#     new_sources_hmtl.change(lambda x : x, inputs = [new_sources_hmtl], outputs = [sources_textbox])
#     new_figures.change(process_figures, inputs=[sources_raw, new_figures], outputs=[sources_raw, figures_cards, gallery_component])
#     current_graphs.change(lambda x: x, inputs=[current_graphs], outputs=[graphs_container])
    
#     # Update sources numbers
#     sources_textbox.change(update_sources_number_display, [sources_textbox, figures_cards, current_graphs, papers_html], [tab_recommended_content, tab_sources, tab_figures, tab_graphs, tab_papers])
#     figures_cards.change(update_sources_number_display, [sources_textbox, figures_cards, current_graphs, papers_html], [tab_recommended_content, tab_sources, tab_figures, tab_graphs, tab_papers])
#     current_graphs.change(update_sources_number_display, [sources_textbox, figures_cards, current_graphs, papers_html], [tab_recommended_content, tab_sources, tab_figures, tab_graphs, tab_papers])
#     papers_html.change(update_sources_number_display, [sources_textbox, figures_cards, current_graphs, papers_html], [tab_recommended_content, tab_sources, tab_figures, tab_graphs, tab_papers])

#     # Other questions examples
#     dropdown_samples.change(change_sample_questions, dropdown_samples, samples)

#     # Search for papers
#     textbox.submit(find_papers, [textbox, after, dropdown_external_sources], [papers_html, citations_network, papers_summary])
#     examples_hidden.change(find_papers, [examples_hidden, after, dropdown_external_sources], [papers_html, citations_network, papers_summary])

#     demo.queue()
    
demo = main_ui()
demo.launch(ssr_mode=False)