Merge branch 'dev' of https://bitbucket.org/ekimetrics/climate_qa into dev

app.py

@@ -15,7 +15,8 @@ from climateqa.chat import start_chat, chat_stream, finish_chat
 from climateqa.engine.talk_to_data.main import ask_vanna
 from climateqa.engine.talk_to_data.myVanna import MyVanna
 
-from front.tabs import (create_config_modal, create_examples_tab, create_papers_tab, create_figures_tab, create_chat_interface, create_about_tab)
+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
 
@@ -136,81 +137,14 @@ def create_drias_tab():
 
         vanna_display = gr.Plot()
         vanna_direct_question.submit(ask_vanna_query, [vanna_direct_question], [vanna_sql_query ,vanna_table, vanna_display])
-    
-# # 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"
-                                )
-
-                                
-    return {
-        "chatbot": chatbot,
-        "textbox": textbox,
-        "tabs": tabs,
-        "sources_raw": sources_raw,
-        "new_figures": new_figures,
-        "current_graphs": current_graphs,
-        "examples_hidden": examples_hidden,
-        "sources_textbox": sources_textbox,
-        "figures_cards": figures_cards,
-        "gallery_component": gallery_component,
-        "config_button": config_button,
-        "papers_direct_search" : papers_direct_search,
-        "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,
-        # "vanna_sql_query": vanna_sql_query,
-        # "vanna_table" : vanna_table,
-        # "vanna_display": vanna_display
-    }
+  
                                 
