app.py

# Import necessary libraries
import os
import gradio as gr

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.engine.graph import make_graph_agent,make_graph_agent_poc
from climateqa.engine.chains.retrieve_papers import find_papers
from climateqa.chat import start_chat, chat_stream, finish_chat
from climateqa.engine.talk_to_data.main import ask_drias, DRIAS_MODELS
from climateqa.engine.talk_to_data.myVanna import MyVanna

from front.tabs import (create_config_modal, cqa_tab, create_about_tab)
from front.tabs import (MainTabPanel, ConfigPanel)
from front.utils import process_figures
from gradio_modal import Modal


from utils import create_user_id
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"],
)

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



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

llm = get_llm(provider="openai",max_tokens = 1024,temperature = 0.0)
if os.environ["GRADIO_ENV"] == "local":
    reranker = get_reranker("nano")
else :
    reranker = get_reranker("large")

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

#Vanna object

vn = MyVanna(config = {"temperature": 0, "api_key": os.getenv('THEO_API_KEY'), 'model': os.getenv('VANNA_MODEL'), 'pc_api_key': os.getenv('VANNA_PINECONE_API_KEY'), 'index_name': os.getenv('VANNA_INDEX_NAME'), "top_k" : 4})
db_vanna_path = os.path.join(os.getcwd(), "data/drias/drias.db")
vn.connect_to_sqlite(db_vanna_path)

# def ask_vanna_query(query):
#     return ask_vanna(vn, db_vanna_path, query)

def ask_drias_query(query: str, index_state: int, drias_model: str):
    return ask_drias(db_vanna_path, query, index_state, drias_model)

async def chat(query, history, audience, sources, reports, relevant_content_sources_selection, search_only):
    print("chat cqa - message received")
    async for event in chat_stream(agent, query, history, audience, sources, reports, relevant_content_sources_selection, search_only, share_client, user_id):
        yield event
        
async def chat_poc(query, history, audience, sources, reports, relevant_content_sources_selection, search_only):
    print("chat poc - message received")
    async for event in chat_stream(agent_poc, query, history, audience, sources, reports, relevant_content_sources_selection, search_only, share_client, user_id):
        yield event


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

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

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 create_drias_tab():
#     with gr.Tab("Beta - Talk to DRIAS", elem_id="tab-vanna", id=6) as tab_vanna:
#         vanna_direct_question = gr.Textbox(label="Direct Question", placeholder="You can write direct question here",elem_id="direct-question", interactive=True)
#         with gr.Accordion("Details",elem_id = 'vanna-details', open=False) as vanna_details :
#             vanna_sql_query = gr.Textbox(label="SQL Query Used", elem_id="sql-query", interactive=False)
#             show_vanna_table = gr.Button("Show Table", elem_id="show-table")
#             with Modal(visible=False) as vanna_table_modal:
#                 vanna_table = gr.DataFrame([], elem_id="vanna-table")
#                 close_vanna_modal = gr.Button("Close", elem_id="close-vanna-modal")
#                 close_vanna_modal.click(lambda: Modal(visible=False),None, [vanna_table_modal])
#             show_vanna_table.click(lambda: Modal(visible=True),None ,[vanna_table_modal])

#         vanna_display = gr.Plot()
#         vanna_direct_question.submit(ask_drias_query, [vanna_direct_question], [vanna_sql_query ,vanna_table, vanna_display])

def create_drias_tab():
    with gr.Tab("Beta - Talk to DRIAS", elem_id="tab-vanna", id=6):
        with gr.Row():
            drias_direct_question = gr.Textbox(label="Direct Question", placeholder="You can write direct question here", elem_id="direct-question", interactive=True)
            model_selection = gr.Dropdown(label="Model", choices=DRIAS_MODELS ,elem_id="drias-model", value="ALL", interactive=True)
        
        with gr.Accordion(label="SQL Query Used"):
            drias_sql_query = gr.Textbox(label="", elem_id="sql-query", interactive=False)

        with gr.Accordion(label='Data used', open=False):
            drias_table = gr.DataFrame([], elem_id="vanna-table")

        with gr.Accordion(label="Chart"):
            drias_display = gr.Plot(elem_id="vanna-plot")
        
        with gr.Row():
            prev_button = gr.Button("Previous")
            next_button = gr.Button("Next")

        sql_queries_state = gr.State([])
        dataframes_state = gr.State([])
        plots_state = gr.State([])
        index_state = gr.State(0)

        drias_direct_question.submit(
            ask_drias_query, 
            inputs=[drias_direct_question, index_state, model_selection], 
            outputs=[drias_sql_query, drias_table, drias_display, sql_queries_state, dataframes_state, plots_state, index_state]
        )

        model_selection.change(
            ask_drias_query,
            inputs=[drias_direct_question, index_state, model_selection],
            outputs=[drias_sql_query, drias_table, drias_display, sql_queries_state, dataframes_state, plots_state, index_state]
        )

        def show_previous(index, sql_queries, dataframes, plots):
            if index > 0:
                index -= 1
            return sql_queries[index], dataframes[index], plots[index], index

        def show_next(index, sql_queries, dataframes, plots):
            if index < len(sql_queries) - 1:
                index += 1
            return sql_queries[index], dataframes[index], plots[index], index

        prev_button.click(
            show_previous, 
            inputs=[index_state, sql_queries_state, dataframes_state, plots_state],
            outputs=[drias_sql_query, drias_table, drias_display, index_state]
        )

        next_button.click(
            show_next, 
            inputs=[index_state, sql_queries_state, dataframes_state, plots_state],
            outputs=[drias_sql_query, drias_table, drias_display, index_state]
        )
    
