small cleaning

app.py

@@ -115,14 +115,13 @@ def update_config_modal_visibility(config_open):
     return gr.update(visible=new_config_visibility_status), new_config_visibility_status
 
 # Main chat function
-# async def chat(query, history, audience, sources, reports, relevant_content_sources, search_only):
 async def chat(
     query: str, 
     history: list[ChatMessage], 
     audience: str, 
     sources: list[str], 
     reports: list[str], 
-    relevant_content_sources: 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.
@@ -133,7 +132,7 @@ async def chat(
         audience (str): Target audience type
         sources (list): Knowledge base sources to search
         reports (list): Specific reports to search within sources
-        relevant_content_sources (list): Types of content to retrieve (figures, papers, etc)
+        relevant_content_sources_selection (list): Types of content to retrieve (figures, papers, etc)
         search_only (bool): Whether to only search without generating answer
         
     Yields:
@@ -158,7 +157,7 @@ async def chat(
         "user_input": query,
         "audience": audience_prompt,
         "sources_input": sources,
-        "relevant_content_sources": relevant_content_sources,
+        "relevant_content_sources_selection": relevant_content_sources_selection,
         "search_only": search_only,
         "reports": reports
     }
@@ -168,7 +167,6 @@ async def chat(
 
     # Initialize state variables
     docs = []
-    used_figures = []
     related_contents = []
     docs_html = ""
     output_query = ""
@@ -176,7 +174,6 @@ async def chat(
     output_keywords = ""
     start_streaming = False
     graphs_html = ""    
-    figures = '<div class="figures-container"><p></p> </div>'
     used_documents = []
     answer_message_content = ""
 
@@ -236,7 +233,7 @@ async def chat(
                         sub_questions = [q["question"] for q in event["data"]["output"]["remaining_questions"]]
                         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 #,output_query,output_keywords
+            yield history, docs_html, output_query, output_language, related_contents, graphs_html 
 
     except Exception as e:
         print(f"Event {event} has failed")

climateqa/engine/chains/retrieve_documents.py

@@ -213,20 +213,10 @@ async def retrieve_documents(state,config, vectorstore,reranker,llm,rerank_by_qu
         dict: The updated state containing the retrieved and reranked documents, related content, and remaining questions.
     """
     print("---- Retrieve documents ----")
+    docs = state.get("documents", [])
+    related_content = state.get("related_content", [])
     
-    # Get the documents from the state
-    if "documents" in state and state["documents"] is not None:
-        docs = state["documents"]
-    else:
-        docs = []
-        
-    # Get the related_content from the state
-    if "related_content" in state and state["related_content"] is not None:
-        related_content = state["related_content"]
-    else:
-        related_content = []
-    
-    search_figures = "Figures (IPCC/IPBES)" in state["relevant_content_sources"]
+    search_figures = "Figures (IPCC/IPBES)" in state["relevant_content_sources_selection"]
     search_only = state["search_only"]
     
     reports = state["reports"]

climateqa/engine/chains/retrieve_papers.py

@@ -32,8 +32,8 @@ def generate_keywords(query):
     return keywords
 
 
-async def find_papers(query,after, relevant_content_sources, reranker= reranker):
-    if "Papers (OpenAlex)" in relevant_content_sources:
+async def find_papers(query,after, relevant_content_sources_selection, reranker= reranker):
+    if "Papers (OpenAlex)" in relevant_content_sources_selection:
         summary = ""
         keywords = generate_keywords(query)
         df_works = oa.search(keywords,after = after)

climateqa/engine/graph.py

@@ -36,12 +36,12 @@ class GraphState(TypedDict):
     answer: str
     audience: str = "experts"
     sources_input: List[str] = ["IPCC","IPBES"]
-    relevant_content_sources: List[str] = ["Figures (IPCC/IPBES)"]
+    relevant_content_sources_selection: List[str] = ["Figures (IPCC/IPBES)"]
     sources_auto: bool = True
     min_year: int = 1960
     max_year: int = None
     documents: List[Document]
-    related_contents : Dict[str,Document]
+    related_contents : List[Document]
     recommended_content : List[Document]
     search_only : bool = False
     reports : List[str] = []
@@ -159,7 +159,7 @@ def make_graph_agent(llm, vectorstore_ipcc, vectorstore_graphs, reranker, thresh
     )
     workflow.add_conditional_edges(
         "transform_query", 
-        lambda state : "retrieve_graphs" if "Graphs (OurWorldInData)" in state["relevant_content_sources"]  else END,
+        lambda state : "retrieve_graphs" if "Graphs (OurWorldInData)" in state["relevant_content_sources_selection"]  else END,
         make_id_dict(["retrieve_graphs", END])
     )