-def config_event_handling(main_tabs_components : list[dict], config_componenets : dict):
-    config_open = config_componenets["config_open"]
-    config_modal = config_componenets["config_modal"]
-    close_config_modal = config_componenets["close_config_modal_button"]
+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]:
+    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],
@@ -218,58 +152,45 @@ def config_event_handling(main_tabs_components : list[dict], config_componenets
         ) 
     
 def event_handling(
-    main_tab_components,
-    config_components,
+    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"]
-    # config_button = main_tab_components["config_button"]
-    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"]
-    # vanna_sql_query = main_tab_components["vanna_sql_query"]
-    # vanna_table = main_tab_components["vanna_table"]
-    # vanna_display = main_tab_components["vanna_display"]
+    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
     
-    
-    # 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"]
-    # close_config_modal = config_components["close_config_modal_button"]
+    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([])
 
-     
-    # for button in [config_button, close_config_modal]:
-    #     button.click(
-    #         fn=update_config_modal_visibility,
-    #         inputs=[config_open],
-    #         outputs=[config_modal, config_open]
-    #     )
     
     if tab_name == "ClimateQ&A":
         print("chat cqa - message sent")
@@ -277,15 +198,20 @@ def event_handling(
         # 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], concurrency_limit=8, api_name=f"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], concurrency_limit=8, api_name=f"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")
@@ -301,6 +227,11 @@ def event_handling(
             .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])

climateqa/chat.py

@@ -101,6 +101,7 @@ async def chat_stream(
     audience_prompt = init_audience(audience)
     sources = sources or ["IPCC", "IPBES"]
     reports = reports or []
+    relevant_history_discussion = history[-2:] if len(history) > 1 else []
 
     # Prepare inputs for agent
     inputs = {
@@ -109,7 +110,8 @@ async def chat_stream(
         "sources_input": sources,
         "relevant_content_sources_selection": relevant_content_sources_selection,
         "search_only": search_only,
-        "reports": reports
+        "reports": reports,
+        "chat_history": relevant_history_discussion,
     }
 
     # Get streaming events from agent
@@ -129,6 +131,7 @@ async def chat_stream(
     retrieved_contents = []
     answer_message_content = ""
     vanna_data = {}
+    follow_up_examples = gr.Dataset(samples=[])
 
     # Define processing steps
     steps_display = {
@@ -200,7 +203,12 @@ async def chat_stream(
                         sub_questions = [q["question"] + "-> relevant sources : " + str(q["sources"]) 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#, vanna_data
+                # Handle follow up questions
+                if event["name"] == "generate_follow_up" and event["event"] == "on_chain_end": 
+                    follow_up_examples = event["data"]["output"].get("follow_up_questions", [])
+                    follow_up_examples = gr.Dataset(samples= [ [question] for question in follow_up_examples ])
+
+            yield history, docs_html, output_query, output_language, related_contents, graphs_html, follow_up_examples#, vanna_data
 
     except Exception as e:
         print(f"Event {event} has failed")
@@ -211,4 +219,4 @@ async def chat_stream(
     # Call the function to log interaction
     log_interaction_to_azure(history, output_query, sources, docs, share_client, user_id)
 
-    yield history, docs_html, output_query, output_language, related_contents, graphs_html#, vanna_data
+    yield history, docs_html, output_query, output_language, related_contents, graphs_html, follow_up_examples#, vanna_data

climateqa/engine/chains/answer_rag.py

@@ -65,6 +65,7 @@ def make_rag_node(llm,with_docs = True):
     async def answer_rag(state,config):
         print("---- Answer RAG ----")
         start_time = time.time()
+        chat_history = state.get("chat_history",[])
         print("Sources used : " +  "\n".join([x.metadata["short_name"] + " - page " + str(x.metadata["page_number"])  for x in state["documents"]]))
 
         answer = await rag_chain.ainvoke(state,config)
@@ -73,9 +74,10 @@ def make_rag_node(llm,with_docs = True):
         elapsed_time = end_time - start_time
         print("RAG elapsed time: ", elapsed_time)
         print("Answer size : ", len(answer))
-        # print(f"\n\nAnswer:\n{answer}")
         
-        return {"answer":answer}
+        chat_history.append({"question":state["query"],"answer":answer})
+        
+        return {"answer":answer,"chat_history": chat_history}
 
     return answer_rag
 

climateqa/engine/chains/follow_up.py

@@ -0,0 +1,32 @@
+from typing import List
+from langchain.prompts import ChatPromptTemplate
+
+
+FOLLOW_UP_TEMPLATE = """Based on the previous question and answer, generate 2-3 relevant follow-up questions that would help explore the topic further.
+
+Previous Question: {user_input}
+Previous Answer: {answer}
+
+Generate short, concise, focused follow-up questions
+You don't need a full question as it will be reformulated later as a standalone question with the context. Eg. "Details the first point"
+"""
+
+def make_follow_up_node(llm):
+    prompt = ChatPromptTemplate.from_template(FOLLOW_UP_TEMPLATE)
+    
+    def generate_follow_up(state):
+        if not state.get("answer"):
+            return state
+            
+        response = llm.invoke(prompt.format(
+            user_input=state["user_input"],
+            answer=state["answer"]
+        ))
+        
+        # Extract questions from response
+        follow_ups = [q.strip() for q in response.content.split("\n") if q.strip()]
+        state["follow_up_questions"] = follow_ups
+        
+        return state
+        
+    return generate_follow_up

climateqa/engine/chains/standalone_question.py

@@ -0,0 +1,39 @@
+from langchain.prompts import ChatPromptTemplate
+
+def make_standalone_question_chain(llm):
+    prompt = ChatPromptTemplate.from_messages([
+        ("system", """You are a helpful assistant that transforms user questions into standalone questions 
+        by incorporating context from the chat history if needed. The output should be a self-contained 
+        question that can be understood without any additional context.
+        
+        Examples:
+        Chat History: "Let's talk about renewable energy"
+        User Input: "What about solar?"
+        Output: "What are the key aspects of solar energy as a renewable energy source?"
+        
+        Chat History: "What causes global warming?"
+        User Input: "And what are its effects?"
+        Output: "What are the effects of global warming on the environment and society?"
+        """),
+        ("user", """Chat History: {chat_history}
+        User Question: {question}
+        
+        Transform this into a standalone question:""")
+    ])
+    
+    chain = prompt | llm
+    return chain
+
+def make_standalone_question_node(llm):
+    standalone_chain = make_standalone_question_chain(llm)
+    
+    def transform_to_standalone(state):
+        chat_history = state.get("chat_history", "")
+        output = standalone_chain.invoke({
+            "chat_history": chat_history,
+            "question": state["user_input"]
+        })
+        state["user_input"] = output.content
+        return state
+        
+    return transform_to_standalone

climateqa/engine/graph.py

@@ -23,13 +23,15 @@ from .chains.retrieve_documents import make_IPx_retriever_node, make_POC_retriev
 from .chains.answer_rag import make_rag_node
 from .chains.graph_retriever import make_graph_retriever_node
 from .chains.chitchat_categorization import make_chitchat_intent_categorization_node
-# from .chains.set_defaults import set_defaults
+from .chains.standalone_question import make_standalone_question_node
+from .chains.follow_up import make_follow_up_node  # Add this import
 
 class GraphState(TypedDict):
     """
     Represents the state of our graph.
     """
     user_input : str
+    chat_history : str
     language : str
     intent : str
     search_graphs_chitchat : bool
@@ -49,6 +51,7 @@ class GraphState(TypedDict):
     recommended_content : List[Document] # OWID Graphs  # TODO merge with related_contents
     search_only : bool = False
     reports : List[str] = []
+    follow_up_questions: List[str] = []
 
 def dummy(state):
     return 
@@ -111,6 +114,11 @@ def route_retrieve_documents(state):
         return END
     return sources_to_retrieve
 
+def route_follow_up(state):
+    if state["follow_up_questions"]:
+        return "process_follow_up"
+    return END
+
 def make_id_dict(values):
     return {k:k for k in values}
 
@@ -119,6 +127,7 @@ def make_graph_agent(llm, vectorstore_ipcc, vectorstore_graphs, vectorstore_regi
     workflow = StateGraph(GraphState)
 
     # Define the node functions
+    standalone_question_node = make_standalone_question_node(llm)
     categorize_intent = make_intent_categorization_node(llm)
     transform_query = make_query_transform_node(llm)
     translate_query = make_translation_node(llm)
@@ -130,9 +139,11 @@ def make_graph_agent(llm, vectorstore_ipcc, vectorstore_graphs, vectorstore_regi
     answer_rag = make_rag_node(llm, with_docs=True)
     answer_rag_no_docs = make_rag_node(llm, with_docs=False)
     chitchat_categorize_intent = make_chitchat_intent_categorization_node(llm)
+    generate_follow_up = make_follow_up_node(llm)
 
     # Define the nodes
     # workflow.add_node("set_defaults", set_defaults)
+    workflow.add_node("standalone_question", standalone_question_node)
     workflow.add_node("categorize_intent", categorize_intent)
     workflow.add_node("answer_climate", dummy)
     workflow.add_node("answer_search", answer_search)
@@ -146,9 +157,11 @@ def make_graph_agent(llm, vectorstore_ipcc, vectorstore_graphs, vectorstore_regi
     workflow.add_node("retrieve_documents", retrieve_documents)
     workflow.add_node("answer_rag", answer_rag)
     workflow.add_node("answer_rag_no_docs", answer_rag_no_docs)
+    workflow.add_node("generate_follow_up", generate_follow_up)
+    # workflow.add_node("process_follow_up", standalone_question_node)
 
     # Entry point
-    workflow.set_entry_point("categorize_intent")
+    workflow.set_entry_point("standalone_question")
 
     # CONDITIONAL EDGES
     workflow.add_conditional_edges(
@@ -180,20 +193,29 @@ def make_graph_agent(llm, vectorstore_ipcc, vectorstore_graphs, vectorstore_regi
         make_id_dict(["retrieve_graphs", END])
     )
 
+    # workflow.add_conditional_edges(
+    #     "generate_follow_up",
+    #     route_follow_up,
+    #     make_id_dict(["process_follow_up", END])
+    # )
+
     # Define the edges
+    workflow.add_edge("standalone_question", "categorize_intent")
     workflow.add_edge("translate_query", "transform_query")
     workflow.add_edge("transform_query", "retrieve_documents") #TODO put back
     # workflow.add_edge("transform_query", "retrieve_local_data")
     # workflow.add_edge("transform_query", END) # TODO remove
 
     workflow.add_edge("retrieve_graphs", END)
-    workflow.add_edge("answer_rag", END)
-    workflow.add_edge("answer_rag_no_docs", END)
+    workflow.add_edge("answer_rag", "generate_follow_up")
+    workflow.add_edge("answer_rag_no_docs", "generate_follow_up")
     workflow.add_edge("answer_chitchat", "chitchat_categorize_intent")
     workflow.add_edge("retrieve_graphs_chitchat", END)
 
     # workflow.add_edge("retrieve_local_data", "answer_search")
     workflow.add_edge("retrieve_documents", "answer_search")
+    workflow.add_edge("generate_follow_up",END)
+    # workflow.add_edge("process_follow_up", "categorize_intent")
 
     # Compile
     app = workflow.compile()
@@ -219,6 +241,8 @@ def make_graph_agent_poc(llm, vectorstore_ipcc, vectorstore_graphs, vectorstore_
     workflow = StateGraph(GraphState)
 
     # Define the node functions
+    standalone_question_node = make_standalone_question_node(llm)
+
     categorize_intent = make_intent_categorization_node(llm)
     transform_query = make_query_transform_node(llm)
     translate_query = make_translation_node(llm)
@@ -231,9 +255,11 @@ def make_graph_agent_poc(llm, vectorstore_ipcc, vectorstore_graphs, vectorstore_
     answer_rag = make_rag_node(llm, with_docs=True)
     answer_rag_no_docs = make_rag_node(llm, with_docs=False)
     chitchat_categorize_intent = make_chitchat_intent_categorization_node(llm)
+    generate_follow_up = make_follow_up_node(llm)
 
     # Define the nodes
     # workflow.add_node("set_defaults", set_defaults)
+    workflow.add_node("standalone_question", standalone_question_node)
     workflow.add_node("categorize_intent", categorize_intent)
     workflow.add_node("answer_climate", dummy)
     workflow.add_node("answer_search", answer_search)
@@ -249,9 +275,10 @@ def make_graph_agent_poc(llm, vectorstore_ipcc, vectorstore_graphs, vectorstore_
     workflow.add_node("retrieve_documents", retrieve_documents)
     workflow.add_node("answer_rag", answer_rag)
     workflow.add_node("answer_rag_no_docs", answer_rag_no_docs)
+    workflow.add_node("generate_follow_up", generate_follow_up)
 
     # Entry point
-    workflow.set_entry_point("categorize_intent")
+    workflow.set_entry_point("standalone_question")
 
     # CONDITIONAL EDGES
     workflow.add_conditional_edges(
@@ -284,22 +311,21 @@ def make_graph_agent_poc(llm, vectorstore_ipcc, vectorstore_graphs, vectorstore_
     )
 
     # Define the edges
+    workflow.add_edge("standalone_question", "categorize_intent")
     workflow.add_edge("translate_query", "transform_query")
     workflow.add_edge("transform_query", "retrieve_documents") #TODO put back
     workflow.add_edge("transform_query", "retrieve_local_data")
     # workflow.add_edge("transform_query", END) # TODO remove
 
     workflow.add_edge("retrieve_graphs", END)
-    workflow.add_edge("answer_rag", END)
-    workflow.add_edge("answer_rag_no_docs", END)
+    workflow.add_edge("answer_rag", "generate_follow_up")
+    workflow.add_edge("answer_rag_no_docs", "generate_follow_up")
     workflow.add_edge("answer_chitchat", "chitchat_categorize_intent")
     workflow.add_edge("retrieve_graphs_chitchat", END)
 
     workflow.add_edge("retrieve_local_data", "answer_search")
     workflow.add_edge("retrieve_documents", "answer_search")
-
-    # workflow.add_edge("transform_query", "retrieve_drias_data")
-    # workflow.add_edge("retrieve_drias_data", END)
+    workflow.add_edge("generate_follow_up",END)
 
 
     # Compile

front/tabs/__init__.py

@@ -3,4 +3,7 @@ from .tab_examples import create_examples_tab
 from .tab_papers import create_papers_tab
 from .tab_figures import create_figures_tab
 from .chat_interface import create_chat_interface
-from .tab_about import create_about_tab
+from .tab_about import create_about_tab
+from .main_tab import MainTabPanel
+from .tab_config import ConfigPanel
+from .main_tab import cqa_tab

front/tabs/chat_interface.py

@@ -54,7 +54,10 @@ def create_chat_interface(tab):
         max_height="80vh",
         height="100vh"
     )
-    
+    with gr.Row(elem_id="follow-up-examples"):
+        follow_up_examples_hidden = gr.Textbox(visible=False, elem_id="follow-up-hidden")
+        follow_up_examples = gr.Examples(examples=[ ], label="Follow up questions", inputs= [follow_up_examples_hidden], elem_id="follow-up-button", run_on_click=False)
+
     with gr.Row(elem_id="input-message"):
         
         textbox = gr.Textbox(
@@ -68,7 +71,7 @@ def create_chat_interface(tab):
         
         config_button = gr.Button("", elem_id="config-button")
     
-    return chatbot, textbox, config_button
+    return chatbot, textbox, config_button, follow_up_examples, follow_up_examples_hidden
 
 
         

front/tabs/main_tab.py

@@ -1,8 +1,37 @@
 import gradio as gr
+from gradio.helpers import Examples
+from typing import TypedDict
 from .chat_interface import create_chat_interface
 from .tab_examples import create_examples_tab
 from .tab_papers import create_papers_tab
 from .tab_figures import create_figures_tab
+from dataclasses import dataclass
+
+@dataclass
+class MainTabPanel:
+    chatbot: gr.Chatbot
+    textbox: gr.Textbox
+    tabs: gr.Tabs
+    sources_raw: gr.State
+    new_figures: gr.State
+    current_graphs: gr.State
+    examples_hidden: gr.State
+    sources_textbox: gr.HTML
+    figures_cards: gr.HTML
+    gallery_component: gr.Gallery
+    config_button: gr.Button
+    papers_direct_search: gr.TextArea
+    papers_html: gr.HTML
+    citations_network: gr.Plot
+    papers_summary: gr.Textbox
+    tab_recommended_content: gr.Tab
+    tab_sources: gr.Tab
+    tab_figures: gr.Tab
+    tab_graphs: gr.Tab
+    tab_papers: gr.Tab
+    graph_container: gr.HTML
+    follow_up_examples : Examples
+    follow_up_examples_hidden : gr.Textbox
 
 def cqa_tab(tab_name):
     # State variables
@@ -11,14 +40,14 @@ def cqa_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)
+                chatbot, textbox, config_button, follow_up_examples, follow_up_examples_hidden = 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, dropdown_samples, samples = create_examples_tab()
+                        examples_hidden = create_examples_tab(tab_name)
 
                     # Sources tab
                     with gr.Tab("Sources", elem_id="tab-sources", id=1) as tab_sources:
@@ -34,7 +63,7 @@ def cqa_tab(tab_name):
 
                             # 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()
+                                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:
@@ -42,27 +71,30 @@ def cqa_tab(tab_name):
                                     "<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
-    }
+
+                                
+    return MainTabPanel(
+        chatbot=chatbot,
+        textbox=textbox,
+        tabs=tabs,
+        sources_raw=sources_raw,
+        new_figures=new_figures,
+        current_graphs=current_graphs,
+        examples_hidden=examples_hidden,
+        sources_textbox=sources_textbox,
+        figures_cards=figures_cards,
+        gallery_component=gallery_component,
+        config_button=config_button,
+        papers_direct_search=papers_direct_search,
+        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,
+        follow_up_examples= follow_up_examples,
+        follow_up_examples_hidden = follow_up_examples_hidden
+    )

front/tabs/tab_config.py

@@ -2,8 +2,10 @@ import gradio as gr
 from gradio_modal import Modal
 from climateqa.constants import POSSIBLE_REPORTS
 from typing import TypedDict
+from dataclasses import dataclass
 
-class ConfigPanel(TypedDict):
+@dataclass
+class ConfigPanel:
     config_open: gr.State
     config_modal: Modal
     dropdown_sources: gr.CheckboxGroup
@@ -14,6 +16,7 @@ class ConfigPanel(TypedDict):
     after: gr.Slider
     output_query: gr.Textbox
     output_language: gr.Textbox
+    close_config_modal_button: gr.Button
     
 
 def create_config_modal():
@@ -95,29 +98,16 @@ def create_config_modal():
 
         close_config_modal_button = gr.Button("Validate and Close", elem_id="close-config-modal")
 
-        
-        # return ConfigPanel(
-        #     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
-        # )
-        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,
-            "close_config_modal_button": close_config_modal_button
-        }
+        return ConfigPanel(
+            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,
+            close_config_modal_button=close_config_modal_button
+        )

style.css

@@ -115,6 +115,11 @@ main.flex.flex-1.flex-col {
     border-radius: 40px;
     padding-left: 30px;
     resize: none;
+    background-color: #f0f8ff; /* Light blue background */
+    border: 2px solid #4b8ec3; /* Blue border */
+    font-size: 16px; /* Increase font size */
+    color: #333; /* Text color */
+    box-shadow: 0 4px 8px rgba(0, 0, 0, 0.1); /* Add shadow */
 }
 
 #input-message > div {
@@ -479,11 +484,6 @@ a {
     height: 15vh;
     overflow-y: auto;
     padding: 10px 0;
-    display: none; /* Hide by default */
-}
-
-#follow-up-examples:not(:empty) {
-    display: block; /* Show only when there's content */
 }
 
 #follow-up-button {
@@ -518,17 +518,14 @@ a {
     }
 
     div#chatbot {
-        height: 80vh !important; /* Increased height when no follow-ups */
-        max-height: 80vh !important;
-        transition: height 0.3s ease;
-    }
-
-    /* Adjust chatbot height when follow-up examples are present */
-    div#chatbot:has(+ #follow-up-examples:not(:empty)) {
         height: 65vh !important;
         max-height: 65vh !important;
     }
 
+    div#chatbot-row {
+        max-height: calc(100vh - 90px) !important;
+    }
+
     div#graphs-container {
         height: calc(100vh - 210px) !important;
         overflow-y: scroll !important;
@@ -543,13 +540,7 @@ a {
 /* Mobile Media Query */
 @media screen and (max-width: 767px) {
     div#chatbot {
-        height: 500px !important; /* Increased height when no follow-ups */
-        transition: height 0.3s ease;
-    }
-
-    /* Adjust chatbot height when follow-up examples are present */
-    div#chatbot:has(+ #follow-up-examples:not(:empty)) {
-        height: 400px !important;
+        height: 400px !important;  /* Reduced from 500px */
     }
 
     #follow-up-examples {