# # UI Layout Components
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(tab_name)

            # 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 = create_examples_tab(tab_name)

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

                                
def config_event_handling(main_tabs_components : list[MainTabPanel], config_componenets : ConfigPanel):
    config_open = config_componenets.config_open
    config_modal = config_componenets.config_modal
    close_config_modal = config_componenets.close_config_modal_button
    
    for button in [close_config_modal] + [main_tab_component.config_button for main_tab_component in main_tabs_components]:
        button.click(
            fn=update_config_modal_visibility,
            inputs=[config_open],
            outputs=[config_modal, config_open]
        ) 
    
def event_handling(
    main_tab_components : MainTabPanel,
    config_components : ConfigPanel,
    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
    sources_textbox = main_tab_components.sources_textbox
    figures_cards = main_tab_components.figures_cards
    gallery_component = main_tab_components.gallery_component
    papers_direct_search = main_tab_components.papers_direct_search
    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
    follow_up_examples = main_tab_components.follow_up_examples
    follow_up_examples_hidden = main_tab_components.follow_up_examples_hidden
    
    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([])
    ttd_data = gr.State([])

    
    if tab_name == "ClimateQ&A":
        print("chat cqa - message sent")

        # Event for textbox
        (textbox
            .submit(start_chat, [textbox, chatbot, search_only], [textbox, tabs, chatbot, sources_raw], queue=False, api_name=f"start_chat_{textbox.elem_id}")
            .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, follow_up_examples.dataset], concurrency_limit=8, api_name=f"chat_{textbox.elem_id}")
            .then(finish_chat, None, [textbox], api_name=f"finish_chat_{textbox.elem_id}")
        )
        # Event for examples_hidden
        (examples_hidden
            .change(start_chat, [examples_hidden, chatbot, search_only], [examples_hidden, tabs, chatbot, sources_raw], queue=False, api_name=f"start_chat_{examples_hidden.elem_id}")
            .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,follow_up_examples.dataset], concurrency_limit=8, api_name=f"chat_{examples_hidden.elem_id}")
            .then(finish_chat, None, [textbox], api_name=f"finish_chat_{examples_hidden.elem_id}")
        )
        (follow_up_examples_hidden
            .change(start_chat, [follow_up_examples_hidden, chatbot, search_only], [follow_up_examples_hidden, tabs, chatbot, sources_raw], queue=False, api_name=f"start_chat_{examples_hidden.elem_id}")
            .then(chat, [follow_up_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,follow_up_examples.dataset], concurrency_limit=8, api_name=f"chat_{examples_hidden.elem_id}")
            .then(finish_chat, None, [textbox], api_name=f"finish_chat_{follow_up_examples_hidden.elem_id}")
        )
    
    elif tab_name == "Beta - POC Adapt'Action":
        print("chat poc - message sent")
        # Event for textbox
        (textbox
            .submit(start_chat, [textbox, chatbot, search_only], [textbox, tabs, chatbot, sources_raw], queue=False, api_name=f"start_chat_{textbox.elem_id}")
            .then(chat_poc, [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=f"chat_{textbox.elem_id}")
            .then(finish_chat, None, [textbox], api_name=f"finish_chat_{textbox.elem_id}")
        )
        # Event for examples_hidden
        (examples_hidden
            .change(start_chat, [examples_hidden, chatbot, search_only], [examples_hidden, tabs, chatbot, sources_raw], queue=False, api_name=f"start_chat_{examples_hidden.elem_id}")
            .then(chat_poc, [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=f"chat_{examples_hidden.elem_id}")
            .then(finish_chat, None, [textbox], api_name=f"finish_chat_{examples_hidden.elem_id}")
        )
        (follow_up_examples_hidden
            .change(start_chat, [follow_up_examples_hidden, chatbot, search_only], [follow_up_examples_hidden, tabs, chatbot, sources_raw], queue=False, api_name=f"start_chat_{examples_hidden.elem_id}")
            .then(chat, [follow_up_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,follow_up_examples.dataset], concurrency_limit=8, api_name=f"chat_{examples_hidden.elem_id}")
            .then(finish_chat, None, [textbox], api_name=f"finish_chat_{follow_up_examples_hidden.elem_id}")
        )
    
    
    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
    for component in [sources_textbox, figures_cards, current_graphs, papers_html]:
        component.change(update_sources_number_display, [sources_textbox, figures_cards, current_graphs, papers_html], [tab_recommended_content, tab_sources, tab_figures, tab_graphs, tab_papers])
    
    # Search for papers
    for component in [textbox, examples_hidden, papers_direct_search]:
        component.submit(find_papers, [component, after, dropdown_external_sources], [papers_html, citations_network, papers_summary])
        

    # if tab_name == "Beta - POC Adapt'Action": # Not untill results are good enough
    #     # Drias search
    #     textbox.submit(ask_vanna, [textbox], [vanna_sql_query ,vanna_table, vanna_display])

def main_ui():
    # config_open = gr.State(True)
    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()  
             
        with gr.Tabs():
            cqa_components = cqa_tab(tab_name = "ClimateQ&A")
            local_cqa_components = cqa_tab(tab_name = "Beta - POC Adapt'Action")
            create_drias_tab()
        
            create_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")
        
        config_event_handling([cqa_components,local_cqa_components] ,config_components)
        
        demo.queue()
        
    return demo

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